Upload folder using huggingface_hub

chain_injectors/pid_injector.py

@@ -0,0 +1,278 @@
+import os
+import yaml
+import random
+
+def load_pid_config():
+    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    pid_path = os.path.join(project_root, 'yaml', 'pid.yaml')
+    with open(pid_path, 'r', encoding='utf-8') as f:
+        return yaml.safe_load(f) or {}
+
+def load_model_config():
+    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    model_list_path = os.path.join(project_root, 'yaml', 'model_list.yaml')
+    with open(model_list_path, 'r', encoding='utf-8') as f:
+        return yaml.safe_load(f) or {}
+
+def inject(assembler, chain_definition, chain_items):
+    if not chain_items:
+        return
+
+    pid_config = {}
+    try:
+        pid_config = load_pid_config() or {}
+    except Exception as e:
+        print(f"Error loading PiD config: {e}")
+
+    pid_items = pid_config.get("PiD", [])
+    architectures_settings = {}
+    default_settings = {"unet_name": "pid_flux1_1024_to_4096_4step_mxfp8.safetensors", "latent_format": "flux"}
+    
+    for item in pid_items:
+        unet_name = item.get("filepath")
+        latent_format = item.get("latent_format")
+        archs = item.get("architectures", [])
+        for arch in archs:
+            architectures_settings[arch] = {
+                "unet_name": unet_name,
+                "latent_format": latent_format
+            }
+            if arch == "flux1":
+                default_settings = {
+                    "unet_name": unet_name,
+                    "latent_format": latent_format
+                }
+
+    ksampler_name = chain_definition.get('ksampler_node', 'ksampler')
+    if ksampler_name not in assembler.node_map:
+        print(f"Warning: [PiD Injector] KSampler node '{ksampler_name}' not found. Skipping.")
+        return
+    
+    original_ksampler_id = assembler.node_map[ksampler_name]
+    
+    original_vae_loader_id = assembler.node_map.get('vae_loader')
+    original_vae_decode_id = assembler.node_map.get('vae_decode')
+    original_pos_prompt_id = assembler.node_map.get('pos_prompt')
+    original_neg_prompt_id = assembler.node_map.get('neg_prompt')
+    
+    if not original_vae_loader_id:
+        for node_id, node_data in assembler.workflow.items():
+            if node_data.get('class_type') == 'VAELoader':
+                original_vae_loader_id = node_id
+                break
+                
+    if not original_vae_decode_id:
+        for node_id, node_data in assembler.workflow.items():
+            if node_data.get('class_type') == 'VAEDecode':
+                original_vae_decode_id = node_id
+                break
+
+    if not original_pos_prompt_id or not original_neg_prompt_id:
+        for node_id, node_data in assembler.workflow.items():
+            if node_data.get('class_type') == 'CLIPTextEncode':
+                title = node_data.get('_meta', {}).get('title', '')
+                if 'Positive' in title:
+                    if not original_pos_prompt_id:
+                        original_pos_prompt_id = node_id
+                elif 'Negative' in title:
+                    if not original_neg_prompt_id:
+                        original_neg_prompt_id = node_id
+                    
+    pos_text = ""
+    if original_pos_prompt_id and original_pos_prompt_id in assembler.workflow:
+        pos_text = assembler.workflow[original_pos_prompt_id]['inputs'].get('text', '')
+
+    neg_text = ""
+    if original_neg_prompt_id and original_neg_prompt_id in assembler.workflow:
+        neg_text = assembler.workflow[original_neg_prompt_id]['inputs'].get('text', '')
+
+    clip_loader_id = assembler._get_unique_id()
+    clip_loader_node = assembler._get_node_template("CLIPLoader")
+    clip_loader_node['inputs']['clip_name'] = "gemma_2_2b_it_elm_fp8_scaled.safetensors"
+    clip_loader_node['inputs']['type'] = "pixeldit"
+    clip_loader_node['inputs']['device'] = "default"
+    assembler.workflow[clip_loader_id] = clip_loader_node
+
+    pos_text_encode_id = assembler._get_unique_id()
+    pos_text_encode_node = assembler._get_node_template("CLIPTextEncode")
+    pos_text_encode_node['inputs']['text'] = pos_text
+    pos_text_encode_node['inputs']['clip'] = [clip_loader_id, 0]
+    assembler.workflow[pos_text_encode_id] = pos_text_encode_node
+
+    neg_text_encode_id = assembler._get_unique_id()
+    neg_text_encode_node = assembler._get_node_template("CLIPTextEncode")
+    neg_text_encode_node['inputs']['text'] = neg_text
+    neg_text_encode_node['inputs']['clip'] = [clip_loader_id, 0]
+    assembler.workflow[neg_text_encode_id] = neg_text_encode_node
+
+    active_model_file = None
+    for node_id, node_data in assembler.workflow.items():
+        class_type = node_data.get('class_type')
+        if class_type == 'UNETLoader':
+            active_model_file = node_data.get('inputs', {}).get('unet_name')
+            if active_model_file:
+                break
+        elif class_type == 'CheckpointLoaderSimple':
+            active_model_file = node_data.get('inputs', {}).get('ckpt_name')
+            if active_model_file:
+                break
+
+    architecture = None
+    if active_model_file:
+        try:
+            model_config = load_model_config()
+            checkpoints = model_config.get("Checkpoints", {})
+            for arch_name, arch_data in checkpoints.items():
+                models_list = arch_data.get("models", [])
+                for model_entry in models_list:
+                    if model_entry.get('path') == active_model_file:
+                        architecture = arch_name
+                        break
+                    components_dict = model_entry.get('components', {})
+                    if active_model_file in components_dict.values():
+                        architecture = arch_name
+                        break
+                if architecture:
+                    break
+        except Exception as e:
+            print(f"Error looking up model architecture in PiD injector: {e}")
+
+        if architecture:
+            architecture = architecture.lower().replace(" ", "-").replace(".", "")
+        else:
+            file_lower = active_model_file.lower().replace("-", "").replace("_", "").replace(".", "")
+            for arch in sorted(architectures_settings.keys(), key=len, reverse=True):
+                candidates = [arch]
+                if "-image" in arch:
+                    candidates.append(arch.replace("-image", ""))
+                if "-i1" in arch:
+                    candidates.append(arch.replace("-i1", ""))
+                if "-kv" in arch:
+                    candidates.append(arch.replace("-kv", ""))
+                
+                matched = False
+                for cand in candidates:
+                    if cand.replace("-", "").replace(".", "") in file_lower:
+                        architecture = arch
+                        matched = True
+                        break
+                if matched:
+                    break
+
+    unet_name = default_settings.get("unet_name")
+    latent_format = default_settings.get("latent_format")
+
+    if architecture in architectures_settings:
+        arch_config = architectures_settings[architecture]
+        unet_name = arch_config.get("unet_name", unet_name)
+        latent_format = arch_config.get("latent_format", latent_format)
+    else:
+        print(f"[PiD Injector] Warning: Model architecture '{architecture}' (file: '{active_model_file}') not explicitly mapped. Using default settings.")
+
+    pid_pos_id = assembler._get_unique_id()
+    pid_pos_node = assembler._get_node_template("PiDConditioning")
+    pid_pos_node['inputs']['latent_format'] = latent_format
+    pid_pos_node['inputs']['degrade_sigma'] = 0
+    pid_pos_node['inputs']['positive'] = [pos_text_encode_id, 0]
+    pid_pos_node['inputs']['latent'] = [original_ksampler_id, 0]
+    assembler.workflow[pid_pos_id] = pid_pos_node
+
+    pid_neg_id = assembler._get_unique_id()
+    pid_neg_node = assembler._get_node_template("PiDConditioning")
+    pid_neg_node['inputs']['latent_format'] = latent_format
+    pid_neg_node['inputs']['degrade_sigma'] = 0
+    pid_neg_node['inputs']['positive'] = [neg_text_encode_id, 0]
+    pid_neg_node['inputs']['latent'] = [original_ksampler_id, 0]
+    assembler.workflow[pid_neg_id] = pid_neg_node
+
+    pid_unet_loader_id = assembler._get_unique_id()
+    pid_unet_loader_node = assembler._get_node_template("UNETLoader")
+    pid_unet_loader_node['inputs']['unet_name'] = unet_name
+    pid_unet_loader_node['inputs']['weight_dtype'] = "default"
+    assembler.workflow[pid_unet_loader_id] = pid_unet_loader_node
+
+    orig_width = 1024
+    orig_height = 1024
+    original_latent_source_id = assembler.node_map.get('latent_source')
+    if original_latent_source_id in assembler.workflow:
+        node_inputs = assembler.workflow[original_latent_source_id].get('inputs', {})
+        if 'width' in node_inputs and 'height' in node_inputs:
+            orig_width = node_inputs['width']
+            orig_height = node_inputs['height']
+        else:
+            for node_data in assembler.workflow.values():
+                inputs = node_data.get('inputs', {})
+                if 'width' in inputs and 'height' in inputs and isinstance(inputs['width'], (int, float)) and isinstance(inputs['height'], (int, float)):
+                    if 256 <= inputs['width'] <= 4096 and 256 <= inputs['height'] <= 4096:
+                        orig_width = inputs['width']
+                        orig_height = inputs['height']
+                        break
+    else:
+        for node_data in assembler.workflow.values():
+            inputs = node_data.get('inputs', {})
+            if 'width' in inputs and 'height' in inputs and isinstance(inputs['width'], (int, float)) and isinstance(inputs['height'], (int, float)):
+                if 256 <= inputs['width'] <= 4096 and 256 <= inputs['height'] <= 4096:
+                    orig_width = inputs['width']
+                    orig_height = inputs['height']
+                    break
+
+    empty_latent_id = assembler._get_unique_id()
+    empty_latent_node = assembler._get_node_template("EmptyChromaRadianceLatentImage")
+    empty_latent_node['inputs']['width'] = int(orig_width) * 4
+    empty_latent_node['inputs']['height'] = int(orig_height) * 4
+    empty_latent_node['inputs']['batch_size'] = 1
+    
+    if original_latent_source_id in assembler.workflow:
+        orig_batch_size = assembler.workflow[original_latent_source_id]['inputs'].get('batch_size') or assembler.workflow[original_latent_source_id]['inputs'].get('amount')
+        if orig_batch_size:
+            empty_latent_node['inputs']['batch_size'] = orig_batch_size
+            
+    assembler.workflow[empty_latent_id] = empty_latent_node
+
+    orig_seed = 0
+    if original_ksampler_id in assembler.workflow:
+        orig_seed = assembler.workflow[original_ksampler_id]['inputs'].get('seed', 0)
+        if orig_seed == -1:
+            orig_seed = random.randint(0, 2**32 - 1)
+        else:
+            orig_seed = (orig_seed + 1) % (2**32)
+
+    new_ksampler_id = assembler._get_unique_id()
+    new_ksampler_node = assembler._get_node_template("KSampler")
+    new_ksampler_node['inputs']['seed'] = orig_seed
+    new_ksampler_node['inputs']['steps'] = 4
+    new_ksampler_node['inputs']['cfg'] = 1
+    new_ksampler_node['inputs']['sampler_name'] = "lcm"
+    new_ksampler_node['inputs']['scheduler'] = "simple"
+    new_ksampler_node['inputs']['denoise'] = 1.0
+    new_ksampler_node['inputs']['model'] = [pid_unet_loader_id, 0]
+    new_ksampler_node['inputs']['positive'] = [pid_pos_id, 0]
+    new_ksampler_node['inputs']['negative'] = [pid_neg_id, 0]
+    new_ksampler_node['inputs']['latent_image'] = [empty_latent_id, 0]
+    assembler.workflow[new_ksampler_id] = new_ksampler_node
+
+    pid_vae_loader_id = assembler._get_unique_id()
+    pid_vae_loader_node = assembler._get_node_template("VAELoader")
+    pid_vae_loader_node['inputs']['vae_name'] = "pixel_space"
+    assembler.workflow[pid_vae_loader_id] = pid_vae_loader_node
+
+    pid_vae_decode_id = assembler._get_unique_id()
+    pid_vae_decode_node = assembler._get_node_template("VAEDecode")
+    pid_vae_decode_node['inputs']['samples'] = [new_ksampler_id, 0]
+    pid_vae_decode_node['inputs']['vae'] = [pid_vae_loader_id, 0]
+    assembler.workflow[pid_vae_decode_id] = pid_vae_decode_node
+
+    if original_vae_decode_id:
+        for node_id, node_data in assembler.workflow.items():
+            if 'inputs' in node_data:
+                for input_name, input_val in list(node_data['inputs'].items()):
+                    if isinstance(input_val, list) and len(input_val) == 2:
+                        if input_val[0] == original_vae_decode_id:
+                            node_data['inputs'][input_name] = [pid_vae_decode_id, 0]
+
+    if original_vae_loader_id in assembler.workflow:
+        del assembler.workflow[original_vae_loader_id]
+    if original_vae_decode_id in assembler.workflow:
+        del assembler.workflow[original_vae_decode_id]
+
+    print("[PiD Injector] Successfully injected PiD pipeline and replaced VAE decode/loader.")

core/pipelines/pipeline_input_processor.py

@@ -0,0 +1,334 @@
+import os
+import random
+import numpy as np
+import gradio as gr
+from PIL import Image, ImageChops
+from typing import Dict, Any, List
+
+from core.settings import INPUT_DIR
+from utils.app_utils import (
+    sanitize_filename,
+    get_lora_path,
+    get_embedding_path,
+    ensure_controlnet_model_downloaded,
+    ensure_ipadapter_models_downloaded,
+    _ensure_model_downloaded,
+    ensure_sd3_ipadapter_models_downloaded,
+    get_vae_path,
+)
+
+def process_pipeline_inputs(ui_inputs: Dict[str, Any], progress: gr.Progress, workflow_model_type: str) -> Dict[str, Any]:
+    task_type = ui_inputs['task_type']
+    temp_files_to_clean = []
+    
+    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)
+    
+    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":
+                _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:
+                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:
+                    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:
+                _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))
+                
+    hidream_o1_reference_data = ui_inputs.get('hidream_o1_reference_data', [])
+    active_hidream_o1_reference = []
+    if hidream_o1_reference_data:
+        for img in hidream_o1_reference_data:
+            if img:
+                if not os.path.exists(INPUT_DIR): os.makedirs(INPUT_DIR)
+                temp_path = os.path.join(INPUT_DIR, f"temp_ho1_ref_{random.randint(1000, 9999)}.png")
+                img.save(temp_path, "PNG")
+                temp_files_to_clean.append(temp_path)
+                active_hidream_o1_reference.append(os.path.basename(temp_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])
+                })
+
+    return {
+        "active_loras_for_gpu": active_loras_for_gpu,
+        "active_loras_for_meta": active_loras_for_meta,
+        "active_controlnets": active_controlnets,
+        "active_anima_controlnets": active_anima_controlnets,
+        "active_diffsynth_controlnets": active_diffsynth_controlnets,
+        "active_ipadapters": active_ipadapters,
+        "active_flux1_ipadapters": active_flux1_ipadapters,
+        "active_sd3_ipadapters": active_sd3_ipadapters,
+        "active_styles": active_styles,
+        "active_reference_latents": active_reference_latents,
+        "active_hidream_o1_reference": active_hidream_o1_reference,
+        "active_conditioning": active_conditioning,
+        "temp_files_to_clean": temp_files_to_clean
+    }

core/pipelines/sd_image_pipeline.py

@@ -1,619 +1,254 @@
-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]
-                        actual_value = get_value_at_index(source_output_tuple, output_index)
-                    else:
-                        actual_value = param_value
-
-                    if '.' in param_name:
-                        parent_key, child_key = param_name.split('.', 1)
-                        if parent_key not in kwargs or not isinstance(kwargs[parent_key], dict):
-                            kwargs[parent_key] = {}
-                        kwargs[parent_key][child_key] = actual_value
-                    else:
-                        kwargs[param_name] = actual_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))
-                    
-        hidream_o1_reference_data = ui_inputs.get('hidream_o1_reference_data', [])
-        active_hidream_o1_reference = []
-        if hidream_o1_reference_data:
-            for img in hidream_o1_reference_data:
-                if img:
-                    if not os.path.exists(INPUT_DIR): os.makedirs(INPUT_DIR)
-                    temp_path = os.path.join(INPUT_DIR, f"temp_ho1_ref_{random.randint(1000, 9999)}.png")
-                    img.save(temp_path, "PNG")
-                    temp_files_to_clean.append(temp_path)
-                    active_hidream_o1_reference.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_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,
-        }
-
-        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}")
-
+import os
+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
+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)):
+        model_display_name = ui_inputs['model_display_name']
+        
+        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
+
+    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 [],
+            }
+
+            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...")
+
+            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

core/pipelines/workflow_executor.py

@@ -0,0 +1,110 @@
+import torch
+from collections import defaultdict, deque
+from typing import Dict, Any, List
+from comfy_integration.nodes import NODE_CLASS_MAPPINGS
+from utils.app_utils import get_value_at_index
+
+class WorkflowExecutor:
+    @staticmethod
+    def topological_sort(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
+
+    @staticmethod
+    def execute_workflow(workflow: Dict[str, Any], initial_objects: Dict[str, Any]):
+        with torch.no_grad():
+            computed_outputs = initial_objects
+            
+            try:
+                sorted_node_ids = WorkflowExecutor.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]
+                        actual_value = get_value_at_index(source_output_tuple, output_index)
+                    else:
+                        actual_value = param_value
+
+                    if '.' in param_name:
+                        parent_key, child_key = param_name.split('.', 1)
+                        if parent_key not in kwargs or not isinstance(kwargs[parent_key], dict):
+                            kwargs[parent_key] = {}
+                        kwargs[parent_key][child_key] = actual_value
+                    else:
+                        kwargs[param_name] = actual_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)

core/pipelines/workflow_recipes/_partials/conditioning/anima.yaml

@@ -52,6 +52,10 @@ dynamic_conditioning_chains:
     ksampler_node: "ksampler"
     clip_source: "clip_loader:0"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   unet_name: "unet_loader:unet_name"
   vae_name: "vae_loader:vae_name"

core/pipelines/workflow_recipes/_partials/conditioning/chroma1.yaml

@@ -55,6 +55,10 @@ dynamic_conditioning_chains:
     ksampler_node: "ksampler"
     clip_source: "t5_tokenizer:0"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   unet_name: "unet_loader:unet_name"
   vae_name: "vae_loader:vae_name"

core/pipelines/workflow_recipes/_partials/conditioning/ernie-image.yaml

@@ -48,6 +48,10 @@ dynamic_conditioning_chains:
     ksampler_node: "ksampler"
     clip_source: "clip_loader:0"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   unet_name: "unet_loader:unet_name"
   clip_name: "clip_loader:clip_name"

core/pipelines/workflow_recipes/_partials/conditioning/flux1.yaml

@@ -56,6 +56,10 @@ dynamic_conditioning_chains:
     ksampler_node: "ksampler"
     clip_source: "clip_loader:0"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   unet_name: "unet_loader:unet_name"
   vae_name: "vae_loader:vae_name"

core/pipelines/workflow_recipes/_partials/conditioning/flux2-kv.yaml

@@ -86,6 +86,10 @@ dynamic_reference_latent_chains:
     ksampler_node: "ksampler"
     vae_node: "vae_loader"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   unet_name: "unet_loader:unet_name"
   clip_name: "clip_loader:clip_name"

core/pipelines/workflow_recipes/_partials/conditioning/flux2.yaml

@@ -78,6 +78,10 @@ dynamic_reference_latent_chains:
     ksampler_node: "ksampler"
     vae_node: "vae_loader"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   unet_name: "unet_loader:unet_name"
   clip_name: "clip_loader:clip_name"

core/pipelines/workflow_recipes/_partials/conditioning/hidream-i1.yaml

@@ -44,6 +44,10 @@ dynamic_conditioning_chains:
     ksampler_node: "ksampler"
     clip_source: "clip_loader:0"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   unet_name: "unet_loader:unet_name"
   vae_name: "vae_loader:vae_name"

core/pipelines/workflow_recipes/_partials/conditioning/lens.yaml

@@ -48,6 +48,10 @@ dynamic_conditioning_chains:
     ksampler_node: "ksampler"
     clip_source: "clip_loader:0"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   unet_name: "unet_loader:unet_name"
   clip_name: "clip_loader:clip_name"

core/pipelines/workflow_recipes/_partials/conditioning/longcat-image.yaml

@@ -77,6 +77,10 @@ dynamic_conditioning_chains:
     ksampler_node: "ksampler"
     clip_source: "clip_loader:0"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   unet_name: "unet_loader:unet_name"
   vae_name: "vae_loader:vae_name"

core/pipelines/workflow_recipes/_partials/conditioning/lumina.yaml

@@ -53,5 +53,9 @@ dynamic_conditioning_chains:
     ksampler_node: "ksampler"
     clip_source: "ckpt_loader:1"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   model_name: "ckpt_loader:ckpt_name"

core/pipelines/workflow_recipes/_partials/conditioning/newbie-image.yaml

@@ -58,6 +58,10 @@ dynamic_conditioning_chains:
     ksampler_node: "ksampler"
     clip_source: "clip_loader:0"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   unet_name: "unet_loader:unet_name"
   vae_name: "vae_loader:vae_name"

core/pipelines/workflow_recipes/_partials/conditioning/omnigen2.yaml

@@ -53,6 +53,10 @@ dynamic_reference_latent_chains:
     ksampler_node: "ksampler"
     vae_node: "vae_loader"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   unet_name: "unet_loader:unet_name"
   vae_name: "vae_loader:vae_name"

core/pipelines/workflow_recipes/_partials/conditioning/ovis-image.yaml

@@ -44,6 +44,10 @@ dynamic_conditioning_chains:
     ksampler_node: "ksampler"
     clip_source: "clip_loader:0"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   unet_name: "unet_loader:unet_name"
   vae_name: "vae_loader:vae_name"

core/pipelines/workflow_recipes/_partials/conditioning/qwen-image.yaml

@@ -73,6 +73,10 @@ dynamic_conditioning_chains:
     ksampler_node: "ksampler"
     clip_source: "clip_loader:0"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   unet_name: "unet_loader:unet_name"
   vae_name: "vae_loader:vae_name"

core/pipelines/workflow_recipes/_partials/conditioning/sd35.yaml

@@ -54,5 +54,9 @@ dynamic_conditioning_chains:
     ksampler_node: "ksampler"
     clip_source: "ckpt_loader:1"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   model_name: "ckpt_loader:ckpt_name"

core/pipelines/workflow_recipes/_partials/conditioning/sdxl.yaml

@@ -59,5 +59,9 @@ dynamic_conditioning_chains:
     ksampler_node: "ksampler"
     clip_source: "ckpt_loader:1"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   model_name: "ckpt_loader:ckpt_name"

core/pipelines/workflow_recipes/_partials/conditioning/z-image.yaml

@@ -59,6 +59,10 @@ dynamic_diffsynth_controlnet_chains:
     ksampler_node: "ksampler"
     vae_source: "vae_loader:0"
 
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
 ui_map:
   unet_name: "unet_loader:unet_name"
   vae_name: "vae_loader:vae_name"

requirements.txt

@@ -1,5 +1,5 @@
-comfyui-frontend-package==1.44.19
-comfyui-workflow-templates==0.9.92
+comfyui-frontend-package==1.45.15
+comfyui-workflow-templates==0.9.98
 comfyui-embedded-docs==0.5.2
 torch
 torchsde

ui/events/change_handlers.py

@@ -16,7 +16,7 @@ from .config_loaders import (
     load_ipadapter_config
 )
 
-def make_update_fn(m_comp, cat_comp, cs_comp, ar_comp, width_comp, height_comp, cn_types, cn_series, cn_filepaths, anima_cn_types, anima_cn_series, anima_cn_filepaths, diffsynth_cn_types, diffsynth_cn_series, diffsynth_cn_filepaths, ipa_preset, lora_acc, cn_acc, anima_cn_acc, diffsynth_cn_acc, ipa_acc, sd3_ipa_acc, flux1_ipa_acc, style_acc, embed_acc, cond_acc, ref_latent_acc, hidream_o1_ref_acc, guidance_comp, prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp):
+def make_update_fn(m_comp, cat_comp, cs_comp, ar_comp, width_comp, height_comp, cn_types, cn_series, cn_filepaths, anima_cn_types, anima_cn_series, anima_cn_filepaths, diffsynth_cn_types, diffsynth_cn_series, diffsynth_cn_filepaths, ipa_preset, lora_acc, cn_acc, anima_cn_acc, diffsynth_cn_acc, ipa_acc, sd3_ipa_acc, flux1_ipa_acc, style_acc, embed_acc, cond_acc, ref_latent_acc, hidream_o1_ref_acc, guidance_comp, prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp, pid_acc=None):
     def update_fn(*args):
         arch = args[0]
         category = args[1]
@@ -66,6 +66,7 @@ def make_update_fn(m_comp, cat_comp, cs_comp, ar_comp, width_comp, height_comp,
         if cond_acc: updates[cond_acc] = gr.update(visible=('conditioning' in enabled_chains))
         if ref_latent_acc: updates[ref_latent_acc] = gr.update(visible=('reference_latent' in enabled_chains))
         if hidream_o1_ref_acc: updates[hidream_o1_ref_acc] = gr.update(visible=('hidream_o1_reference' in enabled_chains))
+        if pid_acc: updates[pid_acc] = gr.update(visible=('pid' in enabled_chains))
 
         if cs_comp:
             updates[cs_comp] = gr.update(visible=(arch_model_type == "sd15"))
@@ -130,7 +131,7 @@ def make_update_fn(m_comp, cat_comp, cs_comp, ar_comp, width_comp, height_comp,
     return update_fn
 
 
-def make_model_change_fn(cat_comp_ref, cs_comp, ar_comp, width_comp, height_comp, cn_types, cn_series, cn_filepaths, anima_cn_types, anima_cn_series, anima_cn_filepaths, diffsynth_cn_types, diffsynth_cn_series, diffsynth_cn_filepaths, arch_comp_ref, ipa_preset, lora_acc, cn_acc, anima_cn_acc, diffsynth_cn_acc, ipa_acc, sd3_ipa_acc, flux1_ipa_acc, style_acc, embed_acc, cond_acc, ref_latent_acc, hidream_o1_ref_acc, guidance_comp, prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp):
+def make_model_change_fn(cat_comp_ref, cs_comp, ar_comp, width_comp, height_comp, cn_types, cn_series, cn_filepaths, anima_cn_types, anima_cn_series, anima_cn_filepaths, diffsynth_cn_types, diffsynth_cn_series, diffsynth_cn_filepaths, arch_comp_ref, ipa_preset, lora_acc, cn_acc, anima_cn_acc, diffsynth_cn_acc, ipa_acc, sd3_ipa_acc, flux1_ipa_acc, style_acc, embed_acc, cond_acc, ref_latent_acc, hidream_o1_ref_acc, guidance_comp, prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp, pid_acc=None):
     def change_fn(*args):
         model_name = args[0]
         idx = 1
@@ -185,6 +186,7 @@ def make_model_change_fn(cat_comp_ref, cs_comp, ar_comp, width_comp, height_comp
         if cond_acc: updates[cond_acc] = gr.update(visible=('conditioning' in enabled_chains))
         if ref_latent_acc: updates[ref_latent_acc] = gr.update(visible=('reference_latent' in enabled_chains))
         if hidream_o1_ref_acc: updates[hidream_o1_ref_acc] = gr.update(visible=('hidream_o1_reference' in enabled_chains))
+        if pid_acc: updates[pid_acc] = gr.update(visible=('pid' in enabled_chains))
 
         if cs_comp:
             updates[cs_comp] = gr.update(visible=(arch_model_type == "sd15"))
@@ -330,4 +332,4 @@ def on_aspect_ratio_change(ratio_key, model_display_name):
     
     res_map = RESOLUTION_MAP.get(arch_model_type, RESOLUTION_MAP.get("sdxl", {}))
     w, h = res_map.get(ratio_key, (1024, 1024))
-    return w, h
+    return w, h

ui/events/main.py

@@ -60,6 +60,7 @@ def attach_event_handlers(ui_components, demo):
         conditioning_accordion = ui_components.get(f'conditioning_accordion_{prefix}')
         ref_latent_accordion = ui_components.get(f'reference_latent_accordion_{prefix}')
         hidream_o1_ref_accordion = ui_components.get(f'hidream_o1_reference_accordion_{prefix}')
+        pid_accordion = ui_components.get(f'pid_accordion_{prefix}')
         
         ipa_preset_list = ui_components.get(f'ipadapter_final_preset_{prefix}')
         
@@ -93,6 +94,7 @@ def attach_event_handlers(ui_components, demo):
             if conditioning_accordion: outputs.append(conditioning_accordion)
             if ref_latent_accordion: outputs.append(ref_latent_accordion)
             if hidream_o1_ref_accordion: outputs.append(hidream_o1_ref_accordion)
+            if pid_accordion: outputs.append(pid_accordion)
             if ipa_preset_list: outputs.append(ipa_preset_list)
             
             outputs.extend(valid_extra_comps)
@@ -103,7 +105,8 @@ def attach_event_handlers(ui_components, demo):
                 anima_cn_types_list, anima_cn_series_list, anima_cn_filepaths_list,
                 diffsynth_cn_types_list, diffsynth_cn_series_list, diffsynth_cn_filepaths_list, 
                 ipa_preset_list, lora_accordion, cn_accordion, anima_cn_accordion, diffsynth_cn_accordion, ipa_accordion, sd3_ipa_accordion, flux1_ipa_accordion, style_accordion, embedding_accordion, conditioning_accordion,
-                ref_latent_accordion, hidream_o1_ref_accordion, guidance_comp, prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp
+                ref_latent_accordion, hidream_o1_ref_accordion, guidance_comp, prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp,
+                pid_acc=pid_accordion
             )
             inputs = [arch_comp, cat_comp]
             if aspect_ratio_comp:
@@ -133,6 +136,7 @@ def attach_event_handlers(ui_components, demo):
             if conditioning_accordion: outputs2.append(conditioning_accordion)
             if ref_latent_accordion: outputs2.append(ref_latent_accordion)
             if hidream_o1_ref_accordion: outputs2.append(hidream_o1_ref_accordion)
+            if pid_accordion: outputs2.append(pid_accordion)
             if ipa_preset_list: outputs2.append(ipa_preset_list)
             
             outputs2.extend(valid_extra_comps)
@@ -148,7 +152,8 @@ def attach_event_handlers(ui_components, demo):
                     anima_cn_types_list, anima_cn_series_list, anima_cn_filepaths_list,
                     diffsynth_cn_types_list, diffsynth_cn_series_list, diffsynth_cn_filepaths_list, 
                     arch_comp, ipa_preset_list, lora_accordion, cn_accordion, anima_cn_accordion, diffsynth_cn_accordion, ipa_accordion, sd3_ipa_accordion, flux1_ipa_accordion, style_accordion, embedding_accordion, conditioning_accordion,
-                    ref_latent_accordion, hidream_o1_ref_accordion, guidance_comp, prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp
+                    ref_latent_accordion, hidream_o1_ref_accordion, guidance_comp, prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp,
+                    pid_acc=pid_accordion
                 )
                 model_comp.change(fn=change_fn, inputs=inputs2, outputs=outputs2)
 
@@ -237,4 +242,4 @@ def attach_event_handlers(ui_components, demo):
         height_component = ui_components.get(f'height_{prefix}') or ui_components.get(f'{prefix}_height')
         model_dropdown = ui_components.get(f'base_model_{prefix}')
         if aspect_ratio_dropdown and width_component and height_component and model_dropdown:
-            aspect_ratio_dropdown.change(fn=on_aspect_ratio_change, inputs=[aspect_ratio_dropdown, model_dropdown], outputs=[width_component, height_component], show_progress=False)
+            aspect_ratio_dropdown.change(fn=on_aspect_ratio_change, inputs=[aspect_ratio_dropdown, model_dropdown], outputs=[width_component, height_component], show_progress=False)

ui/events/run_handlers.py

@@ -19,6 +19,9 @@ def create_run_event(prefix: str, task_type: str, ui_components: dict):
         'task_type': gr.State(task_type)
     }
     
+    if ui_components.get(f'pid_settings_{prefix}'):
+        run_inputs_map['pid_settings'] = ui_components[f'pid_settings_{prefix}']
+    
     if task_type not in ['img2img', 'inpaint']:
         run_inputs_map.update({
             'width': ui_components.get(f'width_{prefix}') or ui_components.get(f'{prefix}_width'), 
@@ -97,4 +100,4 @@ def create_run_event(prefix: str, task_type: str, ui_components: dict):
             fn=lambda *args, progress=gr.Progress(track_tqdm=True): generate_image_wrapper(create_ui_inputs_dict(*args), progress),
             inputs=input_list_flat,
             outputs=[res_gal]
-        )
+        )

ui/shared/txt2img_ui.py

@@ -6,7 +6,8 @@ from .ui_components import (
     create_conditioning_ui, create_vae_override_ui,
     create_model_architecture_filter_ui, create_category_filter_ui,
     create_sd3_ipadapter_ui, create_flux1_ipadapter_ui, create_style_ui,
-    create_reference_latent_ui, create_hidream_o1_reference_ui
+    create_reference_latent_ui, create_hidream_o1_reference_ui,
+    create_pid_ui
 )
 
 def create_ui():
@@ -53,5 +54,6 @@ def create_ui():
         components.update(create_reference_latent_ui(prefix))
         components.update(create_hidream_o1_reference_ui(prefix))
         components.update(create_vae_override_ui(prefix))
+        components.update(create_pid_ui(prefix))
         
     return components

ui/shared/ui_components.py

@@ -548,8 +548,8 @@ def create_vae_override_ui(prefix: str):
     key = lambda name: f"{name}_{prefix}"
     source_choices = ["None"] + LORA_SOURCE_CHOICES
 
-    with gr.Accordion("VAE Settings (Override)", open=False) as accordion:
-        components[key('vae_accordion')] = accordion
+    with gr.Accordion("VAE Settings (Override)", open=False) as vae_accordion:
+        components[key('vae_accordion')] = vae_accordion
         gr.Markdown("💡 **Tip:** When downloading from Civitai, please use the **Version ID**, not the Model ID. You can find the Version ID in the URL (e.g., `civitai.com/models/123?modelVersionId=456`) or under the model's download button.")
         with gr.Row():
             components[key('vae_source')] = gr.Dropdown(
@@ -638,4 +638,21 @@ def create_hidream_o1_reference_ui(prefix: str, max_units=10):
         
         components[key('all_hidream_o1_reference_components_flat')] = ref_image_inputs
 
+    return components
+
+def create_pid_ui(prefix: str):
+    components = {}
+    key = lambda name: f"{name}_{prefix}"
+
+    with gr.Accordion("PiD Settings", open=False, visible=('pid' in default_enabled_chains)) as pid_accordion:
+        components[key('pid_accordion')] = pid_accordion
+        gr.Markdown("💡 **Tip:** Use PiD (Pixel Diffusion Decoder) instead of the VAE Decoder for 4x decoding.")
+        with gr.Row():
+            components[key('pid_settings')] = gr.Dropdown(
+                label="PiD Mode", 
+                choices=["OFF", "ON"], 
+                value="OFF", 
+                interactive=True
+            )
+            
     return components

yaml/file_list.yaml

@@ -408,6 +408,27 @@ file:
       source: "hf"
       repo_id: "Comfy-Org/PixelDiT"
       repository_file_path: "diffusion_models/pixeldit_1300m_1024px_mxfp8.safetensors"
+    # PiD
+    - filename: "pid_sdxl_1024_to_4096_4step_bf16.safetensors"
+      source: "hf"
+      repo_id: "Comfy-Org/PixelDiT"
+      repository_file_path: "diffusion_models/pid_sdxl_1024_to_4096_4step_bf16.safetensors"
+    - filename: "pid_sd3_1024_to_4096_4step_bf16.safetensors"
+      source: "hf"
+      repo_id: "Comfy-Org/PixelDiT"
+      repository_file_path: "diffusion_models/pid_sd3_1024_to_4096_4step_bf16.safetensors"
+    - filename: "pid_flux1_1024_to_4096_4step_mxfp8.safetensors"
+      source: "hf"
+      repo_id: "Comfy-Org/PixelDiT"
+      repository_file_path: "diffusion_models/pid_flux1_1024_to_4096_4step_mxfp8.safetensors"
+    - filename: "pid_qwenimage_1024_to_4096_4step_bf16.safetensors"
+      source: "hf"
+      repo_id: "Comfy-Org/PixelDiT"
+      repository_file_path: "diffusion_models/pid_qwenimage_1024_to_4096_4step_bf16.safetensors"
+    - filename: "pid_flux2_1024_to_4096_4step_mxfp8.safetensors"
+      source: "hf"
+      repo_id: "Comfy-Org/PixelDiT"
+      repository_file_path: "diffusion_models/pid_flux2_1024_to_4096_4step_mxfp8.safetensors"
     # Lens
     - filename: "lens_mxfp8.safetensors"
       source: "hf"

yaml/image_gen_features.yaml

@@ -1,11 +1,12 @@
 default:
   enabled_chains:
-    - lora
-    - controlnet
-    - ipadapter
-    - embedding
-    - style
-    - conditioning
+  - lora
+  - controlnet
+  - ipadapter
+  - embedding
+  - style
+  - conditioning
+  - vae
 
 pixeldit:
   enabled_chains:
@@ -14,119 +15,129 @@ pixeldit:
 lens:
   enabled_chains:
   - conditioning
+  - pid
 
 ernie-image:
   enabled_chains:
-    - lora
-    - conditioning
-
-flux2:
-  enabled_chains:
-    - lora
-    - conditioning
-    - reference_latent
-
-flux2-kv:
-  enabled_chains:
-    - lora
-    - conditioning
-    - reference_latent
-
-z-image:
-  enabled_chains:
-    - lora
-    - conditioning
-    - controlnet_model_patch
-
-qwen-image:
+  - conditioning
+  - pid
+anima:
   enabled_chains:
-    - lora
-    - controlnet
-    - conditioning
-
+  - lora
+  - anima_controlnet_lllite
+  - conditioning
+  - vae
+  - pid
 longcat-image:
   enabled_chains:
-    - lora
-    - conditioning
-
-anima:
-  enabled_chains:
-    - lora
-    - anima_controlnet_lllite
-    - conditioning
-
+  - lora
+  - conditioning
+  - pid
 newbie-image:
   enabled_chains:
-    - lora
-    - embedding
-    - conditioning
-
-omnigen2:
+  - lora
+  - embedding
+  - conditioning
+  - vae
+  - pid
+z-image:
   enabled_chains:
-    - conditioning
-    - reference_latent
-
+  - lora
+  - conditioning
+  - controlnet_model_patch
+  - vae
+  - pid
 lumina:
   enabled_chains:
-    - lora
-    - embedding
-    - conditioning
-
-ovis-image:
-  enabled_chains:
-    - conditioning
-
+  - lora
+  - embedding
+  - conditioning
+  - vae
+  - pid
 sd35:
   enabled_chains:
-    - lora
-    - controlnet
-    - embedding
-    - conditioning
-    - sd3_ipadapter
-
+  - lora
+  - controlnet
+  - embedding
+  - conditioning
+  - sd3_ipadapter
+  - vae
+  - pid
 sdxl:
   enabled_chains:
-    - lora
-    - controlnet
-    - ipadapter
-    - embedding
-    - conditioning
-
+  - lora
+  - controlnet
+  - ipadapter
+  - embedding
+  - conditioning
+  - vae
+  - pid
 sd15:
   enabled_chains:
-    - lora
-    - controlnet
-    - ipadapter
-    - embedding
-    - conditioning
-
+  - lora
+  - controlnet
+  - ipadapter
+  - embedding
+  - conditioning
+  - vae
+flux2:
+  enabled_chains:
+  - lora
+  - conditioning
+  - reference_latent
+  - vae
+  - pid
+flux2-kv:
+  enabled_chains:
+  - lora
+  - conditioning
+  - reference_latent
+  - vae
+  - pid
 flux1:
   enabled_chains:
-    - lora
-    - controlnet
-    - style
-    - conditioning
-    - flux1_ipadapter
-
+  - lora
+  - controlnet
+  - style
+  - conditioning
+  - flux1_ipadapter
+  - vae
+  - pid
+omnigen2:
+  enabled_chains:
+  - conditioning
+  - reference_latent
+  - pid
+qwen-image:
+  enabled_chains:
+  - lora
+  - controlnet
+  - conditioning
+  - vae
+  - pid
 hidream-o1:
   enabled_chains:
   - lora
   - conditioning
   - hidream_o1_reference
-  
 hidream-i1:
   enabled_chains:
   - lora
   - conditioning
-
-chroma1:
+  - pid
+hunyuanimage:
   enabled_chains:
-    - conditioning
+  - conditioning
+  - vae
 
-chroma1-radiance:
+ovis-image:
   enabled_chains:
-    - conditioning 
+  - conditioning
+  - vae
+  - pid
 
-hunyuanimage:
+chroma1:
   enabled_chains:
-    - conditioning
+  - conditioning
+  - vae
+  - pid

yaml/injectors.yaml

@@ -27,6 +27,8 @@ injector_definitions:
     module: "chain_injectors.hidream_o1_smoothing_injector"
   dynamic_hidream_o1_reference_chains:
     module: "chain_injectors.hidream_o1_reference_injector"
+  dynamic_pid_chains:
+    module: "chain_injectors.pid_injector"
   
 injector_order:
   - dynamic_vae_chains
@@ -42,4 +44,5 @@ injector_order:
   - dynamic_controlnet_chains
   - dynamic_anima_controlnet_lllite_chains
   - dynamic_hidream_o1_smoothing_chains
-  - dynamic_hidream_o1_reference_chains
+  - dynamic_hidream_o1_reference_chains
+  - dynamic_pid_chains

yaml/pid.yaml

@@ -0,0 +1,16 @@
+PiD:
+  - filepath: "pid_flux2_1024_to_4096_4step_mxfp8.safetensors"
+    latent_format: "flux"
+    architectures: ["flux2", "flux2-kv", "lens", "ernie-image"]
+  - filepath: "pid_qwenimage_1024_to_4096_4step_bf16.safetensors"
+    latent_format: "qwenimage"
+    architectures: ["anima", "qwen-image"]
+  - filepath: "pid_flux1_1024_to_4096_4step_mxfp8.safetensors"
+    latent_format: "flux"
+    architectures: ["longcat-image", "newbie-image", "z-image", "ovis-image", "omnigen2", "chroma1", "hidream-i1", "flux1"]
+  - filepath: "pid_sd3_1024_to_4096_4step_bf16.safetensors"
+    latent_format: "sd3"
+    architectures: ["sd35"]
+  - filepath: "pid_sdxl_1024_to_4096_4step_bf16.safetensors"
+    latent_format: "sdxl"
+    architectures: ["sdxl"]