Upload folder using huggingface_hub

README.md

@@ -1,13 +1,10 @@
----
-title: ImageGen FLUX.2
-emoji: 🏆
-colorFrom: indigo
-colorTo: blue
-sdk: gradio
-sdk_version: 6.5.0
-python_version: '3.12'
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: ImageGen - FLUX.2
+emoji: 🖼
+colorFrom: purple
+colorTo: red
+sdk: gradio
+sdk_version: "5.50.0"
+app_file: app.py
+short_description: Multi-task image generator with dynamic, chainable workflows
+---

app.py

@@ -0,0 +1,223 @@
+import spaces
+import os
+import sys
+import requests
+import site
+
+APP_DIR = os.path.dirname(os.path.abspath(__file__))
+if APP_DIR not in sys.path:
+    sys.path.insert(0, APP_DIR)
+    print(f"✅ Added project root '{APP_DIR}' to sys.path.")
+
+SAGE_PATCH_APPLIED = False
+
+def apply_sage_attention_patch():
+    global SAGE_PATCH_APPLIED
+    if SAGE_PATCH_APPLIED:
+        return "SageAttention patch already applied."
+
+    try:
+        from comfy import model_management
+        import sageattention
+        
+        print("--- [Runtime Patch] sageattention package found. Applying patch... ---")
+        model_management.sage_attention_enabled = lambda: True
+        model_management.pytorch_attention_enabled = lambda: False
+        
+        SAGE_PATCH_APPLIED = True
+        return "✅ Successfully enabled SageAttention."
+    except ImportError:
+        SAGE_PATCH_APPLIED = False
+        msg = "--- [Runtime Patch] ⚠️ sageattention package not found. Continuing with default attention. ---"
+        print(msg)
+        return msg
+    except Exception as e:
+        SAGE_PATCH_APPLIED = False
+        msg = f"--- [Runtime Patch] ❌ An error occurred while applying SageAttention patch: {e} ---"
+        print(msg)
+        return msg
+
+@spaces.GPU
+def dummy_gpu_for_startup():
+    print("--- [GPU Startup] Dummy function for startup check initiated. ---")
+    patch_result = apply_sage_attention_patch()
+    print(f"--- [GPU Startup] {patch_result} ---")
+    print("--- [GPU Startup] Startup check passed. ---")
+    return "Startup check passed."
+
+def handle_private_downloads():
+    """
+    Checks for a private_file_list.yaml, downloads required models using HF_TOKEN,
+    and then clears the token from the environment.
+    """
+    import yaml
+    from huggingface_hub import hf_hub_download
+    from core.settings import (
+        DIFFUSION_MODELS_DIR, TEXT_ENCODERS_DIR, VAE_DIR, CHECKPOINT_DIR, 
+        LORA_DIR, CONTROLNET_DIR, MODEL_PATCHES_DIR, EMBEDDING_DIR
+    )
+
+    print("--- [Startup] Checking for private models to download... ---")
+    private_list_path = os.path.join(APP_DIR, 'yaml', 'private_file_list.yaml')
+
+    if not os.path.exists(private_list_path):
+        print("--- [Startup] No private model list found. Skipping. ---")
+        if 'HF_TOKEN' in os.environ:
+             del os.environ['HF_TOKEN']
+             print("--- [Startup] Cleared HF_TOKEN environment variable as it is no longer needed. ---")
+             print(f"--- [Startup] Verifying HF_TOKEN after clearing: {os.environ.get('HF_TOKEN')}")
+        return
+
+    try:
+        with open(private_list_path, 'r', encoding='utf-8') as f:
+            private_files_config = yaml.safe_load(f)
+
+        if not private_files_config or 'file' not in private_files_config:
+            print("--- [Startup] Private model list is empty or malformed. Skipping. ---")
+            return
+        
+        category_to_dir_map = {
+            "diffusion_models": DIFFUSION_MODELS_DIR,
+            "text_encoders": TEXT_ENCODERS_DIR,
+            "vae": VAE_DIR,
+            "checkpoints": CHECKPOINT_DIR,
+            "loras": LORA_DIR,
+            "controlnet": CONTROLNET_DIR,
+            "model_patches": MODEL_PATCHES_DIR,
+            "embeddings": EMBEDDING_DIR,
+        }
+
+        files_to_download = []
+        for category, files in private_files_config.get('file', {}).items():
+            dest_dir = category_to_dir_map.get(category)
+            if not dest_dir:
+                print(f"--- [Startup] ⚠️ Unknown category '{category}' in private_file_list.yaml. Skipping. ---")
+                continue
+            
+            if isinstance(files, list):
+                for file_info in files:
+                    files_to_download.append((file_info, dest_dir))
+
+        if not files_to_download:
+            print("--- [Startup] No private models configured for download. ---")
+            return
+
+        print(f"--- [Startup] Found {len(files_to_download)} private model(s) to download. Using HF_TOKEN if available. ---")
+
+        for file_info, dest_dir in files_to_download:
+            filename = file_info.get("filename")
+            repo_id = file_info.get("repo_id")
+            repo_path = file_info.get("repository_file_path", filename)
+            
+            if not all([filename, repo_id]):
+                print(f"--- [Startup] ⚠️ Skipping malformed entry in private_file_list.yaml: {file_info} ---")
+                continue
+
+            dest_path = os.path.join(dest_dir, filename)
+            if os.path.lexists(dest_path):
+                print(f"--- [Startup] ✅ Model '{filename}' already exists. Skipping download. ---")
+                continue
+            
+            print(f"--- [Startup] ⏳ Downloading '{filename}' from repo '{repo_id}'... ---")
+            try:
+                cached_path = hf_hub_download(repo_id=repo_id, filename=repo_path)
+                os.makedirs(dest_dir, exist_ok=True)
+                os.symlink(cached_path, dest_path)
+                print(f"--- [Startup] ✅ Successfully downloaded and linked '{filename}'. ---")
+            except Exception as e:
+                print(f"--- [Startup] ❌ ERROR: Failed to download '{filename}': {e}")
+                print("--- [Startup] ❌ Please ensure your HF_TOKEN is set correctly and has access to the repository. ---")
+
+    finally:
+        if 'HF_TOKEN' in os.environ:
+            del os.environ['HF_TOKEN']
+            print("--- [Startup] ✅ Cleared HF_TOKEN environment variable. ---")
+            print(f"--- [Startup] Verifying HF_TOKEN after clearing: {os.environ.get('HF_TOKEN')}")
+        else:
+            print("--- [Startup] Note: HF_TOKEN environment variable was not set. Private downloads may fail without it. ---")
+
+def main():
+    from utils.app_utils import print_welcome_message
+    from scripts import build_sage_attention
+
+    print_welcome_message()
+    
+    # Handle downloads that require authentication first.
+    handle_private_downloads()
+    
+    print("--- [Setup] Attempting to build and install SageAttention... ---")
+    try:
+        build_sage_attention.install_sage_attention()
+        print("--- [Setup] ✅ SageAttention installation process finished. ---")
+    except Exception as e:
+        print(f"--- [Setup] ❌ SageAttention installation failed: {e}. Continuing with default attention. ---")
+
+
+    print("--- [Setup] Reloading site-packages to detect newly installed packages... ---")
+    try:
+        site.main()
+        print("--- [Setup] ✅ Site-packages reloaded. ---")
+    except Exception as e:
+        print(f"--- [Setup] ⚠️  Warning: Could not fully reload site-packages: {e} ---")
+
+    from comfy_integration import setup as setup_comfyui
+    from utils.app_utils import (
+        build_preprocessor_model_map,
+        build_preprocessor_parameter_map
+    )
+    from core import shared_state
+    from core.settings import ALL_MODEL_MAP, ALL_FILE_DOWNLOAD_MAP
+
+    def check_all_model_urls_on_startup():
+        print("--- [Setup] Checking all model URL validity (one-time check) ---")
+        for display_name, model_info in ALL_MODEL_MAP.items():
+            _, components, _, _ = model_info
+            if not components: continue
+            
+            for filename in components.values():
+                download_info = ALL_FILE_DOWNLOAD_MAP.get(filename, {})
+                repo_id = download_info.get('repo_id')
+                if not repo_id: continue
+                
+                repo_file_path = download_info.get('repository_file_path', filename)
+                url = f"https://huggingface.co/{repo_id}/resolve/main/{repo_file_path}"
+                
+                try:
+                    response = requests.head(url, timeout=5, allow_redirects=True)
+                    if response.status_code >= 400:
+                        print(f"❌ Invalid URL for '{display_name}' component '{filename}': {url} (Status: {response.status_code})")
+                        shared_state.INVALID_MODEL_URLS[display_name] = True
+                        break 
+                except requests.RequestException as e:
+                    print(f"❌ URL check failed for '{display_name}' component '{filename}': {e}")
+                    shared_state.INVALID_MODEL_URLS[display_name] = True
+                    break
+        print("--- [Setup] ✅ Finished checking model URLs. ---")
+
+    print("--- Starting Application Setup ---")
+    
+    setup_comfyui.initialize_comfyui()
+
+    check_all_model_urls_on_startup()
+    
+    print("--- Building ControlNet preprocessor maps ---")
+    from core.generation_logic import build_reverse_map
+    build_reverse_map()
+    build_preprocessor_model_map()
+    build_preprocessor_parameter_map()
+    print("--- ✅ ControlNet preprocessor setup complete. ---")
+
+    print("--- Environment configured. Proceeding with module imports. ---")
+    from ui.layout import build_ui
+    from ui.events import attach_event_handlers
+
+    print(f"✅ Working directory is stable: {os.getcwd()}")
+
+    demo = build_ui(attach_event_handlers)
+        
+    print("--- Launching Gradio Interface ---")
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860)
+
+
+if __name__ == "__main__":
+    main()

chain_injectors/conditioning_injector.py

@@ -0,0 +1,81 @@
+def inject(assembler, chain_definition, chain_items):
+    if not chain_items:
+        return
+
+    ksampler_name = chain_definition.get('ksampler_node', 'ksampler')
+    
+    target_node_id = None
+    target_input_name = None
+
+    if ksampler_name in assembler.node_map:
+        ksampler_id = assembler.node_map[ksampler_name]
+        if 'positive' in assembler.workflow[ksampler_id]['inputs']:
+            target_node_id = ksampler_id
+            target_input_name = 'positive'
+            print(f"Conditioning injector targeting KSampler node '{ksampler_name}'.")
+    else:
+        print(f"Warning: KSampler node '{ksampler_name}' for Conditioning chain not found. Skipping.")
+        return
+    
+    if not target_node_id:
+        print("Warning: Conditioning chain could not find a valid injection point (KSampler may be missing 'positive' input). Skipping.")
+        return
+
+    clip_source_str = chain_definition.get('clip_source')
+    if not clip_source_str:
+        print("Warning: 'clip_source' definition missing in the recipe for the Conditioning chain. Skipping.")
+        return
+    clip_node_name, clip_idx_str = clip_source_str.split(':')
+    if clip_node_name not in assembler.node_map:
+        print(f"Warning: CLIP source node '{clip_node_name}' for Conditioning chain not found. Skipping.")
+        return
+    clip_connection = [assembler.node_map[clip_node_name], int(clip_idx_str)]
+    
+    original_positive_connection = assembler.workflow[target_node_id]['inputs'][target_input_name]
+    
+    area_conditioning_outputs = []
+    
+    for item_data in chain_items:
+        prompt = item_data.get('prompt', '')
+        if not prompt or not prompt.strip():
+            continue
+
+        text_encode_id = assembler._get_unique_id()
+        text_encode_node = assembler._get_node_template("CLIPTextEncode")
+        text_encode_node['inputs']['text'] = prompt
+        text_encode_node['inputs']['clip'] = clip_connection
+        assembler.workflow[text_encode_id] = text_encode_node
+
+        set_area_id = assembler._get_unique_id()
+        set_area_node = assembler._get_node_template("ConditioningSetArea")
+        set_area_node['inputs']['width'] = item_data.get('width', 1024)
+        set_area_node['inputs']['height'] = item_data.get('height', 1024)
+        set_area_node['inputs']['x'] = item_data.get('x', 0)
+        set_area_node['inputs']['y'] = item_data.get('y', 0)
+        set_area_node['inputs']['strength'] = item_data.get('strength', 1.0)
+        set_area_node['inputs']['conditioning'] = [text_encode_id, 0]
+        assembler.workflow[set_area_id] = set_area_node
+        
+        area_conditioning_outputs.append([set_area_id, 0])
+
+    if not area_conditioning_outputs:
+        return
+
+    current_combined_conditioning = area_conditioning_outputs[0]
+    if len(area_conditioning_outputs) > 1:
+        for i in range(1, len(area_conditioning_outputs)):
+            combine_id = assembler._get_unique_id()
+            combine_node = assembler._get_node_template("ConditioningCombine")
+            combine_node['inputs']['conditioning_1'] = current_combined_conditioning
+            combine_node['inputs']['conditioning_2'] = area_conditioning_outputs[i]
+            assembler.workflow[combine_id] = combine_node
+            current_combined_conditioning = [combine_id, 0]
+
+    final_combine_id = assembler._get_unique_id()
+    final_combine_node = assembler._get_node_template("ConditioningCombine")
+    final_combine_node['inputs']['conditioning_1'] = original_positive_connection
+    final_combine_node['inputs']['conditioning_2'] = current_combined_conditioning
+    assembler.workflow[final_combine_id] = final_combine_node
+
+    assembler.workflow[target_node_id]['inputs'][target_input_name] = [final_combine_id, 0]
+    print(f"Conditioning injector applied. Redirected '{target_input_name}' input with {len(area_conditioning_outputs)} regional prompts.")

chain_injectors/reference_latent_injector.py

@@ -0,0 +1,66 @@
+def inject(assembler, chain_definition, chain_items):
+    if not chain_items:
+        return
+
+    ksampler_name = chain_definition.get('ksampler_node', 'ksampler')
+    flux_guidance_name = chain_definition.get('flux_guidance_node')
+    vae_node_name = chain_definition.get('vae_node', 'vae_loader')
+
+    if ksampler_name not in assembler.node_map:
+        print(f"Warning: [ReferenceLatent] KSampler node '{ksampler_name}' not found. Skipping.")
+        return
+    if vae_node_name not in assembler.node_map:
+        print(f"Warning: [ReferenceLatent] VAE loader node '{vae_node_name}' not found. Skipping.")
+        return
+
+    ksampler_id = assembler.node_map[ksampler_name]
+    vae_node_id = assembler.node_map[vae_node_name]
+
+    pos_target_node_id = None
+    pos_target_input_name = None
+    if flux_guidance_name and flux_guidance_name in assembler.node_map:
+        flux_guidance_id = assembler.node_map[flux_guidance_name]
+        if 'conditioning' in assembler.workflow[flux_guidance_id]['inputs']:
+            pos_target_node_id = flux_guidance_id
+            pos_target_input_name = 'conditioning'
+            print(f"ReferenceLatent injector targeting FluxGuidance node '{flux_guidance_name}'.")
+    
+    if not pos_target_node_id:
+        if 'positive' in assembler.workflow[ksampler_id]['inputs']:
+            pos_target_node_id = ksampler_id
+            pos_target_input_name = 'positive'
+            print(f"ReferenceLatent injector targeting KSampler node '{ksampler_name}'.")
+        else:
+            print(f"Warning: [ReferenceLatent] Could not find a valid positive injection point. Skipping.")
+            return
+
+    current_pos_conditioning = assembler.workflow[pos_target_node_id]['inputs'][pos_target_input_name]
+
+    for i, img_filename in enumerate(chain_items):
+        if not img_filename or not isinstance(img_filename, str):
+            continue
+
+        load_id = assembler._get_unique_id()
+        load_node = assembler._get_node_template("LoadImage")
+        load_node['inputs']['image'] = img_filename
+        assembler.workflow[load_id] = load_node
+        
+        vae_encode_id = assembler._get_unique_id()
+        vae_encode_node = assembler._get_node_template("VAEEncode")
+        vae_encode_node['inputs']['pixels'] = [load_id, 0]
+        vae_encode_node['inputs']['vae'] = [vae_node_id, 0]
+        assembler.workflow[vae_encode_id] = vae_encode_node
+        
+        latent_conn = [vae_encode_id, 0]
+
+        ref_latent_id = assembler._get_unique_id()
+        ref_latent_node = assembler._get_node_template("ReferenceLatent")
+        ref_latent_node['inputs']['conditioning'] = current_pos_conditioning
+        ref_latent_node['inputs']['latent'] = latent_conn
+        assembler.workflow[ref_latent_id] = ref_latent_node
+        
+        current_pos_conditioning = [ref_latent_id, 0]
+
+    assembler.workflow[pos_target_node_id]['inputs'][pos_target_input_name] = current_pos_conditioning
+    
+    print(f"ReferenceLatent injector applied. Re-routed inputs through {len(chain_items)} reference image(s).")

comfy_integration/nodes.py

@@ -0,0 +1,39 @@
+import asyncio
+import execution
+import server
+from nodes import (
+    init_extra_nodes, CheckpointLoaderSimple, EmptyLatentImage, KSampler, 
+    VAEDecode, SaveImage, NODE_CLASS_MAPPINGS, LoadImage, VAEEncode,
+    VAEEncodeForInpaint, ImagePadForOutpaint, LatentUpscaleBy, RepeatLatentBatch
+)
+
+
+def import_custom_nodes() -> None:
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    server_instance = server.PromptServer(loop)
+    execution.PromptQueue(server_instance)
+    
+    loop.run_until_complete(init_extra_nodes())
+
+import_custom_nodes()
+
+CLIPTextEncode = NODE_CLASS_MAPPINGS['CLIPTextEncode']
+CLIPTextEncodeSDXL = NODE_CLASS_MAPPINGS['CLIPTextEncodeSDXL']
+LoraLoader = NODE_CLASS_MAPPINGS['LoraLoader']
+CLIPSetLastLayer = NODE_CLASS_MAPPINGS['CLIPSetLastLayer']
+
+try:
+    KSamplerNode = NODE_CLASS_MAPPINGS['KSampler']
+    SAMPLER_CHOICES = KSamplerNode.INPUT_TYPES()["required"]["sampler_name"][0]
+    SCHEDULER_CHOICES = KSamplerNode.INPUT_TYPES()["required"]["scheduler"][0]
+except Exception:
+    print("⚠️ Could not dynamically get sampler/scheduler choices, using fallback list.")
+    SAMPLER_CHOICES = ['euler', 'dpmpp_2m_sde_gpu']
+    SCHEDULER_CHOICES = ['normal', 'karras']
+
+checkpointloadersimple = CheckpointLoaderSimple()
+loraloader = LoraLoader()
+
+
+print("✅ ComfyUI custom nodes and class mappings are ready.")

comfy_integration/setup.py

@@ -0,0 +1,67 @@
+import os
+import sys
+import shutil
+
+from core.settings import *
+
+def move_and_overwrite(src, dst):
+    if os.path.isdir(src):
+        if os.path.exists(dst):
+            shutil.rmtree(dst)
+        shutil.move(src, dst)
+    elif os.path.isfile(src):
+        if os.path.exists(dst):
+            os.remove(dst)
+        shutil.move(src, dst)
+
+def initialize_comfyui():
+    APP_DIR = sys.path[0]
+    COMFYUI_TEMP_DIR = "ComfyUI_temp"
+
+    print("--- Cloning ComfyUI Repository ---")
+    if not os.path.exists(COMFYUI_TEMP_DIR):
+        os.system(f"git clone https://github.com/comfy-Org/ComfyUI {COMFYUI_TEMP_DIR}")
+        print("✅ ComfyUI repository cloned.")
+    else:
+        print("✅ ComfyUI repository already exists.")
+
+    print(f"--- Merging ComfyUI from '{COMFYUI_TEMP_DIR}' to '{APP_DIR}' ---")
+    for item in os.listdir(COMFYUI_TEMP_DIR):
+        src_path = os.path.join(COMFYUI_TEMP_DIR, item)
+        dst_path = os.path.join(APP_DIR, item)
+        if item == '.git':
+            continue
+        move_and_overwrite(src_path, dst_path)
+    
+    try:
+        shutil.rmtree(COMFYUI_TEMP_DIR)
+        print("✅ ComfyUI merged and temporary directory removed.")
+    except OSError as e:
+        print(f"⚠️ Could not remove temporary directory '{COMFYUI_TEMP_DIR}': {e}")
+
+    print("--- Cloning third-party extensions for ComfyUI ---")
+    controlnet_aux_path = os.path.join(APP_DIR, "custom_nodes", "comfyui_controlnet_aux")
+    if not os.path.exists(controlnet_aux_path):
+        os.system(f"git clone https://github.com/Fannovel16/comfyui_controlnet_aux.git {controlnet_aux_path}")
+        print("✅ comfyui_controlnet_aux extension cloned.")
+    else:
+        print("✅ comfyui_controlnet_aux extension already exists.")
+
+
+    print(f"✅ Current working directory is: {os.getcwd()}")
+    
+    import comfy.model_management
+    print("--- Environment Ready ---")
+
+    print("✅ ComfyUI initialized with default attention mechanism.")
+
+    os.makedirs(os.path.join(APP_DIR, CHECKPOINT_DIR), exist_ok=True)
+    os.makedirs(os.path.join(APP_DIR, LORA_DIR), exist_ok=True)
+    os.makedirs(os.path.join(APP_DIR, EMBEDDING_DIR), exist_ok=True)
+    os.makedirs(os.path.join(APP_DIR, CONTROLNET_DIR), exist_ok=True)
+    os.makedirs(os.path.join(APP_DIR, MODEL_PATCHES_DIR), exist_ok=True)
+    os.makedirs(os.path.join(APP_DIR, DIFFUSION_MODELS_DIR), exist_ok=True)
+    os.makedirs(os.path.join(APP_DIR, VAE_DIR), exist_ok=True)
+    os.makedirs(os.path.join(APP_DIR, TEXT_ENCODERS_DIR), exist_ok=True)
+    os.makedirs(os.path.join(APP_DIR, INPUT_DIR), exist_ok=True)
+    print("✅ All required model directories are present.")

core/generation_logic.py

@@ -0,0 +1,25 @@
+from typing import Any, Dict
+import gradio as gr
+
+from core.pipelines.controlnet_preprocessor import ControlNetPreprocessorPipeline
+from core.pipelines.sd_image_pipeline import SdImagePipeline
+
+controlnet_preprocessor_pipeline = ControlNetPreprocessorPipeline()
+sd_image_pipeline = SdImagePipeline()
+
+
+def build_reverse_map():
+    from nodes import NODE_DISPLAY_NAME_MAPPINGS
+    import core.pipelines.controlnet_preprocessor as cn_module
+    
+    if cn_module.REVERSE_DISPLAY_NAME_MAP is None:
+        cn_module.REVERSE_DISPLAY_NAME_MAP = {v: k for k, v in NODE_DISPLAY_NAME_MAPPINGS.items()}
+        if "Semantic Segmentor (legacy, alias for UniFormer)" not in cn_module.REVERSE_DISPLAY_NAME_MAP:
+             cn_module.REVERSE_DISPLAY_NAME_MAP["Semantic Segmentor (legacy, alias for UniFormer)"] = "SemSegPreprocessor"
+
+
+def run_cn_preprocessor_entry(*args, **kwargs):
+    return controlnet_preprocessor_pipeline.run(*args, **kwargs)
+
+def generate_image_wrapper(ui_inputs: dict, progress=gr.Progress(track_tqdm=True)):
+    return sd_image_pipeline.run(ui_inputs=ui_inputs, progress=progress)

core/model_manager.py

@@ -0,0 +1,168 @@
+import gc
+from typing import Dict, List, Any, Set
+
+import torch
+import gradio as gr
+from comfy import model_management
+
+from core.settings import ALL_MODEL_MAP, CHECKPOINT_DIR, LORA_DIR, DIFFUSION_MODELS_DIR, VAE_DIR, TEXT_ENCODERS_DIR
+from comfy_integration.nodes import LoraLoader
+from nodes import NODE_CLASS_MAPPINGS
+from utils.app_utils import get_value_at_index, _ensure_model_downloaded
+
+
+class ModelManager:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if not cls._instance:
+            cls._instance = super(ModelManager, cls).__new__(cls, *args, **kwargs)
+        return cls._instance
+
+    def __init__(self):
+        if hasattr(self, 'initialized'):
+            return
+        self.loaded_models: Dict[str, Any] = {}
+        self.last_active_loras: List[Dict[str, Any]] = []
+        self.initialized = True
+        print("✅ ModelManager initialized.")
+
+    def get_loaded_model_names(self) -> Set[str]:
+        return set(self.loaded_models.keys())
+
+    def _load_model_combo(self, display_name: str, active_loras: List[Dict[str, Any]], progress) -> Dict[str, Any]:
+        print(f"--- [ModelManager] Loading model combo: '{display_name}' ---")
+        
+        if display_name not in ALL_MODEL_MAP:
+            raise ValueError(f"Model '{display_name}' not found in configuration.")
+            
+        _, components, _, _ = ALL_MODEL_MAP[display_name]
+        
+        unet_filename = components.get('unet')
+        clip_filename = components.get('clip')
+        vae_filename = components.get('vae')
+        
+        if not all([unet_filename, clip_filename, vae_filename]):
+             raise ValueError(f"Model '{display_name}' is missing required components (unet, clip, or vae) in model_list.yaml.")
+
+        unet_loader = NODE_CLASS_MAPPINGS["UNETLoader"]()
+        clip_loader = NODE_CLASS_MAPPINGS["CLIPLoader"]()
+        vae_loader = NODE_CLASS_MAPPINGS["VAELoader"]()
+
+        print("  - Loading UNET...")
+        unet_tuple = unet_loader.load_unet(unet_name=unet_filename, weight_dtype="default")
+        
+        print("  - Loading CLIP...")
+        clip_tuple = clip_loader.load_clip(clip_name=clip_filename, type="flux2", device="default")
+        
+        print("  - Loading VAE...")
+        vae_tuple = vae_loader.load_vae(vae_name=vae_filename)
+        
+        unet_object = get_value_at_index(unet_tuple, 0)
+        clip_object = get_value_at_index(clip_tuple, 0)
+
+        if active_loras:
+            print(f"--- [ModelManager] Applying {len(active_loras)} LoRAs on CPU... ---")
+            lora_loader = LoraLoader()
+            patched_unet, patched_clip = unet_object, clip_object
+
+            for lora_info in active_loras:
+                patched_unet, patched_clip = lora_loader.load_lora(
+                    model=patched_unet,
+                    clip=patched_clip,
+                    lora_name=lora_info["lora_name"],
+                    strength_model=lora_info["strength_model"],
+                    strength_clip=lora_info["strength_clip"]
+                )
+            
+            unet_object = patched_unet
+            clip_object = patched_clip
+            print(f"--- [ModelManager] ✅ All LoRAs merged into the model on CPU. ---")
+            
+        loaded_combo = {
+            "unet": (unet_object,),
+            "clip": (clip_object,),
+            "vae": vae_tuple,
+        }
+        
+        print(f"--- [ModelManager] ✅ Successfully loaded combo '{display_name}' to CPU/RAM ---")
+        return loaded_combo
+
+    def move_models_to_gpu(self, required_models: List[str]):
+        print(f"--- [ModelManager] Moving models to GPU: {required_models} ---")
+        models_to_load_gpu = []
+        for name in required_models:
+            if name in self.loaded_models:
+                model_combo = self.loaded_models[name]
+                models_to_load_gpu.append(get_value_at_index(model_combo.get("unet"), 0))
+
+        if models_to_load_gpu:
+            model_management.load_models_gpu(models_to_load_gpu)
+            print("--- [ModelManager] ✅ Models successfully moved to GPU. ---")
+        else:
+            print("--- [ModelManager] ⚠️ No component models found to move to GPU. ---")
+
+    def ensure_models_downloaded(self, required_models: List[str], progress):
+        print(f"--- [ModelManager] Ensuring models are downloaded: {required_models} ---")
+        
+        files_to_download = set()
+        for display_name in required_models:
+            if display_name in ALL_MODEL_MAP:
+                _, components, _, _ = ALL_MODEL_MAP[display_name]
+                for component_file in components.values():
+                    files_to_download.add(component_file)
+        
+        files_to_download = list(files_to_download)
+        total_files = len(files_to_download)
+        
+        for i, filename in enumerate(files_to_download):
+            if progress and hasattr(progress, '__call__'):
+                progress(i / total_files, desc=f"Checking file: {filename}")
+            try:
+                _ensure_model_downloaded(filename, progress)
+            except Exception as e:
+                raise gr.Error(f"Failed to download model component '{filename}'. Reason: {e}")
+
+        print(f"--- [ModelManager] ✅ All required models are present on disk. ---")
+    
+    def load_managed_models(self, required_models: List[str], active_loras: List[Dict[str, Any]], progress) -> Dict[str, Any]:
+        required_set = set(required_models)
+        current_set = set(self.loaded_models.keys())
+
+        loras_changed = active_loras != self.last_active_loras
+
+        models_to_unload = current_set - required_set
+        if models_to_unload or loras_changed:
+            if models_to_unload:
+                print(f"--- [ModelManager] Models to unload: {models_to_unload} ---")
+            if loras_changed and not models_to_unload:
+                models_to_unload = current_set.intersection(required_set)
+                if active_loras:
+                    print(f"--- [ModelManager] LoRA configuration changed. Reloading base model(s): {models_to_unload} ---")
+                else:
+                    print(f"--- [ModelManager] LoRAs removed. Reloading base model(s) to clear patches: {models_to_unload} ---")
+                
+            model_management.unload_all_models()
+            self.loaded_models.clear()
+            gc.collect()
+            torch.cuda.empty_cache()
+            print("--- [ModelManager] All models unloaded to free RAM. ---")
+        
+        models_to_load = required_set if (models_to_unload or loras_changed) else required_set - current_set
+        
+        if models_to_load:
+            print(f"--- [ModelManager] Models to load: {models_to_load} ---")
+            for i, display_name in enumerate(models_to_load):
+                progress(i / len(models_to_load), desc=f"Loading model: {display_name}")
+                try:
+                    loaded_model_data = self._load_model_combo(display_name, active_loras, progress)
+                    self.loaded_models[display_name] = loaded_model_data
+                except Exception as e:
+                    raise gr.Error(f"Failed to load model combo or apply LoRA for '{display_name}'. Reason: {e}")
+        else:
+             print(f"--- [ModelManager] All required models are already loaded. ---")
+
+        self.last_active_loras = active_loras
+        return {name: self.loaded_models[name] for name in required_models}
+
+model_manager = ModelManager()

core/pipelines/base_pipeline.py

@@ -0,0 +1,53 @@
+from abc import ABC, abstractmethod
+from typing import List, Any, Dict
+import gradio as gr
+import spaces
+import tempfile
+import imageio
+import numpy as np
+
+class BasePipeline(ABC):
+    def __init__(self):
+        from core.model_manager import model_manager
+        self.model_manager = model_manager
+
+    @abstractmethod
+    def get_required_models(self, **kwargs) -> List[str]:
+        pass
+
+    @abstractmethod
+    def run(self, *args, progress: gr.Progress, **kwargs) -> Any:
+        pass
+
+    def _ensure_models_downloaded(self, progress: gr.Progress, **kwargs):
+        """Ensures model files are downloaded before requesting GPU."""
+        required_models = self.get_required_models(**kwargs)
+        self.model_manager.ensure_models_downloaded(required_models, progress=progress)
+
+    def _execute_gpu_logic(self, gpu_function: callable, duration: int, default_duration: int, task_name: str, *args, **kwargs):
+        final_duration = default_duration
+        try:
+            if duration is not None and int(duration) > 0:
+                final_duration = int(duration)
+        except (ValueError, TypeError):
+            print(f"Invalid ZeroGPU duration input for {task_name}. Using default {default_duration}s.")
+            pass
+        
+        print(f"Requesting ZeroGPU for {task_name} with duration: {final_duration} seconds.")
+        gpu_runner = spaces.GPU(duration=final_duration)(gpu_function)
+        
+        return gpu_runner(*args, **kwargs)
+
+    def _encode_video_from_frames(self, frames_tensor_cpu: 'torch.Tensor', fps: int, progress: gr.Progress) -> str:
+        progress(0.9, desc="Encoding video on CPU...")
+        frames_np = (frames_tensor_cpu.numpy() * 255.0).astype(np.uint8)
+        
+        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_video_file:
+            video_path = temp_video_file.name
+            writer = imageio.get_writer(video_path, fps=fps, codec='libx264', quality=8)
+            for frame in frames_np:
+                writer.append_data(frame)
+            writer.close()
+            
+        progress(1.0, desc="Done!")
+        return video_path

core/pipelines/controlnet_preprocessor.py

@@ -0,0 +1,143 @@
+from typing import Dict, Any, List
+import imageio
+import tempfile
+import numpy as np
+import torch
+import gradio as gr
+from PIL import Image
+import spaces
+
+from .base_pipeline import BasePipeline
+from comfy_integration.nodes import NODE_CLASS_MAPPINGS
+from nodes import NODE_DISPLAY_NAME_MAPPINGS
+from utils.app_utils import get_value_at_index
+
+REVERSE_DISPLAY_NAME_MAP = None
+CPU_ONLY_PREPROCESSORS = {
+    "Binary Lines", "Canny Edge", "Color Pallete", "Fake Scribble Lines (aka scribble_hed)",
+    "Image Intensity", "Image Luminance", "Inpaint Preprocessor", "PyraCanny", "Scribble Lines",
+    "Scribble XDoG Lines", "Standard Lineart", "Content Shuffle", "Tile"
+}
+
+def run_node_by_function_name(node_instance: Any, **kwargs) -> Any:
+    node_class = type(node_instance)
+    function_name = getattr(node_class, 'FUNCTION', None)
+    if not function_name:
+        raise AttributeError(f"Node class '{node_class.__name__}' is missing the required 'FUNCTION' attribute.")
+    execution_method = getattr(node_instance, function_name, None)
+    if not callable(execution_method):
+        raise AttributeError(f"Method '{function_name}' not found or not callable on node '{node_class.__name__}'.")
+    return execution_method(**kwargs)
+
+class ControlNetPreprocessorPipeline(BasePipeline):
+    def get_required_models(self, **kwargs) -> List[str]:
+        return []
+
+    def _gpu_logic(
+        self, pil_images: List[Image.Image], preprocessor_name: str, model_name: str,
+        params: Dict[str, Any], progress=gr.Progress(track_tqdm=True)
+    ) -> List[Image.Image]:
+        global REVERSE_DISPLAY_NAME_MAP
+        if REVERSE_DISPLAY_NAME_MAP is None:
+            raise RuntimeError("REVERSE_DISPLAY_NAME_MAP has not been initialized. `build_reverse_map` must be called on startup.")
+            
+        class_name = REVERSE_DISPLAY_NAME_MAP.get(preprocessor_name)
+        if not class_name or class_name not in NODE_CLASS_MAPPINGS:
+            raise ValueError(f"Preprocessor '{preprocessor_name}' not found.")
+            
+        preprocessor_instance = NODE_CLASS_MAPPINGS[class_name]()
+        call_args = {**params, 'ckpt_name': model_name}
+        
+        processed_pil_images = []
+        total_frames = len(pil_images)
+
+        for i, frame_pil in enumerate(pil_images):
+            progress(i / total_frames, desc=f"Processing frame {i+1}/{total_frames} with {preprocessor_name}...")
+            
+            frame_tensor = torch.from_numpy(np.array(frame_pil).astype(np.float32) / 255.0).unsqueeze(0)
+            
+            resolution_arg = {'resolution': max(frame_tensor.shape[2], frame_tensor.shape[3])}
+
+            result_tuple = run_node_by_function_name(
+                preprocessor_instance, 
+                image=frame_tensor, 
+                **resolution_arg,
+                **call_args
+            )
+            
+            processed_tensor = get_value_at_index(result_tuple, 0)
+            processed_np = (processed_tensor.squeeze(0).cpu().numpy().clip(0, 1) * 255.0).astype(np.uint8)
+            processed_pil_images.append(Image.fromarray(processed_np))
+            
+        return processed_pil_images
+
+    def run(self, input_type, image_input, video_input, preprocessor_name, model_name, zero_gpu_duration, *args, progress=gr.Progress(track_tqdm=True)):
+        from utils import app_utils
+        pil_images, is_video, fps = [], False, 30
+        
+        progress(0, desc="Reading input file...")
+        if input_type == "Image":
+            if image_input is None: raise gr.Error("Please provide an input image.")
+            pil_images = [image_input]
+        elif input_type == "Video":
+            if video_input is None: raise gr.Error("Please provide an input video.")
+            try:
+                video_reader = imageio.get_reader(video_input)
+                meta = video_reader.get_meta_data()
+                fps = meta.get('fps', 30)
+                pil_images = [Image.fromarray(frame) for frame in video_reader]
+                is_video = True
+                video_reader.close()
+            except Exception as e: raise gr.Error(f"Failed to read video file: {e}")
+        else:
+            raise gr.Error("Invalid input type selected.")
+
+        if not pil_images: raise gr.Error("Could not extract any frames from the input.")
+        
+        if app_utils.PREPROCESSOR_PARAMETER_MAP is None:
+            raise RuntimeError("Preprocessor parameter map is not built. Check startup logs.")
+            
+        params_config = app_utils.PREPROCESSOR_PARAMETER_MAP.get(preprocessor_name, [])
+        sliders_params = [p for p in params_config if p['type'] in ["INT", "FLOAT"]]
+        dropdown_params = [p for p in params_config if isinstance(p['type'], list)]
+        checkbox_params = [p for p in params_config if p['type'] == "BOOLEAN"]
+        ordered_params_config = sliders_params + dropdown_params + checkbox_params
+        param_names = [p['name'] for p in ordered_params_config]
+        provided_params = {param_names[i]: args[i] for i in range(len(param_names))}
+        
+        if preprocessor_name not in CPU_ONLY_PREPROCESSORS:
+            print(f"--- '{preprocessor_name}' requires GPU, requesting ZeroGPU. ---")
+            try:
+                processed_pil_images = self._execute_gpu_logic(
+                    self._gpu_logic,
+                    duration=zero_gpu_duration,
+                    default_duration=60,
+                    task_name=f"Preprocessor '{preprocessor_name}'",
+                    pil_images=pil_images,
+                    preprocessor_name=preprocessor_name,
+                    model_name=model_name,
+                    params=provided_params,
+                    progress=progress
+                )
+            except Exception as e:
+                import traceback; traceback.print_exc()
+                raise gr.Error(f"Failed to run preprocessor '{preprocessor_name}' on GPU: {e}")
+        else:
+            print(f"--- Running '{preprocessor_name}' on CPU, no ZeroGPU requested. ---")
+            try:
+                processed_pil_images = self._gpu_logic(pil_images, preprocessor_name, model_name, provided_params, progress=progress)
+            except Exception as e:
+                import traceback; traceback.print_exc()
+                raise gr.Error(f"Failed to run preprocessor '{preprocessor_name}' on CPU: {e}")
+
+        if not processed_pil_images: raise gr.Error("Processing returned no frames.")
+        
+        progress(0.9, desc="Finalizing output...")
+        if is_video:
+            frames_np = [np.array(img) for img in processed_pil_images]
+            frames_tensor = torch.from_numpy(np.stack(frames_np)).to(torch.float32) / 255.0
+            video_path = self._encode_video_from_frames(frames_tensor, fps, progress)
+            return [video_path]
+        else:
+            progress(1.0, desc="Done!")
+            return processed_pil_images

core/pipelines/sd_image_pipeline.py

@@ -0,0 +1,389 @@
+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, sanitize_filename
+from core.workflow_assembler import WorkflowAssembler
+
+class SdImagePipeline(BasePipeline):
+    def get_required_models(self, model_display_name: str, **kwargs) -> List[str]:
+        return [model_display_name]
+
+    def _topological_sort(self, workflow: Dict[str, Any]) -> List[str]:
+        graph = defaultdict(list)
+        in_degree = {node_id: 0 for node_id in workflow}
+
+        for node_id, node_info in workflow.items():
+            for input_value in node_info.get('inputs', {}).values():
+                if isinstance(input_value, list) and len(input_value) == 2 and isinstance(input_value[0], str):
+                    source_node_id = input_value[0]
+                    if source_node_id in workflow:
+                        graph[source_node_id].append(node_id)
+                        in_degree[node_id] += 1
+
+        queue = deque([node_id for node_id, degree in in_degree.items() if degree == 0])
+        
+        sorted_nodes = []
+        while queue:
+            current_node_id = queue.popleft()
+            sorted_nodes.append(current_node_id)
+
+            for neighbor_node_id in graph[current_node_id]:
+                in_degree[neighbor_node_id] -= 1
+                if in_degree[neighbor_node_id] == 0:
+                    queue.append(neighbor_node_id)
+        
+        if len(sorted_nodes) != len(workflow):
+            raise RuntimeError("Workflow contains a cycle and cannot be executed.")
+            
+        return sorted_nodes
+
+
+    def _execute_workflow(self, workflow: Dict[str, Any], initial_objects: Dict[str, Any]):
+        with torch.no_grad():
+            computed_outputs = initial_objects
+            
+            try:
+                sorted_node_ids = self._topological_sort(workflow)
+                print(f"--- [Workflow Executor] Execution order: {sorted_node_ids}")
+            except RuntimeError as e:
+                print("--- [Workflow Executor] ERROR: Failed to sort workflow. Dumping graph details. ---")
+                for node_id, node_info in workflow.items():
+                    print(f"  Node {node_id} ({node_info['class_type']}):")
+                    for input_name, input_value in node_info['inputs'].items():
+                         if isinstance(input_value, list) and len(input_value) == 2 and isinstance(input_value[0], str):
+                             print(f"    - {input_name} <- [{input_value[0]}, {input_value[1]}]")
+                raise e
+
+            for node_id in sorted_node_ids:
+                if node_id in computed_outputs:
+                    continue
+                    
+                node_info = workflow[node_id]
+                class_type = node_info['class_type']
+                
+                is_loader_with_filename = 'Loader' in class_type and any(key.endswith('_name') for key in node_info['inputs'])
+                if node_id in initial_objects and is_loader_with_filename:
+                    continue
+                
+                node_class = NODE_CLASS_MAPPINGS.get(class_type)
+                if node_class is None:
+                     raise RuntimeError(f"Could not find node class '{class_type}'. Is it imported in comfy_integration/nodes.py?")
+                
+                node_instance = node_class()
+                
+                kwargs = {}
+                for param_name, param_value in node_info['inputs'].items():
+                    if isinstance(param_value, list) and len(param_value) == 2 and isinstance(param_value[0], str):
+                        source_node_id, output_index = param_value
+                        if source_node_id not in computed_outputs:
+                            raise RuntimeError(f"Workflow integrity error: Output of node {source_node_id} needed for {node_id} but not yet computed.")
+                        
+                        source_output_tuple = computed_outputs[source_node_id]
+                        kwargs[param_name] = get_value_at_index(source_output_tuple, output_index)
+                    else:
+                        kwargs[param_name] = param_value
+
+                function_name = getattr(node_class, 'FUNCTION')
+                execution_method = getattr(node_instance, function_name)
+                
+                result = execution_method(**kwargs)
+                computed_outputs[node_id] = result
+            
+            final_node_id = None
+            for node_id in reversed(sorted_node_ids):
+                 if workflow[node_id]['class_type'] == 'SaveImage':
+                     final_node_id = node_id
+                     break
+            
+            if not final_node_id:
+                raise RuntimeError("Workflow does not contain a 'SaveImage' node as the output.")
+
+            save_image_inputs = workflow[final_node_id]['inputs']
+            image_source_node_id, image_source_index = save_image_inputs['images']
+            
+            return get_value_at_index(computed_outputs[image_source_node_id], image_source_index)
+
+    def _gpu_logic(self, ui_inputs: Dict, loras_string: str, required_models_for_gpu: List[str], workflow: Dict[str, Any], assembler: WorkflowAssembler, progress=gr.Progress(track_tqdm=True)):
+        model_display_name = ui_inputs['model_display_name']
+        
+        progress(0.1, desc="Moving models to GPU...")
+        self.model_manager.move_models_to_gpu(required_models_for_gpu)
+        
+        progress(0.4, desc="Executing workflow...")
+        
+        loaded_model_combo = self.model_manager.loaded_models[model_display_name]
+        
+        initial_objects = {}
+        
+        unet_loader_id = assembler.node_map.get("unet_loader")
+        clip_loader_id = assembler.node_map.get("clip_loader")
+        vae_loader_id = assembler.node_map.get("vae_loader")
+
+        if unet_loader_id: initial_objects[unet_loader_id] = loaded_model_combo.get("unet")
+        if clip_loader_id: initial_objects[clip_loader_id] = loaded_model_combo.get("clip")
+        if vae_loader_id: initial_objects[vae_loader_id] = loaded_model_combo.get("vae")
+
+        if not all([unet_loader_id, clip_loader_id, vae_loader_id]):
+            raise RuntimeError("Workflow is missing one or more required loaders (unet_loader, clip_loader, vae_loader).")
+            
+        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 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']
+        
+        ui_inputs['positive_prompt'] = sanitize_prompt(ui_inputs.get('positive_prompt', ''))
+        ui_inputs['negative_prompt'] = sanitize_prompt(ui_inputs.get('negative_prompt', ''))
+        
+        required_models = self.get_required_models(model_display_name=ui_inputs['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 = [], []
+        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, ui_inputs['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}")
+        
+        progress(0.1, desc="Loading models into RAM...")
+        self.model_manager.load_managed_models(required_models, active_loras=active_loras_for_gpu, progress=progress)
+
+        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 input_image_pil:
+                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['inpaint_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 input_image_pil:
+                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)
+
+        elif task_type == 'hires_fix':
+            input_image_pil = ui_inputs.get('hires_image')
+            if input_image_pil:
+                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, ui_inputs['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
+
+        from utils.app_utils import get_vae_path
+        vae_source = ui_inputs.get('vae_source')
+        vae_id = ui_inputs.get('vae_id')
+        vae_file = ui_inputs.get('vae_file')
+        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, ui_inputs.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   = conditioning_data[0*num_units : 1*num_units]
+            widths    = conditioning_data[1*num_units : 2*num_units]
+            heights   = conditioning_data[2*num_units : 3*num_units]
+            xs        = conditioning_data[3*num_units : 4*num_units]
+            ys        = conditioning_data[4*num_units : 5*num_units]
+            strengths = conditioning_data[5*num_units : 6*num_units]
+
+            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])
+                    })
+
+        reference_latent_data = ui_inputs.get('reference_latent_data', [])
+        active_reference_latents = []
+        if reference_latent_data:
+            for img_pil in reference_latent_data:
+                if img_pil is not None:
+                    temp_file_path = os.path.join(INPUT_DIR, f"temp_ref_{random.randint(1000, 9999)}.png")
+                    img_pil.save(temp_file_path, "PNG")
+                    active_reference_latents.append(os.path.basename(temp_file_path))
+                    temp_files_to_clean.append(temp_file_path)
+
+        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)
+
+        dynamic_values = {'task_type': ui_inputs['task_type']}
+        
+        recipe_path = os.path.join(os.path.dirname(__file__), "workflow_recipes", "sd_unified_recipe.yaml")
+        assembler = WorkflowAssembler(recipe_path, dynamic_values=dynamic_values)
+        
+        model_display_name = ui_inputs['model_display_name']
+        if model_display_name not in ALL_MODEL_MAP:
+            raise gr.Error(f"Model '{model_display_name}' is not configured in model_list.yaml.")
+        
+        _, components, _, _ = ALL_MODEL_MAP[model_display_name]
+        
+        workflow_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'],
+            "denoise": ui_inputs['denoise'], 
+            "input_image": ui_inputs.get('input_image'),
+            "inpaint_image": ui_inputs.get('inpaint_image'),
+            "inpaint_mask": ui_inputs.get('inpaint_mask'),
+            "left": ui_inputs.get('outpaint_left'), "top": ui_inputs.get('outpaint_top'),
+            "right": ui_inputs.get('outpaint_right'), "bottom": ui_inputs.get('outpaint_bottom'),
+            "hires_upscaler": ui_inputs.get('hires_upscaler'), "hires_scale_by": ui_inputs.get('hires_scale_by'),
+            "unet_name": components['unet'],
+            "clip_name": components['clip'],
+            "vae_name": ui_inputs.get('vae_name', components['vae']),
+            "conditioning_chain": active_conditioning,
+            "reference_latent_chain": active_reference_latents,
+        }
+        
+        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,
+                required_models_for_gpu=required_models,
+                workflow=workflow,
+                assembler=assembler, 
+                progress=progress
+            )
+        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_recipes/_partials/_base_sampler.yaml

@@ -0,0 +1,23 @@
+nodes:
+  ksampler:
+    class_type: KSampler
+
+  vae_decode:
+    class_type: VAEDecode
+  save_image:
+    class_type: SaveImage
+    params: {}
+
+connections:
+  - from: "ksampler:0"
+    to: "vae_decode:samples"
+  - from: "vae_decode:0"
+    to: "save_image:images"
+
+ui_map:
+  seed: "ksampler:seed"
+  steps: "ksampler:steps"
+  cfg: "ksampler:cfg"
+  sampler_name: "ksampler:sampler_name"
+  scheduler: "ksampler:scheduler"
+  denoise: "ksampler:denoise"

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

@@ -0,0 +1,56 @@
+nodes:
+  unet_loader:
+    class_type: UNETLoader
+    title: "Load Diffusion Model"
+    params:
+      weight_dtype: "default"
+  clip_loader:
+    class_type: CLIPLoader
+    title: "Load CLIP"
+    params:
+      type: "flux2"
+      device: "default"
+  vae_loader:
+    class_type: VAELoader
+    title: "Load VAE"
+  
+  pos_prompt:
+    class_type: CLIPTextEncode
+    title: "CLIP Text Encode (Positive)"
+  neg_prompt:
+    class_type: CLIPTextEncode
+    title: "CLIP Text Encode (Negative)"
+
+connections:
+  - from: "unet_loader:0"
+    to: "ksampler:model"
+  - from: "clip_loader:0"
+    to: "pos_prompt:clip"
+  - from: "clip_loader:0"
+    to: "neg_prompt:clip"
+  - from: "vae_loader:0"
+    to: "vae_decode:vae"
+  - from: "vae_loader:0"
+    to: "vae_encode:vae"
+  
+  - from: "pos_prompt:0"
+    to: "ksampler:positive"
+  - from: "neg_prompt:0"
+    to: "ksampler:negative"
+
+dynamic_conditioning_chains:
+  conditioning_chain:
+    ksampler_node: "ksampler"
+    clip_source: "clip_loader:0"
+
+dynamic_reference_latent_chains:
+  reference_latent_chain:
+    ksampler_node: "ksampler"
+    vae_node: "vae_loader"
+
+ui_map:
+  unet_name: "unet_loader:unet_name"
+  clip_name: "clip_loader:clip_name"
+  vae_name: "vae_loader:vae_name"
+  positive_prompt: "pos_prompt:text"
+  negative_prompt: "neg_prompt:text"

core/pipelines/workflow_recipes/_partials/input/hires_fix.yaml

@@ -0,0 +1,26 @@
+nodes:
+  input_image_loader:
+    class_type: LoadImage
+    
+  vae_encode:
+    class_type: VAEEncode
+    
+  latent_upscaler:
+    class_type: LatentUpscaleBy
+
+  latent_source:
+    class_type: RepeatLatentBatch
+
+connections:
+  - from: "input_image_loader:0"
+    to: "vae_encode:pixels"
+  - from: "vae_encode:0"
+    to: "latent_upscaler:samples"
+  - from: "latent_upscaler:0"
+    to: "latent_source:samples"
+
+ui_map:
+  input_image: "input_image_loader:image"
+  hires_upscaler: "latent_upscaler:upscale_method"
+  hires_scale_by: "latent_upscaler:scale_by"
+  batch_size: "latent_source:amount"

core/pipelines/workflow_recipes/_partials/input/img2img.yaml

@@ -0,0 +1,19 @@
+nodes:
+  input_image_loader:
+    class_type: LoadImage
+  
+  vae_encode:
+    class_type: VAEEncode
+
+  latent_source:
+    class_type: RepeatLatentBatch
+
+connections:
+  - from: "input_image_loader:0"
+    to: "vae_encode:pixels"
+  - from: "vae_encode:0"
+    to: "latent_source:samples"
+    
+ui_map:
+  input_image: "input_image_loader:image"
+  batch_size: "latent_source:amount"

core/pipelines/workflow_recipes/_partials/input/inpaint.yaml

@@ -0,0 +1,25 @@
+nodes:
+  inpaint_loader:
+    class_type: LoadImage
+    title: "Load Inpaint Image+Mask"
+      
+  vae_encode:
+    class_type: VAEEncodeForInpaint
+    params:
+      grow_mask_by: 6
+      
+  latent_source:
+    class_type: RepeatLatentBatch
+      
+connections:
+  - from: "inpaint_loader:0"
+    to: "vae_encode:pixels"
+  - from: "inpaint_loader:1"
+    to: "vae_encode:mask"
+    
+  - from: "vae_encode:0"
+    to: "latent_source:samples"
+
+ui_map:
+  inpaint_image: "inpaint_loader:image"
+  batch_size: "latent_source:amount"

core/pipelines/workflow_recipes/_partials/input/outpaint.yaml

@@ -0,0 +1,38 @@
+nodes:
+  input_image_loader:
+    class_type: LoadImage
+    
+  pad_image:
+    class_type: ImagePadForOutpaint
+    params:
+      feathering: 10
+      
+  vae_encode:
+    class_type: VAEEncodeForInpaint
+    params:
+      grow_mask_by: 6
+
+  latent_source:
+    class_type: RepeatLatentBatch
+
+connections:
+  - from: "input_image_loader:0"
+    to: "pad_image:image"
+    
+  - from: "pad_image:0"
+    to: "vae_encode:pixels"
+  - from: "pad_image:1"
+    to: "vae_encode:mask"
+
+  - from: "vae_encode:0"
+    to: "latent_source:samples"
+
+ui_map:
+  input_image: "input_image_loader:image"
+  
+  left: "pad_image:left"
+  top: "pad_image:top"
+  right: "pad_image:right"
+  bottom: "pad_image:bottom"
+  
+  batch_size: "latent_source:amount"

core/pipelines/workflow_recipes/_partials/input/txt2img.yaml

@@ -0,0 +1,8 @@
+nodes:
+  latent_source:
+    class_type: EmptyFlux2LatentImage
+
+ui_map:
+  width: "latent_source:width"
+  height: "latent_source:height"
+  batch_size: "latent_source:batch_size"

core/pipelines/workflow_recipes/sd_unified_recipe.yaml

@@ -0,0 +1,8 @@
+imports:
+  - "_partials/_base_sampler.yaml"
+  - "_partials/input/{{ task_type }}.yaml"
+  - "_partials/conditioning/flux2.yaml"
+
+connections:
+  - from: "latent_source:0"
+    to: "ksampler:latent_image"

core/settings.py

@@ -0,0 +1,125 @@
+import yaml
+import os
+from collections import OrderedDict
+
+CHECKPOINT_DIR = "models/checkpoints"
+LORA_DIR = "models/loras"
+EMBEDDING_DIR = "models/embeddings" 
+CONTROLNET_DIR = "models/controlnet"
+MODEL_PATCHES_DIR = "models/model_patches"
+DIFFUSION_MODELS_DIR = "models/diffusion_models"
+VAE_DIR = "models/vae"
+TEXT_ENCODERS_DIR = "models/text_encoders"
+INPUT_DIR = "input"
+OUTPUT_DIR = "output"
+
+_PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+_MODEL_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'model_list.yaml')
+_FILE_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'file_list.yaml')
+_CONSTANTS_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'constants.yaml')
+_MODEL_DEFAULTS_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'model_defaults.yaml')
+
+
+def load_constants_from_yaml(filepath=_CONSTANTS_PATH):
+    if not os.path.exists(filepath):
+        print(f"Warning: Constants file not found at {filepath}. Using fallback values.")
+        return {}
+    with open(filepath, 'r', encoding='utf-8') as f:
+        return yaml.safe_load(f)
+
+def load_file_download_map(filepath=_FILE_LIST_PATH):
+    if not os.path.exists(filepath):
+        raise FileNotFoundError(f"The file list (for downloads) was not found at: {filepath}")
+    
+    with open(filepath, 'r', encoding='utf-8') as f:
+        file_list_data = yaml.safe_load(f)
+    
+    download_info_map = {}
+    for category, files in file_list_data.get('file', {}).items():
+        if isinstance(files, list):
+            for file_info in files:
+                if 'filename' in file_info:
+                    file_info['category'] = category
+                    download_info_map[file_info['filename']] = file_info
+    return download_info_map
+
+
+def load_models_from_yaml(model_list_filepath=_MODEL_LIST_PATH, download_map=None):
+    if not os.path.exists(model_list_filepath):
+        raise FileNotFoundError(f"The model list file was not found at: {model_list_filepath}")
+    if download_map is None:
+        raise ValueError("download_map must be provided to load_models_from_yaml")
+
+    with open(model_list_filepath, 'r', encoding='utf-8') as f:
+        model_data = yaml.safe_load(f)
+
+    model_maps = {
+        "MODEL_MAP_CHECKPOINT": OrderedDict(),
+        "ALL_MODEL_MAP": OrderedDict(),
+    }
+    category_map_names = {
+        "Checkpoint": "MODEL_MAP_CHECKPOINT",
+    }
+    
+    for category, models in model_data.items():
+        if category in category_map_names:
+            map_name = category_map_names[category]
+            if not isinstance(models, list): continue
+            for model in models:
+                display_name = model['display_name']
+                components = model.get('components', {})
+
+                model_tuple = (
+                    None,
+                    components,
+                    "SDXL",
+                    None
+                )
+                model_maps[map_name][display_name] = model_tuple
+                model_maps["ALL_MODEL_MAP"][display_name] = model_tuple
+    
+    return model_maps
+
+def load_model_defaults(filepath=_MODEL_DEFAULTS_PATH):
+    if not os.path.exists(filepath):
+        print(f"Warning: Model defaults file not found at {filepath}. Using empty defaults.")
+        return {}
+    with open(filepath, 'r', encoding='utf-8') as f:
+        return yaml.safe_load(f)
+
+try:
+    ALL_FILE_DOWNLOAD_MAP = load_file_download_map()
+    loaded_maps = load_models_from_yaml(download_map=ALL_FILE_DOWNLOAD_MAP)
+    MODEL_MAP_CHECKPOINT = loaded_maps["MODEL_MAP_CHECKPOINT"]
+    ALL_MODEL_MAP = loaded_maps["ALL_MODEL_MAP"]
+    
+    MODEL_TYPE_MAP = {k: v[2] for k, v in ALL_MODEL_MAP.items()}
+    
+    ALL_MODEL_DEFAULTS = load_model_defaults()
+
+except Exception as e:
+    print(f"FATAL: Could not load model configuration from YAML. Error: {e}")
+    ALL_FILE_DOWNLOAD_MAP = {}
+    MODEL_MAP_CHECKPOINT, ALL_MODEL_MAP = {}, {}
+    MODEL_TYPE_MAP = {}
+    ALL_MODEL_DEFAULTS = {}
+
+
+try:
+    _constants = load_constants_from_yaml()
+    MAX_LORAS = _constants.get('MAX_LORAS', 5)
+    MAX_EMBEDDINGS = _constants.get('MAX_EMBEDDINGS', 5)
+    MAX_CONDITIONINGS = _constants.get('MAX_CONDITIONINGS', 10)
+    MAX_CONTROLNETS = _constants.get('MAX_CONTROLNETS', 5)
+    MAX_REFERENCE_LATENTS = _constants.get('MAX_REFERENCE_LATENTS', 10)
+    LORA_SOURCE_CHOICES = _constants.get('LORA_SOURCE_CHOICES', ["Civitai", "File"])
+    RESOLUTION_MAP = _constants.get('RESOLUTION_MAP', {})
+except Exception as e:
+    print(f"FATAL: Could not load constants from YAML. Error: {e}")
+    MAX_LORAS, MAX_EMBEDDINGS, MAX_CONDITIONINGS, MAX_CONTROLNETS = 5, 5, 10, 5
+    MAX_REFERENCE_LATENTS = 10
+    LORA_SOURCE_CHOICES = ["Civitai", "File"]
+    RESOLUTION_MAP = {}
+
+
+DEFAULT_NEGATIVE_PROMPT = ""

core/shared_state.py

@@ -0,0 +1 @@
+INVALID_MODEL_URLS = {}

core/workflow_assembler.py

@@ -0,0 +1,179 @@
+import yaml
+import os
+import importlib
+from copy import deepcopy
+from comfy_integration.nodes import NODE_CLASS_MAPPINGS
+
+class WorkflowAssembler:
+    def __init__(self, recipe_path, dynamic_values=None):
+        self.base_path = os.path.dirname(recipe_path)
+        self.node_counter = 0
+        self.workflow = {}
+        self.node_map = {}
+        
+        self._load_injector_config()
+        
+        self.recipe = self._load_and_merge_recipe(os.path.basename(recipe_path), dynamic_values or {})
+
+    def _load_injector_config(self):
+        try:
+            project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+            injectors_path = os.path.join(project_root, 'yaml', 'injectors.yaml')
+            
+            with open(injectors_path, 'r', encoding='utf-8') as f:
+                injector_config = yaml.safe_load(f)
+
+            definitions = injector_config.get("injector_definitions", {})
+            self.injector_order = injector_config.get("injector_order", [])
+            self.global_injectors = {}
+
+            for chain_type, config in definitions.items():
+                module_path = config.get("module")
+                if not module_path:
+                    print(f"Warning: Injector '{chain_type}' in injectors.yaml is missing 'module' path.")
+                    continue
+                try:
+                    module = importlib.import_module(module_path)
+                    if hasattr(module, 'inject'):
+                        self.global_injectors[chain_type] = module.inject
+                        print(f"✅ Successfully registered global injector: {chain_type} from {module_path}")
+                    else:
+                        print(f"⚠️ Warning: Module '{module_path}' for injector '{chain_type}' does not have an 'inject' function.")
+                except ImportError as e:
+                    print(f"❌ Error importing module '{module_path}' for injector '{chain_type}': {e}")
+            
+            if not self.injector_order:
+                 print("⚠️ Warning: 'injector_order' is not defined in injectors.yaml. Using definition order.")
+                 self.injector_order = list(definitions.keys())
+
+        except FileNotFoundError:
+            print(f"❌ FATAL: Could not find injectors.yaml at {injectors_path}. Dynamic chains will not work.")
+            self.injector_order = []
+            self.global_injectors = {}
+        except Exception as e:
+            print(f"❌ FATAL: Could not load or parse injectors.yaml. Dynamic chains will not work. Error: {e}")
+            self.injector_order = []
+            self.global_injectors = {}
+
+    def _get_unique_id(self):
+        self.node_counter += 1
+        return str(self.node_counter)
+
+    def _get_node_template(self, class_type):
+        if class_type not in NODE_CLASS_MAPPINGS:
+            raise ValueError(f"Node class '{class_type}' not found. Ensure it's correctly imported in comfy_integration/nodes.py.")
+        
+        node_class = NODE_CLASS_MAPPINGS[class_type]
+        input_types = node_class.INPUT_TYPES()
+        
+        template = {
+            "inputs": {},
+            "class_type": class_type,
+            "_meta": {"title": node_class.NODE_NAME if hasattr(node_class, 'NODE_NAME') else class_type}
+        }
+        
+        all_inputs = {**input_types.get('required', {}), **input_types.get('optional', {})}
+        for name, details in all_inputs.items():
+            config = details[1] if len(details) > 1 and isinstance(details[1], dict) else {}
+            template["inputs"][name] = config.get("default")
+            
+        return template
+
+    def _load_and_merge_recipe(self, recipe_filename, dynamic_values, search_context_dir=None):
+        search_path = search_context_dir or self.base_path
+        recipe_path_to_use = os.path.join(search_path, recipe_filename)
+
+        if not os.path.exists(recipe_path_to_use):
+            raise FileNotFoundError(f"Recipe file not found: {recipe_path_to_use}")
+
+        with open(recipe_path_to_use, 'r', encoding='utf-8') as f:
+            content = f.read()
+
+        for key, value in dynamic_values.items():
+            if value is not None:
+                content = content.replace(f"{{{{ {key} }}}}", str(value))
+        
+        main_recipe = yaml.safe_load(content)
+        
+        merged_recipe = {'nodes': {}, 'connections': [], 'ui_map': {}}
+        for key in self.injector_order:
+            if key.startswith('dynamic_'):
+                merged_recipe[key] = {}
+
+        parent_recipe_dir = os.path.dirname(recipe_path_to_use)
+        for import_path_template in main_recipe.get('imports', []):
+            import_path = import_path_template
+            for key, value in dynamic_values.items():
+                if value is not None:
+                    import_path = import_path.replace(f"{{{{ {key} }}}}", str(value))
+            
+            try:
+                imported_recipe = self._load_and_merge_recipe(import_path, dynamic_values, search_context_dir=parent_recipe_dir)
+                merged_recipe['nodes'].update(imported_recipe.get('nodes', {}))
+                merged_recipe['connections'].extend(imported_recipe.get('connections', []))
+                merged_recipe['ui_map'].update(imported_recipe.get('ui_map', {}))
+                for key in self.injector_order:
+                    if key in imported_recipe and key.startswith('dynamic_'):
+                        merged_recipe[key].update(imported_recipe.get(key, {}))
+            except FileNotFoundError:
+                print(f"Warning: Optional recipe partial '{import_path}' not found. Skipping.")
+
+        merged_recipe['nodes'].update(main_recipe.get('nodes', {}))
+        merged_recipe['connections'].extend(main_recipe.get('connections', []))
+        merged_recipe['ui_map'].update(main_recipe.get('ui_map', {}))
+        for key in self.injector_order:
+            if key in main_recipe and key.startswith('dynamic_'):
+                merged_recipe[key].update(main_recipe.get(key, {}))
+        
+        return merged_recipe
+
+    def assemble(self, ui_values):
+        for name, details in self.recipe['nodes'].items():
+            class_type = details['class_type']
+            template = self._get_node_template(class_type)
+            node_data = deepcopy(template)
+            
+            unique_id = self._get_unique_id()
+            self.node_map[name] = unique_id
+            
+            if 'params' in details:
+                for param, value in details['params'].items():
+                    if param in node_data['inputs']:
+                        node_data['inputs'][param] = value
+            
+            self.workflow[unique_id] = node_data
+
+        for ui_key, target in self.recipe.get('ui_map', {}).items():
+            if ui_key in ui_values and ui_values[ui_key] is not None:
+                target_list = target if isinstance(target, list) else [target]
+                for t in target_list:
+                    target_name, target_param = t.split(':')
+                    if target_name in self.node_map:
+                        self.workflow[self.node_map[target_name]]['inputs'][target_param] = ui_values[ui_key]
+
+        for conn in self.recipe.get('connections', []):
+            from_name, from_output_idx = conn['from'].split(':')
+            to_name, to_input_name = conn['to'].split(':')
+            
+            from_id = self.node_map.get(from_name)
+            to_id = self.node_map.get(to_name)
+            
+            if from_id and to_id:
+                self.workflow[to_id]['inputs'][to_input_name] = [from_id, int(from_output_idx)]
+        
+        print("--- [Assembler] Applying dynamic injectors ---")
+        recipe_chain_types = {key for key in self.recipe if key.startswith('dynamic_')}
+        processing_order = [key for key in self.injector_order if key in recipe_chain_types]
+
+        for chain_type in processing_order:
+            injector_func = self.global_injectors.get(chain_type)
+            if injector_func:
+                for chain_key, chain_def in self.recipe.get(chain_type, {}).items():
+                    if chain_key in ui_values and ui_values[chain_key]:
+                        print(f"  -> Injecting '{chain_type}' for '{chain_key}'...")
+                        chain_items = ui_values[chain_key]
+                        injector_func(self, chain_def, chain_items)
+        
+        print("--- [Assembler] Finished applying injectors ---")
+        
+        return self.workflow

requirements.txt

@@ -0,0 +1,58 @@
+comfyui-frontend-package==1.37.11
+comfyui-workflow-templates==0.8.24
+comfyui-embedded-docs==0.4.0
+torch==2.9.1
+torchsde
+torchvision
+torchaudio
+numpy>=1.25.0
+einops
+transformers>=4.50.3
+tokenizers>=0.13.3
+sentencepiece
+safetensors>=0.4.2
+aiohttp>=3.11.8
+yarl>=1.18.0
+pyyaml
+Pillow
+scipy
+tqdm
+psutil
+alembic
+SQLAlchemy
+av>=14.2.0
+comfy-kitchen>=0.2.7
+requests
+
+#non essential dependencies:
+kornia>=0.7.1
+spandrel
+pydantic~=2.0
+pydantic-settings~=2.0
+
+
+addict
+albumentations
+filelock
+ftfy
+fvcore
+huggingface-hub
+imageio[ffmpeg]
+importlib_metadata
+matplotlib
+mediapipe
+ninja
+omegaconf
+opencv-python>=4.7.0.72
+python-dateutil
+requests
+scikit-image
+scikit-learn
+soundfile
+spaces
+svglib
+torchsde
+trimesh[easy]
+yacs
+yapf
+onnxruntime-gpu

scripts/build_sage_attention.py

@@ -0,0 +1,99 @@
+import os
+import subprocess
+import sys
+import textwrap
+
+REPO_URL = "https://github.com/thu-ml/SageAttention.git"
+REPO_DIR = "SageAttention"
+
+def run_command(command, cwd=None, env=None):
+    print(f"🚀 Running command: {' '.join(command)}")
+    result = subprocess.run(
+        command,
+        cwd=cwd,
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True
+    )
+    
+    if result.returncode != 0:
+        print(result.stdout)
+        raise subprocess.CalledProcessError(result.returncode, command)
+
+def patch_setup_py(setup_py_path):
+    print(f"--- [SageAttention Build] Applying patches to {setup_py_path} ---")
+
+    with open(setup_py_path, 'r', encoding='utf-8') as f:
+        content = f.read()
+
+    original_cxx_flags = 'CXX_FLAGS = ["-g", "-O3", "-fopenmp", "-lgomp", "-std=c++17", "-DENABLE_BF16"]'
+    modified_cxx_flags = 'CXX_FLAGS = ["-g", "-O3", "-std=c++17", "-DENABLE_BF16"]'
+    
+    if original_cxx_flags in content:
+        content = content.replace(original_cxx_flags, modified_cxx_flags)
+        print("🔧 Patch 1/1: Removed '-fopenmp' and '-lgomp' from CXX_FLAGS.")
+    else:
+        print("⚠️ Patch 1/1: CXX_FLAGS line not found as expected. It might have been changed upstream. Skipping.")
+
+    with open(setup_py_path, 'w', encoding='utf-8') as f:
+        f.write(content)
+    
+    print("✅ Patches applied successfully.")
+
+
+def install_sage_attention():
+    print("--- [SageAttention Build] Checking environment ---")
+    
+    if os.path.isdir(REPO_DIR):
+        print(f"✅ Directory '{REPO_DIR}' already exists, assuming SageAttention is installed. Skipping build.")
+        return
+
+    print(f"⏳ Directory '{REPO_DIR}' not found. Starting a fresh installation of SageAttention.")
+    
+    try:
+        print(f"--- [SageAttention Build] Step 1/3: Cloning repository ---")
+        run_command(["git", "clone", REPO_URL])
+        print("✅ Repository cloned successfully.")
+
+        print(f"--- [SageAttention Build] Step 2/3: Patching setup.py ---")
+        setup_py_path = os.path.join(REPO_DIR, "setup.py")
+        patch_setup_py(setup_py_path)
+
+        print(f"--- [SageAttention Build] Step 3/3: Compiling and installing ---")
+        
+        build_env = os.environ.copy()
+        build_env.update({
+            "TORCH_CUDA_ARCH_LIST": "9.0",
+            "EXT_PARALLEL": "4",
+            "NVCC_APPEND_FLAGS": "--threads 8",
+            "MAX_JOBS": "32"
+        })
+        print("🔧 Setting build environment variables:")
+        print(f"   - TORCH_CUDA_ARCH_LIST='{build_env['TORCH_CUDA_ARCH_LIST']}'")
+        print(f"   - EXT_PARALLEL={build_env['EXT_PARALLEL']}")
+        print(f"   - NVCC_APPEND_FLAGS='{build_env['NVCC_APPEND_FLAGS']}'")
+        print(f"   - MAX_JOBS={build_env['MAX_JOBS']}")
+
+        install_command = [sys.executable, "setup.py", "install"]
+        
+        run_command(install_command, cwd=REPO_DIR, env=build_env)
+
+        print("🎉 SageAttention compiled and installed successfully! ---")
+
+    except FileNotFoundError:
+        print("❌ ERROR: 'git' command not found. Please ensure Git is installed in your environment.")
+        sys.exit(1)
+    except subprocess.CalledProcessError as e:
+        print(f"❌ Command failed with return code: {e.returncode}")
+        print(f"❌ Command: {' '.join(e.cmd)}")
+        print("❌ SageAttention installation failed. Please check the logs above for details.")
+        sys.exit(1)
+    except Exception as e:
+        print(f"❌ An unknown error occurred: {e}")
+        sys.exit(1)
+
+if __name__ == "__main__":
+    if os.path.isdir(REPO_DIR):
+         print(f"Note: To force a rebuild, please delete the '{REPO_DIR}' directory first.")
+    install_sage_attention()

ui/events.py

@@ -0,0 +1,500 @@
+import gradio as gr
+import yaml
+import os
+import shutil
+from functools import lru_cache
+from core.settings import *
+from utils.app_utils import *
+from core.generation_logic import *
+from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
+
+from core.pipelines.controlnet_preprocessor import CPU_ONLY_PREPROCESSORS
+from utils.app_utils import PREPROCESSOR_MODEL_MAP, PREPROCESSOR_PARAMETER_MAP, save_uploaded_file_with_hash
+from ui.shared.ui_components import RESOLUTION_MAP, MAX_CONTROLNETS, MAX_EMBEDDINGS, MAX_CONDITIONINGS, MAX_LORAS, MAX_REFERENCE_LATENTS
+
+
+def on_model_change(model_display_name):
+    """
+    Callback function to update UI elements when the base model changes.
+    It loads default values for steps and cfg from model_defaults.yaml.
+    """
+    defaults = ALL_MODEL_DEFAULTS.get('Default', {}).copy()
+    
+    model_found = False
+    for category, models_in_category in ALL_MODEL_DEFAULTS.items():
+        if category == 'Default' or not isinstance(models_in_category, dict):
+            continue
+            
+        if model_display_name in models_in_category:
+            if '_defaults' in models_in_category:
+                defaults.update(models_in_category['_defaults'])
+            defaults.update(models_in_category[model_display_name])
+            model_found = True
+            break
+    
+    if not model_found:
+        print(f"No specific defaults found for '{model_display_name}'. Using category or global defaults.")
+
+    steps_update = gr.update(value=defaults.get('steps'))
+    cfg_update = gr.update(value=defaults.get('cfg'))
+    
+    return steps_update, cfg_update
+
+def attach_event_handlers(ui_components, demo):
+    def update_cn_input_visibility(choice):
+        return {
+            ui_components["cn_image_input"]: gr.update(visible=choice == "Image"),
+            ui_components["cn_video_input"]: gr.update(visible=choice == "Video")
+        }
+    ui_components["cn_input_type"].change(
+        fn=update_cn_input_visibility,
+        inputs=[ui_components["cn_input_type"]],
+        outputs=[ui_components["cn_image_input"], ui_components["cn_video_input"]]
+    )
+    
+    def update_preprocessor_models_dropdown(preprocessor_name):
+        models = PREPROCESSOR_MODEL_MAP.get(preprocessor_name)
+        if models:
+            model_filenames = [m[1] for m in models]
+            return gr.update(choices=model_filenames, value=model_filenames[0], visible=True)
+        else:
+            return gr.update(choices=[], value=None, visible=False)
+
+    def update_preprocessor_settings_ui(preprocessor_name):
+        from ui.layout import MAX_DYNAMIC_CONTROLS
+        params = PREPROCESSOR_PARAMETER_MAP.get(preprocessor_name, [])
+        
+        slider_updates, dropdown_updates, checkbox_updates = [], [], []
+        
+        s_idx, d_idx, c_idx = 0, 0, 0
+
+        for param in params:
+            if s_idx + d_idx + c_idx >= MAX_DYNAMIC_CONTROLS: break
+            
+            name = param["name"]
+            ptype = param["type"]
+            config = param["config"]
+            label = name.replace('_', ' ').title()
+
+            if ptype == "INT" or ptype == "FLOAT":
+                if s_idx < MAX_DYNAMIC_CONTROLS:
+                    slider_updates.append(gr.update(
+                        label=label,
+                        minimum=config.get('min', 0),
+                        maximum=config.get('max', 255),
+                        step=config.get('step', 0.1 if ptype == "FLOAT" else 1),
+                        value=config.get('default', 0),
+                        visible=True
+                    ))
+                    s_idx += 1
+            elif isinstance(ptype, list):
+                if d_idx < MAX_DYNAMIC_CONTROLS:
+                    dropdown_updates.append(gr.update(
+                        label=label,
+                        choices=ptype,
+                        value=config.get('default', ptype[0] if ptype else None),
+                        visible=True
+                    ))
+                    d_idx += 1
+            elif ptype == "BOOLEAN":
+                if c_idx < MAX_DYNAMIC_CONTROLS:
+                    checkbox_updates.append(gr.update(
+                        label=label,
+                        value=config.get('default', False),
+                        visible=True
+                    ))
+                    c_idx += 1
+
+        for _ in range(s_idx, MAX_DYNAMIC_CONTROLS): slider_updates.append(gr.update(visible=False))
+        for _ in range(d_idx, MAX_DYNAMIC_CONTROLS): dropdown_updates.append(gr.update(visible=False))
+        for _ in range(c_idx, MAX_DYNAMIC_CONTROLS): checkbox_updates.append(gr.update(visible=False))
+
+        return slider_updates + dropdown_updates + checkbox_updates
+
+    def update_run_button_for_cpu(preprocessor_name):
+        if preprocessor_name in CPU_ONLY_PREPROCESSORS:
+            return gr.update(value="Run Preprocessor CPU Only", variant="primary"), gr.update(visible=False)
+        else:
+            return gr.update(value="Run Preprocessor", variant="primary"), gr.update(visible=True)
+
+    ui_components["preprocessor_cn"].change(
+        fn=update_preprocessor_models_dropdown,
+        inputs=[ui_components["preprocessor_cn"]],
+        outputs=[ui_components["preprocessor_model_cn"]]
+    ).then(
+        fn=update_preprocessor_settings_ui,
+        inputs=[ui_components["preprocessor_cn"]],
+        outputs=ui_components["cn_sliders"] + ui_components["cn_dropdowns"] + ui_components["cn_checkboxes"]
+    ).then(
+        fn=update_run_button_for_cpu,
+        inputs=[ui_components["preprocessor_cn"]],
+        outputs=[ui_components["run_cn"], ui_components["zero_gpu_cn"]]
+    )
+
+    all_dynamic_inputs = (
+        ui_components["cn_sliders"] + 
+        ui_components["cn_dropdowns"] + 
+        ui_components["cn_checkboxes"]
+    )
+
+    ui_components["run_cn"].click(
+        fn=run_cn_preprocessor_entry,
+        inputs=[
+            ui_components["cn_input_type"],
+            ui_components["cn_image_input"],
+            ui_components["cn_video_input"],
+            ui_components["preprocessor_cn"],
+            ui_components["preprocessor_model_cn"],
+            ui_components["zero_gpu_cn"],
+        ] + all_dynamic_inputs,
+        outputs=[ui_components["output_gallery_cn"]]
+    )
+
+    def create_lora_event_handlers(prefix):
+        lora_rows = ui_components[f'lora_rows_{prefix}']
+        lora_ids = ui_components[f'lora_ids_{prefix}']
+        lora_scales = ui_components[f'lora_scales_{prefix}']
+        lora_uploads = ui_components[f'lora_uploads_{prefix}']
+        count_state = ui_components[f'lora_count_state_{prefix}']
+        add_button = ui_components[f'add_lora_button_{prefix}']
+        del_button = ui_components[f'delete_lora_button_{prefix}']
+
+        def add_lora_row(c):
+            updates = {}
+            if c < MAX_LORAS:
+                c += 1
+                updates[lora_rows[c - 1]] = gr.update(visible=True)
+            
+            updates[count_state] = c
+            updates[add_button] = gr.update(visible=c < MAX_LORAS)
+            updates[del_button] = gr.update(visible=c > 1)
+            return updates
+
+        def del_lora_row(c):
+            updates = {}
+            if c > 1:
+                updates[lora_rows[c - 1]] = gr.update(visible=False)
+                updates[lora_ids[c - 1]] = ""
+                updates[lora_scales[c - 1]] = 0.0
+                updates[lora_uploads[c - 1]] = None
+                c -= 1
+
+            updates[count_state] = c
+            updates[add_button] = gr.update(visible=True)
+            updates[del_button] = gr.update(visible=c > 1)
+            return updates
+        
+        add_outputs = [count_state, add_button, del_button] + lora_rows
+        del_outputs = [count_state, add_button, del_button] + lora_rows + lora_ids + lora_scales + lora_uploads
+
+        add_button.click(add_lora_row, [count_state], add_outputs, show_progress=False)
+        del_button.click(del_lora_row, [count_state], del_outputs, show_progress=False)
+        
+    def create_embedding_event_handlers(prefix):
+        rows = ui_components[f'embedding_rows_{prefix}']
+        ids = ui_components[f'embeddings_ids_{prefix}']
+        files = ui_components[f'embeddings_files_{prefix}']
+        count_state = ui_components[f'embedding_count_state_{prefix}']
+        add_button = ui_components[f'add_embedding_button_{prefix}']
+        del_button = ui_components[f'delete_embedding_button_{prefix}']
+
+        def add_row(c):
+            c += 1
+            return {
+                count_state: c,
+                rows[c - 1]: gr.update(visible=True),
+                add_button: gr.update(visible=c < MAX_EMBEDDINGS),
+                del_button: gr.update(visible=True)
+            }
+
+        def del_row(c):
+            c -= 1
+            return {
+                count_state: c,
+                rows[c]: gr.update(visible=False),
+                ids[c]: "",
+                files[c]: None,
+                add_button: gr.update(visible=True),
+                del_button: gr.update(visible=c > 0)
+            }
+        
+        add_outputs = [count_state, add_button, del_button] + rows
+        del_outputs = [count_state, add_button, del_button] + rows + ids + files
+        add_button.click(fn=add_row, inputs=[count_state], outputs=add_outputs, show_progress=False)
+        del_button.click(fn=del_row, inputs=[count_state], outputs=del_outputs, show_progress=False)
+    
+    def create_conditioning_event_handlers(prefix):
+        rows = ui_components[f'conditioning_rows_{prefix}']
+        prompts = ui_components[f'conditioning_prompts_{prefix}']
+        count_state = ui_components[f'conditioning_count_state_{prefix}']
+        add_button = ui_components[f'add_conditioning_button_{prefix}']
+        del_button = ui_components[f'delete_conditioning_button_{prefix}']
+        
+        def add_row(c):
+            c += 1
+            return {
+                count_state: c,
+                rows[c - 1]: gr.update(visible=True),
+                add_button: gr.update(visible=c < MAX_CONDITIONINGS),
+                del_button: gr.update(visible=True),
+            }
+
+        def del_row(c):
+            c -= 1
+            return {
+                count_state: c,
+                rows[c]: gr.update(visible=False),
+                prompts[c]: "",
+                add_button: gr.update(visible=True),
+                del_button: gr.update(visible=c > 0),
+            }
+
+        add_outputs = [count_state, add_button, del_button] + rows
+        del_outputs = [count_state, add_button, del_button] + rows + prompts
+        add_button.click(fn=add_row, inputs=[count_state], outputs=add_outputs, show_progress=False)
+        del_button.click(fn=del_row, inputs=[count_state], outputs=del_outputs, show_progress=False)
+
+    def create_reference_latent_event_handlers(prefix):
+        rows = ui_components[f'reference_latent_rows_{prefix}']
+        images = ui_components[f'reference_latent_images_{prefix}']
+        count_state = ui_components[f'reference_latent_count_state_{prefix}']
+        add_button = ui_components[f'add_reference_latent_button_{prefix}']
+        del_button = ui_components[f'delete_reference_latent_button_{prefix}']
+
+        def add_row(c):
+            c += 1
+            return {
+                count_state: c,
+                rows[c - 1]: gr.update(visible=True),
+                add_button: gr.update(visible=c < MAX_REFERENCE_LATENTS),
+                del_button: gr.update(visible=True),
+            }
+
+        def del_row(c):
+            c -= 1
+            return {
+                count_state: c,
+                rows[c]: gr.update(visible=False),
+                images[c]: None,
+                add_button: gr.update(visible=True),
+                del_button: gr.update(visible=c > 0),
+            }
+
+        add_outputs = [count_state, add_button, del_button] + rows
+        del_outputs = [count_state, add_button, del_button] + rows + images
+        add_button.click(fn=add_row, inputs=[count_state], outputs=add_outputs, show_progress=False)
+        del_button.click(fn=del_row, inputs=[count_state], outputs=del_outputs, show_progress=False)
+    
+    def on_vae_upload(file_obj):
+        if not file_obj:
+            return gr.update(), gr.update(), None
+        
+        hashed_filename = save_uploaded_file_with_hash(file_obj, VAE_DIR)
+        return hashed_filename, "File", file_obj
+    
+    def on_lora_upload(file_obj):
+        if not file_obj:
+            return gr.update(), gr.update()
+        
+        hashed_filename = save_uploaded_file_with_hash(file_obj, LORA_DIR)
+        return hashed_filename, "File"
+
+    def on_embedding_upload(file_obj):
+        if not file_obj:
+            return gr.update(), gr.update(), None
+        
+        hashed_filename = save_uploaded_file_with_hash(file_obj, EMBEDDING_DIR)
+        return hashed_filename, "File", file_obj
+
+
+    def create_run_event(prefix: str, task_type: str):
+        run_inputs_map = {
+            'model_display_name': ui_components[f'base_model_{prefix}'],
+            'positive_prompt': ui_components[f'prompt_{prefix}'],
+            'negative_prompt': ui_components[f'neg_prompt_{prefix}'],
+            'seed': ui_components[f'seed_{prefix}'],
+            'batch_size': ui_components[f'batch_size_{prefix}'],
+            'guidance_scale': ui_components[f'cfg_{prefix}'],
+            'num_inference_steps': ui_components[f'steps_{prefix}'],
+            'sampler': ui_components[f'sampler_{prefix}'],
+            'scheduler': ui_components[f'scheduler_{prefix}'],
+            'zero_gpu_duration': ui_components[f'zero_gpu_{prefix}'],
+            'civitai_api_key': ui_components.get(f'civitai_api_key_{prefix}'),
+            'clip_skip': ui_components[f'clip_skip_{prefix}'],
+            'task_type': gr.State(task_type)
+        }
+        
+        if task_type not in ['img2img', 'inpaint']:
+            run_inputs_map.update({'width': ui_components[f'width_{prefix}'], 'height': ui_components[f'height_{prefix}']})
+        
+        task_specific_map = {
+            'img2img': {'img2img_image': f'input_image_{prefix}', 'img2img_denoise': f'denoise_{prefix}'},
+            'inpaint': {'inpaint_image_dict': f'input_image_dict_{prefix}'},
+            'outpaint': {'outpaint_image': f'input_image_{prefix}', 'outpaint_left': f'outpaint_left_{prefix}', 'outpaint_top': f'outpaint_top_{prefix}', 'outpaint_right': f'outpaint_right_{prefix}', 'outpaint_bottom': f'outpaint_bottom_{prefix}'},
+            'hires_fix': {'hires_image': f'input_image_{prefix}', 'hires_upscaler': f'hires_upscaler_{prefix}', 'hires_scale_by': f'hires_scale_by_{prefix}', 'hires_denoise': f'denoise_{prefix}'}
+        }
+        if task_type in task_specific_map:
+            for key, comp_name in task_specific_map[task_type].items():
+                run_inputs_map[key] = ui_components[comp_name]
+        
+        lora_data_components = ui_components.get(f'all_lora_components_flat_{prefix}', [])
+        embedding_data_components = ui_components.get(f'all_embedding_components_flat_{prefix}', [])
+        conditioning_data_components = ui_components.get(f'all_conditioning_components_flat_{prefix}', [])
+        reference_latent_components = ui_components.get(f'all_reference_latent_components_flat_{prefix}', [])
+        
+        run_inputs_map['vae_source'] = ui_components.get(f'vae_source_{prefix}')
+        run_inputs_map['vae_id'] = ui_components.get(f'vae_id_{prefix}')
+        run_inputs_map['vae_file'] = ui_components.get(f'vae_file_{prefix}')
+
+        input_keys = list(run_inputs_map.keys())
+        input_list_flat = [v for v in run_inputs_map.values() if v is not None]
+        input_list_flat += lora_data_components + embedding_data_components + conditioning_data_components + reference_latent_components
+
+        def create_ui_inputs_dict(*args):
+            valid_keys = [k for k in input_keys if run_inputs_map[k] is not None]
+            ui_dict = dict(zip(valid_keys, args[:len(valid_keys)]))
+            arg_idx = len(valid_keys)
+            
+            ui_dict['lora_data'] = list(args[arg_idx : arg_idx + len(lora_data_components)])
+            arg_idx += len(lora_data_components)
+            ui_dict['embedding_data'] = list(args[arg_idx : arg_idx + len(embedding_data_components)])
+            arg_idx += len(embedding_data_components)
+            ui_dict['conditioning_data'] = list(args[arg_idx : arg_idx + len(conditioning_data_components)])
+            arg_idx += len(conditioning_data_components)
+            ui_dict['reference_latent_data'] = list(args[arg_idx : arg_idx + len(reference_latent_components)])
+
+
+            return ui_dict
+
+        ui_components[f'run_{prefix}'].click(
+            fn=lambda *args, progress=gr.Progress(track_tqdm=True): generate_image_wrapper(create_ui_inputs_dict(*args), progress),
+            inputs=input_list_flat,
+            outputs=[ui_components[f'result_{prefix}']]
+        )
+
+
+    for prefix, task_type in [
+        ("txt2img", "txt2img"), ("img2img", "img2img"), ("inpaint", "inpaint"),
+        ("outpaint", "outpaint"), ("hires_fix", "hires_fix"),
+    ]:
+        model_dropdown = ui_components.get(f'base_model_{prefix}')
+        steps_slider = ui_components.get(f'steps_{prefix}')
+        cfg_slider = ui_components.get(f'cfg_{prefix}')
+        if all([model_dropdown, steps_slider, cfg_slider]):
+            model_dropdown.change(
+                fn=on_model_change,
+                inputs=[model_dropdown],
+                outputs=[steps_slider, cfg_slider],
+                show_progress=False
+            )
+
+        if f'add_lora_button_{prefix}' in ui_components: 
+            create_lora_event_handlers(prefix)
+            lora_uploads = ui_components[f'lora_uploads_{prefix}']
+            lora_ids = ui_components[f'lora_ids_{prefix}']
+            lora_sources = ui_components[f'lora_sources_{prefix}']
+            for i in range(MAX_LORAS):
+                lora_uploads[i].upload(
+                    fn=on_lora_upload,
+                    inputs=[lora_uploads[i]],
+                    outputs=[lora_ids[i], lora_sources[i]],
+                    show_progress=False
+                )
+
+        if f'add_embedding_button_{prefix}' in ui_components: 
+            create_embedding_event_handlers(prefix)
+            if f'embeddings_uploads_{prefix}' in ui_components:
+                emb_uploads = ui_components[f'embeddings_uploads_{prefix}']
+                emb_ids = ui_components[f'embeddings_ids_{prefix}']
+                emb_sources = ui_components[f'embeddings_sources_{prefix}']
+                emb_files = ui_components[f'embeddings_files_{prefix}']
+                for i in range(MAX_EMBEDDINGS):
+                    emb_uploads[i].upload(
+                        fn=on_embedding_upload,
+                        inputs=[emb_uploads[i]],
+                        outputs=[emb_ids[i], emb_sources[i], emb_files[i]],
+                        show_progress=False
+                    )
+        if f'add_conditioning_button_{prefix}' in ui_components: create_conditioning_event_handlers(prefix)
+        if f'add_reference_latent_button_{prefix}' in ui_components: create_reference_latent_event_handlers(prefix)
+        if f'vae_source_{prefix}' in ui_components:
+            upload_button = ui_components.get(f'vae_upload_button_{prefix}')
+            if upload_button:
+                 upload_button.upload(
+                    fn=on_vae_upload, 
+                    inputs=[upload_button], 
+                    outputs=[
+                        ui_components[f'vae_id_{prefix}'], 
+                        ui_components[f'vae_source_{prefix}'], 
+                        ui_components[f'vae_file_{prefix}']
+                    ]
+                )
+
+        create_run_event(prefix, task_type)
+
+    def on_aspect_ratio_change(ratio_key, model_display_name):
+        model_type = MODEL_TYPE_MAP.get(model_display_name, 'sdxl').lower()
+        res_map = RESOLUTION_MAP.get(model_type, RESOLUTION_MAP.get("sdxl", {}))
+        w, h = res_map.get(ratio_key, (1024, 1024))
+        return w, h
+
+    for prefix in ["txt2img", "img2img", "inpaint", "outpaint", "hires_fix"]:
+        if f'aspect_ratio_{prefix}' in ui_components:
+            aspect_ratio_dropdown = ui_components[f'aspect_ratio_{prefix}']
+            width_component = ui_components[f'width_{prefix}']
+            height_component = ui_components[f'height_{prefix}']
+            model_dropdown = ui_components[f'base_model_{prefix}']
+            aspect_ratio_dropdown.change(fn=on_aspect_ratio_change, inputs=[aspect_ratio_dropdown, model_dropdown], outputs=[width_component, height_component], show_progress=False)
+            
+    if 'view_mode_inpaint' in ui_components:
+        def toggle_inpaint_fullscreen_view(view_mode):
+            is_fullscreen = (view_mode == "Fullscreen View")
+            other_elements_visible = not is_fullscreen
+            editor_height = 800 if is_fullscreen else 272
+            return {
+                ui_components['model_and_run_row_inpaint']: gr.update(visible=other_elements_visible),
+                ui_components['prompts_column_inpaint']: gr.update(visible=other_elements_visible),
+                ui_components['params_and_gallery_row_inpaint']: gr.update(visible=other_elements_visible),
+                ui_components['accordion_wrapper_inpaint']: gr.update(visible=other_elements_visible),
+                ui_components['input_image_dict_inpaint']: gr.update(height=editor_height),
+            }
+        
+        output_components = [
+            ui_components['model_and_run_row_inpaint'], ui_components['prompts_column_inpaint'],
+            ui_components['params_and_gallery_row_inpaint'], ui_components['accordion_wrapper_inpaint'],
+            ui_components['input_image_dict_inpaint']
+        ]
+        ui_components['view_mode_inpaint'].change(fn=toggle_inpaint_fullscreen_view, inputs=[ui_components['view_mode_inpaint']], outputs=output_components, show_progress=False)
+
+    def run_on_load():
+        all_updates = {}
+        
+        default_preprocessor = "Canny Edge" 
+        model_update = update_preprocessor_models_dropdown(default_preprocessor)
+        all_updates[ui_components["preprocessor_model_cn"]] = model_update
+        
+        settings_outputs = update_preprocessor_settings_ui(default_preprocessor)
+        dynamic_outputs = ui_components["cn_sliders"] + ui_components["cn_dropdowns"] + ui_components["cn_checkboxes"]
+        for i, comp in enumerate(dynamic_outputs):
+            all_updates[comp] = settings_outputs[i]
+
+        run_button_update, zero_gpu_update = update_run_button_for_cpu(default_preprocessor)
+        all_updates[ui_components["run_cn"]] = run_button_update
+        all_updates[ui_components["zero_gpu_cn"]] = zero_gpu_update
+        
+        return all_updates
+    
+    all_load_outputs = [
+        ui_components["preprocessor_model_cn"],
+        *ui_components["cn_sliders"],
+        *ui_components["cn_dropdowns"],
+        *ui_components["cn_checkboxes"],
+        ui_components["run_cn"],
+        ui_components["zero_gpu_cn"]
+    ]
+
+    if all_load_outputs:
+        demo.load(
+            fn=run_on_load,
+            outputs=all_load_outputs
+        )

ui/layout.py

@@ -0,0 +1,92 @@
+import os
+import gradio as gr
+from core.settings import *
+
+from .shared import txt2img_ui, img2img_ui, inpaint_ui, outpaint_ui, hires_fix_ui
+
+MAX_DYNAMIC_CONTROLS = 10
+
+def get_preprocessor_choices():
+    from nodes import NODE_DISPLAY_NAME_MAPPINGS
+    
+    preprocessor_names = [
+        display_name for class_name, display_name in NODE_DISPLAY_NAME_MAPPINGS.items()
+        if "Preprocessor" in class_name or "Segmentor" in class_name or 
+           "Estimator" in class_name or "Detector" in class_name
+    ]
+    return sorted(list(set(preprocessor_names)))
+
+
+def build_ui(event_handler_function):
+    ui_components = {}
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# ImageGen - FLUX.2")
+        gr.Markdown(
+            "This demo is a streamlined version of the [Comfy web UI](https://github.com/RioShiina47/comfy-webui)'s ImgGen functionality. "
+            "Other versions are also available: "
+            "[Z-Image](https://huggingface.co/spaces/RioShiina/ImageGen-Z-Image), "
+            "[Qwen-Image](https://huggingface.co/spaces/RioShiina/ImageGen-Qwen-Image), "
+            "[NewBie-Image](https://huggingface.co/spaces/RioShiina/ImageGen-NewBie-Image), "
+            "[Illstrious](https://huggingface.co/spaces/RioShiina/ImageGen-Illstrious), "
+            "[NoobAI](https://huggingface.co/spaces/RioShiina/ImageGen-NoobAI), "
+            "[Pony](https://huggingface.co/spaces/RioShiina/ImageGen-Pony1), "
+            "[SDXL](https://huggingface.co/spaces/RioShiina/ImageGen-SDXL), "
+            "[SD1.5](https://huggingface.co/spaces/RioShiina/ImageGen-SD15)"
+        )
+        with gr.Tabs(elem_id="tabs_container") as tabs:
+            with gr.TabItem("FLUX.2", id=0):
+                with gr.Tabs(elem_id="image_gen_tabs") as image_gen_tabs:
+                    with gr.TabItem("Txt2Img", id=0):
+                        ui_components.update(txt2img_ui.create_ui())
+                    
+                    with gr.TabItem("Img2Img", id=1):
+                        ui_components.update(img2img_ui.create_ui())
+
+                    with gr.TabItem("Inpaint", id=2):
+                        ui_components.update(inpaint_ui.create_ui())
+
+                    with gr.TabItem("Outpaint", id=3):
+                        ui_components.update(outpaint_ui.create_ui())
+
+                    with gr.TabItem("Hires. Fix", id=4):
+                        ui_components.update(hires_fix_ui.create_ui())
+                
+                ui_components['image_gen_tabs'] = image_gen_tabs
+
+            with gr.TabItem("Controlnet Preprocessors", id=1):
+                gr.Markdown("## ControlNet Auxiliary Preprocessors")
+                gr.Markdown("Powered by [Fannovel16/comfyui_controlnet_aux](https://github.com/Fannovel16/comfyui_controlnet_aux).")
+                gr.Markdown("Upload an image or video to process it with a ControlNet preprocessor.")
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        cn_input_type = gr.Radio(["Image", "Video"], label="Input Type", value="Image")
+                        cn_image_input = gr.Image(type="pil", label="Input Image", visible=True, height=384)
+                        cn_video_input = gr.Video(label="Input Video", visible=False)
+                        preprocessor_cn = gr.Dropdown(label="Preprocessor", choices=get_preprocessor_choices(), value="Canny Edge")
+                        preprocessor_model_cn = gr.Dropdown(label="Preprocessor Model", choices=[], value=None, visible=False)
+                        with gr.Column() as preprocessor_settings_ui:
+                            cn_sliders, cn_dropdowns, cn_checkboxes = [], [], []
+                            for i in range(MAX_DYNAMIC_CONTROLS):
+                                cn_sliders.append(gr.Slider(visible=False, label=f"dyn_slider_{i}"))
+                                cn_dropdowns.append(gr.Dropdown(visible=False, label=f"dyn_dropdown_{i}"))
+                                cn_checkboxes.append(gr.Checkbox(visible=False, label=f"dyn_checkbox_{i}"))
+                        run_cn = gr.Button("Run Preprocessor", variant="primary")
+                    with gr.Column(scale=1):
+                        output_gallery_cn = gr.Gallery(label="Output", show_label=False, object_fit="contain", height=512)
+                        zero_gpu_cn = gr.Number(label="ZeroGPU Duration (s)", value=None, placeholder="Default: 60, Max: 120", info="Optional")
+                ui_components.update({
+                    "cn_input_type": cn_input_type, "cn_image_input": cn_image_input, "cn_video_input": cn_video_input,
+                    "preprocessor_cn": preprocessor_cn, "preprocessor_model_cn": preprocessor_model_cn, "run_cn": run_cn,
+                    "zero_gpu_cn": zero_gpu_cn, "output_gallery_cn": output_gallery_cn,
+                    "preprocessor_settings_ui": preprocessor_settings_ui, "cn_sliders": cn_sliders,
+                    "cn_dropdowns": cn_dropdowns, "cn_checkboxes": cn_checkboxes
+                })
+            
+            ui_components["tabs"] = tabs
+        
+        gr.Markdown("<div style='text-align: center; margin-top: 20px;'>Made by RioShiina with ❤️<br><a href='https://github.com/RioShiina47' target='_blank'>GitHub</a> | <a href='https://huggingface.co/RioShiina' target='_blank'>Hugging Face</a> | <a href='https://civitai.com/user/RioShiina' target='_blank'>Civitai</a></div>")
+        
+        event_handler_function(ui_components, demo)
+
+    return demo

ui/shared/hires_fix_ui.py

@@ -0,0 +1,76 @@
+import gradio as gr
+from core.settings import MODEL_MAP_CHECKPOINT
+from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
+from .ui_components import (
+    create_lora_settings_ui,
+    create_embedding_ui,
+    create_conditioning_ui, create_vae_override_ui, create_api_key_ui,
+    create_reference_latent_ui
+)
+
+def create_ui():
+    prefix = "hires_fix"
+    components = {}
+    
+    with gr.Column():
+        with gr.Row():
+            components[f'base_model_{prefix}'] = gr.Dropdown(
+                label="Base Model", 
+                choices=list(MODEL_MAP_CHECKPOINT.keys()), 
+                value=list(MODEL_MAP_CHECKPOINT.keys())[0], 
+                scale=3
+            )
+            with gr.Column(scale=1): 
+                components[f'run_{prefix}'] = gr.Button("Run Hires. Fix", variant="primary")
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                components[f'input_image_{prefix}'] = gr.Image(type="pil", label="Input Image", height=255)
+            with gr.Column(scale=2):
+                components[f'prompt_{prefix}'] = gr.Text(label="Prompt", lines=3, placeholder="Describe the final image...")
+                components[f'neg_prompt_{prefix}'] = gr.Text(label="Negative prompt", lines=3, value="")
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                with gr.Row():
+                    components[f'hires_upscaler_{prefix}'] = gr.Dropdown(
+                        label="Upscaler", 
+                        choices=["nearest-exact", "bilinear", "area", "bicubic", "bislerp"], 
+                        value="nearest-exact"
+                    )
+                    components[f'hires_scale_by_{prefix}'] = gr.Slider(
+                        label="Upscale by", minimum=1.0, maximum=4.0, step=0.1, value=1.5
+                    )
+                
+                with gr.Row():
+                    components[f'denoise_{prefix}'] = gr.Slider(label="Denoise Strength", minimum=0.0, maximum=1.0, step=0.01, value=0.55)
+                
+                with gr.Row():
+                    components[f'sampler_{prefix}'] = gr.Dropdown(label="Sampler", choices=SAMPLER_CHOICES, value="euler")
+                    components[f'scheduler_{prefix}'] = gr.Dropdown(label="Scheduler", choices=SCHEDULER_CHOICES, value="simple")
+                with gr.Row():
+                    components[f'steps_{prefix}'] = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=4)
+                    components[f'cfg_{prefix}'] = gr.Slider(label="CFG Scale", minimum=1.0, maximum=20.0, step=0.1, value=1.0)
+                with gr.Row():
+                    components[f'seed_{prefix}'] = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
+                    components[f'batch_size_{prefix}'] = gr.Slider(label="Batch Size", minimum=1, maximum=16, step=1, value=1)
+                with gr.Row():
+                    components[f'zero_gpu_{prefix}'] = gr.Number(label="ZeroGPU Duration (s)", value=None, placeholder="Default: 60, Max: 120", info="Optional: Set how long to reserve the GPU.")
+
+                components[f'clip_skip_{prefix}'] = gr.State(value=1)
+                components[f'width_{prefix}'] = gr.State(value=512)
+                components[f'height_{prefix}'] = gr.State(value=512)
+
+            with gr.Column(scale=1):
+                components[f'result_{prefix}'] = gr.Gallery(label="Result", show_label=False, columns=1, object_fit="contain", height=610)
+
+        components.update(create_api_key_ui(prefix))
+        components.update(create_lora_settings_ui(prefix))
+        # components.update(create_diffsynth_controlnet_ui(prefix))
+        # components.update(create_controlnet_ui(prefix))
+        # components.update(create_embedding_ui(prefix))
+        components.update(create_reference_latent_ui(prefix))
+        components.update(create_conditioning_ui(prefix))
+        # components.update(create_vae_override_ui(prefix))
+            
+    return components

ui/shared/img2img_ui.py

@@ -0,0 +1,59 @@
+import gradio as gr
+from core.settings import MODEL_MAP_CHECKPOINT
+from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
+from .ui_components import (
+    create_lora_settings_ui,
+    create_embedding_ui,
+    create_conditioning_ui, create_vae_override_ui, create_api_key_ui,
+    create_reference_latent_ui
+)
+
+def create_ui():
+    prefix = "img2img"
+    components = {}
+    
+    with gr.Column():
+        with gr.Row():
+            components[f'base_model_{prefix}'] = gr.Dropdown(label="Base Model", choices=list(MODEL_MAP_CHECKPOINT.keys()), value=list(MODEL_MAP_CHECKPOINT.keys())[0], scale=3)
+            with gr.Column(scale=1): 
+                components[f'run_{prefix}'] = gr.Button("Run", variant="primary")
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                components[f'input_image_{prefix}'] = gr.Image(type="pil", label="Input Image", height=255)
+
+            with gr.Column(scale=2):
+                components[f'prompt_{prefix}'] = gr.Text(label="Prompt", lines=3, placeholder="Enter your prompt")
+                components[f'neg_prompt_{prefix}'] = gr.Text(label="Negative prompt", lines=3, value="")
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                components[f'denoise_{prefix}'] = gr.Slider(label="Denoise Strength", minimum=0.0, maximum=1.0, step=0.01, value=0.7)
+                
+                with gr.Row():
+                    components[f'sampler_{prefix}'] = gr.Dropdown(label="Sampler", choices=SAMPLER_CHOICES, value="euler")
+                    components[f'scheduler_{prefix}'] = gr.Dropdown(label="Scheduler", choices=SCHEDULER_CHOICES, value="simple")
+                with gr.Row():
+                    components[f'steps_{prefix}'] = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=4)
+                    components[f'cfg_{prefix}'] = gr.Slider(label="CFG Scale", minimum=1.0, maximum=20.0, step=0.1, value=1.0)
+                with gr.Row():
+                    components[f'seed_{prefix}'] = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
+                    components[f'batch_size_{prefix}'] = gr.Slider(label="Batch Size", minimum=1, maximum=16, step=1, value=1)
+                with gr.Row():
+                    components[f'zero_gpu_{prefix}'] = gr.Number(label="ZeroGPU Duration (s)", value=None, placeholder="Default: 60, Max: 120", info="Optional: Set how long to reserve the GPU. Longer jobs may need more time.")
+                
+                components[f'clip_skip_{prefix}'] = gr.State(value=1)
+                
+            with gr.Column(scale=1):
+                components[f'result_{prefix}'] = gr.Gallery(label="Result", show_label=False, columns=1, object_fit="contain", height=505)
+        
+        components.update(create_api_key_ui(prefix))
+        components.update(create_lora_settings_ui(prefix))
+        # components.update(create_diffsynth_controlnet_ui(prefix))
+        # components.update(create_controlnet_ui(prefix))
+        # components.update(create_embedding_ui(prefix))
+        components.update(create_reference_latent_ui(prefix))
+        components.update(create_conditioning_ui(prefix))
+        # components.update(create_vae_override_ui(prefix))
+            
+    return components

ui/shared/inpaint_ui.py

@@ -0,0 +1,83 @@
+import gradio as gr
+from core.settings import MODEL_MAP_CHECKPOINT
+from .ui_components import (
+    create_base_parameter_ui, create_lora_settings_ui,
+    create_embedding_ui,
+    create_conditioning_ui, create_vae_override_ui, create_api_key_ui,
+    create_reference_latent_ui
+)
+
+def create_ui():
+    prefix = "inpaint"
+    components = {}
+    
+    with gr.Column():
+        with gr.Row() as model_and_run_row:
+            components[f'base_model_{prefix}'] = gr.Dropdown(
+                label="Base Model", 
+                choices=list(MODEL_MAP_CHECKPOINT.keys()), 
+                value=list(MODEL_MAP_CHECKPOINT.keys())[0], 
+                scale=3
+            )
+            with gr.Column(scale=1): 
+                components[f'run_{prefix}'] = gr.Button("Run Inpaint", variant="primary")
+        
+        components[f'model_and_run_row_{prefix}'] = model_and_run_row
+
+        with gr.Row() as main_content_row:
+            with gr.Column(scale=1) as editor_column:
+                components[f'view_mode_{prefix}'] = gr.Radio(
+                    ["Normal View", "Fullscreen View"], 
+                    label="Editor View", 
+                    value="Normal View", 
+                    interactive=True
+                )
+                components[f'input_image_dict_{prefix}'] = gr.ImageEditor(
+                    type="pil", 
+                    label="Image & Mask", 
+                    height=272
+                )
+            components[f'editor_column_{prefix}'] = editor_column
+
+            with gr.Column(scale=2) as prompts_column:
+                components[f'prompt_{prefix}'] = gr.Text(label="Prompt", lines=6, placeholder="Describe what to fill in the mask...")
+                components[f'neg_prompt_{prefix}'] = gr.Text(label="Negative prompt", lines=6, value="")
+            components[f'prompts_column_{prefix}'] = prompts_column
+
+        with gr.Row() as params_and_gallery_row:
+            with gr.Column(scale=1):
+                param_defaults = {'w': 1024, 'h': 1024, 'cs_vis': False, 'cs_val': 1}
+                from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
+                with gr.Row():
+                    components[f'sampler_{prefix}'] = gr.Dropdown(label="Sampler", choices=SAMPLER_CHOICES, value="euler")
+                    components[f'scheduler_{prefix}'] = gr.Dropdown(label="Scheduler", choices=SCHEDULER_CHOICES, value="simple")
+                with gr.Row():
+                    components[f'steps_{prefix}'] = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=4)
+                    components[f'cfg_{prefix}'] = gr.Slider(label="CFG Scale", minimum=1.0, maximum=20.0, step=0.1, value=1.0)
+                with gr.Row():
+                    components[f'seed_{prefix}'] = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
+                    components[f'batch_size_{prefix}'] = gr.Slider(label="Batch Size", minimum=1, maximum=16, step=1, value=1)
+                with gr.Row():
+                    components[f'zero_gpu_{prefix}'] = gr.Number(label="ZeroGPU Duration (s)", value=None, placeholder="Default: 60, Max: 120", info="Optional: Set how long to reserve the GPU.")
+
+                components[f'clip_skip_{prefix}'] = gr.State(value=1)
+                components[f'width_{prefix}'] = gr.State(value=512)
+                components[f'height_{prefix}'] = gr.State(value=512)
+
+            with gr.Column(scale=1):
+                components[f'result_{prefix}'] = gr.Gallery(label="Result", show_label=False, columns=1, object_fit="contain", height=414)
+        
+        components[f'params_and_gallery_row_{prefix}'] = params_and_gallery_row
+            
+        with gr.Column() as accordion_wrapper:
+            components.update(create_api_key_ui(prefix))
+            components.update(create_lora_settings_ui(prefix))
+            # components.update(create_diffsynth_controlnet_ui(prefix))
+            # components.update(create_controlnet_ui(prefix))
+            # components.update(create_embedding_ui(prefix))
+            components.update(create_reference_latent_ui(prefix))
+            components.update(create_conditioning_ui(prefix))
+            # components.update(create_vae_override_ui(prefix))
+        components[f'accordion_wrapper_{prefix}'] = accordion_wrapper
+            
+    return components

ui/shared/outpaint_ui.py

@@ -0,0 +1,70 @@
+import gradio as gr
+from core.settings import MODEL_MAP_CHECKPOINT
+from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
+from .ui_components import (
+    create_lora_settings_ui,
+    create_embedding_ui,
+    create_conditioning_ui, create_vae_override_ui, create_api_key_ui,
+    create_reference_latent_ui
+)
+
+def create_ui():
+    prefix = "outpaint"
+    components = {}
+    
+    with gr.Column():
+        with gr.Row():
+            components[f'base_model_{prefix}'] = gr.Dropdown(
+                label="Base Model", 
+                choices=list(MODEL_MAP_CHECKPOINT.keys()), 
+                value=list(MODEL_MAP_CHECKPOINT.keys())[0], 
+                scale=3
+            )
+            with gr.Column(scale=1): 
+                components[f'run_{prefix}'] = gr.Button("Run Outpaint", variant="primary")
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                components[f'input_image_{prefix}'] = gr.Image(type="pil", label="Input Image", height=255)
+            with gr.Column(scale=2):
+                components[f'prompt_{prefix}'] = gr.Text(label="Prompt", lines=3, placeholder="Describe the content for the expanded areas...")
+                components[f'neg_prompt_{prefix}'] = gr.Text(label="Negative prompt", lines=3, value="")
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                with gr.Row():
+                    components[f'outpaint_left_{prefix}'] = gr.Slider(label="Pad Left", minimum=0, maximum=512, step=64, value=0)
+                    components[f'outpaint_right_{prefix}'] = gr.Slider(label="Pad Right", minimum=0, maximum=512, step=64, value=256)
+                with gr.Row():
+                    components[f'outpaint_top_{prefix}'] = gr.Slider(label="Pad Top", minimum=0, maximum=512, step=64, value=0)
+                    components[f'outpaint_bottom_{prefix}'] = gr.Slider(label="Pad Bottom", minimum=0, maximum=512, step=64, value=0)
+                
+                with gr.Row():
+                    components[f'sampler_{prefix}'] = gr.Dropdown(label="Sampler", choices=SAMPLER_CHOICES, value="euler")
+                    components[f'scheduler_{prefix}'] = gr.Dropdown(label="Scheduler", choices=SCHEDULER_CHOICES, value="simple")
+                with gr.Row():
+                    components[f'steps_{prefix}'] = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=4)
+                    components[f'cfg_{prefix}'] = gr.Slider(label="CFG Scale", minimum=1.0, maximum=20.0, step=0.1, value=1.0)
+                with gr.Row():
+                    components[f'seed_{prefix}'] = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
+                    components[f'batch_size_{prefix}'] = gr.Slider(label="Batch Size", minimum=1, maximum=16, step=1, value=1)
+                with gr.Row():
+                    components[f'zero_gpu_{prefix}'] = gr.Number(label="ZeroGPU Duration (s)", value=None, placeholder="Default: 60, Max: 120", info="Optional: Set how long to reserve the GPU.")
+
+                components[f'clip_skip_{prefix}'] = gr.State(value=1)
+                components[f'width_{prefix}'] = gr.State(value=512)
+                components[f'height_{prefix}'] = gr.State(value=512)
+
+            with gr.Column(scale=1):
+                components[f'result_{prefix}'] = gr.Gallery(label="Result", show_label=False, columns=1, object_fit="contain", height=595)
+        
+        components.update(create_api_key_ui(prefix))
+        components.update(create_lora_settings_ui(prefix))
+        # components.update(create_diffsynth_controlnet_ui(prefix))
+        # components.update(create_controlnet_ui(prefix))
+        # components.update(create_embedding_ui(prefix))
+        components.update(create_reference_latent_ui(prefix))
+        components.update(create_conditioning_ui(prefix))
+        # components.update(create_vae_override_ui(prefix))
+            
+    return components

ui/shared/txt2img_ui.py

@@ -0,0 +1,40 @@
+import gradio as gr
+from core.settings import MODEL_MAP_CHECKPOINT
+from .ui_components import (
+    create_base_parameter_ui, create_lora_settings_ui,
+    create_embedding_ui,
+    create_conditioning_ui, create_vae_override_ui, create_api_key_ui,
+    create_reference_latent_ui
+)
+
+def create_ui():
+    """Creates the UI components for the Txt2Img tab."""
+    prefix = "txt2img"
+    components = {}
+    
+    with gr.Column():
+        with gr.Row():
+            components[f'base_model_{prefix}'] = gr.Dropdown(label="Base Model", choices=list(MODEL_MAP_CHECKPOINT.keys()), value=list(MODEL_MAP_CHECKPOINT.keys())[0], scale=3)
+            with gr.Column(scale=1): 
+                components[f'run_{prefix}'] = gr.Button("Run", variant="primary")
+        
+        components[f'prompt_{prefix}'] = gr.Text(label="Prompt", lines=3, placeholder="Enter your prompt")
+        components[f'neg_prompt_{prefix}'] = gr.Text(label="Negative prompt", lines=3, value="")
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                param_defaults = {'w': 1024, 'h': 1024, 'cs_vis': False, 'cs_val': 1}
+                components.update(create_base_parameter_ui(prefix, param_defaults))
+            with gr.Column(scale=1):
+                components[f'result_{prefix}'] = gr.Gallery(label="Result", show_label=False, columns=2, object_fit="contain", height=627)
+        
+        components.update(create_api_key_ui(prefix))
+        components.update(create_lora_settings_ui(prefix))
+        # components.update(create_diffsynth_controlnet_ui(prefix))
+        # components.update(create_controlnet_ui(prefix))
+        # components.update(create_embedding_ui(prefix))
+        components.update(create_reference_latent_ui(prefix))
+        components.update(create_conditioning_ui(prefix))
+        # components.update(create_vae_override_ui(prefix))
+        
+    return components

ui/shared/ui_components.py

@@ -0,0 +1,236 @@
+import gradio as gr
+from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
+from core.settings import (
+    MAX_LORAS, LORA_SOURCE_CHOICES, MAX_EMBEDDINGS, MAX_CONDITIONINGS,
+    MAX_CONTROLNETS, RESOLUTION_MAP, MAX_REFERENCE_LATENTS
+)
+import yaml
+import os
+from functools import lru_cache
+
+def create_base_parameter_ui(prefix, defaults=None):
+    if defaults is None:
+        defaults = {}
+    
+    components = {}
+    
+    with gr.Row():
+        components[f'aspect_ratio_{prefix}'] = gr.Dropdown(
+            label="Aspect Ratio",
+            choices=list(RESOLUTION_MAP['sdxl'].keys()),
+            value="1:1 (Square)",
+            interactive=True
+        )
+    with gr.Row():
+        components[f'width_{prefix}'] = gr.Number(label="Width", value=defaults.get('w', 1024), interactive=True)
+        components[f'height_{prefix}'] = gr.Number(label="Height", value=defaults.get('h', 1024), interactive=True)
+    with gr.Row():
+        components[f'sampler_{prefix}'] = gr.Dropdown(label="Sampler", choices=SAMPLER_CHOICES, value="euler")
+        components[f'scheduler_{prefix}'] = gr.Dropdown(label="Scheduler", choices=SCHEDULER_CHOICES, value="simple")
+    with gr.Row():
+        components[f'steps_{prefix}'] = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=4)
+        components[f'cfg_{prefix}'] = gr.Slider(label="CFG Scale", minimum=1.0, maximum=20.0, step=0.1, value=1.0)
+    with gr.Row():
+        components[f'seed_{prefix}'] = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
+        components[f'batch_size_{prefix}'] = gr.Slider(label="Batch Size", minimum=1, maximum=16, step=1, value=1)
+    with gr.Row():
+        components[f'zero_gpu_{prefix}'] = gr.Number(label="ZeroGPU Duration (s)", value=None, placeholder="Default: 60, Max: 120", info="Optional: Set how long to reserve the GPU. Longer jobs may need more time.")
+    
+    components[f'clip_skip_{prefix}'] = gr.State(value=1)
+            
+    return components
+
+
+def create_api_key_ui(prefix: str):
+    components = {}
+    with gr.Accordion("API Key Settings", open=False) as api_key_accordion:
+        components[f'api_key_accordion_{prefix}'] = api_key_accordion
+        gr.Markdown("💡 **Tip:** Enter API key (optional). An API key is required for resources that need a login to download. The key will be used for all Civitai downloads on this tab. You can also manually upload the corresponding files to avoid API Key leakage caused by potential vulnerabilities.")
+        with gr.Row():
+             components[f'civitai_api_key_{prefix}'] = gr.Textbox(
+                 label="Civitai API Key", 
+                 type="password", 
+                 placeholder="Enter your Civitai API key here (optional)"
+            )
+    return components
+
+
+def create_lora_settings_ui(prefix: str):
+    components = {}
+    
+    lora_rows, lora_sources, lora_ids, lora_scales, lora_uploads = [], [], [], [], []
+
+    with gr.Accordion("LoRA Settings", open=False) as lora_accordion:
+        components[f'lora_accordion_{prefix}'] = lora_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.")
+        components[f'lora_count_state_{prefix}'] = gr.State(1)
+
+        for i in range(MAX_LORAS):
+            with gr.Row(visible=i==0) as row:
+                source = gr.Dropdown(label=f"LoRA Source {i+1}", choices=LORA_SOURCE_CHOICES, value=LORA_SOURCE_CHOICES[0], scale=1)
+                lora_id = gr.Textbox(label=f"Civitai Version ID / File", placeholder="Civitai Version ID or Filename", scale=2, type="text")
+                scale = gr.Slider(label=f"Scale", minimum=-2.0, maximum=2.0, step=0.05, value=0.8, scale=1)
+                upload = gr.UploadButton(label="Upload", file_types=[".safetensors"], scale=1)
+            
+            lora_rows.append(row)
+            lora_sources.append(source)
+            lora_ids.append(lora_id)
+            lora_scales.append(scale)
+            lora_uploads.append(upload)
+            
+        with gr.Row():
+            components[f'add_lora_button_{prefix}'] = gr.Button("Add LoRA", variant="secondary")
+            components[f'delete_lora_button_{prefix}'] = gr.Button("Remove LoRA", variant="secondary", visible=False)
+
+    components[f'lora_rows_{prefix}'] = lora_rows
+    components[f'lora_sources_{prefix}'] = lora_sources
+    components[f'lora_ids_{prefix}'] = lora_ids
+    components[f'lora_scales_{prefix}'] = lora_scales
+    components[f'lora_uploads_{prefix}'] = lora_uploads
+    
+    all_lora_components_flat = []
+    for i in range(MAX_LORAS):
+        all_lora_components_flat.extend([lora_sources[i], lora_ids[i], lora_scales[i], lora_uploads[i]])
+    components[f'all_lora_components_flat_{prefix}'] = all_lora_components_flat
+    
+    return components
+
+def create_embedding_ui(prefix: str):
+    components = {}
+    key = lambda name: f"{name}_{prefix}"
+
+    with gr.Accordion("Embedding Settings", open=False, visible=True) as accordion:
+        components[key('embedding_accordion')] = accordion
+        gr.Markdown("💡 **Tip:** Embeddings are automatically added to your prompt using `embedding:filename` syntax. 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. For instance, using the Version ID `456` from the example above would automatically append `embedding:civitai_456` to your positive prompt.")
+        
+        embedding_rows, sources, ids, files, upload_buttons = [], [], [], [], []
+        components.update({
+            key('embedding_rows'): embedding_rows, 
+            key('embeddings_sources'): sources, 
+            key('embeddings_ids'): ids, 
+            key('embeddings_files'): files,
+            key('embeddings_uploads'): upload_buttons
+        })
+        
+        for i in range(MAX_EMBEDDINGS):
+            with gr.Row(visible=(i < 1)) as row:
+                sources.append(gr.Dropdown(label=f"Embedding Source {i+1}", choices=LORA_SOURCE_CHOICES, value="Civitai", scale=1, interactive=True))
+                ids.append(gr.Textbox(label="Civitai Version ID / File", placeholder="Civitai Version ID or Filename", scale=3, interactive=True, type="text"))
+                upload_btn = gr.UploadButton("Upload", file_types=[".safetensors"], scale=1)
+                files.append(gr.State(None))
+                upload_buttons.append(upload_btn)
+                embedding_rows.append(row)
+
+        with gr.Row():
+            components[key('add_embedding_button')] = gr.Button("✚ Add Embedding")
+            components[key('delete_embedding_button')] = gr.Button("➖ Delete Embedding", visible=False)
+        components[key('embedding_count_state')] = gr.State(1)
+
+        all_embedding_components_flat = []
+        for i in range(MAX_EMBEDDINGS):
+            all_embedding_components_flat.extend([sources[i], ids[i], files[i]])
+        components[key('all_embedding_components_flat')] = all_embedding_components_flat
+
+    return components
+
+def create_conditioning_ui(prefix: str):
+    components = {}
+    key = lambda name: f"{name}_{prefix}"
+    
+    with gr.Accordion("Conditioning Settings", open=False) as accordion:
+        components[key('conditioning_accordion')] = accordion
+        gr.Markdown("💡 **Tip:** Define rectangular areas and assign specific prompts to them. Coordinates (X, Y) start from the top-left corner.")
+        
+        cond_rows, prompts, widths, heights, xs, ys, strengths = [], [], [], [], [], [], []
+        components.update({
+            key('conditioning_rows'): cond_rows,
+            key('conditioning_prompts'): prompts,
+            key('conditioning_widths'): widths,
+            key('conditioning_heights'): heights,
+            key('conditioning_xs'): xs,
+            key('conditioning_ys'): ys,
+            key('conditioning_strengths'): strengths
+        })
+        
+        for i in range(MAX_CONDITIONINGS):
+            with gr.Column(visible=(i < 1)) as row_wrapper:
+                prompts.append(gr.Textbox(label=f"Area Prompt {i+1}", lines=2, interactive=True))
+                with gr.Row():
+                    xs.append(gr.Number(label="X", value=0, interactive=True, step=8, scale=1))
+                    ys.append(gr.Number(label="Y", value=0, interactive=True, step=8, scale=1))
+                    widths.append(gr.Number(label="Width", value=512, interactive=True, step=8, scale=1))
+                    heights.append(gr.Number(label="Height", value=512, interactive=True, step=8, scale=1))
+                    strengths.append(gr.Slider(label="Strength", minimum=0.1, maximum=2.0, step=0.05, value=1.0, interactive=True, scale=2))
+                cond_rows.append(row_wrapper)
+
+        with gr.Row():
+            components[key('add_conditioning_button')] = gr.Button("✚ Add Area")
+            components[key('delete_conditioning_button')] = gr.Button("➖ Delete Area", visible=False)
+        components[key('conditioning_count_state')] = gr.State(1)
+
+        all_cond_components_flat = prompts + widths + heights + xs + ys + strengths
+        components[key('all_conditioning_components_flat')] = all_cond_components_flat
+    
+    return components
+    
+def create_reference_latent_ui(prefix: str):
+    components = {}
+    key = lambda name: f"{name}_{prefix}"
+
+    with gr.Accordion("Reference Edit", open=False) as accordion:
+        components[key('reference_latent_accordion')] = accordion
+        gr.Markdown("💡 **Tip:** For multimodal models (like FLUX.2), this feature enables powerful editing and combining capabilities. In txt2img mode, adding a single reference image performs an **Image Edit**, while adding multiple images performs an **Image Combine**.")
+        
+        ref_rows, ref_images = [], []
+        components.update({
+            key('reference_latent_rows'): ref_rows,
+            key('reference_latent_images'): ref_images,
+        })
+        
+        with gr.Row():
+            for i in range(MAX_REFERENCE_LATENTS):
+                with gr.Column(visible=(i < 1), min_width=160) as row_wrapper:
+                    ref_images.append(gr.Image(type="pil", label=f"Reference {i+1}", sources=["upload"], height=150))
+                    ref_rows.append(row_wrapper)
+
+        with gr.Row():
+            components[key('add_reference_latent_button')] = gr.Button("✚ Add Reference Image")
+            components[key('delete_reference_latent_button')] = gr.Button("➖ Delete Reference Image", visible=False)
+        components[key('reference_latent_count_state')] = gr.State(1)
+
+        components[key('all_reference_latent_components_flat')] = ref_images
+
+    return components
+
+def create_vae_override_ui(prefix: str):
+    components = {}
+    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
+        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(
+                label="VAE Source", 
+                choices=source_choices, 
+                value="None", 
+                scale=1, 
+                interactive=True
+            )
+            components[key('vae_id')] = gr.Textbox(
+                label="Civitai Version ID / File", 
+                placeholder="Civitai Version ID or Filename", 
+                scale=3, 
+                interactive=True,
+                type="text"
+            )
+            upload_btn = gr.UploadButton(
+                "Upload", 
+                file_types=[".safetensors"], 
+                scale=1
+            )
+            components[key('vae_upload_button')] = upload_btn
+            components[key('vae_file')] = gr.State(None)
+            
+    return components

utils/app_utils.py

@@ -0,0 +1,462 @@
+import os
+import requests
+import hashlib
+import re
+from typing import Sequence, Mapping, Any, Union, Set
+from pathlib import Path
+import shutil
+
+import gradio as gr
+from huggingface_hub import hf_hub_download, constants as hf_constants
+import torch
+import numpy as np
+from PIL import Image, ImageChops
+
+
+from core.settings import *
+
+DISK_LIMIT_GB = 120
+MODELS_ROOT_DIR = "ComfyUI/models"
+
+PREPROCESSOR_MODEL_MAP = None
+PREPROCESSOR_PARAMETER_MAP = None
+
+
+def save_uploaded_file_with_hash(file_obj: gr.File, target_dir: str) -> str:
+    if not file_obj:
+        return ""
+    
+    temp_path = file_obj.name
+    
+    sha256 = hashlib.sha256()
+    with open(temp_path, 'rb') as f:
+        for block in iter(lambda: f.read(65536), b''):
+            sha256.update(block)
+    
+    file_hash = sha256.hexdigest()
+    _, extension = os.path.splitext(temp_path)
+    hashed_filename = f"{file_hash}{extension.lower()}"
+    
+    dest_path = os.path.join(target_dir, hashed_filename)
+    
+    os.makedirs(target_dir, exist_ok=True)
+    if not os.path.exists(dest_path):
+        shutil.copy(temp_path, dest_path)
+        print(f"✅ Saved uploaded file as: {dest_path}")
+    else:
+        print(f"ℹ️ File already exists (deduplicated): {dest_path}")
+        
+    return hashed_filename
+
+
+def bytes_to_gb(byte_size: int) -> float:
+    if byte_size is None or byte_size == 0:
+        return 0.0
+    return round(byte_size / (1024 ** 3), 2)
+
+def get_directory_size(path: str) -> int:
+    total_size = 0
+    if not os.path.exists(path):
+        return 0
+    try:
+        for dirpath, _, filenames in os.walk(path):
+            for f in filenames:
+                fp = os.path.join(dirpath, f)
+                if os.path.isfile(fp) and not os.path.islink(fp):
+                    total_size += os.path.getsize(fp)
+    except OSError as e:
+        print(f"Warning: Could not access {path} to calculate size: {e}")
+    return total_size
+
+def enforce_disk_limit():
+    disk_limit_bytes = DISK_LIMIT_GB * (1024 ** 3)
+    cache_dir = hf_constants.HF_HUB_CACHE
+
+    if not os.path.exists(cache_dir):
+        return
+
+    print(f"--- [Storage Manager] Checking disk usage in '{cache_dir}' (Limit: {DISK_LIMIT_GB} GB) ---")
+
+    try:
+        all_files = []
+        current_size_bytes = 0
+        for dirpath, _, filenames in os.walk(cache_dir):
+            for f in filenames:
+                if f.endswith(".incomplete") or f.endswith(".lock"):
+                    continue
+                file_path = os.path.join(dirpath, f)
+                if os.path.isfile(file_path) and not os.path.islink(file_path):
+                    try:
+                        file_size = os.path.getsize(file_path)
+                        creation_time = os.path.getctime(file_path)
+                        all_files.append((creation_time, file_path, file_size))
+                        current_size_bytes += file_size
+                    except OSError:
+                        continue
+
+        print(f"--- [Storage Manager] Current usage: {bytes_to_gb(current_size_bytes)} GB ---")
+
+        if current_size_bytes > disk_limit_bytes:
+            print(f"--- [Storage Manager] Usage exceeds limit. Starting cleanup... ---")
+            all_files.sort(key=lambda x: x[0])
+
+            while current_size_bytes > disk_limit_bytes and all_files:
+                oldest_file_time, oldest_file_path, oldest_file_size = all_files.pop(0)
+                try:
+                    os.remove(oldest_file_path)
+                    current_size_bytes -= oldest_file_size
+                    print(f"--- [Storage Manager] Deleted oldest file: {os.path.basename(oldest_file_path)} ({bytes_to_gb(oldest_file_size)} GB freed) ---")
+                except OSError as e:
+                    print(f"--- [Storage Manager] Error deleting file {oldest_file_path}: {e} ---")
+
+            print(f"--- [Storage Manager] Cleanup finished. New usage: {bytes_to_gb(current_size_bytes)} GB ---")
+        else:
+            print("--- [Storage Manager] Disk usage is within the limit. No action needed. ---")
+
+    except Exception as e:
+        print(f"--- [Storage Manager] An unexpected error occurred: {e} ---")
+
+
+def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
+    try:
+        return obj[index]
+    except (KeyError, IndexError):
+        try:
+            return obj["result"][index]
+        except (KeyError, IndexError):
+            return None
+
+def sanitize_prompt(prompt: str) -> str:
+    if not isinstance(prompt, str):
+        return ""
+    return "".join(char for char in prompt if char.isprintable() or char in ('\n', '\t'))
+
+def sanitize_id(input_id: str) -> str:
+    if not isinstance(input_id, str):
+        return ""
+    return re.sub(r'[^0-9]', '', input_id)
+
+def sanitize_url(url: str) -> str:
+    if not isinstance(url, str):
+        raise ValueError("URL must be a string.")
+    url = url.strip()
+    if not re.match(r'^https?://[^\s/$.?#].[^\s]*$', url):
+        raise ValueError("Invalid URL format or scheme. Only HTTP and HTTPS are allowed.")
+    return url
+
+def sanitize_filename(filename: str) -> str:
+    if not isinstance(filename, str):
+        return ""
+    sanitized = filename.replace('..', '')
+    sanitized = re.sub(r'[^\w\.\-]', '_', sanitized)
+    return sanitized.lstrip('/\\')
+
+
+def get_civitai_file_info(version_id: str) -> dict | None:
+    api_url = f"https://civitai.com/api/v1/model-versions/{version_id}"
+    try:
+        response = requests.get(api_url, timeout=10)
+        response.raise_for_status()
+        data = response.json()
+        
+        for file_data in data.get('files', []):
+            if file_data.get('type') == 'Model' and file_data['name'].endswith(('.safetensors', '.pt', '.bin')):
+                return file_data
+        
+        if data.get('files'):
+            return data['files'][0]
+    except Exception:
+        return None
+
+
+def download_file(url: str, save_path: str, api_key: str = None, progress=None, desc: str = "") -> str:
+    enforce_disk_limit()
+    
+    if os.path.exists(save_path):
+        return f"File already exists: {os.path.basename(save_path)}"
+    
+    headers = {'Authorization': f'Bearer {api_key}'} if api_key and api_key.strip() else {}
+    try:
+        if progress:
+            progress(0, desc=desc)
+        
+        response = requests.get(url, stream=True, headers=headers, timeout=15)
+        response.raise_for_status()
+        total_size = int(response.headers.get('content-length', 0))
+        
+        with open(save_path, "wb") as f:
+            downloaded = 0
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+                if progress and total_size > 0:
+                    downloaded += len(chunk)
+                    progress(downloaded / total_size, desc=desc)
+        return f"Successfully downloaded: {os.path.basename(save_path)}"
+    except Exception as e:
+        if os.path.exists(save_path):
+            os.remove(save_path)
+        return f"Download failed for {os.path.basename(save_path)}: {e}"
+
+
+def get_lora_path(source: str, id_or_url: str, civitai_key: str, progress) -> tuple[str | None, str]:
+    if not id_or_url or not id_or_url.strip():
+        return None, "No ID/URL provided."
+
+    try:
+        if source == "Civitai":
+            version_id = sanitize_id(id_or_url)
+            if not version_id:
+                return None, "Invalid Civitai ID provided. Must be numeric."
+            filename = sanitize_filename(f"civitai_{version_id}.safetensors")
+            local_path = os.path.join(LORA_DIR, filename)
+            file_info = get_civitai_file_info(version_id)
+            api_key_to_use = civitai_key
+            source_name = f"Civitai ID {version_id}"
+        else:
+            return None, "Invalid source."
+
+    except ValueError as e:
+        return None, f"Input validation failed: {e}"
+
+    if os.path.exists(local_path):
+        return local_path, "File already exists."
+
+    if not file_info or not file_info.get('downloadUrl'):
+        return None, f"Could not get download link for {source_name}."
+
+    status = download_file(file_info['downloadUrl'], local_path, api_key_to_use, progress=progress, desc=f"Downloading {source_name}")
+    
+    return (local_path, status) if "Successfully" in status else (None, status)
+
+def get_embedding_path(source: str, id_or_url: str, civitai_key: str, progress) -> tuple[str | None, str]:
+    if not id_or_url or not id_or_url.strip():
+        return None, "No ID/URL provided."
+
+    try:
+        file_ext = ".safetensors"
+
+        if source == "Civitai":
+            version_id = sanitize_id(id_or_url)
+            if not version_id:
+                return None, "Invalid Civitai ID. Must be numeric."
+            
+            file_info = get_civitai_file_info(version_id)
+            if file_info and file_info['name'].lower().endswith(('.pt', '.bin')):
+                file_ext = os.path.splitext(file_info['name'])[1]
+
+            filename = sanitize_filename(f"civitai_{version_id}{file_ext}")
+            local_path = os.path.join(EMBEDDING_DIR, filename)
+            api_key_to_use = civitai_key
+            source_name = f"Embedding Civitai ID {version_id}"
+        else:
+            return None, "Invalid source."
+
+    except ValueError as e:
+        return None, f"Input validation failed: {e}"
+
+    if os.path.exists(local_path):
+        return local_path, "File already exists."
+
+    if not file_info or not file_info.get('downloadUrl'):
+        return None, f"Could not get download link for {source_name}."
+
+    status = download_file(file_info['downloadUrl'], local_path, api_key_to_use, progress=progress, desc=f"Downloading {source_name}")
+    
+    return (local_path, status) if "Successfully" in status else (None, status)
+
+def get_vae_path(source: str, id_or_url: str, civitai_key: str, progress) -> tuple[str | None, str]:
+    if not id_or_url or not id_or_url.strip():
+        return None, "No ID/URL provided."
+
+    try:
+        file_ext = ".safetensors"
+
+        if source == "Civitai":
+            version_id = sanitize_id(id_or_url)
+            if not version_id:
+                return None, "Invalid Civitai ID. Must be numeric."
+            
+            file_info = get_civitai_file_info(version_id)
+            if file_info and file_info['name'].lower().endswith(('.pt', '.bin')):
+                file_ext = os.path.splitext(file_info['name'])[1]
+
+            filename = sanitize_filename(f"civitai_{version_id}{file_ext}")
+            local_path = os.path.join(VAE_DIR, filename)
+            api_key_to_use = civitai_key
+            source_name = f"VAE Civitai ID {version_id}"
+        else:
+            return None, "Invalid source."
+
+    except ValueError as e:
+        return None, f"Input validation failed: {e}"
+
+    if os.path.exists(local_path):
+        return local_path, "File already exists."
+
+    if not file_info or not file_info.get('downloadUrl'):
+        return None, f"Could not get download link for {source_name}."
+
+    status = download_file(file_info['downloadUrl'], local_path, api_key_to_use, progress=progress, desc=f"Downloading {source_name}")
+    
+    return (local_path, status) if "Successfully" in status else (None, status)
+
+
+def _ensure_model_downloaded(filename: str, progress=gr.Progress()):
+    download_info = ALL_FILE_DOWNLOAD_MAP.get(filename)
+    if not download_info:
+        raise gr.Error(f"Model component '{filename}' not found in file_list.yaml. Cannot download.")
+        
+    category_to_dir_map = {
+        "diffusion_models": DIFFUSION_MODELS_DIR,
+        "text_encoders": TEXT_ENCODERS_DIR,
+        "vae": VAE_DIR,
+        "checkpoints": CHECKPOINT_DIR,
+        "loras": LORA_DIR,
+        "controlnet": CONTROLNET_DIR,
+        "model_patches": MODEL_PATCHES_DIR,
+        "clip_vision": os.path.join(os.path.dirname(LORA_DIR), "clip_vision")
+    }
+    
+    category = download_info.get('category')
+    dest_dir = category_to_dir_map.get(category)
+    if not dest_dir:
+        raise ValueError(f"Unknown model category '{category}' for file '{filename}'.")
+    
+    dest_path = os.path.join(dest_dir, filename)
+
+    if os.path.lexists(dest_path):
+        if not os.path.exists(dest_path):
+            print(f"⚠️ Found and removed broken symlink: {dest_path}")
+            os.remove(dest_path)
+        else:
+            return filename
+
+    source = download_info.get("source")
+    try:
+        progress(0, desc=f"Downloading: {filename}")
+        
+        if source == "hf":
+            repo_id = download_info.get("repo_id")
+            hf_filename = download_info.get("repository_file_path", filename)
+            if not repo_id:
+                raise ValueError(f"repo_id is missing for HF model '{filename}'")
+            
+            cached_path = hf_hub_download(repo_id=repo_id, filename=hf_filename)
+            os.makedirs(dest_dir, exist_ok=True)
+            os.symlink(cached_path, dest_path)
+            print(f"✅ Symlinked '{cached_path}' to '{dest_path}'")
+
+        elif source == "civitai":
+            model_version_id = download_info.get("model_version_id")
+            if not model_version_id:
+                raise ValueError(f"model_version_id is missing for Civitai model '{filename}'")
+            
+            file_info = get_civitai_file_info(model_version_id)
+            if not file_info or not file_info.get('downloadUrl'):
+                raise ConnectionError(f"Could not get download URL for Civitai model version ID {model_version_id}")
+            
+            status = download_file(
+                file_info['downloadUrl'], dest_path, progress=progress, desc=f"Downloading: {filename}"
+            )
+            if "Failed" in status:
+                raise ConnectionError(status)
+        else:
+            raise NotImplementedError(f"Download source '{source}' is not implemented for '{filename}'")
+            
+        progress(1.0, desc=f"Downloaded: {filename}")
+
+    except Exception as e:
+        if os.path.lexists(dest_path):
+            try:
+                os.remove(dest_path)
+            except OSError: pass
+        raise gr.Error(f"Failed to download and link '{filename}': {e}")
+    
+    return filename
+
+def ensure_controlnet_model_downloaded(filename: str, progress):
+    if not filename or filename == "None":
+        return
+    _ensure_model_downloaded(filename, progress)
+
+
+def build_preprocessor_model_map():
+    global PREPROCESSOR_MODEL_MAP
+    if PREPROCESSOR_MODEL_MAP is not None: return PREPROCESSOR_MODEL_MAP
+    print("--- Building ControlNet Preprocessor model map ---")
+    manual_map = {
+        "dwpose": [("yzd-v/DWPose", "yolox_l.onnx"), ("yzd-v/DWPose", "dw-ll_ucoco_384.onnx"), ("hr16/UnJIT-DWPose", "dw-ll_ucoco.onnx"), ("hr16/DWPose-TorchScript-BatchSize5", "dw-ll_ucoco_384_bs5.torchscript.pt"), ("hr16/DWPose-TorchScript-BatchSize5", "rtmpose-m_ap10k_256_bs5.torchscript.pt"), ("hr16/yolo-nas-fp16", "yolo_nas_l_fp16.onnx"), ("hr16/yolo-nas-fp16", "yolo_nas_m_fp16.onnx"), ("hr16/yolo-nas-fp16", "yolo_nas_s_fp16.onnx")],
+        "densepose": [("LayerNorm/DensePose-TorchScript-with-hint-image", "densepose_r50_fpn_dl.torchscript"), ("LayerNorm/DensePose-TorchScript-with-hint-image", "densepose_r101_fpn_dl.torchscript")]
+    }
+    temp_map = {}
+    from nodes import NODE_DISPLAY_NAME_MAPPINGS
+    wrappers_dir = Path("./custom_nodes/comfyui_controlnet_aux/node_wrappers/")
+    if not wrappers_dir.exists():
+        print("⚠️ ControlNet AUX wrappers directory not found. Cannot build model map.")
+        PREPROCESSOR_MODEL_MAP = {}; return PREPROCESSOR_MODEL_MAP
+    for wrapper_file in wrappers_dir.glob("*.py"):
+        if wrapper_file.name == "__init__.py": continue
+        with open(wrapper_file, 'r', encoding='utf-8') as f:
+            content = f.read()
+            display_name_matches = re.findall(r'NODE_DISPLAY_NAME_MAPPINGS\s*=\s*{(?:.|\n)*?["\'](.*?)["\']\s*:\s*["\'](.*?)["\']', content)
+            for _, display_name in display_name_matches:
+                if display_name not in temp_map: temp_map[display_name] = []
+                manual_key = wrapper_file.stem
+                if manual_key in manual_map: temp_map[display_name].extend(manual_map[manual_key])
+                matches = re.findall(r"from_pretrained\s*\(\s*(?:filename=)?\s*f?[\"']([^\"']+)[\"']", content)
+                for model_filename in matches:
+                    repo_id = "lllyasviel/Annotators"
+                    if "depth_anything" in model_filename and "v2" in model_filename: repo_id = "LiheYoung/Depth-Anything-V2"
+                    elif "depth_anything" in model_filename: repo_id = "LiheYoung/Depth-Anything"
+                    elif "diffusion_edge" in model_filename: repo_id = "hr16/Diffusion-Edge"
+                    temp_map[display_name].append((repo_id, model_filename))
+    final_map = {name: sorted(list(set(models))) for name, models in temp_map.items() if models}
+    PREPROCESSOR_MODEL_MAP = final_map
+    print("✅ ControlNet Preprocessor model map built."); return PREPROCESSOR_MODEL_MAP
+
+def build_preprocessor_parameter_map():
+    global PREPROCESSOR_PARAMETER_MAP
+    if PREPROCESSOR_PARAMETER_MAP is not None: return
+    print("--- Building ControlNet Preprocessor parameter map ---")
+    param_map = {}
+    from nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+    for class_name, node_class in NODE_CLASS_MAPPINGS.items():
+        if not hasattr(node_class, "INPUT_TYPES"): continue
+        if hasattr(node_class, '__module__') and 'comfyui_controlnet_aux.node_wrappers' not in node_class.__module__: continue
+        display_name = NODE_DISPLAY_NAME_MAPPINGS.get(class_name)
+        if not display_name: continue
+        try:
+            input_types = node_class.INPUT_TYPES()
+            all_inputs = {**input_types.get('required', {}), **input_types.get('optional', {})}
+            params = []
+            for name, details in all_inputs.items():
+                if name in ['image', 'resolution', 'pose_kps']: continue
+                if not isinstance(details, (list, tuple)) or not details: continue
+                param_type = details[0]
+                param_config = details[1] if len(details) > 1 and isinstance(details[1], dict) else {}
+                param_info = {"name": name, "type": param_type, "config": param_config}
+                params.append(param_info)
+            if params: param_map[display_name] = params
+        except Exception as e:
+            print(f"⚠️ Could not parse parameters for {display_name}: {e}")
+    PREPROCESSOR_PARAMETER_MAP = param_map
+    print("✅ ControlNet Preprocessor parameter map built.")
+
+def print_welcome_message():
+    author_name = "RioShiina"
+    project_url = "https://huggingface.co/RioShiina"
+    border = "=" * 72
+    
+    message = (
+        f"\n{border}\n\n"
+        f"  Thank you for using this project!\n\n"
+        f"  **Author:** {author_name}\n"
+        f"  **Find more from the author:** {project_url}\n\n"
+        f"  This project is open-source under the GNU General Public License v3.0 (GPL-3.0).\n"
+        f"  As it's built upon GPL-3.0 components (like ComfyUI), any modifications you\n"
+        f"  distribute must also be open-sourced under the same license.\n\n"
+        f"  Your respect for the principles of free software is greatly appreciated!\n\n"
+        f"{border}\n"
+    )
+    
+    print(message)

yaml/constants.yaml

@@ -0,0 +1,16 @@
+MAX_LORAS: 5
+MAX_CONTROLNETS: 5
+MAX_EMBEDDINGS: 5
+MAX_CONDITIONINGS: 10
+MAX_REFERENCE_LATENTS: 10
+LORA_SOURCE_CHOICES: ["Civitai", "File"]
+
+RESOLUTION_MAP:
+  sdxl:
+    "1:1 (Square)": [1024, 1024]
+    "16:9 (Landscape)": [1344, 768]
+    "9:16 (Portrait)": [768, 1344]
+    "4:3 (Classic)": [1152, 896]
+    "3:4 (Classic Portrait)": [896, 1152]
+    "3:2 (Photography)": [1216, 832]
+    "2:3 (Photography Portrait)": [832, 1216]

yaml/file_list.yaml

@@ -0,0 +1,50 @@
+file:
+  diffusion_models: 
+    # FLUX.2-klein-4B
+    - filename: "flux-2-klein-4b-fp8.safetensors"
+      source: "hf"
+      repo_id: "black-forest-labs/FLUX.2-klein-4b-fp8"
+      repository_file_path: "flux-2-klein-4b-fp8.safetensors"
+    # FLUX.2-klein-base-4B
+    - filename: "flux-2-klein-base-4b-fp8.safetensors"
+      source: "hf"
+      repo_id: "black-forest-labs/FLUX.2-klein-base-4b-fp8"
+      repository_file_path: "flux-2-klein-base-4b-fp8.safetensors"
+    # FLUX.2-klein-9B
+    - filename: "flux-2-klein-9b-fp8.safetensors"
+      source: "hf"
+      repo_id: "black-forest-labs/FLUX.2-klein-9b-fp8"
+      repository_file_path: "flux-2-klein-9b-fp8.safetensors"
+    # FLUX.2-klein-base-9B
+    - filename: "flux-2-klein-base-9b-fp8.safetensors"
+      source: "hf"
+      repo_id: "black-forest-labs/FLUX.2-klein-base-9b-fp8"
+      repository_file_path: "flux-2-klein-base-9b-fp8.safetensors"
+    # FLUX.2-dev
+    - filename: "flux2_dev_fp8mixed.safetensors"
+      source: "hf"
+      repo_id: "Comfy-Org/flux2-dev"
+      repository_file_path: "split_files/diffusion_models/flux2_dev_fp8mixed.safetensors"
+
+  text_encoders: 
+    # FLUX.2-klein-4B & base
+    - filename: "qwen_3_4b.safetensors"
+      source: "hf"
+      repo_id: "Comfy-Org/vae-text-encorder-for-flux-klein-4b"
+      repository_file_path: "split_files/text_encoders/qwen_3_4b.safetensors"
+    # FLUX.2-klein-9B & base
+    - filename: "qwen_3_8b_fp8mixed.safetensors"
+      source: "hf"
+      repo_id: "Comfy-Org/vae-text-encorder-for-flux-klein-9b"
+      repository_file_path: "split_files/text_encoders/qwen_3_8b_fp8mixed.safetensors"
+    # FLUX.2-dev
+    - filename: "mistral_3_small_flux2_fp8.safetensors"
+      source: "hf"
+      repo_id: "Comfy-Org/flux2-dev"
+      repository_file_path: "split_files/text_encoders/mistral_3_small_flux2_fp8.safetensors"
+
+  vae: 
+    - filename: "flux2-vae.safetensors"
+      source: "hf"
+      repo_id: "Comfy-Org/flux2-dev"
+      repository_file_path: "split_files/vae/flux2-vae.safetensors"

yaml/injectors.yaml

@@ -0,0 +1,9 @@
+injector_definitions:
+  dynamic_conditioning_chains:
+    module: "chain_injectors.conditioning_injector"
+  dynamic_reference_latent_chains:
+    module: "chain_injectors.reference_latent_injector"
+  
+injector_order:
+  - dynamic_reference_latent_chains
+  - dynamic_conditioning_chains

yaml/model_defaults.yaml

@@ -0,0 +1,18 @@
+Default:
+  steps: 20
+  cfg: 4.0
+  sampler_name: "euler"
+  scheduler: "simple"
+
+FLUX.2:
+  _defaults:
+    steps: 20
+    cfg: 4.0
+    sampler_name: "euler"
+    scheduler: "simple"
+  "black-forest-labs/FLUX.2-klein-4B":
+    steps: 4
+    cfg: 1.0
+  "black-forest-labs/FLUX.2-klein-9B":
+    steps: 4
+    cfg: 1.0

yaml/model_list.yaml

@@ -0,0 +1,26 @@
+Checkpoint:
+  - display_name: "black-forest-labs/FLUX.2-klein-4B"
+    components:
+      unet: "flux-2-klein-4b-fp8.safetensors"
+      clip: "qwen_3_4b.safetensors"
+      vae: "flux2-vae.safetensors"
+  - display_name: "black-forest-labs/FLUX.2-klein-9B"
+    components:
+      unet: "flux-2-klein-9b-fp8.safetensors"
+      clip: "qwen_3_8b_fp8mixed.safetensors"
+      vae: "flux2-vae.safetensors"
+  - display_name: "black-forest-labs/FLUX.2-klein-base-4B"
+    components:
+      unet: "flux-2-klein-base-4b-fp8.safetensors"
+      clip: "qwen_3_4b.safetensors"
+      vae: "flux2-vae.safetensors"
+  - display_name: "black-forest-labs/FLUX.2-klein-base-9B"
+    components:
+      unet: "flux-2-klein-base-9b-fp8.safetensors"
+      clip: "qwen_3_8b_fp8mixed.safetensors"
+      vae: "flux2-vae.safetensors"
+  - display_name: "black-forest-labs/FLUX.2-dev (Need to set ZeroGPU Duration to 120)"
+    components:
+      unet: "flux2_dev_fp8mixed.safetensors"
+      clip: "mistral_3_small_flux2_fp8.safetensors"
+      vae: "flux2-vae.safetensors"

yaml/private_file_list.yaml

@@ -0,0 +1,12 @@
+file:
+  diffusion_models: 
+    # FLUX.2-klein-9B
+    - filename: "flux-2-klein-9b-fp8.safetensors"
+      source: "hf"
+      repo_id: "black-forest-labs/FLUX.2-klein-9b-fp8"
+      repository_file_path: "flux-2-klein-9b-fp8.safetensors"
+    # FLUX.2-klein-base-9B
+    - filename: "flux-2-klein-base-9b-fp8.safetensors"
+      source: "hf"
+      repo_id: "black-forest-labs/FLUX.2-klein-base-9b-fp8"
+      repository_file_path: "flux-2-klein-base-9b-fp8.safetensors"