nodes/NyaFuP-ComfyUI_Preview_Selector
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

1.2.5

Browse Source
This commit is contained in:
NyaFu
2025-08-06 14:44:16 +09:00
parent e964510273
commit bc47a1cbe8
2 changed files with 75 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
js/preview_dialog.js
+2 -1
View File
@@ -884,7 +884,8 @@ class ExpandedImageView {
handleWheelZoom(event) {
event.preventDefault();
const rect = this.imageElement.getBoundingClientRect();
// Use container coordinates for consistency with zoomAtPoint
const rect = this.container.getBoundingClientRect();
const mouseX = event.clientX - rect.left;
const mouseY = event.clientY - rect.top;
nf_preview_selector.py
+73 -14
View File
@@ -95,11 +95,22 @@ class NFPreviewSelector(PreviewImage):
# 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": images,
"latents": latents,
"images": cpu_images,
"latents": cpu_latents,
"review_id": review_id,
"timeout": timeout
"timeout": timeout,
"original_device": original_device
}
# Send review request to frontend via WebSocket
@@ -107,10 +118,18 @@ class NFPreviewSelector(PreviewImage):
"unique_id": unique_id,
"review_id": review_id,
"images": preview_result["ui"]["images"],
"count": len(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)
@@ -136,11 +155,20 @@ class NFPreviewSelector(PreviewImage):
pass
# Wait for response from frontend
result = self.wait_for_response(review_id, timeout)
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:
@@ -150,24 +178,46 @@ class NFPreviewSelector(PreviewImage):
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, 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]))}
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
selected_images = torch.stack([images[i] for i in selected_indices])
# 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 latents is not None:
selected_latents = {"samples": torch.stack([latents["samples"][i] for i in selected_indices])}
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
selected_latents = {"samples": torch.zeros((len(selected_indices), 4, 64, 64))}
# 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)
@@ -178,7 +228,7 @@ class NFPreviewSelector(PreviewImage):
"ui": preview_result.get("ui", {})
}
def wait_for_response(self, review_id, timeout):
def wait_for_response(self, review_id, timeout, unique_id):
"""Wait for user response from frontend"""
import time
@@ -198,6 +248,15 @@ class NFPreviewSelector(PreviewImage):
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