118 lines
3.5 KiB
Python
118 lines
3.5 KiB
Python
"""
|
|
ComfyUI Preview Image with Pause
|
|
Version: 1.0
|
|
Author: Cordux (with help from Grok and Claude)
|
|
"""
|
|
import comfy
|
|
from server import PromptServer
|
|
from aiohttp import web
|
|
import time
|
|
import os
|
|
import random
|
|
import numpy as np
|
|
from PIL import Image as PILImage
|
|
import folder_paths
|
|
from comfy.model_management import InterruptProcessingException
|
|
|
|
|
|
class ImagePreviewPause:
|
|
status_by_id = {}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
},
|
|
"hidden": {
|
|
"id": "UNIQUE_ID",
|
|
"prompt": "PROMPT",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "image"
|
|
OUTPUT_NODE = True
|
|
|
|
def execute(self, images, id=None, prompt=None):
|
|
# Save images for preview
|
|
filename_prefix = "preview_pause_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for _ in range(5))
|
|
|
|
full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
|
|
filename_prefix,
|
|
folder_paths.get_temp_directory(),
|
|
images[0].shape[1],
|
|
images[0].shape[0]
|
|
)
|
|
|
|
results = []
|
|
for batch_number, image in enumerate(images):
|
|
i = 255. * image.cpu().numpy()
|
|
img = PILImage.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
|
|
filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
|
|
file = f"{filename_with_batch_num}_{counter:05}_.png"
|
|
img.save(os.path.join(full_output_folder, file), compress_level=1)
|
|
|
|
results.append({
|
|
"filename": file,
|
|
"subfolder": subfolder,
|
|
"type": "temp"
|
|
})
|
|
counter += 1
|
|
|
|
# Send preview to UI immediately using PromptServer
|
|
PromptServer.instance.send_sync("executing", {
|
|
"node": id,
|
|
"prompt_id": None # You might need to track this
|
|
})
|
|
|
|
# Send the preview images
|
|
PromptServer.instance.send_sync("executed", {
|
|
"node": id,
|
|
"output": {
|
|
"images": results
|
|
},
|
|
"prompt_id": None
|
|
})
|
|
|
|
# Small delay to ensure UI receives it
|
|
time.sleep(0.3)
|
|
|
|
# Now pause and wait
|
|
self.status_by_id[id] = "paused"
|
|
|
|
while self.status_by_id[id] == "paused":
|
|
time.sleep(0.1)
|
|
|
|
# Check for cancel (key still exists)
|
|
if self.status_by_id[id] == "cancelled":
|
|
raise InterruptProcessingException()
|
|
|
|
# Safe cleanup on continue path
|
|
self.status_by_id.pop(id, None)
|
|
|
|
# Return images to continue workflow
|
|
return (images,)
|
|
# return {
|
|
# "ui": {"images": results},
|
|
# "result": (images,)
|
|
# }
|
|
|
|
|
|
@PromptServer.instance.routes.post("/image_preview_pause/continue/{node_id}")
|
|
async def handle_continue(request):
|
|
node_id = request.match_info["node_id"].strip()
|
|
ImagePreviewPause.status_by_id[node_id] = "continue"
|
|
return web.json_response({"status": "ok"})
|
|
|
|
|
|
@PromptServer.instance.routes.post("/image_preview_pause/cancel")
|
|
async def handle_cancel(request):
|
|
comfy.model_management.interrupt_current_processing()
|
|
for node_id in ImagePreviewPause.status_by_id:
|
|
ImagePreviewPause.status_by_id[node_id] = "cancelled"
|
|
return web.json_response({"status": "ok"})
|