diff --git a/js/snapshot.js b/js/snapshot.js
index ea6bad90..4b765969 100644
--- a/js/snapshot.js
+++ b/js/snapshot.js
@@ -1,9 +1,16 @@
 import { app } from "../../scripts/app.js";
 import { api } from "../../scripts/api.js"
-import { ComfyDialog, $el } from "../../scripts/ui.js";
 import { manager_instance, rebootAPI, show_message, handle403Response, loadCss } from  "./common.js";
 import { buildGuiFrame } from "./comfyui-gui-builder.js";
 
+// === SHIM FOR NEW COMFYUI (removes ui.js warning) ===
+let ComfyDialog, $el;
+if (window?.comfyAPI?.ui) {
+    ({ ComfyDialog, $el } = window.comfyAPI.ui);
+} else {
+    ({ ComfyDialog, $el } = await import("../../scripts/ui.js"));
+}
+
 loadCss("./snapshot.css");
 
 async function restore_snapshot(target) {