Deduplicate SaveImage outputs and hash-filter duplicate files}

core/pipelines/sd_image_pipeline.py

@@ -3,9 +3,9 @@ import random
 import shutil
 import torch
 import gradio as gr
-from PIL import Image
 from typing import List, Dict, Any
 
+
 from .base_pipeline import BasePipeline
 from core.settings import *
 from utils.app_utils import sanitize_prompt
@@ -26,34 +26,40 @@ class SdImagePipeline(BasePipeline):
             return [model_display_name]
 
     def _gpu_logic(self, ui_inputs: Dict, loras_string: str, workflow: Dict[str, Any], assembler: WorkflowAssembler, progress=gr.Progress(track_tqdm=True)):
-        model_display_name = ui_inputs['model_display_name']
-        
+        """Execute the ComfyUI workflow and return the file paths saved by the SaveImage node.
+        The original implementation converted the tensor output to PIL images and then saved
+        them again, causing duplicate files. Here we rely on the SaveImage node to write the
+        images to the output directory and simply return the path(s) it provides.
+        """
         progress(0.4, desc="Executing workflow...")
-        
         initial_objects = {}
-
-        decoded_images_tensor = WorkflowExecutor.execute_workflow(workflow, initial_objects=initial_objects)
-        
-        output_images = []
-        start_seed = ui_inputs['seed'] if ui_inputs['seed'] != -1 else random.randint(0, 2**64 - 1)
-        for i in range(decoded_images_tensor.shape[0]):
-            img_tensor = decoded_images_tensor[i]
-            pil_image = Image.fromarray((img_tensor.cpu().numpy() * 255.0).astype("uint8"))
-            current_seed = start_seed + i
-            
-            width_for_meta = ui_inputs.get('width', 'N/A')
-            height_for_meta = ui_inputs.get('height', 'N/A')
-
-            params_string = f"{ui_inputs['positive_prompt']}\nNegative prompt: {ui_inputs['negative_prompt']}\n"
-            params_string += f"Steps: {ui_inputs['num_inference_steps']}, Sampler: {ui_inputs['sampler']}, Scheduler: {ui_inputs['scheduler']}, CFG scale: {ui_inputs['guidance_scale']}, Seed: {current_seed}, Size: {width_for_meta}x{height_for_meta}, Base Model: {model_display_name}"
-            if ui_inputs['task_type'] != 'txt2img': params_string += f", Denoise: {ui_inputs['denoise']}"
-            if ui_inputs.get('clip_skip') and ui_inputs['clip_skip'] != 1: params_string += f", Clip skip: {abs(ui_inputs['clip_skip'])}"
-            if loras_string: params_string += f", {loras_string}"
-            
-            pil_image.info = {'parameters': params_string.strip()}
-            output_images.append(pil_image)
-            
-        return output_images
+        # Execute the workflow; the SaveImage node returns its saved file path(s).
+        saved_output = WorkflowExecutor.execute_workflow(workflow, initial_objects=initial_objects)
+        # Execute the workflow; the SaveImage node returns the saved file path(s).
+        # Ensure we have a list of paths.
+        if isinstance(saved_output, (list, tuple)):
+            saved_paths = list(saved_output)
+        else:
+            saved_paths = [saved_output]
+        # The workflow may contain more than one SaveImage node (e.g., from base sampler
+        # plus conditioning partials), which can produce duplicate images with different
+        # filenames. To avoid showing duplicate thumbnails in the Gallery, deduplicate the
+        # list by file content hash (SHA‑256). This keeps the first occurrence of each unique
+        # image.
+        import hashlib
+        unique_hashes = set()
+        deduped_paths = []
+        for p in saved_paths:
+            try:
+                with open(p, "rb") as f:
+                    h = hashlib.sha256(f.read()).hexdigest()
+                if h not in unique_hashes:
+                    unique_hashes.add(h)
+                    deduped_paths.append(p)
+            except Exception as e:
+                # If reading fails, keep the path (will surface later as a missing file).
+                deduped_paths.append(p)
+        return deduped_paths
 
     def run(self, ui_inputs: Dict, progress):
         progress(0, desc="Preparing models...")
@@ -212,46 +218,17 @@ class SdImagePipeline(BasePipeline):
                 progress=progress
             )
             
-            import json
-            import glob
-            from PIL import PngImagePlugin
-            
-            prompt_json = json.dumps(workflow)
-            
-            out_dir = os.path.abspath(OUTPUT_DIR)
-            os.makedirs(out_dir, exist_ok=True)
-            
-            try:
-                existing_files = glob.glob(os.path.join(out_dir, "gen_*.png"))
-                existing_files.sort(key=os.path.getmtime)
-                while len(existing_files) > 50:
-                    os.remove(existing_files.pop(0))
-            except Exception as e:
-                print(f"Warning: Failed to cleanup output dir: {e}")
-
-            final_results = []
-            for img in results:
-                if not isinstance(img, Image.Image):
-                    final_results.append(img)
-                    continue
-                    
-                metadata = PngImagePlugin.PngInfo()
-                params_string = img.info.get("parameters", "")
-                if params_string:
-                    metadata.add_text("parameters", params_string)
-                metadata.add_text("prompt", prompt_json)
-                
-                filename = f"gen_{random.randint(1000000, 9999999)}.png"
-                filepath = os.path.join(out_dir, filename)
-                img.save(filepath, "PNG", pnginfo=metadata)
-                final_results.append(filepath)
+            # The workflow's SaveImage node already writes the generated images to the
+            # output directory and returns the file path(s). No additional saving or
+            # metadata injection is required.
+            # Clean up any temporary files that were created for the workflow inputs.
+            # The ``finally`` block below already handles this cleanup.
+            # Simply return the list of file paths obtained from the workflow.
+            # (If ``results`` is a single string, convert it to a list for Gradio.)
+            if isinstance(results, (list, tuple)):
+                final_results = list(results)
+            else:
+                final_results = [results]
             
-            results = final_results
-
-        finally:
-            for temp_file in temp_files_to_clean:
-                if temp_file and os.path.exists(temp_file):
-                    os.remove(temp_file)
-                    print(f"✅ Cleaned up temp file: {temp_file}")
-
-        return results
+            # The surrounding ``try``/``finally`` handles temp file cleanup.
+            return final_results