core/pipelines/sd_image_pipeline.py

import os
import random
import shutil
import torch
import gradio as gr
from PIL import Image, ImageChops
from typing import List, Dict, Any
from collections import defaultdict, deque
import numpy as np

from .base_pipeline import BasePipeline
from core.settings import *
from comfy_integration.nodes import *
from utils.app_utils import get_value_at_index, sanitize_prompt, get_lora_path, get_embedding_path, ensure_controlnet_model_downloaded, ensure_ipadapter_models_downloaded, sanitize_filename
from core.workflow_assembler import WorkflowAssembler

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 _topological_sort(self, workflow: Dict[str, Any]) -> List[str]:
        graph = defaultdict(list)
        in_degree = {node_id: 0 for node_id in workflow}

        for node_id, node_info in workflow.items():
            for input_value in node_info.get('inputs', {}).values():
                if isinstance(input_value, list) and len(input_value) == 2 and isinstance(input_value[0], str):
                    source_node_id = input_value[0]
                    if source_node_id in workflow:
                        graph[source_node_id].append(node_id)
                        in_degree[node_id] += 1

        queue = deque([node_id for node_id, degree in in_degree.items() if degree == 0])
        
        sorted_nodes = []
        while queue:
            current_node_id = queue.popleft()
            sorted_nodes.append(current_node_id)

            for neighbor_node_id in graph[current_node_id]:
                in_degree[neighbor_node_id] -= 1
                if in_degree[neighbor_node_id] == 0:
                    queue.append(neighbor_node_id)
        
        if len(sorted_nodes) != len(workflow):
            raise RuntimeError("Workflow contains a cycle and cannot be executed.")
            
        return sorted_nodes

    def _execute_workflow(self, workflow: Dict[str, Any], initial_objects: Dict[str, Any]):
        with torch.no_grad():
            computed_outputs = initial_objects
            
            try:
                sorted_node_ids = self._topological_sort(workflow)
                print(f"--- [Workflow Executor] Execution order: {sorted_node_ids}")
            except RuntimeError as e:
                print("--- [Workflow Executor] ERROR: Failed to sort workflow. Dumping graph details. ---")
                for node_id, node_info in workflow.items():
                    print(f"  Node {node_id} ({node_info['class_type']}):")
                    for input_name, input_value in node_info['inputs'].items():
                         if isinstance(input_value, list) and len(input_value) == 2 and isinstance(input_value[0], str):
                             print(f"    - {input_name} <- [{input_value[0]}, {input_value[1]}]")
                raise e

            for node_id in sorted_node_ids:
                if node_id in computed_outputs:
                    continue
                    
                node_info = workflow[node_id]
                class_type = node_info['class_type']
                
                is_loader_with_filename = 'Loader' in class_type and any(key.endswith('_name') for key in node_info['inputs'])
                if node_id in initial_objects and is_loader_with_filename:
                    continue
                
                node_class = NODE_CLASS_MAPPINGS.get(class_type)
                if node_class is None:
                     raise RuntimeError(f"Could not find node class '{class_type}'. Is it imported in comfy_integration/nodes.py?")
                
                node_instance = node_class()
                
                kwargs = {}
                for param_name, param_value in node_info['inputs'].items():
                    if isinstance(param_value, list) and len(param_value) == 2 and isinstance(param_value[0], str):
                        source_node_id, output_index = param_value
                        if source_node_id not in computed_outputs:
                            raise RuntimeError(f"Workflow integrity error: Output of node {source_node_id} needed for {node_id} but not yet computed.")
                        
                        source_output_tuple = computed_outputs[source_node_id]
                        kwargs[param_name] = get_value_at_index(source_output_tuple, output_index)
                    else:
                        kwargs[param_name] = param_value

                function_name = getattr(node_class, 'FUNCTION')
                execution_method = getattr(node_instance, function_name)
                
                result = execution_method(**kwargs)
                computed_outputs[node_id] = result
            
            final_node_id = None
            for node_id in reversed(sorted_node_ids):
                 if workflow[node_id]['class_type'] == 'SaveImage':
                     final_node_id = node_id
                     break
            
            if not final_node_id:
                raise RuntimeError("Workflow does not contain a 'SaveImage' node as the output.")

            save_image_inputs = workflow[final_node_id]['inputs']
            image_source_node_id, image_source_index = save_image_inputs['images']
            
            return get_value_at_index(computed_outputs[image_source_node_id], image_source_index)

    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']
        
        progress(0.4, desc="Executing workflow...")
        
        initial_objects = {}

        decoded_images_tensor = self._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

    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", "sdxl")
        
        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)
        self.model_manager.ensure_models_downloaded(required_models, progress=progress)
        
        lora_data = ui_inputs.get('lora_data', [])
        active_loras_for_gpu, active_loras_for_meta = [], []
        if lora_data:
            sources, ids, scales, files = lora_data[0::4], lora_data[1::4], lora_data[2::4], lora_data[3::4]
            for i, (source, lora_id, scale, _) in enumerate(zip(sources, ids, scales, files)):
                if scale > 0 and lora_id and lora_id.strip():
                    lora_filename = None
                    if source == "File":
                        lora_filename = sanitize_filename(lora_id)
                    elif source == "Civitai":
                        local_path, status = get_lora_path(source, lora_id, os.environ.get("CIVITAI_API_KEY", ""), progress)
                        if local_path: lora_filename = os.path.basename(local_path)
                        else: raise gr.Error(f"Failed to prepare LoRA {lora_id}: {status}")
                    
                    if lora_filename:
                        active_loras_for_gpu.append({"lora_name": lora_filename, "strength_model": scale, "strength_clip": scale})
                        active_loras_for_meta.append(f"{source} {lora_id}:{scale}")

        ui_inputs['denoise'] = 1.0
        if task_type == 'img2img': ui_inputs['denoise'] = ui_inputs.get('img2img_denoise', 0.7)
        elif task_type == 'hires_fix': ui_inputs['denoise'] = ui_inputs.get('hires_denoise', 0.55)
        
        temp_files_to_clean = []
        if not os.path.exists(INPUT_DIR): os.makedirs(INPUT_DIR)

        if task_type == 'img2img':
            input_image_pil = ui_inputs.get('img2img_image')
            if not input_image_pil:
                raise gr.Error("Please upload an image for Image-to-Image.")
            temp_file_path = os.path.join(INPUT_DIR, f"temp_input_{random.randint(1000, 9999)}.png")
            input_image_pil.save(temp_file_path, "PNG")
            ui_inputs['input_image'] = os.path.basename(temp_file_path)
            temp_files_to_clean.append(temp_file_path)
            ui_inputs['width'] = input_image_pil.width
            ui_inputs['height'] = input_image_pil.height
        
        elif task_type == 'inpaint':
            inpaint_dict = ui_inputs.get('inpaint_image_dict')
            if not inpaint_dict or not inpaint_dict.get('background') or not inpaint_dict.get('layers'):
                raise gr.Error("Inpainting requires an input image and a drawn mask.")
            
            background_img = inpaint_dict['background'].convert("RGBA")
            composite_mask_pil = Image.new('L', background_img.size, 0)
            for layer in inpaint_dict['layers']:
                if layer:
                    layer_alpha = layer.split()[-1]
                    composite_mask_pil = ImageChops.lighter(composite_mask_pil, layer_alpha)

            inverted_mask_alpha = Image.fromarray(255 - np.array(composite_mask_pil), mode='L')
            r, g, b, _ = background_img.split()
            composite_image_with_mask = Image.merge('RGBA', [r, g, b, inverted_mask_alpha])
            
            temp_file_path = os.path.join(INPUT_DIR, f"temp_inpaint_composite_{random.randint(1000, 9999)}.png")
            composite_image_with_mask.save(temp_file_path, "PNG")
            
            ui_inputs['input_image'] = os.path.basename(temp_file_path)
            temp_files_to_clean.append(temp_file_path)
            ui_inputs.pop('inpaint_mask', None)

        elif task_type == 'outpaint':
            input_image_pil = ui_inputs.get('outpaint_image')
            if not input_image_pil:
                raise gr.Error("Please upload an image for Outpainting.")
            temp_file_path = os.path.join(INPUT_DIR, f"temp_input_{random.randint(1000, 9999)}.png")
            input_image_pil.save(temp_file_path, "PNG")
            ui_inputs['input_image'] = os.path.basename(temp_file_path)
            temp_files_to_clean.append(temp_file_path)
            
            ui_inputs['megapixels'] = 0.25
            ui_inputs['grow_mask_by'] = ui_inputs.get('feathering', 10)

        elif task_type == 'hires_fix':
            input_image_pil = ui_inputs.get('hires_image')
            if not input_image_pil:
                raise gr.Error("Please upload an image for Hires Fix.")
            temp_file_path = os.path.join(INPUT_DIR, f"temp_input_{random.randint(1000, 9999)}.png")
            input_image_pil.save(temp_file_path, "PNG")
            ui_inputs['input_image'] = os.path.basename(temp_file_path)
            temp_files_to_clean.append(temp_file_path)

        embedding_data = ui_inputs.get('embedding_data', [])
        embedding_filenames = []
        if embedding_data:
            emb_sources, emb_ids, emb_files = embedding_data[0::3], embedding_data[1::3], embedding_data[2::3]
            for i, (source, emb_id, _) in enumerate(zip(emb_sources, emb_ids, emb_files)):
                if emb_id and emb_id.strip():
                    emb_filename = None
                    if source == "File":
                        emb_filename = sanitize_filename(emb_id)
                    elif source == "Civitai":
                        local_path, status = get_embedding_path(source, emb_id, os.environ.get("CIVITAI_API_KEY", ""), progress)
                        if local_path: emb_filename = os.path.basename(local_path)
                        else: raise gr.Error(f"Failed to prepare Embedding {emb_id}: {status}")
                    
                    if emb_filename:
                        embedding_filenames.append(emb_filename)
        
        if embedding_filenames:
            embedding_prompt_text = " ".join([f"embedding:{f}" for f in embedding_filenames])
            if ui_inputs['positive_prompt']:
                ui_inputs['positive_prompt'] = f"{ui_inputs['positive_prompt']}, {embedding_prompt_text}"
            else:
                ui_inputs['positive_prompt'] = embedding_prompt_text

        controlnet_data = ui_inputs.get('controlnet_data', [])
        active_controlnets = []
        if controlnet_data:
            (cn_images, _, _, cn_strengths, cn_filepaths) = [controlnet_data[i::5] for i in range(5)]
            for i in range(len(cn_images)):
                if cn_images[i] and cn_strengths[i] > 0 and cn_filepaths[i] and cn_filepaths[i] != "None":
                    ensure_controlnet_model_downloaded(cn_filepaths[i], progress)
                    if not os.path.exists(INPUT_DIR): os.makedirs(INPUT_DIR)
                    cn_temp_path = os.path.join(INPUT_DIR, f"temp_cn_{i}_{random.randint(1000, 9999)}.png")
                    cn_images[i].save(cn_temp_path, "PNG")
                    temp_files_to_clean.append(cn_temp_path)
                    active_controlnets.append({
                        "image": os.path.basename(cn_temp_path), "strength": cn_strengths[i],
                        "start_percent": 0.0, "end_percent": 1.0, "control_net_name": cn_filepaths[i]
                    })
        
        anima_controlnet_lllite_data = ui_inputs.get('anima_controlnet_lllite_data', [])
        active_anima_controlnets = []
        if anima_controlnet_lllite_data:
            (cn_images, _, _, cn_strengths, cn_filepaths, cn_starts, cn_ends) = [anima_controlnet_lllite_data[i::7] for i in range(7)]
            for i in range(len(cn_images)):
                if cn_images[i] and cn_strengths[i] > 0 and cn_filepaths[i] and cn_filepaths[i] != "None":
                    from utils.app_utils import _ensure_model_downloaded
                    _ensure_model_downloaded(cn_filepaths[i], progress)
                    if not os.path.exists(INPUT_DIR): os.makedirs(INPUT_DIR)
                    cn_temp_path = os.path.join(INPUT_DIR, f"temp_anima_cn_{i}_{random.randint(1000, 9999)}.png")
                    cn_images[i].save(cn_temp_path, "PNG")
                    temp_files_to_clean.append(cn_temp_path)
                    active_anima_controlnets.append({
                        "image": os.path.basename(cn_temp_path), "strength": cn_strengths[i],
                        "start_percent": cn_starts[i], "end_percent": cn_ends[i], "control_net_name": cn_filepaths[i]
                    })

        diffsynth_controlnet_data = ui_inputs.get('diffsynth_controlnet_data', [])
        active_diffsynth_controlnets = []
        if diffsynth_controlnet_data:
            (cn_images, _, _, cn_strengths, cn_filepaths) = [diffsynth_controlnet_data[i::5] for i in range(5)]
            for i in range(len(cn_images)):
                if cn_images[i] and cn_strengths[i] > 0 and cn_filepaths[i] and cn_filepaths[i] != "None":
                    ensure_controlnet_model_downloaded(cn_filepaths[i], progress)

                    if not os.path.exists(INPUT_DIR): os.makedirs(INPUT_DIR)
                    cn_temp_path = os.path.join(INPUT_DIR, f"temp_diffsynth_cn_{i}_{random.randint(1000, 9999)}.png")
                    cn_images[i].save(cn_temp_path, "PNG")
                    temp_files_to_clean.append(cn_temp_path)
                    active_diffsynth_controlnets.append({
                        "image": os.path.basename(cn_temp_path), "strength": cn_strengths[i],
                        "control_net_name": cn_filepaths[i]
                    })
        
        ipadapter_data = ui_inputs.get('ipadapter_data', [])
        active_ipadapters = []
        if ipadapter_data:
            num_ipa_units = (len(ipadapter_data) - 5) // 3
            final_preset, final_weight, final_lora_strength, final_embeds_scaling, final_combine_method = ipadapter_data[-5:]
            ipa_images, ipa_weights, ipa_lora_strengths = [ipadapter_data[i*num_ipa_units:(i+1)*num_ipa_units] for i in range(3)]
            all_presets_to_download = set()
            for i in range(num_ipa_units):
                if ipa_images[i] and ipa_weights[i] > 0 and final_preset:
                    all_presets_to_download.add(final_preset)
                    if not os.path.exists(INPUT_DIR): os.makedirs(INPUT_DIR)
                    ipa_temp_path = os.path.join(INPUT_DIR, f"temp_ipa_{i}_{random.randint(1000, 9999)}.png")
                    ipa_images[i].save(ipa_temp_path, "PNG")
                    temp_files_to_clean.append(ipa_temp_path)
                    active_ipadapters.append({
                        "image": os.path.basename(ipa_temp_path), "preset": final_preset,
                        "weight": ipa_weights[i], "lora_strength": ipa_lora_strengths[i]
                    })
            if active_ipadapters and final_preset:
                all_presets_to_download.add(final_preset)
            for preset in all_presets_to_download:
                ensure_ipadapter_models_downloaded(preset, progress)
                
            model_type_key = 'sd15' if workflow_model_type == 'sd15' else 'sdxl'
            if active_ipadapters:
                active_ipadapters.append({
                    'is_final_settings': True, 'model_type': model_type_key, 'final_preset': final_preset,
                    'final_weight': final_weight, 'final_lora_strength': final_lora_strength,
                    'final_embeds_scaling': final_embeds_scaling, 'final_combine_method': final_combine_method
                })
                
        flux1_ipadapter_data = ui_inputs.get('flux1_ipadapter_data', [])
        active_flux1_ipadapters = []
        if flux1_ipadapter_data:
            num_units = len(flux1_ipadapter_data) // 4
            f_images = flux1_ipadapter_data[0*num_units : 1*num_units]
            f_weights = flux1_ipadapter_data[1*num_units : 2*num_units]
            f_starts = flux1_ipadapter_data[2*num_units : 3*num_units]
            f_ends = flux1_ipadapter_data[3*num_units : 4*num_units]
            for i in range(len(f_images)):
                if f_images[i] and f_weights[i] > 0:
                    from utils.app_utils import _ensure_model_downloaded
                    for filename in ["ip-adapter.bin"]:
                        _ensure_model_downloaded(filename, progress)
                    
                    from huggingface_hub import snapshot_download
                    progress(0.5, desc="Caching HF SigLIP model...")
                    snapshot_download(
                        repo_id="google/siglip-so400m-patch14-384", 
                        allow_patterns=["*.json", "*.safetensors", "*.txt"], 
                        ignore_patterns=["*.msgpack", "*.h5", "*.bin"]
                    )

                    temp_path = os.path.join(INPUT_DIR, f"temp_fipa_{i}_{random.randint(1000, 9999)}.png")
                    f_images[i].save(temp_path, "PNG")
                    temp_files_to_clean.append(temp_path)
                    active_flux1_ipadapters.append({
                        "image": os.path.basename(temp_path),
                        "weight": f_weights[i], "start_percent": f_starts[i], "end_percent": f_ends[i]
                    })
                    
        sd3_ipadapter_data = ui_inputs.get('sd3_ipadapter_chain', [])
        active_sd3_ipadapters = []
        if sd3_ipadapter_data:
            num_units = len(sd3_ipadapter_data) // 4
            s_images = sd3_ipadapter_data[0*num_units : 1*num_units]
            s_weights = sd3_ipadapter_data[1*num_units : 2*num_units]
            s_starts = sd3_ipadapter_data[2*num_units : 3*num_units]
            s_ends = sd3_ipadapter_data[3*num_units : 4*num_units]
            sd3_ipa_downloaded = False
            for i in range(len(s_images)):
                if s_images[i] and s_weights[i] > 0:
                    if not sd3_ipa_downloaded:
                        from utils.app_utils import ensure_sd3_ipadapter_models_downloaded
                        ensure_sd3_ipadapter_models_downloaded(progress)
                        sd3_ipa_downloaded = True
                    temp_path = os.path.join(INPUT_DIR, f"temp_s3ipa_{i}_{random.randint(1000, 9999)}.png")
                    s_images[i].save(temp_path, "PNG")
                    temp_files_to_clean.append(temp_path)
                    active_sd3_ipadapters.append({
                        "image": os.path.basename(temp_path),
                        "weight": s_weights[i], "start_percent": s_starts[i], "end_percent": s_ends[i]
                    })
                    
        style_data = ui_inputs.get('style_data', [])
        active_styles = []
        if style_data:
            num_units = len(style_data) // 2
            st_images = style_data[0*num_units : 1*num_units]
            st_strengths = style_data[1*num_units : 2*num_units]
            for i in range(len(st_images)):
                if st_images[i] and st_strengths[i] > 0:
                    from utils.app_utils import _ensure_model_downloaded
                    _ensure_model_downloaded("sigclip_vision_patch14_384.safetensors", progress)
                    temp_path = os.path.join(INPUT_DIR, f"temp_style_{i}_{random.randint(1000, 9999)}.png")
                    st_images[i].save(temp_path, "PNG")
                    temp_files_to_clean.append(temp_path)
                    active_styles.append({
                        "image": os.path.basename(temp_path), "strength": st_strengths[i]
                    })
        
        reference_latent_data = ui_inputs.get('reference_latent_data', [])
        active_reference_latents = []
        if reference_latent_data:
            for img in reference_latent_data:
                if img:
                    if not os.path.exists(INPUT_DIR): os.makedirs(INPUT_DIR)
                    temp_path = os.path.join(INPUT_DIR, f"temp_ref_{random.randint(1000, 9999)}.png")
                    img.save(temp_path, "PNG")
                    temp_files_to_clean.append(temp_path)
                    active_reference_latents.append(os.path.basename(temp_path))
        
        from utils.app_utils import get_vae_path
        vae_source = ui_inputs.get('vae_source')
        vae_id = ui_inputs.get('vae_id')
        vae_name_override = None
        if vae_source and vae_source != "None":
            if vae_source == "File":
                vae_name_override = sanitize_filename(vae_id)
            elif vae_source == "Civitai" and vae_id and vae_id.strip():
                local_path, status = get_vae_path(vae_source, vae_id, os.environ.get("CIVITAI_API_KEY", ""), progress)
                if local_path: vae_name_override = os.path.basename(local_path)
                else: raise gr.Error(f"Failed to prepare VAE {vae_id}: {status}")
        if vae_name_override:
            ui_inputs['vae_name'] = vae_name_override

        conditioning_data = ui_inputs.get('conditioning_data', [])
        active_conditioning = []
        if conditioning_data:
            num_units = len(conditioning_data) // 6
            prompts, widths, heights, xs, ys, strengths = [conditioning_data[i*num_units : (i+1)*num_units] for i in range(6)]
            for i in range(num_units):
                if prompts[i] and prompts[i].strip():
                    active_conditioning.append({
                        "prompt": prompts[i], "width": int(widths[i]), "height": int(heights[i]),
                        "x": int(xs[i]), "y": int(ys[i]), "strength": float(strengths[i])
                    })

        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_smoothing_data = []
        if workflow_model_type == 'hidream-o1' and model_display_name == "HiDream-O1-Image":
            hidream_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,
            "vae_chain": [ui_inputs.get('vae_name')] if ui_inputs.get('vae_name') else [],
            "hidream_smoothing_chain": hidream_smoothing_data,
        }

        if isinstance(path_or_components, dict):
            workflow_inputs.update({
                'unet_name': path_or_components.get('unet'),
                '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...")

        try:
            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
            )
            
            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)
            
            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