Restore image saving from tensor, dedupe duplicates, simplify result handling

core/pipelines/sd_image_pipeline.py

@@ -3,6 +3,7 @@ import random
 import shutil
 import torch
 import gradio as gr
+from PIL import Image
 from typing import List, Dict, Any
 
 
@@ -33,32 +34,53 @@ class SdImagePipeline(BasePipeline):
         """
         progress(0.4, desc="Executing workflow...")
         initial_objects = {}
-        # 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.
+        # Execute the workflow; the SaveImage node returns the image tensor(s).
+        decoded_images_tensor = WorkflowExecutor.execute_workflow(workflow, initial_objects=initial_objects)
+
+        # Convert tensors to PIL images, attach metadata, and save them to the output folder.
+        out_dir = os.path.abspath(OUTPUT_DIR)
+        os.makedirs(out_dir, exist_ok=True)
+        saved_file_paths = []
+        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: {ui_inputs['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()}
+            filename = f"gen_{random.randint(1000000, 9999999)}.png"
+            filepath = os.path.join(out_dir, filename)
+            pil_image.save(filepath, "PNG")
+            saved_file_paths.append(filepath)
+
+        # After both the internal SaveImage node (which already wrote files) and the above
+        # manual save, the output directory may contain duplicate images. Deduplicate by
+        # content hash (SHA‑256) and return a list of unique file paths.
         import hashlib
         unique_hashes = set()
         deduped_paths = []
-        for p in saved_paths:
+        for p in sorted(os.listdir(out_dir)):
+            full_path = os.path.join(out_dir, p)
             try:
-                with open(p, "rb") as f:
+                with open(full_path, "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)
+                    deduped_paths.append(full_path)
+            except Exception:
+                continue
         return deduped_paths
 
     def run(self, ui_inputs: Dict, progress):
@@ -218,20 +240,8 @@ class SdImagePipeline(BasePipeline):
                 progress=progress
             )
             
-            # 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]
-            
-            # The surrounding ``try``/``finally`` handles temp file cleanup.
-            return final_results
+            # The workflow already saved images and returned a deduplicated list of file paths.
+            return results
         finally:
             for temp_file in temp_files_to_clean:
                 if temp_file and os.path.exists(temp_file):