102 lines
3.9 KiB
Python
102 lines
3.9 KiB
Python
"""ComfyUI Encrypt Preview: a single output node and its binary transport."""
|
|
|
|
import asyncio
|
|
from contextlib import suppress
|
|
import re
|
|
|
|
from aiohttp import web
|
|
import torch
|
|
from comfy.model_management import throw_exception_if_processing_interrupted
|
|
from comfy_execution.utils import get_executing_context
|
|
from server import PromptServer
|
|
|
|
from .preview_core import PreviewExpired, PreviewStore, process_batch
|
|
|
|
WEB_DIRECTORY = "./web"
|
|
STORE = PreviewStore()
|
|
|
|
|
|
class EncryptPreviewImage:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"images": ("IMAGE",)}, "hidden": {"unique_id": "UNIQUE_ID"}}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "preview_image"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "image/preview"
|
|
DESCRIPTION = "Fast in-memory obfuscated preview. Click a thumbnail to view or download the original PNG."
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
# Refresh expired preview handles without invalidating upstream generation.
|
|
return float("nan")
|
|
|
|
def preview_image(self, images, unique_id):
|
|
if not isinstance(images, torch.Tensor) or not images.is_floating_point():
|
|
raise ValueError("Encrypt Preview requires a floating point IMAGE tensor.")
|
|
server = PromptServer.instance
|
|
context = get_executing_context()
|
|
node_id = str(unique_id)
|
|
prompt_id = context.prompt_id if context else None
|
|
client_id = server.client_id
|
|
|
|
def notify(data):
|
|
if client_id is not None:
|
|
server.send_sync("encrypt-preview", {**data, "node_id": node_id, "prompt_id": prompt_id}, client_id)
|
|
|
|
result = process_batch(
|
|
images, STORE,
|
|
lambda image: image.detach().to(device="cpu", dtype=torch.float32).numpy(),
|
|
notify, throw_exception_if_processing_interrupted,
|
|
)
|
|
return {"ui": {"encrypt_preview": [{**result, "node_id": node_id, "prompt_id": prompt_id}]}}
|
|
|
|
|
|
async def read_preview(request):
|
|
batch_id = request.match_info["batch_id"]
|
|
index_text = request.match_info["index"]
|
|
kind = request.match_info["kind"]
|
|
headers = {"Cache-Control": "no-store", "X-Content-Type-Options": "nosniff"}
|
|
if not re.fullmatch(r"[0-9a-f]{32}", batch_id) or not re.fullmatch(r"[0-9]{1,3}", index_text) or kind not in ("thumbnail", "original"):
|
|
raise web.HTTPBadRequest(text="Invalid preview request.", headers=headers)
|
|
try:
|
|
packet = STORE.get(batch_id, int(index_text), kind)
|
|
except PreviewExpired:
|
|
raise web.HTTPGone(text="Preview expired. Run the output node again.", headers=headers) from None
|
|
return web.Response(body=packet, content_type="application/octet-stream", headers=headers)
|
|
|
|
|
|
def disable_native_preview(payload):
|
|
"""Per-job override: the native sampler preview is a separate plaintext path."""
|
|
prompt = payload.get("prompt", {})
|
|
if isinstance(prompt, dict) and any(isinstance(node, dict) and node.get("class_type") == "EncryptPreviewImage" for node in prompt.values()):
|
|
extra_data = payload.setdefault("extra_data", {})
|
|
if isinstance(extra_data, dict):
|
|
extra_data["preview_method"] = "none"
|
|
return payload
|
|
|
|
|
|
async def cache_lifecycle(app):
|
|
async def sweep():
|
|
while True:
|
|
await asyncio.sleep(60)
|
|
STORE.expire()
|
|
|
|
task = asyncio.create_task(sweep())
|
|
try:
|
|
yield
|
|
finally:
|
|
task.cancel()
|
|
with suppress(asyncio.CancelledError):
|
|
await task
|
|
|
|
|
|
PromptServer.instance.routes.get("/encrypt-preview/v1/{batch_id}/{index}/{kind}")(read_preview)
|
|
PromptServer.instance.add_on_prompt_handler(disable_native_preview)
|
|
PromptServer.instance.app.cleanup_ctx.append(cache_lifecycle)
|
|
|
|
NODE_CLASS_MAPPINGS = {"EncryptPreviewImage": EncryptPreviewImage}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"EncryptPreviewImage": "Encrypt Preview Image"}
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|