From 8ad0d806e996cf59e5c9a35a67b998ca193584ed Mon Sep 17 00:00:00 2001 From: NyaFuP Date: Tue, 15 Jul 2025 15:05:01 +0900 Subject: [PATCH] [Add] nf_preview_selector.py --- README.md | 72 +++++++ __init__.py | 12 ++ js/floating_window.css | 165 ++++++++++++++++ js/floating_window.js | 113 +++++++++++ js/nf_preview_selector.js | 87 ++++++++ js/preview_dialog.js | 405 ++++++++++++++++++++++++++++++++++++++ nf_preview_selector.py | 210 ++++++++++++++++++++ 7 files changed, 1064 insertions(+) create mode 100644 README.md create mode 100644 __init__.py create mode 100644 js/floating_window.css create mode 100644 js/floating_window.js create mode 100644 js/nf_preview_selector.js create mode 100644 js/preview_dialog.js create mode 100644 nf_preview_selector.py diff --git a/README.md b/README.md new file mode 100644 index 0000000..0d97dd8 --- /dev/null +++ b/README.md @@ -0,0 +1,72 @@ +# NF Review Selector + +A floating dialog-based image review and selection system for ComfyUI. Inspired by cg-image-filter but designed as a simpler, more focused solution for image review workflows. + +## Features + +- **Floating Dialog Interface**: Non-intrusive dialog window instead of full-screen overlay +- **Drag & Drop Positioning**: Move the review window anywhere on screen +- **Multi-Selection**: Click images to select/deselect multiple images +- **Timeout Support**: Configurable timeout with visual countdown +- **Multiple Modes**: Review mode, pass-through, take first/last +- **Latent Support**: Optional latent passthrough for efficiency +- **Memory**: Window position and settings are remembered + +## Usage + +1. Add the "NF Review Selector" node to your workflow +2. Connect image input (and optionally latent input) +3. Choose a mode: + - **review_and_select**: Opens floating dialog for manual selection + - **pass_through**: Passes all images without review + - **take_first**: Automatically selects first image + - **take_last**: Automatically selects last image +4. Set timeout value (10-300 seconds) + +## Dialog Interface + +When in review mode, a floating dialog will appear showing: +- Grid of images to review +- Click images to select/deselect (green border indicates selection) +- Selection counter +- Countdown timer +- Confirm/Cancel buttons + +## Settings + +Access via ComfyUI Settings menu: +- **Remember window position**: Saves dialog position between uses +- **Use compact mode**: Smaller layout for limited screen space +- **Auto-select first image**: Automatically selects first image when dialog opens + +## Outputs + +- **selected_images**: IMAGE batch of selected images +- **selected_latents**: LATENT batch of corresponding latents (if input provided) +- **selection_indices**: STRING of selected indices (e.g., "0,2,4") + +## Installation + +1. Place this folder in `ComfyUI/custom_nodes/` +2. Restart ComfyUI +3. The "NF Review Selector" node will appear in the "image/review" category + +## Technical Details + +- Built with modern JavaScript ES6 modules +- CSS-styled floating window system +- WebSocket communication with ComfyUI backend +- Supports both image and latent workflows +- Minimal dependencies (uses ComfyUI's built-in APIs) + +## Comparison to Alternatives + +Unlike full-screen solutions, NF Review Selector: +- Doesn't take over the entire interface +- Allows interaction with other parts of ComfyUI while reviewing +- Provides a cleaner, more focused review experience +- Maintains workflow visibility + +## License + +This implementation is provided as-is for educational and development purposes. \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..8ce889d --- /dev/null +++ b/__init__.py @@ -0,0 +1,12 @@ +""" +NF Preview Selector - A floating dialog image review and selection system for ComfyUI +Based on cg-image-filter architecture but simplified for review purposes +""" + +from .nf_preview_selector import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +WEB_DIRECTORY = "./js" + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] + +print("NF Preview Selector: Loaded with floating dialog implementation") \ No newline at end of file diff --git a/js/floating_window.css b/js/floating_window.css new file mode 100644 index 0000000..5d5183e --- /dev/null +++ b/js/floating_window.css @@ -0,0 +1,165 @@ +/** + * NF Preview Selector - Floating Window Styles + */ + +.nf-float { + position: absolute; + z-index: 999999; + border: 2px solid #333; + background-color: rgba(45, 45, 45, 0.95); + border-radius: 8px; + box-shadow: 0 4px 12px rgba(0, 0, 0, 0.3); + min-width: 300px; + max-width: 800px; + min-height: 300px; /* Minimum height for proper layout */ + font-family: Arial, sans-serif; + resize: both; + overflow: hidden; + display: flex; + flex-direction: column; +} + +.nf-float.hidden { + display: none; +} + +.nf-float-header { + width: 100%; + text-align: center; + background-color: rgba(70, 130, 180, 0.9); + color: white; + border-bottom: 1px solid #333; + padding: 8px 12px; + font-size: 14px; + font-weight: bold; + cursor: move; + border-radius: 6px 6px 0 0; + user-select: none; + flex-shrink: 0; /* Don't shrink header */ +} + +.nf-float-header:hover { + background-color: rgba(70, 130, 180, 1); +} + +.nf-float-body { + padding: 12px; + color: white; + display: flex; + flex-direction: column; + gap: 8px; + flex: 1; /* Take remaining space */ + min-height: 0; /* Allow shrinking */ +} + +/* Image Grid */ +.nf-image-grid { + display: grid; + gap: 8px; + grid-template-columns: repeat(auto-fit, minmax(150px, 1fr)); + flex: 1; /* Take up remaining space */ + overflow-y: auto; + min-height: 0; /* Allow shrinking */ +} + +.nf-image-container { + position: relative; + border: 2px solid transparent; + border-radius: 4px; + cursor: pointer; + transition: border-color 0.2s; +} + +.nf-image-container:hover { + border-color: #70a0d0; +} + +.nf-image-container.selected { + border-color: #4CAF50; + background-color: rgba(76, 175, 80, 0.1); +} + +.nf-image { + width: 100%; + height: auto; + display: block; + border-radius: 2px; +} + +.nf-image-overlay { + position: absolute; + top: 4px; + right: 4px; + background-color: rgba(0, 0, 0, 0.7); + color: white; + padding: 2px 6px; + border-radius: 3px; + font-size: 11px; +} + +.nf-image-container.selected .nf-image-overlay { + background-color: rgba(76, 175, 80, 0.9); +} + +/* Controls */ +.nf-controls { + display: flex; + justify-content: space-between; + align-items: center; + padding: 8px 0; + border-top: 1px solid #555; + margin-top: auto; /* Push to bottom */ + flex-shrink: 0; /* Don't shrink */ +} + +.nf-selection-info { + font-size: 12px; + color: #ccc; +} + +.nf-buttons { + display: flex; + gap: 8px; +} + +.nf-button { + padding: 6px 12px; + border: none; + border-radius: 4px; + cursor: pointer; + font-size: 12px; + transition: background-color 0.2s; +} + +.nf-button.primary { + background-color: #4CAF50; + color: white; +} + +.nf-button.primary:hover { + background-color: #45a049; +} + +.nf-button.primary:disabled { + background-color: #666; + cursor: not-allowed; +} + +.nf-button.secondary { + background-color: #666; + color: white; +} + +.nf-button.secondary:hover { + background-color: #777; +} + +/* Countdown */ +.nf-countdown { + font-size: 11px; + color: #ff9800; + text-align: center; + padding: 4px; +} + +/* Max width is now configurable via settings */ \ No newline at end of file diff --git a/js/floating_window.js b/js/floating_window.js new file mode 100644 index 0000000..a54c153 --- /dev/null +++ b/js/floating_window.js @@ -0,0 +1,113 @@ +/** + * NF Preview Selector - Floating Window Component + * Based on cg-image-filter FloatingWindow but simplified + */ + +export class NFFloatingWindow extends HTMLDivElement { + constructor(title, x = 100, y = 100, parent = null, moveCallback = null, resizeCallback = null) { + super(); + this.moveCallback = moveCallback; + this.resizeCallback = resizeCallback; + + // Setup window structure + this.classList.add('nf-float'); + + // Header + this.header = document.createElement('div'); + this.header.classList.add('nf-float-header'); + this.header.innerText = title; + this.appendChild(this.header); + + // Body + this.body = document.createElement('div'); + this.body.classList.add('nf-float-body'); + this.appendChild(this.body); + + // Setup dragging + this.dragging = false; + this.position = { x: x, y: y }; + + this.header.addEventListener('mousedown', this.startDrag.bind(this)); + document.addEventListener('mouseup', this.stopDrag.bind(this)); + document.addEventListener('mousemove', this.drag.bind(this)); + document.addEventListener('mouseleave', this.stopDrag.bind(this)); + + // Position window + this.moveTo(x, y); + + // Add to parent or body + if (parent) { + parent.appendChild(this); + } else { + document.body.appendChild(this); + } + + // Setup resize observer for size changes + this.setupResizeObserver(); + } + + show() { + this.style.display = 'block'; + } + + hide() { + this.style.display = 'none'; + } + + setTitle(title) { + this.header.innerText = title; + } + + moveTo(x, y) { + this.position = { x: x, y: y }; + this.style.left = `${this.position.x}px`; + this.style.top = `${this.position.y}px`; + + if (this.moveCallback) { + this.moveCallback(x, y); + } + } + + startDrag(e) { + this.dragging = true; + e.stopPropagation(); + e.preventDefault(); + } + + stopDrag(e) { + this.dragging = false; + } + + drag(e) { + if (this.dragging) { + this.moveTo( + this.position.x + e.movementX, + this.position.y + e.movementY + ); + e.stopPropagation(); + e.preventDefault(); + } + } + + setupResizeObserver() { + if (this.resizeCallback && window.ResizeObserver) { + this.resizeObserver = new ResizeObserver((entries) => { + for (let entry of entries) { + const { width, height } = entry.contentRect; + this.resizeCallback(width, height); + } + }); + this.resizeObserver.observe(this); + } + } + + destroy() { + if (this.resizeObserver) { + this.resizeObserver.disconnect(); + } + this.remove(); + } +} + +// Register custom element +customElements.define('nf-floating-window', NFFloatingWindow, { extends: 'div' }); \ No newline at end of file diff --git a/js/nf_preview_selector.js b/js/nf_preview_selector.js new file mode 100644 index 0000000..27eb08f --- /dev/null +++ b/js/nf_preview_selector.js @@ -0,0 +1,87 @@ +/** + * NF Preview Selector - Main Extension + * A floating dialog-based image review and selection system for ComfyUI + */ + +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; +import { nfPreviewDialog } from './preview_dialog.js'; + +console.log("NF Preview Selector: Loading extension..."); + +app.registerExtension({ + name: "nf.preview_selector", + + async setup() { + console.log("NF Preview Selector: Setting up extension"); + console.log("NF Preview Selector: API object:", api); + console.log("NF Preview Selector: API addEventListener available:", typeof api.addEventListener); + + // Register message handler using ComfyUI's system + if (api.addEventListener) { + // Register for our custom event type + api.addEventListener("nf_preview_request", (event) => { + console.log("NF Preview Selector: Received review request via event", event.detail); + nfPreviewDialog.show(event.detail); + }); + + console.log("NF Preview Selector: Event listener registered for nf_preview_request"); + } + + // Event system is working correctly + + // Add settings + app.ui.settings.addSetting({ + id: "NFPreviewSelector.UI.WindowPosition", + name: "Remember window position and size", + type: "boolean", + defaultValue: true, + tooltip: "Remember the position and size of the preview dialog window between sessions" + }); + + app.ui.settings.addSetting({ + id: "NFPreviewSelector.UI.MaxWidth.v2", + name: "Dialog maximum width (px)", + type: "number", + defaultValue: 800, + tooltip: "Maximum width of the preview dialog in pixels (400-1800)" + }); + + // Auto-select first image setting removed - this is available in node UI + + console.log("NF Preview Selector: Setup complete"); + }, + + async nodeCreated(node) { + if (node.comfyClass === "NFPreviewSelector") { + console.log("NF Preview Selector: Node created", node.id); + + // Set node title + node.title = "NF Preview Selector"; + + // Add any node-specific setup here + node.onResize = function() { + // Handle node resize if needed + }; + } + }, + + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeData.name === "NFPreviewSelector") { + console.log("NF Preview Selector: Registering node definition"); + + // Add any node definition modifications here + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function() { + if (onNodeCreated) { + onNodeCreated.apply(this, arguments); + } + + // Node-specific initialization + this.setProperty("review_position", { x: 100, y: 100 }); + }; + } + } +}); + +console.log("NF Preview Selector: Extension loaded successfully"); \ No newline at end of file diff --git a/js/preview_dialog.js b/js/preview_dialog.js new file mode 100644 index 0000000..7e2160c --- /dev/null +++ b/js/preview_dialog.js @@ -0,0 +1,405 @@ +/** + * NF Preview Selector - Preview Dialog Component + */ + +import { NFFloatingWindow } from './floating_window.js'; +import { api } from "../../scripts/api.js"; +import { app } from "../../scripts/app.js"; + +export class NFPreviewDialog { + constructor() { + this.window = null; + this.images = []; + this.selectedIndices = new Set(); + this.reviewId = null; + this.uniqueId = null; + this.timeout = 60; + this.countdownInterval = null; + this.remainingTime = 0; + + // Load CSS + this.loadCSS(); + + // Apply max width setting + this.applyMaxWidthSetting(); + } + + loadCSS() { + const link = document.createElement('link'); + link.rel = 'stylesheet'; + link.href = new URL("./floating_window.css", import.meta.url).href; + document.head.appendChild(link); + } + + applyMaxWidthSetting() { + // Get max width setting from ComfyUI settings (try new ID first, fallback to old) + let maxWidth = app.ui.settings.getSettingValue("NFPreviewSelector.UI.MaxWidth.v2", + app.ui.settings.getSettingValue("NFPreviewSelector.UI.MaxWidth", 800)); + + // Validate and clamp the value between 400 and 1800 + maxWidth = Math.max(400, Math.min(1800, maxWidth)); + + // Create or update CSS rule for max width + const style = document.createElement('style'); + style.id = 'nf-preview-max-width-style'; + style.textContent = ` + .nf-float { + max-width: ${maxWidth}px !important; + } + `; + + // Remove existing style if it exists + const existingStyle = document.getElementById('nf-preview-max-width-style'); + if (existingStyle) { + existingStyle.remove(); + } + + document.head.appendChild(style); + } + + show(data) { + console.log("NF Preview Dialog: Showing review dialog", data); + + // Apply latest max width setting + this.applyMaxWidthSetting(); + + // If dialog is already open for this review, don't create another + if (this.reviewId === data.review_id && this.window) { + console.log("NF Preview Dialog: Dialog already open for this review"); + return; + } + + // Stop any existing countdown + this.stopCountdown(); + + this.images = data.images || []; + this.reviewId = data.review_id; + this.uniqueId = data.unique_id; + this.timeout = data.timeout || 60; + this.remainingTime = this.timeout; + this.selectedIndices.clear(); + + this.createWindow(); + this.renderImages(); + this.startCountdown(); + } + + createWindow() { + if (this.window) { + this.window.destroy(); + } + + // Get saved position or calculate center position + const rememberPosition = app.ui.settings.getSettingValue("NFPreviewSelector.UI.WindowPosition", true); + const savedPosition = rememberPosition ? this.getSavedPosition() : null; + let x, y; + + if (savedPosition) { + // Use saved position from previous session + x = savedPosition.x; + y = savedPosition.y; + } else { + // First time or position memory disabled: calculate center position + x = Math.max(0, (window.innerWidth - 600) / 2); // 600px is approximate dialog width + y = Math.max(0, (window.innerHeight - 400) / 2); // 400px is approximate dialog height + } + + this.window = new NFFloatingWindow( + `NF Preview - Select Images (${this.images.length})`, + x, + y, + null, + this.onWindowMove.bind(this), + this.onWindowResize.bind(this) + ); + + this.setupWindowContent(); + + // Apply saved size if available and position memory is enabled + if (rememberPosition) { + const savedSize = this.getSavedSize(); + if (savedSize) { + this.window.style.width = savedSize.width + 'px'; + this.window.style.height = savedSize.height + 'px'; + } + } + + // Center the dialog only if no saved position (first time) + if (!savedPosition) { + this.centerDialog(); + } + } + + setupWindowContent() { + // Image grid container + this.imageGrid = document.createElement('div'); + this.imageGrid.className = 'nf-image-grid'; + this.window.body.appendChild(this.imageGrid); + + // Controls + this.controls = document.createElement('div'); + this.controls.className = 'nf-controls'; + + // Selection info + this.selectionInfo = document.createElement('div'); + this.selectionInfo.className = 'nf-selection-info'; + this.updateSelectionInfo(); + + // Countdown + this.countdown = document.createElement('div'); + this.countdown.className = 'nf-countdown'; + + // Buttons + this.buttons = document.createElement('div'); + this.buttons.className = 'nf-buttons'; + + this.confirmButton = document.createElement('button'); + this.confirmButton.className = 'nf-button primary'; + this.confirmButton.textContent = 'Confirm Selection'; + this.confirmButton.onclick = this.confirmSelection.bind(this); + this.updateConfirmButton(); + + this.cancelButton = document.createElement('button'); + this.cancelButton.className = 'nf-button secondary'; + this.cancelButton.textContent = 'Cancel'; + this.cancelButton.onclick = this.cancelSelection.bind(this); + + this.buttons.appendChild(this.confirmButton); + this.buttons.appendChild(this.cancelButton); + + this.controls.appendChild(this.selectionInfo); + this.controls.appendChild(this.countdown); + this.controls.appendChild(this.buttons); + + this.window.body.appendChild(this.controls); + } + + renderImages() { + this.imageGrid.innerHTML = ''; + + this.images.forEach((imageInfo, index) => { + const container = document.createElement('div'); + container.className = 'nf-image-container'; + container.dataset.index = index; + container.onclick = () => this.toggleImage(index); + + const img = document.createElement('img'); + img.className = 'nf-image'; + img.src = this.getImageUrl(imageInfo); + img.alt = `Image ${index + 1}`; + + const overlay = document.createElement('div'); + overlay.className = 'nf-image-overlay'; + overlay.textContent = index + 1; + + container.appendChild(img); + container.appendChild(overlay); + this.imageGrid.appendChild(container); + }); + } + + getImageUrl(imageInfo) { + return `/view?filename=${encodeURIComponent(imageInfo.filename)}&type=${imageInfo.type}&subfolder=${encodeURIComponent(imageInfo.subfolder || '')}`; + } + + toggleImage(index) { + const container = this.imageGrid.querySelector(`[data-index="${index}"]`); + + if (this.selectedIndices.has(index)) { + this.selectedIndices.delete(index); + container.classList.remove('selected'); + } else { + this.selectedIndices.add(index); + container.classList.add('selected'); + } + + this.updateSelectionInfo(); + this.updateConfirmButton(); + + console.log("NF Preview Dialog: Selected indices:", Array.from(this.selectedIndices)); + } + + updateSelectionInfo() { + const count = this.selectedIndices.size; + const total = this.images.length; + this.selectionInfo.textContent = `Selected: ${count} of ${total}`; + } + + updateConfirmButton() { + const hasSelection = this.selectedIndices.size > 0; + this.confirmButton.disabled = !hasSelection; + this.confirmButton.textContent = hasSelection + ? `Confirm (${this.selectedIndices.size} selected)` + : 'Select at least one image'; + } + + startCountdown() { + this.updateCountdown(); + this.countdownInterval = setInterval(() => { + this.remainingTime--; + this.updateCountdown(); + + if (this.remainingTime <= 0) { + console.log("NF Preview Dialog: Timeout reached"); + this.timeoutSelection(); + } + }, 1000); + } + + updateCountdown() { + const minutes = Math.floor(this.remainingTime / 60); + const seconds = this.remainingTime % 60; + this.countdown.textContent = `Time remaining: ${minutes}:${seconds.toString().padStart(2, '0')}`; + } + + stopCountdown() { + if (this.countdownInterval) { + clearInterval(this.countdownInterval); + this.countdownInterval = null; + } + } + + confirmSelection() { + this.stopCountdown(); + + const selection = Array.from(this.selectedIndices).sort((a, b) => a - b); + console.log("NF Preview Dialog: Confirming selection:", selection); + + this.sendSelection(selection, false); + this.hide(); + } + + cancelSelection() { + this.stopCountdown(); + + console.log("NF Preview Dialog: Cancelling selection"); + this.sendSelection([], true); + this.hide(); + } + + timeoutSelection() { + this.stopCountdown(); + + // On timeout, select all images or none based on setting + const selection = []; // Default to empty selection on timeout + console.log("NF Preview Dialog: Timeout - returning empty selection"); + + this.sendSelection(selection, false); + this.hide(); + } + + sendSelection(indices, cancelled) { + const response = { + review_id: this.reviewId, + unique_id: this.uniqueId, + selection: indices, + cancelled: cancelled + }; + + // Send response to backend + fetch('/nf_preview_response', { + method: 'POST', + headers: { + 'Content-Type': 'application/json' + }, + body: JSON.stringify(response) + }).catch(error => { + console.error("NF Preview Dialog: Error sending response:", error); + }); + } + + hide() { + this.stopCountdown(); + + if (this.window) { + this.window.destroy(); + this.window = null; + } + + this.selectedIndices.clear(); + } + + centerDialog() { + if (!this.window) return; + + // Wait for the DOM to update and get actual dimensions + setTimeout(() => { + const dialogRect = this.window.getBoundingClientRect(); + const centerX = Math.max(0, (window.innerWidth - dialogRect.width) / 2); + const centerY = Math.max(0, (window.innerHeight - dialogRect.height) / 2); + + this.window.moveTo(centerX, centerY); + }, 10); + } + + getSavedPosition() { + try { + const saved = localStorage.getItem('nf-preview-dialog-position'); + if (saved) { + const position = JSON.parse(saved); + // Validate position is within screen bounds + if (position.x >= 0 && position.y >= 0 && + position.x < window.innerWidth && position.y < window.innerHeight) { + return position; + } + } + } catch (e) { + console.warn("NF Preview Dialog: Error loading saved position:", e); + } + return null; + } + + savePosition(x, y) { + try { + localStorage.setItem('nf-preview-dialog-position', JSON.stringify({ x, y })); + } catch (e) { + console.warn("NF Preview Dialog: Error saving position:", e); + } + } + + getSavedSize() { + try { + const saved = localStorage.getItem('nf-preview-dialog-size'); + if (saved) { + const size = JSON.parse(saved); + // Validate size is reasonable + if (size.width > 200 && size.height > 150 && + size.width < window.innerWidth + 100 && size.height < window.innerHeight + 100) { + return size; + } + } + } catch (e) { + console.warn("NF Preview Dialog: Error loading saved size:", e); + } + return null; + } + + saveSize(width, height) { + try { + localStorage.setItem('nf-preview-dialog-size', JSON.stringify({ width, height })); + } catch (e) { + console.warn("NF Preview Dialog: Error saving size:", e); + } + } + + onWindowMove(x, y) { + // Save position when user moves the dialog (only if setting is enabled) + const rememberPosition = app.ui.settings.getSettingValue("NFPreviewSelector.UI.WindowPosition", true); + if (rememberPosition) { + this.savePosition(x, y); + } + console.log("NF Preview Dialog: Window moved to", x, y); + } + + onWindowResize(width, height) { + // Save size when user resizes the dialog (only if setting is enabled) + const rememberPosition = app.ui.settings.getSettingValue("NFPreviewSelector.UI.WindowPosition", true); + if (rememberPosition) { + this.saveSize(width, height); + } + console.log("NF Preview Dialog: Window resized to", width, height); + } +} + +// Global instance +export const nfPreviewDialog = new NFPreviewDialog(); \ No newline at end of file diff --git a/nf_preview_selector.py b/nf_preview_selector.py new file mode 100644 index 0000000..812effc --- /dev/null +++ b/nf_preview_selector.py @@ -0,0 +1,210 @@ +""" +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) + + print(f"NF Preview Selector: Received response for review {review_id}: {selection}") + + # Store response + review_responses[review_id] = { + 'selection': selection, + 'cancelled': cancelled + } + + return web.json_response({'status': 'success'}) + except Exception as e: + print(f"NF Preview Selector: Error handling response: {e}") + 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""" + + print(f"NF Preview Selector: Processing {len(images)} images with mode '{mode}'") + + # 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()) + pending_reviews[unique_id] = { + "images": images, + "latents": latents, + "review_id": review_id, + "timeout": timeout + } + + # Send review request to frontend via WebSocket + message_data = { + "unique_id": unique_id, + "review_id": review_id, + "images": preview_result["ui"]["images"], + "count": len(images), + "timeout": timeout + } + + # Use ComfyUI's message dispatch system + try: + PromptServer.instance.send_sync("nf_preview_request", message_data) + print(f"NF Preview Selector: Sent review request via send_sync") + except Exception as e: + print(f"NF Preview Selector: send_sync failed, trying direct WebSocket: {e}") + # 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: + print(f"Failed to send message to client: {e}") + + if PromptServer.instance.loop: + asyncio.run_coroutine_threadsafe(send_message(), PromptServer.instance.loop) + + print(f"NF Preview Selector: Sent review request for {len(images)} images") + + # Wait for response from frontend + result = self.wait_for_response(review_id, timeout) + if result == "CANCELLED": + print(f"NF Preview Selector: Workflow cancelled by user") + raise InterruptProcessingException() + elif result == "TIMEOUT": + print(f"NF Preview Selector: Workflow timed out") + raise InterruptProcessingException() + else: + selected_indices = result + else: + selected_indices = [0] if len(images) > 0 else [] + else: + selected_indices = [0] if len(images) > 0 else [] + + # Process selection + if not selected_indices: + # For timeout or other cases, return empty tensors + print(f"NF Preview Selector: No selection made, returning empty tensors") + empty_image = torch.zeros((1, images.shape[1], images.shape[2], images.shape[3])) + empty_latent = {"samples": torch.zeros((1, 4, 64, 64))} if latents is None else {"samples": torch.zeros((1, latents["samples"].shape[1], latents["samples"].shape[2], latents["samples"].shape[3]))} + return (empty_image, empty_latent, "") + + # Get selected images + selected_images = torch.stack([images[i] for i in selected_indices]) + + # Handle latents + selected_latents = None + if latents is not None: + selected_latents = {"samples": torch.stack([latents["samples"][i] for i in selected_indices])} + else: + # Create placeholder latents + selected_latents = {"samples": torch.zeros((len(selected_indices), 4, 64, 64))} + + indices_str = ",".join(str(i) for i in selected_indices) + + print(f"NF Preview Selector: Selected {len(selected_indices)} images: [{indices_str}]") + + return { + "result": (selected_images, selected_latents, indices_str), + "ui": preview_result.get("ui", {}) + } + + def wait_for_response(self, review_id, timeout): + """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']: + print(f"NF Preview Selector: Review {review_id} was cancelled") + return "CANCELLED" + else: + selection = response['selection'] + print(f"NF Preview Selector: Review {review_id} completed with selection: {selection}") + return selection + + time.sleep(0.1) # Poll every 100ms + + print(f"NF Preview Selector: Review {review_id} timed out") + return "TIMEOUT" # Indicate timeout occurred + +# Node registration +NODE_CLASS_MAPPINGS = { + "NFPreviewSelector": NFPreviewSelector, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "NFPreviewSelector": "NF Preview Selector", +} \ No newline at end of file