From bdf506b94ea994e2337bba4afee9ea3de91adcf0 Mon Sep 17 00:00:00 2001 From: XiaoYan Date: Fri, 17 Jul 2026 03:58:10 +0800 Subject: [PATCH] Initial commit --- __init__.py | 5 + node/hd_preview.py | 57 +++++++++ node/sampler.py | 52 ++++++++ web/high_preview.js | 287 ++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 401 insertions(+) create mode 100644 __init__.py create mode 100644 node/hd_preview.py create mode 100644 node/sampler.py create mode 100644 web/high_preview.js diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..db7e16e --- /dev/null +++ b/__init__.py @@ -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"] \ No newline at end of file diff --git a/node/hd_preview.py b/node/hd_preview.py new file mode 100644 index 0000000..bf60286 --- /dev/null +++ b/node/hd_preview.py @@ -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 \ No newline at end of file diff --git a/node/sampler.py b/node/sampler.py new file mode 100644 index 0000000..5ad5f39 --- /dev/null +++ b/node/sampler.py @@ -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"} \ No newline at end of file diff --git a/web/high_preview.js b/web/high_preview.js new file mode 100644 index 0000000..40a6748 --- /dev/null +++ b/web/high_preview.js @@ -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"; + } + }); + } +}); \ No newline at end of file