Initial commit

2026-07-17 04:02:11 +08:00
parent bc211dbbbf
commit bdf506b94e
4 changed files with 401 additions and 0 deletions
@@ -0,0 +1,5 @@
+from .node.sampler import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+WEB_DIRECTORY = "./web"
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
@@ -0,0 +1,57 @@
+import io
+import base64
+import time
+
+import torch.nn.functional as tnnf
+import comfy.utils
+from server import PromptServer
+from PIL import Image
+
+def scale_latent(latent, max_pixels=1000000):
+    h = latent.shape[2]
+    w = latent.shape[3]
+    pixels = h * 8 * w * 8
+    if pixels <= max_pixels:
+        return latent
+    scale = (max_pixels / pixels) ** 0.5
+    new_h = int(h * scale)
+    new_w = int(w * scale)
+    new_h = max(1,new_h)
+    new_w = max(1,new_w)
+    return tnnf.interpolate(latent, size = (new_h, new_w), mode = "bilinear",
+    align_corners = False)
+
+def tensor_to_base64(image):
+    # Receive a vae image and change it to base64, then send it to front-end.
+    image = image.cpu().clamp(0, 1)
+    image = (image[0].numpy() * 255).astype("uint8")
+    img = Image.fromarray(image)
+    buffer = io.BytesIO()
+    img.save(buffer, format = "PNG")
+    return base64.b64encode(buffer.getvalue()).decode()
+
+def create_preview_callback(model, vae, steps, preview_decode = "Scale", interval = 1, sync_preview = True):
+    # Sampler's callback, it has two mode.
+    # Sync: The HD preview will sync with latent's step, one step = one preview image.
+    # Fast: The HD preview will be rendered once at intervals.
+    pbar = comfy.utils.ProgressBar(steps)
+    def inner_callback(step, x0, x, total_steps):
+        if(steps % interval == 0 or steps == total_steps):
+            latent = (model.model.latent_format.process_out(x0))
+            if preview_decode == "Scale":
+                latent = scale_latent(latent)
+            image = vae.decode(latent)
+            data = tensor_to_base64(image)
+            PromptServer.instance.send_sync(
+            "high_preview", {
+                "image": data,
+                "step": step + 1,
+                "total":total_steps
+                }
+            )
+            if sync_preview:
+                # Give browser some time to flash image.
+                time.sleep(0.15)
+        # Not use default preview window, so there is None.
+        pbar.update_absolute(step + 1, total_steps, None)
+    return inner_callback
@@ -0,0 +1,52 @@
+import comfy.samplers
+import comfy.sample
+
+from .hd_preview import create_preview_callback
+
+def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative,
+latent, preview_decode, vae, preview_mode,denoise = 1.0):
+    latent_image = latent["samples"]
+    latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image)
+    noise = comfy.sample.prepare_noise(latent_image, seed, None)
+
+    if preview_mode == "Sync":
+        callback = create_preview_callback(model, vae, steps, preview_decode, interval = 1, sync_preview = True)
+    else:
+        callback = create_preview_callback(model, vae, steps, preview_decode, interval = 5, sync_preview = False)
+    samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler,
+    positive, negative, latent_image, denoise = denoise, callback = callback, seed = seed)
+    out = latent.copy()
+    out["samples"] = samples
+    return (out, )
+
+class HighPreviewKSampler:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {
+            "model": ("MODEL", ),
+            "seed": ("INT", {"default" : 0}),
+            "steps": ("INT", {"default": 20}),
+            "cfg": ("FLOAT", {"default": 8}),
+            "sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
+            "scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
+            "positive": ("CONDITIONING", ),
+            "negative": ("CONDITIONING", ),
+            "latent_image": ("LATENT", ),
+            "vae": ("VAE", ),
+            "preview_mode": (["Sync", "Fast"], {"default": "Sync"}),
+            "preview_decode": (["Full", "Scale"], {"default": "Scale"}),
+            "denoise": ("FLOAT", {"default": 1.0, "min": 0, "max":1})
+            }
+        }
+    RETURN_TYPES = ("LATENT", )
+    FUNCTION = "sample"
+    CATEGORY = "High_Preview"
+
+    def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative,
+    latent_image, preview_decode, vae, preview_mode, denoise):
+        return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
+            negative, latent_image, preview_decode, vae, preview_mode, denoise,)
+
+NODE_CLASS_MAPPINGS = {"HighPreviewKSampler": HighPreviewKSampler}
+NODE_DISPLAY_NAME_MAPPINGS = {"HighPreviewKSampler": "High Preview KSampler"}
@@ -0,0 +1,287 @@
+import { app } from "../../scripts/app.js";
+
+let panel;
+let img;
+let info;
+let mini;
+
+let executing = false;
+let userClosed = false;
+
+let scale = 1;
+let offsetX = 0;
+let offsetY = 0;
+
+// Create the preview window
+function createPreview() {
+  panel = document.createElement("div");
+
+  Object.assign(panel.style, {
+    position: "fixed",
+    left: "120px",
+    top: "120px",
+    width: "800px",
+    height: "700px",
+    background: "#111",
+    border: "1px solid #777",
+    zIndex: 99999,
+    resize: "both",
+    overflow: "hidden",
+    display: "none",
+    flexDirection: "column"
+  });
+
+  // Title bar
+  let header = document.createElement("div");
+
+  Object.assign(header.style, {
+    height: "35px",
+    background: "#222",
+    color: "white",
+    cursor: "move",
+    padding: "5px",
+    userSelect: "none"
+  });
+
+  header.innerHTML = "";
+  header.appendChild(document.createTextNode("HD Preview"));
+
+  let close = document.createElement("button");
+  close.innerText = "×";
+
+  Object.assign(close.style, {
+    float: "right",
+    width: "30px"
+  });
+
+  close.onclick = (e) => {
+    e.stopPropagation();
+    panel.style.display = "none";
+    mini.style.display = "block";
+    userClosed = true;
+  };
+
+  header.appendChild(close);
+
+  // Content area
+  let content = document.createElement("div");
+
+  Object.assign(content.style, {
+    flex: "1",
+    overflow: "hidden",
+    display: "flex",
+    justifyContent: "center",
+    alignItems: "center"
+  });
+
+  img = document.createElement("img");
+
+  Object.assign(img.style, {
+    maxWidth: "100%",
+    maxHeight: "100%",
+    cursor: "grab",
+    transformOrigin: "center"
+  });
+
+  content.appendChild(img);
+
+  info = document.createElement("div");
+
+  Object.assign(info.style, {
+    height: "25px",
+    color: "white"
+  });
+
+  panel.appendChild(header);
+  panel.appendChild(content);
+  panel.appendChild(info);
+
+  document.body.appendChild(panel);
+
+  // Window dragging
+  let moving = false;
+  let moveStartX = 0;
+  let moveStartY = 0;
+  let panelStartX = 0;
+  let panelStartY = 0;
+
+  header.addEventListener("pointerdown", (e) => {
+    if (e.target === close) return;
+
+    moving = true;
+    moveStartX = e.clientX;
+    moveStartY = e.clientY;
+    panelStartX = panel.offsetLeft;
+    panelStartY = panel.offsetTop;
+    header.setPointerCapture(e.pointerId);
+    header.style.cursor = "grabbing";
+
+    e.preventDefault();
+    e.stopPropagation();
+  });
+
+  header.addEventListener("pointermove", (e) => {
+    if (!moving) return;
+
+    let dx = e.clientX - moveStartX;
+    let dy = e.clientY - moveStartY;
+
+    panel.style.left = panelStartX + dx + "px";
+    panel.style.top = panelStartY + dy + "px";
+  });
+
+  header.addEventListener("pointerup", (e) => {
+    moving = false;
+    header.style.cursor = "move";
+
+    try {
+      header.releasePointerCapture(e.pointerId);
+    } catch (err) {}
+  });
+
+  // Image zoom
+  img.addEventListener("wheel", (e) => {
+    e.preventDefault();
+
+    scale *= e.deltaY < 0 ? 1.1 : 0.9;
+    scale = Math.max(0.2, Math.min(scale, 5));
+
+    updateImage();
+  });
+
+  // Image drag
+  let dragImg = false;
+  let lastX;
+  let lastY;
+
+  img.addEventListener("mousedown", (e) => {
+    dragImg = true;
+    lastX = e.clientX;
+    lastY = e.clientY;
+    img.style.cursor = "grabbing";
+  });
+
+  document.addEventListener("mousemove", (e) => {
+    if (!dragImg) return;
+
+    offsetX += e.clientX - lastX;
+    offsetY += e.clientY - lastY;
+    lastX = e.clientX;
+    lastY = e.clientY;
+
+    updateImage();
+  });
+
+  document.addEventListener("mouseup", () => {
+    dragImg = false;
+    if (img) img.style.cursor = "grab";
+  });
+
+  img.ondblclick = () => {
+    scale = 1;
+    offsetX = 0;
+    offsetY = 0;
+    updateImage();
+  };
+
+  // Mini floating button
+  mini = document.createElement("button");
+  mini.innerText = "HD";
+
+  Object.assign(mini.style, {
+    position: "fixed",
+    left: "30px",
+    top: "30px",
+    width: "60px",
+    height: "60px",
+    borderRadius: "50%",
+    zIndex: 100000
+  });
+
+  // Mini button interactions
+  let miniMoving = false;
+  let miniDragged = false;
+  let miniStartX = 0;
+  let miniStartY = 0;
+  let miniStartLeft = 0;
+  let miniStartTop = 0;
+
+  mini.addEventListener("click", () => {
+    if (miniDragged) {
+      miniDragged = false;
+      return;
+    }
+
+    panel.style.display = "flex";
+    mini.style.display = "none";
+    userClosed = false;
+  });
+
+  mini.addEventListener("mousedown", (e) => {
+    miniMoving = true;
+    miniDragged = false;
+    miniStartX = e.clientX;
+    miniStartY = e.clientY;
+    miniStartLeft = mini.offsetLeft;
+    miniStartTop = mini.offsetTop;
+    e.preventDefault();
+  });
+
+  document.addEventListener("mousemove", (e) => {
+    if (!miniMoving) return;
+
+    let dx = e.clientX - miniStartX;
+    let dy = e.clientY - miniStartY;
+
+    if (Math.abs(dx) > 5 || Math.abs(dy) > 5) {
+      miniDragged = true;
+    }
+
+    mini.style.left = miniStartLeft + dx + "px";
+    mini.style.top = miniStartTop + dy + "px";
+  });
+
+  document.addEventListener("mouseup", () => {
+    miniMoving = false;
+  });
+
+  document.body.appendChild(mini);
+}
+
+function updateImage() {
+  img.style.transform = `translate(${offsetX}px, ${offsetY}px) scale(${scale})`;
+}
+
+app.registerExtension({
+  name: "HighPreviewKSampler",
+
+  async setup() {
+    createPreview();
+
+    app.api.addEventListener("execution_start", () => {
+      executing = true;
+      // Do not reset userClosed here
+      panel.style.display = "flex";
+      mini.style.display = "none";
+    });
+
+    app.api.addEventListener("high_preview", (e) => {
+      let data = e.detail;
+
+      if (executing && !userClosed) {
+        panel.style.display = "flex";
+        mini.style.display = "none";
+      }
+
+      img.src = "data:image/png;base64," + data.image;
+      info.innerText = `Step ${data.step}/${data.total}`;
+    });
+
+    app.api.addEventListener("execution_success", () => {
+      executing = false;
+      if (panel.style.display === "none") {
+        mini.style.display = "block";
+      }
+    });
+  }
+});