import os import random import shutil import torch import gradio as gr from typing import List, Dict, Any from .base_pipeline import BasePipeline from core.settings import * from utils.app_utils import sanitize_prompt from core.workflow_assembler import WorkflowAssembler from .workflow_executor import WorkflowExecutor from .pipeline_input_processor import process_pipeline_inputs class SdImagePipeline(BasePipeline): def get_required_models(self, model_display_name: str, **kwargs) -> List[str]: model_info = ALL_MODEL_MAP.get(model_display_name) if not model_info: return [model_display_name] path_or_components = model_info[1] if isinstance(path_or_components, dict): return [v for v in path_or_components.values() if v and v != "pixel_space"] else: 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)): """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 = {} # 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...") task_type = ui_inputs['task_type'] model_display_name = ui_inputs['model_display_name'] model_type = MODEL_TYPE_MAP.get(model_display_name, 'sdxl') architectures_dict = ARCHITECTURES_CONFIG.get('architectures', {}) workflow_model_type = architectures_dict.get(model_type, {}).get("model_type", model_type.lower().replace(" ", "").replace(".", "")) ui_inputs['positive_prompt'] = sanitize_prompt(ui_inputs.get('positive_prompt', '')) ui_inputs['negative_prompt'] = sanitize_prompt(ui_inputs.get('negative_prompt', '')) if 'clip_skip' in ui_inputs and ui_inputs['clip_skip'] is not None: ui_inputs['clip_skip'] = -int(ui_inputs['clip_skip']) else: ui_inputs['clip_skip'] = -1 required_models = self.get_required_models(model_display_name=model_display_name) is_pid_enabled = (ui_inputs.get('pid_settings', 'OFF') == 'ON' and task_type == 'txt2img') if is_pid_enabled: import yaml pid_config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), 'yaml', 'pid.yaml') pid_unet_name = "pid_flux1_1024_to_4096_4step_mxfp8.safetensors" try: with open(pid_config_path, 'r', encoding='utf-8') as f: pid_config = yaml.safe_load(f) or {} pid_items = pid_config.get("PiD", []) for item in pid_items: archs = item.get("architectures", []) if workflow_model_type in archs: pid_unet_name = item.get("filepath") break except Exception as e: print(f"Error loading PiD config for download: {e}") if pid_unet_name not in required_models: required_models.append(pid_unet_name) if "gemma_2_2b_it_elm_fp8_scaled.safetensors" not in required_models: required_models.append("gemma_2_2b_it_elm_fp8_scaled.safetensors") self.model_manager.ensure_models_downloaded(required_models, progress=progress) temp_files_to_clean = [] try: processed = process_pipeline_inputs(ui_inputs, progress, workflow_model_type) temp_files_to_clean.extend(processed["temp_files_to_clean"]) active_loras_for_gpu = processed["active_loras_for_gpu"] active_loras_for_meta = processed["active_loras_for_meta"] active_controlnets = processed["active_controlnets"] active_anima_controlnets = processed["active_anima_controlnets"] active_diffsynth_controlnets = processed["active_diffsynth_controlnets"] active_ipadapters = processed["active_ipadapters"] active_flux1_ipadapters = processed["active_flux1_ipadapters"] active_sd3_ipadapters = processed["active_sd3_ipadapters"] active_styles = processed["active_styles"] active_reference_latents = processed["active_reference_latents"] active_hidream_o1_reference = processed["active_hidream_o1_reference"] active_conditioning = processed["active_conditioning"] loras_string = f"LoRAs: [{', '.join(active_loras_for_meta)}]" if active_loras_for_meta else "" progress(0.8, desc="Assembling workflow...") if ui_inputs.get('seed') == -1: ui_inputs['seed'] = random.randint(0, 2**32 - 1) model_info = ALL_MODEL_MAP[model_display_name] path_or_components = model_info[1] latent_type = model_info[3] if len(model_info) > 3 and model_info[3] else 'latent' latent_generator_template = "EmptyLatentImage" if latent_type == 'sd3_latent': latent_generator_template = "EmptySD3LatentImage" elif latent_type == 'chroma_radiance_latent': latent_generator_template = "EmptyChromaRadianceLatentImage" elif latent_type == 'hunyuan_latent': latent_generator_template = "EmptyHunyuanImageLatent" dynamic_values = { 'task_type': ui_inputs['task_type'], 'model_type': workflow_model_type, 'latent_type': latent_type, 'latent_generator_template': latent_generator_template } recipe_path = os.path.join(os.path.dirname(__file__), "workflow_recipes", "sd_unified_recipe.yaml") assembler = WorkflowAssembler(recipe_path, dynamic_values=dynamic_values) hidream_o1_smoothing_data = [] if workflow_model_type == 'hidream-o1' and model_display_name == "HiDream-O1-Image": hidream_o1_smoothing_data.append({}) workflow_inputs = { **ui_inputs, "positive_prompt": ui_inputs['positive_prompt'], "negative_prompt": ui_inputs['negative_prompt'], "seed": ui_inputs['seed'], "steps": ui_inputs['num_inference_steps'], "cfg": ui_inputs['guidance_scale'], "sampler_name": ui_inputs['sampler'], "scheduler": ui_inputs['scheduler'], "batch_size": ui_inputs['batch_size'], "clip_skip": ui_inputs['clip_skip'], "denoise": ui_inputs['denoise'], "vae_name": ui_inputs.get('vae_name'), "guidance": ui_inputs.get('guidance', 3.5), "lora_chain": active_loras_for_gpu, "controlnet_chain": active_controlnets if not active_anima_controlnets else [], "anima_controlnet_lllite_chain": active_anima_controlnets, "diffsynth_controlnet_chain": active_diffsynth_controlnets, "ipadapter_chain": active_ipadapters, "flux1_ipadapter_chain": active_flux1_ipadapters, "sd3_ipadapter_chain": active_sd3_ipadapters, "style_chain": active_styles, "conditioning_chain": active_conditioning, "reference_latent_chain": active_reference_latents, "hidream_o1_reference_chain": active_hidream_o1_reference, "vae_chain": [ui_inputs.get('vae_name')] if ui_inputs.get('vae_name') else [], "hidream_o1_smoothing_chain": hidream_o1_smoothing_data, "pid_chain": [ui_inputs.get('pid_settings', 'OFF')] if is_pid_enabled else [], "scheduler_width": ui_inputs.get('width', 1024), "scheduler_height": ui_inputs.get('height', 1024), } if isinstance(path_or_components, dict): workflow_inputs.update({ 'unet_name': path_or_components.get('unet'), 'unet_uncond_name': path_or_components.get('unet_uncond'), 'vae_name': ui_inputs.get('vae_name') or path_or_components.get('vae'), 'clip_name': path_or_components.get('clip'), 'clip1_name': path_or_components.get('clip1'), 'clip2_name': path_or_components.get('clip2'), 'clip3_name': path_or_components.get('clip3'), 'clip4_name': path_or_components.get('clip4'), 'lora_name': path_or_components.get('lora'), }) else: workflow_inputs['model_name'] = path_or_components if task_type == 'txt2img': workflow_inputs['width'] = ui_inputs['width'] workflow_inputs['height'] = ui_inputs['height'] workflow = assembler.assemble(workflow_inputs) progress(1.0, desc="All models ready. Requesting GPU for generation...") results = self._execute_gpu_logic( self._gpu_logic, duration=ui_inputs['zero_gpu_duration'], default_duration=60, task_name=f"ImageGen ({task_type})", ui_inputs=ui_inputs, loras_string=loras_string, workflow=workflow, assembler=assembler, 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 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}")