271 lines
10 KiB
Python
271 lines
10 KiB
Python
"""
|
|
NF Preview Selector - Floating dialog image review system
|
|
|
|
"""
|
|
|
|
|
|
import torch
|
|
import folder_paths
|
|
import json
|
|
import uuid
|
|
from server import PromptServer
|
|
from nodes import PreviewImage
|
|
from aiohttp import web
|
|
import asyncio
|
|
from comfy.model_management import InterruptProcessingException
|
|
|
|
# Global storage for pending reviews
|
|
pending_reviews = {}
|
|
review_responses = {}
|
|
|
|
async def handle_review_response(request):
|
|
"""Handle review response from frontend"""
|
|
try:
|
|
data = await request.json()
|
|
review_id = data.get('review_id')
|
|
selection = data.get('selection', [])
|
|
cancelled = data.get('cancelled', False)
|
|
|
|
pass
|
|
|
|
# Store response
|
|
review_responses[review_id] = {
|
|
'selection': selection,
|
|
'cancelled': cancelled
|
|
}
|
|
|
|
return web.json_response({'status': 'success'})
|
|
except Exception as e:
|
|
pass
|
|
return web.json_response({'status': 'error', 'message': str(e)}, status=500)
|
|
|
|
# Register API endpoint
|
|
PromptServer.instance.app.router.add_post('/nf_preview_response', handle_review_response)
|
|
|
|
class NFPreviewSelector(PreviewImage):
|
|
"""
|
|
A node for reviewing and selecting images using a floating dialog interface
|
|
"""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"mode": (["review_and_select", "pass_through", "take_first", "take_last"], {"default": "review_and_select"}),
|
|
"timeout": ("INT", {"default": 60, "min": 10, "max": 300, "step": 1}),
|
|
},
|
|
"optional": {
|
|
"latents": ("LATENT",),
|
|
"selection_indices": ("STRING", {"default": ""}),
|
|
},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO"
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "LATENT", "STRING")
|
|
RETURN_NAMES = ("selected_images", "selected_latents", "selection_indices")
|
|
FUNCTION = "preview_images"
|
|
CATEGORY = "image/preview"
|
|
OUTPUT_NODE = True
|
|
|
|
def preview_images(self, images, mode, timeout, unique_id=None, latents=None, selection_indices="", prompt=None, extra_pnginfo=None):
|
|
"""Main function for image preview and selection"""
|
|
|
|
pass
|
|
|
|
# No flags needed anymore
|
|
|
|
# Save images for UI display using parent class method
|
|
preview_result = self.save_images(images, filename_prefix="NFPreview", prompt=prompt, extra_pnginfo=extra_pnginfo)
|
|
|
|
# Handle different modes
|
|
if mode == "pass_through":
|
|
selected_indices = list(range(len(images)))
|
|
elif mode == "take_first":
|
|
selected_indices = [0] if len(images) > 0 else []
|
|
elif mode == "take_last":
|
|
selected_indices = [len(images) - 1] if len(images) > 0 else []
|
|
elif mode == "review_and_select":
|
|
# Send to frontend for review
|
|
if unique_id:
|
|
review_id = str(uuid.uuid4())
|
|
# Store necessary values before transferring to CPU
|
|
original_device = images.device if images is not None else 'cpu'
|
|
image_count = len(images) if images is not None else 0
|
|
|
|
# Transfer to CPU to save GPU memory during wait
|
|
cpu_images = images.cpu() if images is not None else None
|
|
cpu_latents = None
|
|
if latents is not None:
|
|
cpu_latents = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in latents.items()}
|
|
|
|
pending_reviews[unique_id] = {
|
|
"images": cpu_images,
|
|
"latents": cpu_latents,
|
|
"review_id": review_id,
|
|
"timeout": timeout,
|
|
"original_device": original_device
|
|
}
|
|
|
|
# Send review request to frontend via WebSocket
|
|
message_data = {
|
|
"unique_id": unique_id,
|
|
"review_id": review_id,
|
|
"images": preview_result["ui"]["images"],
|
|
"count": image_count,
|
|
"timeout": timeout
|
|
}
|
|
|
|
# Clear GPU tensor references (reassign to None to allow garbage collection)
|
|
images = None
|
|
latents = None
|
|
|
|
# Force GPU memory cleanup if CUDA is available
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
# Use ComfyUI's message dispatch system
|
|
try:
|
|
PromptServer.instance.send_sync("nf_preview_request", message_data)
|
|
pass
|
|
except Exception as e:
|
|
pass
|
|
# Fallback to direct WebSocket
|
|
message = {
|
|
"type": "nf_preview_request",
|
|
"data": message_data
|
|
}
|
|
|
|
async def send_message():
|
|
for ws in PromptServer.instance.sockets.values():
|
|
try:
|
|
await ws.send_str(json.dumps(message))
|
|
except Exception as e:
|
|
pass
|
|
|
|
if PromptServer.instance.loop:
|
|
asyncio.run_coroutine_threadsafe(send_message(), PromptServer.instance.loop)
|
|
|
|
pass
|
|
|
|
# Wait for response from frontend
|
|
result = self.wait_for_response(review_id, timeout, unique_id)
|
|
if result == "CANCELLED":
|
|
# Clean up pending review data
|
|
if unique_id in pending_reviews:
|
|
del pending_reviews[unique_id]
|
|
# Force garbage collection and GPU memory cleanup
|
|
import gc
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
pass
|
|
raise InterruptProcessingException()
|
|
elif result == "TIMEOUT":
|
|
# Cleanup is already handled in wait_for_response
|
|
pass
|
|
raise InterruptProcessingException()
|
|
else:
|
|
selected_indices = result
|
|
else:
|
|
selected_indices = [0] if len(images) > 0 else []
|
|
else:
|
|
selected_indices = [0] if len(images) > 0 else []
|
|
|
|
# Get images and latents from CPU storage if needed (for review mode)
|
|
if mode == "review_and_select" and unique_id and unique_id in pending_reviews:
|
|
# For review mode, get data from CPU storage
|
|
review_data = pending_reviews[unique_id]
|
|
source_images = review_data["images"]
|
|
source_latents = review_data["latents"]
|
|
target_device = review_data.get("original_device", 'cpu')
|
|
# Clean up pending review data
|
|
del pending_reviews[unique_id]
|
|
|
|
# Clear CPU references after moving to GPU
|
|
import gc
|
|
del review_data
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
else:
|
|
# For other modes, use original data
|
|
source_images = images
|
|
source_latents = latents
|
|
target_device = images.device if images is not None else 'cpu'
|
|
|
|
# Process selection
|
|
if not selected_indices:
|
|
# For timeout or other cases, return empty tensors
|
|
pass
|
|
empty_image = torch.zeros((1, source_images.shape[1], source_images.shape[2], source_images.shape[3]), device=target_device)
|
|
empty_latent = {"samples": torch.zeros((1, 4, 64, 64), device=target_device)} if source_latents is None else {"samples": torch.zeros((1, source_latents["samples"].shape[1], source_latents["samples"].shape[2], source_latents["samples"].shape[3]), device=target_device)}
|
|
return (empty_image, empty_latent, "")
|
|
|
|
# Get selected images and transfer back to original device
|
|
selected_images = torch.stack([source_images[i] for i in selected_indices]).to(target_device)
|
|
|
|
# Handle latents
|
|
selected_latents = None
|
|
if source_latents is not None:
|
|
selected_latents = {"samples": torch.stack([source_latents["samples"][i] for i in selected_indices]).to(target_device)}
|
|
else:
|
|
# Create placeholder latents on target device
|
|
selected_latents = {"samples": torch.zeros((len(selected_indices), 4, 64, 64), device=target_device)}
|
|
|
|
indices_str = ",".join(str(i) for i in selected_indices)
|
|
|
|
pass
|
|
|
|
return {
|
|
"result": (selected_images, selected_latents, indices_str),
|
|
"ui": preview_result.get("ui", {})
|
|
}
|
|
|
|
def wait_for_response(self, review_id, timeout, unique_id):
|
|
"""Wait for user response from frontend"""
|
|
import time
|
|
|
|
start_time = time.time()
|
|
|
|
while time.time() - start_time < timeout:
|
|
if review_id in review_responses:
|
|
response = review_responses.pop(review_id)
|
|
|
|
if response['cancelled']:
|
|
pass
|
|
return "CANCELLED"
|
|
else:
|
|
selection = response['selection']
|
|
pass
|
|
return selection
|
|
|
|
time.sleep(0.1) # Poll every 100ms
|
|
|
|
# Timeout occurred - clean up pending review data to free memory
|
|
if unique_id in pending_reviews:
|
|
del pending_reviews[unique_id]
|
|
# Force garbage collection and GPU memory cleanup
|
|
import gc
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
pass
|
|
return "TIMEOUT" # Indicate timeout occurred
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"NFPreviewSelector": NFPreviewSelector,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"NFPreviewSelector": "NF Preview Selector",
|
|
} |