Upload folder using huggingface_hub

README.md

@@ -1,13 +1,10 @@
----
-title: ImageGen Z Image
-emoji: 🚀
-colorFrom: green
-colorTo: yellow
-sdk: gradio
-sdk_version: 6.0.2
-app_file: app.py
-pinned: false
-license: gpl-3.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: ImageGen - Z-Image
+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,122 @@
+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 main():
+    from utils.app_utils import print_welcome_message
+    from scripts import build_sage_attention
+
+    print_welcome_message()
+    
+    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("--- 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/controlnet_injector.py

@@ -0,0 +1,49 @@
+def inject(assembler, chain_definition, chain_items):
+    if not chain_items:
+        return
+
+    ksampler_name = chain_definition.get('ksampler_node', 'ksampler')
+    if ksampler_name not in assembler.node_map:
+        print(f"Warning: Target node '{ksampler_name}' for ControlNet chain not found. Skipping chain injection.")
+        return
+        
+    ksampler_id = assembler.node_map[ksampler_name]
+
+    if 'positive' not in assembler.workflow[ksampler_id]['inputs'] or \
+       'negative' not in assembler.workflow[ksampler_id]['inputs']:
+        print(f"Warning: KSampler node '{ksampler_name}' is missing 'positive' or 'negative' inputs. Skipping ControlNet chain.")
+        return
+        
+    current_positive_connection = assembler.workflow[ksampler_id]['inputs']['positive']
+    current_negative_connection = assembler.workflow[ksampler_id]['inputs']['negative']
+    
+    for item_data in chain_items:
+        cn_loader_id = assembler._get_unique_id()
+        cn_loader_node = assembler._get_node_template("ControlNetLoader")
+        cn_loader_node['inputs']['control_net_name'] = item_data['control_net_name']
+        assembler.workflow[cn_loader_id] = cn_loader_node
+
+        image_loader_id = assembler._get_unique_id()
+        image_loader_node = assembler._get_node_template("LoadImage")
+        image_loader_node['inputs']['image'] = item_data['image']
+        assembler.workflow[image_loader_id] = image_loader_node
+
+        apply_cn_id = assembler._get_unique_id()
+        apply_cn_node = assembler._get_node_template(chain_definition['template'])
+        
+        apply_cn_node['inputs']['strength'] = item_data['strength']
+        
+        apply_cn_node['inputs']['positive'] = current_positive_connection
+        apply_cn_node['inputs']['negative'] = current_negative_connection
+        apply_cn_node['inputs']['control_net'] = [cn_loader_id, 0]
+        apply_cn_node['inputs']['image'] = [image_loader_id, 0]
+        
+        assembler.workflow[apply_cn_id] = apply_cn_node
+
+        current_positive_connection = [apply_cn_id, 0]
+        current_negative_connection = [apply_cn_id, 1]
+
+    assembler.workflow[ksampler_id]['inputs']['positive'] = current_positive_connection
+    assembler.workflow[ksampler_id]['inputs']['negative'] = current_negative_connection
+    
+    print(f"ControlNet injector applied. KSampler inputs redirected through {len(chain_items)} ControlNet nodes.")

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,58 @@
+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/comfyanonymous/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(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, 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,20 @@
+from typing import Any, Dict
+import gradio as gr
+
+from core.pipelines.sd_image_pipeline import SdImagePipeline
+
+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 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,163 @@
+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.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="lumina2", 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 = len(active_loras) > 0 or len(current_set - required_set) > 0
+
+        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)
+                print(f"--- [ModelManager] LoRA configuration changed. Reloading base model(s): {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. ---")
+
+        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/sd_image_pipeline.py

@@ -0,0 +1,395 @@
+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
+
+        controlnet_data = ui_inputs.get('controlnet_data', [])
+        active_controlnets = []
+        (cn_images, _, _, cn_strengths, cn_filepaths) = [controlnet_data[i::5] for i in range(5)]
+        for i in range(len(cn_images)):
+            if cn_images[i] and cn_strengths[i] > 0 and cn_filepaths[i] and cn_filepaths[i] != "None":
+                ensure_controlnet_model_downloaded(cn_filepaths[i], progress)
+
+                if not os.path.exists(INPUT_DIR): os.makedirs(INPUT_DIR)
+                cn_temp_path = os.path.join(INPUT_DIR, f"temp_cn_{i}_{random.randint(1000, 9999)}.png")
+                cn_images[i].save(cn_temp_path, "PNG")
+                temp_files_to_clean.append(cn_temp_path)
+                active_controlnets.append({
+                    "image": os.path.basename(cn_temp_path), "strength": cn_strengths[i],
+                    "start_percent": 0.0, "end_percent": 1.0, "control_net_name": cn_filepaths[i]
+                })
+        
+        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])
+                    })
+
+        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']),
+            "controlnet_chain": active_controlnets,
+            "conditioning_chain": active_conditioning,
+        }
+        
+        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/z-image.yaml

@@ -0,0 +1,63 @@
+nodes:
+  unet_loader:
+    class_type: UNETLoader
+    title: "Load Diffusion Model"
+    params:
+      weight_dtype: "default"
+  vae_loader:
+    class_type: VAELoader
+    title: "Load VAE"
+  clip_loader:
+    class_type: CLIPLoader
+    title: "Load CLIP"
+    params:
+      type: "lumina2"
+      device: "default"
+  model_sampler:
+    class_type: ModelSamplingAuraFlow
+    params:
+      shift: 3.0
+  pos_prompt:
+    class_type: CLIPTextEncode
+    title: "Positive Prompt Encoder"
+  neg_prompt:
+    class_type: CLIPTextEncode
+    title: "Negative Prompt Encoder"
+
+connections:
+  - from: "unet_loader:0"
+    to: "model_sampler:model"
+  - from: "model_sampler:0"
+    to: "ksampler:model"
+  
+  - from: "clip_loader:0"
+    to: "pos_prompt:clip"
+  - from: "clip_loader:0"
+    to: "neg_prompt:clip"
+    
+  - from: "pos_prompt:0"
+    to: "ksampler:positive"
+  - from: "neg_prompt:0"
+    to: "ksampler:negative"
+  
+  - from: "vae_loader:0"
+    to: "vae_decode:vae"
+  - from: "vae_loader:0"
+    to: "vae_encode:vae"
+
+dynamic_controlnet_chains:
+  controlnet_chain:
+    template: "ControlNetApplyAdvanced"
+    ksampler_node: "ksampler"
+
+dynamic_conditioning_chains:
+  conditioning_chain:
+    ksampler_node: "ksampler"
+    clip_source: "clip_loader:0"
+
+ui_map:
+  unet_name: "unet_loader:unet_name"
+  vae_name: "vae_loader:vae_name"
+  clip_name: "clip_loader:clip_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: EmptySD3LatentImage
+
+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/z-image.yaml"
+
+connections:
+  - from: "latent_source:0"
+    to: "ksampler:latent_image"

core/settings.py

@@ -0,0 +1,111 @@
+import yaml
+import os
+from collections import OrderedDict
+
+CHECKPOINT_DIR = "models/checkpoints"
+LORA_DIR = "models/loras"
+EMBEDDING_DIR = "models/embeddings" 
+CONTROLNET_DIR = "models/controlnet"
+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')
+
+
+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
+
+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()}
+
+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 = {}
+
+
+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)
+    LORA_SOURCE_CHOICES = _constants.get('LORA_SOURCE_CHOICES', ["Civitai", "Custom URL", "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
+    LORA_SOURCE_CHOICES = ["Civitai", "Custom URL", "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,56 @@
+comfyui-frontend-package==1.33.10
+comfyui-workflow-templates==0.7.25
+comfyui-embedded-docs==0.3.1
+torch
+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
+
+#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,66 @@
+import os
+import subprocess
+import sys
+
+def run_command(command, env=None):
+    """
+    Runs a command with a specified environment, prints its output, 
+    and raises an exception on failure.
+    """
+    print(f"🚀 Running command: {' '.join(command)}")
+    result = subprocess.run(
+        command,
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        encoding='utf-8',
+        errors='replace'
+    )
+
+    if result.stdout:
+        print("--- Pip Output ---")
+        print(result.stdout.strip())
+        print("------------------")
+
+    if result.returncode != 0:
+        raise subprocess.CalledProcessError(result.returncode, command)
+
+def install_sage_attention():
+    """
+    Installs the sageattention package from PyPI using pip, ensuring the
+    correct CUDA architecture is set for any potential on-the-fly compilation.
+    """
+    print("--- [SageAttention Install] Starting installation using pip ---")
+
+    build_env = os.environ.copy()
+    build_env["TORCH_CUDA_ARCH_LIST"] = "9.0"
+    print(f"🔧 Setting build environment variable: TORCH_CUDA_ARCH_LIST='{build_env['TORCH_CUDA_ARCH_LIST']}'")
+
+    install_command = [
+        sys.executable,
+        "-m",
+        "pip",
+        "install",
+        "sageattention==2.2.0",
+        "--no-build-isolation",
+    ]
+
+    try:
+        run_command(install_command, env=build_env)
+        print("🎉 SageAttention installed successfully via pip!")
+    except subprocess.CalledProcessError as e:
+        print(f"❌ SageAttention installation via pip failed (Exit Code: {e.returncode}).")
+        print("   The application will continue with the default attention mechanism.")
+        raise e
+    except Exception as e:
+        print(f"❌ An unexpected error occurred during pip installation: {e}")
+        raise e
+
+
+if __name__ == "__main__":
+    try:
+        install_sage_attention()
+    except Exception:
+        print("\nInstallation script finished with an error.")
+        sys.exit(1)

ui/events.py

@@ -0,0 +1,439 @@
+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 utils.app_utils import save_uploaded_file_with_hash
+from ui.shared.ui_components import RESOLUTION_MAP, MAX_CONTROLNETS, MAX_EMBEDDINGS, MAX_CONDITIONINGS, MAX_LORAS
+
+
+@lru_cache(maxsize=1)
+def load_controlnet_config():
+    _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    _CN_MODEL_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'controlnet_models.yaml')
+    try:
+        print("--- Loading controlnet_models.yaml ---")
+        with open(_CN_MODEL_LIST_PATH, 'r', encoding='utf-8') as f:
+            config = yaml.safe_load(f)
+        print("--- ✅ controlnet_models.yaml loaded successfully ---")
+        return config.get("ControlNet", {}).get("SDXL", [])
+    except Exception as e:
+        print(f"Error loading controlnet_models.yaml: {e}")
+        return []
+
+
+def attach_event_handlers(ui_components, demo):
+    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_controlnet_event_handlers(prefix):
+        cn_rows = ui_components[f'controlnet_rows_{prefix}']
+        cn_types = ui_components[f'controlnet_types_{prefix}']
+        cn_series = ui_components[f'controlnet_series_{prefix}']
+        cn_filepaths = ui_components[f'controlnet_filepaths_{prefix}']
+        cn_images = ui_components[f'controlnet_images_{prefix}']
+        cn_strengths = ui_components[f'controlnet_strengths_{prefix}']
+        
+        count_state = ui_components[f'controlnet_count_state_{prefix}']
+        add_button = ui_components[f'add_controlnet_button_{prefix}']
+        del_button = ui_components[f'delete_controlnet_button_{prefix}']
+        accordion = ui_components[f'controlnet_accordion_{prefix}']
+
+        def add_cn_row(c):
+            c += 1
+            updates = {
+                count_state: c,
+                cn_rows[c-1]: gr.update(visible=True),
+                add_button: gr.update(visible=c < MAX_CONTROLNETS),
+                del_button: gr.update(visible=True)
+            }
+            return updates
+
+        def del_cn_row(c):
+            c -= 1
+            updates = {
+                count_state: c,
+                cn_rows[c]: gr.update(visible=False),
+                cn_images[c]: None,
+                cn_strengths[c]: 1.0,
+                add_button: gr.update(visible=True),
+                del_button: gr.update(visible=c > 0)
+            }
+            return updates
+            
+        add_outputs = [count_state, add_button, del_button] + cn_rows
+        del_outputs = [count_state, add_button, del_button] + cn_rows + cn_images + cn_strengths
+        add_button.click(fn=add_cn_row, inputs=[count_state], outputs=add_outputs, show_progress=False)
+        del_button.click(fn=del_cn_row, inputs=[count_state], outputs=del_outputs, show_progress=False)
+        
+        def on_cn_type_change(selected_type):
+            cn_config = load_controlnet_config()
+            series_choices = []
+            if selected_type:
+                series_choices = sorted(list(set(
+                    model.get("Series", "Default") for model in cn_config
+                    if selected_type in model.get("Type", [])
+                )))
+            default_series = series_choices[0] if series_choices else None
+            filepath = "None"
+            if default_series:
+                for model in cn_config:
+                    if model.get("Series") == default_series and selected_type in model.get("Type", []):
+                        filepath = model.get("Filepath")
+                        break
+            return gr.update(choices=series_choices, value=default_series), filepath
+
+        def on_cn_series_change(selected_series, selected_type):
+            cn_config = load_controlnet_config()
+            filepath = "None"
+            if selected_series and selected_type:
+                for model in cn_config:
+                    if model.get("Series") == selected_series and selected_type in model.get("Type", []):
+                        filepath = model.get("Filepath")
+                        break
+            return filepath
+            
+        for i in range(MAX_CONTROLNETS):
+            cn_types[i].change(
+                fn=on_cn_type_change,
+                inputs=[cn_types[i]],
+                outputs=[cn_series[i], cn_filepaths[i]],
+                show_progress=False
+            )
+            cn_series[i].change(
+                fn=on_cn_series_change,
+                inputs=[cn_series[i], cn_types[i]],
+                outputs=[cn_filepaths[i]],
+                show_progress=False
+            )
+
+        def on_accordion_expand(*images):
+            return [gr.update() for _ in images]
+        
+        accordion.expand(
+            fn=on_accordion_expand,
+            inputs=cn_images,
+            outputs=cn_images,
+            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 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}', [])
+        controlnet_data_components = ui_components.get(f'all_controlnet_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}', [])
+        
+        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 + controlnet_data_components + embedding_data_components + conditioning_data_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['controlnet_data'] = list(args[arg_idx : arg_idx + len(controlnet_data_components)])
+            arg_idx += len(controlnet_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)])
+
+            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"),
+    ]:
+        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_controlnet_button_{prefix}' in ui_components: create_controlnet_event_handlers(prefix)
+        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'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 initialize_all_cn_dropdowns():
+        cn_config = load_controlnet_config()
+        if not cn_config: return {}
+
+        all_types = sorted(list(set(t for model in cn_config for t in model.get("Type", []))))
+        default_type = all_types[0] if all_types else None
+        
+        series_choices = []
+        if default_type:
+            series_choices = sorted(list(set(model.get("Series", "Default") for model in cn_config if default_type in model.get("Type", []))))
+        default_series = series_choices[0] if series_choices else None
+        
+        filepath = "None"
+        if default_series and default_type:
+            for model in cn_config:
+                if model.get("Series") == default_series and default_type in model.get("Type", []):
+                    filepath = model.get("Filepath")
+                    break
+
+        updates = {}
+        for prefix in ["txt2img", "img2img", "inpaint", "outpaint", "hires_fix"]:
+            if f'controlnet_types_{prefix}' in ui_components:
+                for type_dd in ui_components[f'controlnet_types_{prefix}']:
+                    updates[type_dd] = gr.update(choices=all_types, value=default_type)
+                for series_dd in ui_components[f'controlnet_series_{prefix}']:
+                    updates[series_dd] = gr.update(choices=series_choices, value=default_series)
+                for filepath_state in ui_components[f'controlnet_filepaths_{prefix}']:
+                    updates[filepath_state] = filepath
+        return updates
+
+    def run_on_load():
+        cn_updates = initialize_all_cn_dropdowns()
+        return cn_updates
+    
+    all_load_outputs = []
+    for prefix in ["txt2img", "img2img", "inpaint", "outpaint", "hires_fix"]:
+        if f'controlnet_types_{prefix}' in ui_components:
+            all_load_outputs.extend(ui_components[f'controlnet_types_{prefix}'])
+            all_load_outputs.extend(ui_components[f'controlnet_series_{prefix}'])
+            all_load_outputs.extend(ui_components[f'controlnet_filepaths_{prefix}'])
+
+    if all_load_outputs:
+        demo.load(
+            fn=run_on_load,
+            outputs=all_load_outputs
+        )

ui/layout.py

@@ -0,0 +1,45 @@
+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 build_ui(event_handler_function):
+    ui_components = {}
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# ImageGen - Z-Image")
+        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: "
+            "[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="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
+        
+        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,73 @@
+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_controlnet_ui, create_embedding_ui,
+    create_conditioning_ui, create_vae_override_ui, create_api_key_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="res_multistep")
+                    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=9)
+                    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: 60s, Max: 120s", 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_controlnet_ui(prefix))
+        # components.update(create_embedding_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,56 @@
+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_controlnet_ui, create_embedding_ui,
+    create_conditioning_ui, create_vae_override_ui, create_api_key_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="res_multistep")
+                    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=9)
+                    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: 60s, Max: 120s", 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_controlnet_ui(prefix))
+        # components.update(create_embedding_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,80 @@
+import gradio as gr
+from core.settings import MODEL_MAP_CHECKPOINT
+from .ui_components import (
+    create_base_parameter_ui, create_lora_settings_ui,
+    create_controlnet_ui, create_embedding_ui,
+    create_conditioning_ui, create_vae_override_ui, create_api_key_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="res_multistep")
+                    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=9)
+                    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: 60s, Max: 120s", 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_controlnet_ui(prefix))
+            # components.update(create_embedding_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,67 @@
+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_controlnet_ui, create_embedding_ui,
+    create_conditioning_ui, create_vae_override_ui, create_api_key_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="res_multistep")
+                    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=9)
+                    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: 60s, Max: 120s", 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_controlnet_ui(prefix))
+        # components.update(create_embedding_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,37 @@
+import gradio as gr
+from core.settings import MODEL_MAP_CHECKPOINT
+from .ui_components import (
+    create_base_parameter_ui, create_lora_settings_ui,
+    create_controlnet_ui, create_embedding_ui,
+    create_conditioning_ui, create_vae_override_ui, create_api_key_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_controlnet_ui(prefix))
+        # components.update(create_embedding_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,249 @@
+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
+)
+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="res_multistep")
+        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=9)
+        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: 60s, Max: 120s", 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_controlnet_ui(prefix: str, max_units=MAX_CONTROLNETS):
+    components = {}
+    key = lambda name: f"{name}_{prefix}"
+
+    with gr.Accordion("ControlNet Settings", open=False) as accordion:
+        components[key('controlnet_accordion')] = accordion
+        
+        cn_rows, images, series, types, strengths, filepaths = [], [], [], [], [], []
+        components.update({
+            key('controlnet_rows'): cn_rows,
+            key('controlnet_images'): images,
+            key('controlnet_series'): series,
+            key('controlnet_types'): types,
+            key('controlnet_strengths'): strengths,
+            key('controlnet_filepaths'): filepaths
+        })
+        
+        for i in range(max_units):
+            with gr.Row(visible=(i < 1)) as row:
+                with gr.Column(scale=1):
+                    images.append(gr.Image(label=f"Control Image {i+1}", type="pil", sources=["upload"], height=256))
+                with gr.Column(scale=2):
+                    types.append(gr.Dropdown(label="Type", choices=[], interactive=True))
+                    series.append(gr.Dropdown(label="Series", choices=[], interactive=True))
+                    strengths.append(gr.Slider(label="Strength", minimum=0.0, maximum=2.0, step=0.05, value=1.0, interactive=True))
+                    filepaths.append(gr.State(None))
+                cn_rows.append(row)
+
+        with gr.Row():
+            components[key('add_controlnet_button')] = gr.Button("✚ Add ControlNet")
+            components[key('delete_controlnet_button')] = gr.Button("➖ Delete ControlNet", visible=False)
+        components[key('controlnet_count_state')] = gr.State(1)
+        
+        all_cn_components_flat = []
+        for i in range(max_units):
+            all_cn_components_flat.extend([
+                images[i], types[i], series[i], strengths[i], filepaths[i]
+            ])
+        components[key('all_controlnet_components_flat')] = all_cn_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_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,413 @@
+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,
+        "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 ---")
+    PREPROCESSOR_MODEL_MAP = {}
+    print("✅ ControlNet Preprocessor model map build skipped as feature is disabled."); 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 ---")
+    PREPROCESSOR_PARAMETER_MAP = {}
+    print("✅ ControlNet Preprocessor parameter map build skipped as feature is disabled.")
+
+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,15 @@
+MAX_LORAS: 5
+MAX_CONTROLNETS: 5
+MAX_EMBEDDINGS: 5
+MAX_CONDITIONINGS: 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/controlnet_models.yaml

@@ -0,0 +1,5 @@
+ControlNet:
+  SDXL:
+    - Filepath: "Z-Image-Turbo-Fun-Controlnet-Union.safetensors"
+      Series: "alibaba-pai Union"
+      Type: ["Canny", "HED", "Depth", "Pose", "MLSD"]

yaml/file_list.yaml

@@ -0,0 +1,24 @@
+file:
+  diffusion_models: 
+    - filename: "z_image_turbo_bf16.safetensors"
+      source: "hf"
+      repo_id: "Comfy-Org/z_image_turbo"
+      repository_file_path: "split_files/diffusion_models/z_image_turbo_bf16.safetensors"
+
+  text_encoders: 
+    - filename: "qwen_3_4b.safetensors"
+      source: "hf"
+      repo_id: "Comfy-Org/z_image_turbo"
+      repository_file_path: "split_files/text_encoders/qwen_3_4b.safetensors"
+
+  vae: 
+    - filename: "ae.safetensors"
+      source: "hf"
+      repo_id: "Comfy-Org/z_image_turbo"
+      repository_file_path: "split_files/vae/ae.safetensors"
+
+  controlnet:
+    - filename: "Z-Image-Turbo-Fun-Controlnet-Union.safetensors"
+      source: "hf"
+      repo_id: "alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union"
+      repository_file_path: "Z-Image-Turbo-Fun-Controlnet-Union.safetensors"

yaml/injectors.yaml

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

yaml/model_list.yaml

@@ -0,0 +1,6 @@
+Checkpoint:
+  - display_name: "Tongyi-MAI/Z Image Turbo"
+    components:
+      unet: "z_image_turbo_bf16.safetensors"
+      vae: "ae.safetensors"
+      clip: "qwen_3_4b.safetensors"