From bc47a1cbe866b4e90cc5163861d4a531f611dbbd Mon Sep 17 00:00:00 2001 From: NyaFu Date: Wed, 6 Aug 2025 14:44:16 +0900 Subject: [PATCH] 1.2.5 --- js/preview_dialog.js | 3 +- nf_preview_selector.py | 87 +++++++++++++++++++++++++++++++++++------- 2 files changed, 75 insertions(+), 15 deletions(-) diff --git a/js/preview_dialog.js b/js/preview_dialog.js index f09f048..4758df6 100644 --- a/js/preview_dialog.js +++ b/js/preview_dialog.js @@ -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; diff --git a/nf_preview_selector.py b/nf_preview_selector.py index 8637c13..8351226 100644 --- a/nf_preview_selector.py +++ b/nf_preview_selector.py @@ -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