Add High-Level API/MCP Tools

.gitattributes

@@ -1,35 +1,35 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

app.py

@@ -84,13 +84,15 @@ def main():
     print("--- Environment configured. Proceeding with module imports. ---")
     from ui.layout import build_ui
     from ui.events import attach_event_handlers
+    import mcp_tools as mcp
+    print(f"✅ Loaded MCP module with tools: {[fn.__name__ for fn in mcp.MCP_FUNCTIONS]}")
 
     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, show_api=False)
+    demo.queue().launch(mcp_server=True)
 
 
 if __name__ == "__main__":

chain_injectors/__init__.py

@@ -1,50 +1,50 @@
-import os
-import importlib
-import pkgutil
-
-def discover_injectors():
-    injectors = {}
-    package_dir = os.path.dirname(__file__)
-
-    for _, module_name, is_pkg in pkgutil.iter_modules([package_dir]):
-        if is_pkg or module_name.startswith('_'):
-            continue
-
-        full_module_name = f"chain_injectors.{module_name}"
-        try:
-            module = importlib.import_module(full_module_name)
-            if hasattr(module, 'inject') and callable(module.inject):
-                feature_name = getattr(module, 'FEATURE_NAME', None)
-                if not feature_name:
-                    feature_name = module_name[:-9] if module_name.endswith('_injector') else module_name
-
-                chain_type = getattr(module, 'CHAIN_TYPE', None)
-                if not chain_type:
-                    chain_type = f"dynamic_{feature_name}_chains"
-
-                injectors[chain_type] = module.inject
-            else:
-                print(f"Warning: Module '{full_module_name}' does not have a callable 'inject' function.")
-        except Exception as e:
-            print(f"Error importing injector module '{full_module_name}': {e}")
-
-    return injectors
-
-def get_registered_features():
-    features = {}
-    package_dir = os.path.dirname(__file__)
-
-    for _, module_name, is_pkg in pkgutil.iter_modules([package_dir]):
-        if is_pkg or module_name.startswith('_'):
-            continue
-
-        feature_name = module_name[:-9] if module_name.endswith('_injector') else module_name
-        full_module_name = f"chain_injectors.{module_name}"
-        chain_type = f"dynamic_{feature_name}_chains"
-
-        features[feature_name] = {
-            'module': full_module_name,
-            'chain_type': chain_type
-        }
-
-    return features
+import os
+import importlib
+import pkgutil
+
+def discover_injectors():
+    injectors = {}
+    package_dir = os.path.dirname(__file__)
+
+    for _, module_name, is_pkg in pkgutil.iter_modules([package_dir]):
+        if is_pkg or module_name.startswith('_'):
+            continue
+
+        full_module_name = f"chain_injectors.{module_name}"
+        try:
+            module = importlib.import_module(full_module_name)
+            if hasattr(module, 'inject') and callable(module.inject):
+                feature_name = getattr(module, 'FEATURE_NAME', None)
+                if not feature_name:
+                    feature_name = module_name[:-9] if module_name.endswith('_injector') else module_name
+
+                chain_type = getattr(module, 'CHAIN_TYPE', None)
+                if not chain_type:
+                    chain_type = f"dynamic_{feature_name}_chains"
+
+                injectors[chain_type] = module.inject
+            else:
+                print(f"Warning: Module '{full_module_name}' does not have a callable 'inject' function.")
+        except Exception as e:
+            print(f"Error importing injector module '{full_module_name}': {e}")
+
+    return injectors
+
+def get_registered_features():
+    features = {}
+    package_dir = os.path.dirname(__file__)
+
+    for _, module_name, is_pkg in pkgutil.iter_modules([package_dir]):
+        if is_pkg or module_name.startswith('_'):
+            continue
+
+        feature_name = module_name[:-9] if module_name.endswith('_injector') else module_name
+        full_module_name = f"chain_injectors.{module_name}"
+        chain_type = f"dynamic_{feature_name}_chains"
+
+        features[feature_name] = {
+            'module': full_module_name,
+            'chain_type': chain_type
+        }
+
+    return features

chain_injectors/boogu_image_edit_injector.py

@@ -1,73 +1,73 @@
-def create_node(assembler, class_type, title):
-    try:
-        node = assembler._get_node_template(class_type)
-    except Exception:
-        node = {
-            "inputs": {},
-            "class_type": class_type,
-            "_meta": {"title": title}
-        }
-    node['_meta']['title'] = title
-    return node
-
-def inject(assembler, chain_definition, chain_items):
-    if not chain_items:
-        return
-
-    valid_images = []
-    for item in chain_items:
-        if not item:
-            continue
-        img_path = item
-        if isinstance(item, dict):
-            img_path = item.get('image') or item.get('filename') or item.get('path')
-        if img_path:
-            valid_images.append(img_path)
-
-    if not valid_images:
-        return
-
-    valid_images = valid_images[:10]
-
-    boogu_prompt_name = chain_definition.get('boogu_prompt_node', 'boogu_prompt')
-    vae_loader_name = chain_definition.get('vae_loader_node', 'vae_loader')
-
-    boogu_prompt_id = assembler.node_map.get(boogu_prompt_name)
-    if not boogu_prompt_id or boogu_prompt_id not in assembler.workflow:
-        for node_id, node in assembler.workflow.items():
-            if isinstance(node, dict) and node.get('class_type') == 'TextEncodeBooguEdit':
-                boogu_prompt_id = node_id
-                break
-
-    if not boogu_prompt_id:
-        print(f"Warning: Target node '{boogu_prompt_name}' (TextEncodeBooguEdit) for Boogu Edit chain not found. Skipping.")
-        return
-
-    vae_id = assembler.node_map.get(vae_loader_name)
-    if not vae_id:
-        for node_id, node in assembler.workflow.items():
-            if isinstance(node, dict) and node.get('class_type') == 'VAELoader':
-                vae_id = node_id
-                break
-
-    if vae_id:
-        assembler.workflow[boogu_prompt_id]['inputs']['vae'] = [vae_id, 0]
-
-    for i, img_filename in enumerate(valid_images):
-        load_id = assembler._get_unique_id()
-        load_node = create_node(assembler, "LoadImage", f"Load Reference Image {i+1}")
-        load_node['inputs']['image'] = img_filename
-        assembler.workflow[load_id] = load_node
-
-        scale_id = assembler._get_unique_id()
-        scale_node = create_node(assembler, "ImageScaleToTotalPixels", f"Scale Reference {i+1}")
-        scale_node['inputs']['upscale_method'] = "nearest-exact"
-        scale_node['inputs']['megapixels'] = 1
-        scale_node['inputs']['resolution_steps'] = 1
-        scale_node['inputs']['image'] = [load_id, 0]
-        assembler.workflow[scale_id] = scale_node
-
-        image_key = f"images.image_{i+1}"
-        assembler.workflow[boogu_prompt_id]['inputs'][image_key] = [scale_id, 0]
-
-    print(f"Boogu Edit injector applied with {len(valid_images)} reference image(s). Connected VAE dynamically.")
+def create_node(assembler, class_type, title):
+    try:
+        node = assembler._get_node_template(class_type)
+    except Exception:
+        node = {
+            "inputs": {},
+            "class_type": class_type,
+            "_meta": {"title": title}
+        }
+    node['_meta']['title'] = title
+    return node
+
+def inject(assembler, chain_definition, chain_items):
+    if not chain_items:
+        return
+
+    valid_images = []
+    for item in chain_items:
+        if not item:
+            continue
+        img_path = item
+        if isinstance(item, dict):
+            img_path = item.get('image') or item.get('filename') or item.get('path')
+        if img_path:
+            valid_images.append(img_path)
+
+    if not valid_images:
+        return
+
+    valid_images = valid_images[:10]
+
+    boogu_prompt_name = chain_definition.get('boogu_prompt_node', 'boogu_prompt')
+    vae_loader_name = chain_definition.get('vae_loader_node', 'vae_loader')
+
+    boogu_prompt_id = assembler.node_map.get(boogu_prompt_name)
+    if not boogu_prompt_id or boogu_prompt_id not in assembler.workflow:
+        for node_id, node in assembler.workflow.items():
+            if isinstance(node, dict) and node.get('class_type') == 'TextEncodeBooguEdit':
+                boogu_prompt_id = node_id
+                break
+
+    if not boogu_prompt_id:
+        print(f"Warning: Target node '{boogu_prompt_name}' (TextEncodeBooguEdit) for Boogu Edit chain not found. Skipping.")
+        return
+
+    vae_id = assembler.node_map.get(vae_loader_name)
+    if not vae_id:
+        for node_id, node in assembler.workflow.items():
+            if isinstance(node, dict) and node.get('class_type') == 'VAELoader':
+                vae_id = node_id
+                break
+
+    if vae_id:
+        assembler.workflow[boogu_prompt_id]['inputs']['vae'] = [vae_id, 0]
+
+    for i, img_filename in enumerate(valid_images):
+        load_id = assembler._get_unique_id()
+        load_node = create_node(assembler, "LoadImage", f"Load Reference Image {i+1}")
+        load_node['inputs']['image'] = img_filename
+        assembler.workflow[load_id] = load_node
+
+        scale_id = assembler._get_unique_id()
+        scale_node = create_node(assembler, "ImageScaleToTotalPixels", f"Scale Reference {i+1}")
+        scale_node['inputs']['upscale_method'] = "nearest-exact"
+        scale_node['inputs']['megapixels'] = 1
+        scale_node['inputs']['resolution_steps'] = 1
+        scale_node['inputs']['image'] = [load_id, 0]
+        assembler.workflow[scale_id] = scale_node
+
+        image_key = f"images.image_{i+1}"
+        assembler.workflow[boogu_prompt_id]['inputs'][image_key] = [scale_id, 0]
+
+    print(f"Boogu Edit injector applied with {len(valid_images)} reference image(s). Connected VAE dynamically.")

chain_injectors/joyai_image_injector.py

@@ -1,63 +1,63 @@
-import os
-
-def inject(assembler, chain_definition, chain_items):
-    if not chain_items:
-        return
-
-    valid_images = []
-    for item in chain_items:
-        if not item:
-            continue
-        img_path = item
-        if isinstance(item, dict):
-            img_path = item.get('image') or item.get('filename') or item.get('path')
-        if img_path:
-            valid_images.append(img_path)
-
-    if not valid_images:
-        return
-
-    valid_images = valid_images[:2]
-
-    pos_prompt_name = chain_definition.get('pos_prompt_node', 'pos_prompt')
-    neg_prompt_name = chain_definition.get('neg_prompt_node', 'neg_prompt')
-    vae_node_name = chain_definition.get('vae_node', 'vae_loader')
-
-    if pos_prompt_name not in assembler.node_map:
-        print(f"Warning: Positive prompt node '{pos_prompt_name}' not found for JoyAI Reference chain. Skipping.")
-        return
-
-    if vae_node_name not in assembler.node_map:
-        print(f"Warning: VAE loader node '{vae_node_name}' not found for JoyAI Reference chain. Skipping.")
-        return
-
-    pos_prompt_id = assembler.node_map[pos_prompt_name]
-    neg_prompt_id = assembler.node_map.get(neg_prompt_name)
-    vae_node_id = assembler.node_map[vae_node_name]
-
-    assembler.workflow[pos_prompt_id]['inputs']['vae'] = [vae_node_id, 0]
-    if neg_prompt_id and neg_prompt_id in assembler.workflow:
-        assembler.workflow[neg_prompt_id]['inputs']['vae'] = [vae_node_id, 0]
-
-    for i, img_filename in enumerate(valid_images):
-        load_id = assembler._get_unique_id()
-        load_node = assembler._get_node_template("LoadImage")
-        load_node['inputs']['image'] = img_filename
-        load_node['_meta']['title'] = f"Load Reference Image {i+1}"
-        assembler.workflow[load_id] = load_node
-
-        scale_id = assembler._get_unique_id()
-        scale_node = assembler._get_node_template("ImageScaleToTotalPixels")
-        scale_node['inputs']['megapixels'] = 1.0
-        scale_node['inputs']['upscale_method'] = "nearest-exact"
-        scale_node['inputs']['resolution_steps'] = 1
-        scale_node['inputs']['image'] = [load_id, 0]
-        scale_node['_meta']['title'] = f"Scale Reference {i+1}"
-        assembler.workflow[scale_id] = scale_node
-
-        input_key = f"images.image{i}"
-        assembler.workflow[pos_prompt_id]['inputs'][input_key] = [scale_id, 0]
-        if neg_prompt_id and neg_prompt_id in assembler.workflow:
-            assembler.workflow[neg_prompt_id]['inputs'][input_key] = [scale_id, 0]
-
-    print(f"JoyAI Reference injector applied. Injected {len(valid_images)} reference images to JoyAI text encoding nodes.")
+import os
+
+def inject(assembler, chain_definition, chain_items):
+    if not chain_items:
+        return
+
+    valid_images = []
+    for item in chain_items:
+        if not item:
+            continue
+        img_path = item
+        if isinstance(item, dict):
+            img_path = item.get('image') or item.get('filename') or item.get('path')
+        if img_path:
+            valid_images.append(img_path)
+
+    if not valid_images:
+        return
+
+    valid_images = valid_images[:2]
+
+    pos_prompt_name = chain_definition.get('pos_prompt_node', 'pos_prompt')
+    neg_prompt_name = chain_definition.get('neg_prompt_node', 'neg_prompt')
+    vae_node_name = chain_definition.get('vae_node', 'vae_loader')
+
+    if pos_prompt_name not in assembler.node_map:
+        print(f"Warning: Positive prompt node '{pos_prompt_name}' not found for JoyAI Reference chain. Skipping.")
+        return
+
+    if vae_node_name not in assembler.node_map:
+        print(f"Warning: VAE loader node '{vae_node_name}' not found for JoyAI Reference chain. Skipping.")
+        return
+
+    pos_prompt_id = assembler.node_map[pos_prompt_name]
+    neg_prompt_id = assembler.node_map.get(neg_prompt_name)
+    vae_node_id = assembler.node_map[vae_node_name]
+
+    assembler.workflow[pos_prompt_id]['inputs']['vae'] = [vae_node_id, 0]
+    if neg_prompt_id and neg_prompt_id in assembler.workflow:
+        assembler.workflow[neg_prompt_id]['inputs']['vae'] = [vae_node_id, 0]
+
+    for i, img_filename in enumerate(valid_images):
+        load_id = assembler._get_unique_id()
+        load_node = assembler._get_node_template("LoadImage")
+        load_node['inputs']['image'] = img_filename
+        load_node['_meta']['title'] = f"Load Reference Image {i+1}"
+        assembler.workflow[load_id] = load_node
+
+        scale_id = assembler._get_unique_id()
+        scale_node = assembler._get_node_template("ImageScaleToTotalPixels")
+        scale_node['inputs']['megapixels'] = 1.0
+        scale_node['inputs']['upscale_method'] = "nearest-exact"
+        scale_node['inputs']['resolution_steps'] = 1
+        scale_node['inputs']['image'] = [load_id, 0]
+        scale_node['_meta']['title'] = f"Scale Reference {i+1}"
+        assembler.workflow[scale_id] = scale_node
+
+        input_key = f"images.image{i}"
+        assembler.workflow[pos_prompt_id]['inputs'][input_key] = [scale_id, 0]
+        if neg_prompt_id and neg_prompt_id in assembler.workflow:
+            assembler.workflow[neg_prompt_id]['inputs'][input_key] = [scale_id, 0]
+
+    print(f"JoyAI Reference injector applied. Injected {len(valid_images)} reference images to JoyAI text encoding nodes.")

chain_injectors/krea2_identity_edit_injector.py

@@ -1,173 +1,173 @@
-import os
-from utils.app_utils import ensure_file_downloaded
-
-def inject(assembler, chain_definition, chain_items):
-    if not chain_items:
-        return
-
-    valid_images = []
-    for item in chain_items:
-        if not item:
-            continue
-        img_path = item
-        if isinstance(item, dict):
-            img_path = item.get('image') or item.get('filename') or item.get('path')
-        if img_path:
-            valid_images.append(img_path)
-
-    if not valid_images:
-        return
-
-    valid_images = valid_images[:2]
-
-    lora_filename = "krea2_identity_edit_v1_2.safetensors"
-    try:
-        ensure_file_downloaded(lora_filename)
-    except Exception as e:
-        print(f"Warning: Failed to ensure '{lora_filename}' downloaded: {e}")
-
-    ksampler_name = chain_definition.get('ksampler_node', 'ksampler')
-    pos_prompt_name = chain_definition.get('pos_prompt_node', 'pos_prompt')
-    neg_prompt_name = chain_definition.get('neg_prompt_node', 'neg_prompt')
-    clip_loader_name = chain_definition.get('clip_loader_node', 'clip_loader')
-    vae_loader_name = chain_definition.get('vae_loader_node', 'vae_loader')
-
-    if ksampler_name not in assembler.node_map:
-        print(f"Warning: Target node '{ksampler_name}' for Krea2 Identity Edit chain not found. Skipping.")
-        return
-
-    ksampler_id = assembler.node_map[ksampler_name]
-
-    if 'model' not in assembler.workflow[ksampler_id]['inputs']:
-        print(f"Warning: KSampler node '{ksampler_name}' is missing 'model' input. Skipping.")
-        return
-
-    latent_connection = assembler.workflow[ksampler_id]['inputs'].get('latent_image')
-    if not latent_connection:
-        print(f"Warning: KSampler node '{ksampler_name}' is missing 'latent_image' input. Skipping.")
-        return
-
-    current_model_connection = assembler.workflow[ksampler_id]['inputs']['model']
-
-    vae_connection = None
-    if vae_loader_name in assembler.node_map:
-        vae_connection = [assembler.node_map[vae_loader_name], 0]
-
-    clip_connection = None
-    if clip_loader_name in assembler.node_map:
-        clip_connection = [assembler.node_map[clip_loader_name], 0]
-    elif pos_prompt_name in assembler.node_map:
-        pos_id = assembler.node_map[pos_prompt_name]
-        clip_connection = assembler.workflow[pos_id]['inputs'].get('clip')
-
-    lora_loader_id = assembler._get_unique_id()
-    lora_loader_node = assembler._get_node_template("LoraLoaderModelOnly")
-    lora_loader_node['inputs']['lora_name'] = lora_filename
-    lora_loader_node['inputs']['strength_model'] = 1.0
-    lora_loader_node['inputs']['model'] = current_model_connection
-    lora_loader_node['_meta']['title'] = "Load LoRA (Krea2 Identity Edit)"
-    assembler.workflow[lora_loader_id] = lora_loader_node
-
-    image_ids = []
-    vae_encode_ids = []
-
-    for i, img_filename in enumerate(valid_images):
-        load_id = assembler._get_unique_id()
-        load_node = assembler._get_node_template("LoadImage")
-        load_node['inputs']['image'] = img_filename
-        load_node['_meta']['title'] = f"Load Image (Ref {i+1})"
-        assembler.workflow[load_id] = load_node
-        image_ids.append(load_id)
-
-        vae_enc_id = assembler._get_unique_id()
-        vae_enc_node = assembler._get_node_template("VAEEncode")
-        vae_enc_node['inputs']['pixels'] = [load_id, 0]
-        if vae_connection:
-            vae_enc_node['inputs']['vae'] = vae_connection
-        vae_enc_node['_meta']['title'] = f"VAE Encode (Ref {i+1})"
-        assembler.workflow[vae_enc_id] = vae_enc_node
-        vae_encode_ids.append(vae_enc_id)
-
-    patch_id = assembler._get_unique_id()
-    patch_node = assembler._get_node_template("Krea2EditModelPatch")
-    patch_node['inputs']['ref_boost'] = 4
-    patch_node['inputs']['ref_boost_a'] = 1
-    patch_node['inputs']['fit_mode'] = "fit"
-    patch_node['inputs']['model'] = [lora_loader_id, 0]
-    patch_node['inputs']['source_latent'] = [vae_encode_ids[0], 0]
-    if vae_connection:
-        patch_node['inputs']['vae'] = vae_connection
-    patch_node['inputs']['source_image'] = [image_ids[0], 0]
-    patch_node['inputs']['target_latent'] = latent_connection
-
-    if len(valid_images) > 1:
-        patch_node['inputs']['source_latent_b'] = [vae_encode_ids[1], 0]
-        patch_node['inputs']['source_image_b'] = [image_ids[1], 0]
-
-    patch_node['_meta']['title'] = "Krea2 Edit (source patch)"
-    assembler.workflow[patch_id] = patch_node
-
-    assembler.workflow[ksampler_id]['inputs']['model'] = [patch_id, 0]
-
-    pos_prompt_id = assembler.node_map.get(pos_prompt_name)
-    neg_prompt_id = assembler.node_map.get(neg_prompt_name)
-
-    pos_text = ""
-    if pos_prompt_id and pos_prompt_id in assembler.workflow:
-        pos_text = assembler.workflow[pos_prompt_id]['inputs'].get('text', '')
-    elif hasattr(assembler, 'ui_values') and isinstance(assembler.ui_values, dict):
-        pos_text = assembler.ui_values.get('positive_prompt') or assembler.ui_values.get('prompt') or ''
-
-    if not pos_text:
-        for node_id, node in assembler.workflow.items():
-            if isinstance(node, dict):
-                cls = node.get('class_type', '')
-                if cls in ['Krea2EditGroundedEncode', 'TextEncodeQwenImageEditPlus', 'CLIPTextEncode']:
-                    t = node.get('inputs', {}).get('prompt') or node.get('inputs', {}).get('text')
-                    if t:
-                        pos_text = t
-                        break
-
-    neg_text = ""
-    if neg_prompt_id and neg_prompt_id in assembler.workflow:
-        neg_text = assembler.workflow[neg_prompt_id]['inputs'].get('text', '')
-    elif hasattr(assembler, 'ui_values') and isinstance(assembler.ui_values, dict):
-        neg_text = assembler.ui_values.get('negative_prompt') or assembler.ui_values.get('neg_prompt') or ''
-
-    pos_grounded_id = assembler._get_unique_id()
-    pos_grounded_node = assembler._get_node_template("Krea2EditGroundedEncode")
-    pos_grounded_node['inputs']['prompt'] = pos_text
-    pos_grounded_node['inputs']['grounding_px'] = 768
-    pos_grounded_node['inputs']['system_prompt'] = ""
-    if clip_connection:
-        pos_grounded_node['inputs']['clip'] = clip_connection
-    pos_grounded_node['inputs']['image'] = [image_ids[0], 0]
-    if len(valid_images) > 1:
-        pos_grounded_node['inputs']['image_b'] = [image_ids[1], 0]
-    pos_grounded_node['_meta']['title'] = "Krea2 Edit (grounded encode positive)"
-    assembler.workflow[pos_grounded_id] = pos_grounded_node
-
-    assembler.workflow[ksampler_id]['inputs']['positive'] = [pos_grounded_id, 0]
-
-    neg_grounded_id = assembler._get_unique_id()
-    neg_grounded_node = assembler._get_node_template("Krea2EditGroundedEncode")
-    neg_grounded_node['inputs']['prompt'] = neg_text
-    neg_grounded_node['inputs']['grounding_px'] = 768
-    neg_grounded_node['inputs']['system_prompt'] = ""
-    if clip_connection:
-        neg_grounded_node['inputs']['clip'] = clip_connection
-    neg_grounded_node['inputs']['image'] = [image_ids[0], 0]
-    if len(valid_images) > 1:
-        neg_grounded_node['inputs']['image_b'] = [image_ids[1], 0]
-    neg_grounded_node['_meta']['title'] = "Krea2 Edit (grounded encode negative)"
-    assembler.workflow[neg_grounded_id] = neg_grounded_node
-
-    assembler.workflow[ksampler_id]['inputs']['negative'] = [neg_grounded_id, 0]
-
-    if pos_prompt_id and pos_prompt_id in assembler.workflow:
-        del assembler.workflow[pos_prompt_id]
-
-    if neg_prompt_id and neg_prompt_id in assembler.workflow:
-        del assembler.workflow[neg_prompt_id]
-
-    print(f"Krea2 Identity Edit injector applied with {len(valid_images)} reference image(s). Original CLIPTextEncode nodes removed.")
+import os
+from utils.app_utils import ensure_file_downloaded
+
+def inject(assembler, chain_definition, chain_items):
+    if not chain_items:
+        return
+
+    valid_images = []
+    for item in chain_items:
+        if not item:
+            continue
+        img_path = item
+        if isinstance(item, dict):
+            img_path = item.get('image') or item.get('filename') or item.get('path')
+        if img_path:
+            valid_images.append(img_path)
+
+    if not valid_images:
+        return
+
+    valid_images = valid_images[:2]
+
+    lora_filename = "krea2_identity_edit_v1_2.safetensors"
+    try:
+        ensure_file_downloaded(lora_filename)
+    except Exception as e:
+        print(f"Warning: Failed to ensure '{lora_filename}' downloaded: {e}")
+
+    ksampler_name = chain_definition.get('ksampler_node', 'ksampler')
+    pos_prompt_name = chain_definition.get('pos_prompt_node', 'pos_prompt')
+    neg_prompt_name = chain_definition.get('neg_prompt_node', 'neg_prompt')
+    clip_loader_name = chain_definition.get('clip_loader_node', 'clip_loader')
+    vae_loader_name = chain_definition.get('vae_loader_node', 'vae_loader')
+
+    if ksampler_name not in assembler.node_map:
+        print(f"Warning: Target node '{ksampler_name}' for Krea2 Identity Edit chain not found. Skipping.")
+        return
+
+    ksampler_id = assembler.node_map[ksampler_name]
+
+    if 'model' not in assembler.workflow[ksampler_id]['inputs']:
+        print(f"Warning: KSampler node '{ksampler_name}' is missing 'model' input. Skipping.")
+        return
+
+    latent_connection = assembler.workflow[ksampler_id]['inputs'].get('latent_image')
+    if not latent_connection:
+        print(f"Warning: KSampler node '{ksampler_name}' is missing 'latent_image' input. Skipping.")
+        return
+
+    current_model_connection = assembler.workflow[ksampler_id]['inputs']['model']
+
+    vae_connection = None
+    if vae_loader_name in assembler.node_map:
+        vae_connection = [assembler.node_map[vae_loader_name], 0]
+
+    clip_connection = None
+    if clip_loader_name in assembler.node_map:
+        clip_connection = [assembler.node_map[clip_loader_name], 0]
+    elif pos_prompt_name in assembler.node_map:
+        pos_id = assembler.node_map[pos_prompt_name]
+        clip_connection = assembler.workflow[pos_id]['inputs'].get('clip')
+
+    lora_loader_id = assembler._get_unique_id()
+    lora_loader_node = assembler._get_node_template("LoraLoaderModelOnly")
+    lora_loader_node['inputs']['lora_name'] = lora_filename
+    lora_loader_node['inputs']['strength_model'] = 1.0
+    lora_loader_node['inputs']['model'] = current_model_connection
+    lora_loader_node['_meta']['title'] = "Load LoRA (Krea2 Identity Edit)"
+    assembler.workflow[lora_loader_id] = lora_loader_node
+
+    image_ids = []
+    vae_encode_ids = []
+
+    for i, img_filename in enumerate(valid_images):
+        load_id = assembler._get_unique_id()
+        load_node = assembler._get_node_template("LoadImage")
+        load_node['inputs']['image'] = img_filename
+        load_node['_meta']['title'] = f"Load Image (Ref {i+1})"
+        assembler.workflow[load_id] = load_node
+        image_ids.append(load_id)
+
+        vae_enc_id = assembler._get_unique_id()
+        vae_enc_node = assembler._get_node_template("VAEEncode")
+        vae_enc_node['inputs']['pixels'] = [load_id, 0]
+        if vae_connection:
+            vae_enc_node['inputs']['vae'] = vae_connection
+        vae_enc_node['_meta']['title'] = f"VAE Encode (Ref {i+1})"
+        assembler.workflow[vae_enc_id] = vae_enc_node
+        vae_encode_ids.append(vae_enc_id)
+
+    patch_id = assembler._get_unique_id()
+    patch_node = assembler._get_node_template("Krea2EditModelPatch")
+    patch_node['inputs']['ref_boost'] = 4
+    patch_node['inputs']['ref_boost_a'] = 1
+    patch_node['inputs']['fit_mode'] = "fit"
+    patch_node['inputs']['model'] = [lora_loader_id, 0]
+    patch_node['inputs']['source_latent'] = [vae_encode_ids[0], 0]
+    if vae_connection:
+        patch_node['inputs']['vae'] = vae_connection
+    patch_node['inputs']['source_image'] = [image_ids[0], 0]
+    patch_node['inputs']['target_latent'] = latent_connection
+
+    if len(valid_images) > 1:
+        patch_node['inputs']['source_latent_b'] = [vae_encode_ids[1], 0]
+        patch_node['inputs']['source_image_b'] = [image_ids[1], 0]
+
+    patch_node['_meta']['title'] = "Krea2 Edit (source patch)"
+    assembler.workflow[patch_id] = patch_node
+
+    assembler.workflow[ksampler_id]['inputs']['model'] = [patch_id, 0]
+
+    pos_prompt_id = assembler.node_map.get(pos_prompt_name)
+    neg_prompt_id = assembler.node_map.get(neg_prompt_name)
+
+    pos_text = ""
+    if pos_prompt_id and pos_prompt_id in assembler.workflow:
+        pos_text = assembler.workflow[pos_prompt_id]['inputs'].get('text', '')
+    elif hasattr(assembler, 'ui_values') and isinstance(assembler.ui_values, dict):
+        pos_text = assembler.ui_values.get('positive_prompt') or assembler.ui_values.get('prompt') or ''
+
+    if not pos_text:
+        for node_id, node in assembler.workflow.items():
+            if isinstance(node, dict):
+                cls = node.get('class_type', '')
+                if cls in ['Krea2EditGroundedEncode', 'TextEncodeQwenImageEditPlus', 'CLIPTextEncode']:
+                    t = node.get('inputs', {}).get('prompt') or node.get('inputs', {}).get('text')
+                    if t:
+                        pos_text = t
+                        break
+
+    neg_text = ""
+    if neg_prompt_id and neg_prompt_id in assembler.workflow:
+        neg_text = assembler.workflow[neg_prompt_id]['inputs'].get('text', '')
+    elif hasattr(assembler, 'ui_values') and isinstance(assembler.ui_values, dict):
+        neg_text = assembler.ui_values.get('negative_prompt') or assembler.ui_values.get('neg_prompt') or ''
+
+    pos_grounded_id = assembler._get_unique_id()
+    pos_grounded_node = assembler._get_node_template("Krea2EditGroundedEncode")
+    pos_grounded_node['inputs']['prompt'] = pos_text
+    pos_grounded_node['inputs']['grounding_px'] = 768
+    pos_grounded_node['inputs']['system_prompt'] = ""
+    if clip_connection:
+        pos_grounded_node['inputs']['clip'] = clip_connection
+    pos_grounded_node['inputs']['image'] = [image_ids[0], 0]
+    if len(valid_images) > 1:
+        pos_grounded_node['inputs']['image_b'] = [image_ids[1], 0]
+    pos_grounded_node['_meta']['title'] = "Krea2 Edit (grounded encode positive)"
+    assembler.workflow[pos_grounded_id] = pos_grounded_node
+
+    assembler.workflow[ksampler_id]['inputs']['positive'] = [pos_grounded_id, 0]
+
+    neg_grounded_id = assembler._get_unique_id()
+    neg_grounded_node = assembler._get_node_template("Krea2EditGroundedEncode")
+    neg_grounded_node['inputs']['prompt'] = neg_text
+    neg_grounded_node['inputs']['grounding_px'] = 768
+    neg_grounded_node['inputs']['system_prompt'] = ""
+    if clip_connection:
+        neg_grounded_node['inputs']['clip'] = clip_connection
+    neg_grounded_node['inputs']['image'] = [image_ids[0], 0]
+    if len(valid_images) > 1:
+        neg_grounded_node['inputs']['image_b'] = [image_ids[1], 0]
+    neg_grounded_node['_meta']['title'] = "Krea2 Edit (grounded encode negative)"
+    assembler.workflow[neg_grounded_id] = neg_grounded_node
+
+    assembler.workflow[ksampler_id]['inputs']['negative'] = [neg_grounded_id, 0]
+
+    if pos_prompt_id and pos_prompt_id in assembler.workflow:
+        del assembler.workflow[pos_prompt_id]
+
+    if neg_prompt_id and neg_prompt_id in assembler.workflow:
+        del assembler.workflow[neg_prompt_id]
+
+    print(f"Krea2 Identity Edit injector applied with {len(valid_images)} reference image(s). Original CLIPTextEncode nodes removed.")

chain_injectors/krea2_style_reference_injector.py

@@ -1,201 +1,168 @@
-import os
-from utils.app_utils import ensure_file_downloaded
-
-def create_node(assembler, class_type, title):
-    try:
-        node = assembler._get_node_template(class_type)
-    except Exception:
-        node = {
-            "inputs": {},
-            "class_type": class_type,
-            "_meta": {"title": title}
-        }
-    node['_meta']['title'] = title
-    return node
-
-def inject(assembler, chain_definition, chain_items):
-    if not chain_items:
-        return
-
-    valid_images = []
-    for item in chain_items:
-        if not item:
-            continue
-        img_path = item
-        if isinstance(item, dict):
-            img_path = item.get('image') or item.get('filename') or item.get('path')
-        if img_path:
-            valid_images.append(img_path)
-
-    if not valid_images:
-        return
-
-    valid_images = valid_images[:3]
-
-    lora_filename = "krea2_style_reference.safetensors"
-    try:
-        ensure_file_downloaded(lora_filename)
-    except Exception as e:
-        print(f"Warning: Failed to ensure '{lora_filename}' downloaded: {e}")
-
-    ksampler_name = chain_definition.get('ksampler_node', 'ksampler')
-    pos_prompt_name = chain_definition.get('pos_prompt_node', 'pos_prompt')
-    neg_prompt_name = chain_definition.get('neg_prompt_node', 'neg_prompt')
-    clip_loader_name = chain_definition.get('clip_loader_node', 'clip_loader')
-    vae_loader_name = chain_definition.get('vae_loader_node', 'vae_loader')
-
-    if ksampler_name not in assembler.node_map:
-        print(f"Warning: Target node '{ksampler_name}' for Krea2 Style Reference Edit chain not found. Skipping.")
-        return
-
-    ksampler_id = assembler.node_map[ksampler_name]
-
-    if 'model' not in assembler.workflow[ksampler_id]['inputs']:
-        print(f"Warning: KSampler node '{ksampler_name}' is missing 'model' input. Skipping.")
-        return
-
-    current_model_connection = assembler.workflow[ksampler_id]['inputs']['model']
-
-    vae_connection = None
-    if vae_loader_name in assembler.node_map:
-        vae_connection = [assembler.node_map[vae_loader_name], 0]
-    else:
-        for node_id, node in assembler.workflow.items():
-            if isinstance(node, dict) and node.get('class_type') == 'VAELoader':
-                vae_connection = [node_id, 0]
-                break
-
-    clip_connection = None
-    if clip_loader_name in assembler.node_map:
-        clip_connection = [assembler.node_map[clip_loader_name], 0]
-    elif pos_prompt_name in assembler.node_map:
-        pos_id = assembler.node_map[pos_prompt_name]
-        clip_connection = assembler.workflow[pos_id]['inputs'].get('clip')
-
-    scaled_image_ids = []
-    for i, img_filename in enumerate(valid_images):
-        load_id = assembler._get_unique_id()
-        load_node = create_node(assembler, "LoadImage", f"Load Reference Image {i+1}")
-        load_node['inputs']['image'] = img_filename
-        assembler.workflow[load_id] = load_node
-
-        scale_id = assembler._get_unique_id()
-        scale_node = create_node(assembler, "ImageScaleToTotalPixels", f"Scale Reference {i+1}")
-        scale_node['inputs']['upscale_method'] = "nearest-exact"
-        scale_node['inputs']['megapixels'] = 1
-        scale_node['inputs']['resolution_steps'] = 1
-        scale_node['inputs']['image'] = [load_id, 0]
-        assembler.workflow[scale_id] = scale_node
-        scaled_image_ids.append(scale_id)
-
-    lora_loader_id = assembler._get_unique_id()
-    lora_loader_node = create_node(assembler, "LoraLoaderModelOnly", "Load LoRA (Krea2 Style Reference)")
-    lora_loader_node['inputs']['lora_name'] = lora_filename
-    lora_loader_node['inputs']['strength_model'] = 1.0
-    lora_loader_node['inputs']['model'] = current_model_connection
-    assembler.workflow[lora_loader_id] = lora_loader_node
-
-    assembler.workflow[ksampler_id]['inputs']['model'] = [lora_loader_id, 0]
-
-    pos_prompt_id = assembler.node_map.get(pos_prompt_name)
-    neg_prompt_id = assembler.node_map.get(neg_prompt_name)
-
-    pos_text = ""
-    if pos_prompt_id and pos_prompt_id in assembler.workflow:
-        pos_text = assembler.workflow[pos_prompt_id]['inputs'].get('text', '')
-    elif hasattr(assembler, 'ui_values') and isinstance(assembler.ui_values, dict):
-        pos_text = assembler.ui_values.get('positive_prompt') or assembler.ui_values.get('prompt') or ''
-
-    if not pos_text:
-        for node_id, node in assembler.workflow.items():
-            if isinstance(node, dict):
-                cls = node.get('class_type', '')
-                if cls in ['Krea2EditGroundedEncode', 'TextEncodeQwenImageEditPlus', 'CLIPTextEncode']:
-                    t = node.get('inputs', {}).get('prompt') or node.get('inputs', {}).get('text')
-                    if t:
-                        pos_text = t
-                        break
-
-    neg_text = ""
-    if neg_prompt_id and neg_prompt_id in assembler.workflow:
-        neg_text = assembler.workflow[neg_prompt_id]['inputs'].get('text', '')
-    elif hasattr(assembler, 'ui_values') and isinstance(assembler.ui_values, dict):
-        neg_text = assembler.ui_values.get('negative_prompt') or assembler.ui_values.get('neg_prompt') or ''
-
-    pos_encode_id = assembler._get_unique_id()
-    pos_encode_node = create_node(assembler, "TextEncodeQwenImageEditPlus", "TextEncodeQwenImageEditPlus (Positive)")
-    pos_encode_node['inputs']['prompt'] = pos_text
-    if clip_connection:
-        pos_encode_node['inputs']['clip'] = clip_connection
-    if vae_connection:
-        pos_encode_node['inputs']['vae'] = vae_connection
-    for idx, s_id in enumerate(scaled_image_ids):
-        pos_encode_node['inputs'][f"image{idx+1}"] = [s_id, 0]
-    assembler.workflow[pos_encode_id] = pos_encode_node
-
-    neg_encode_id = assembler._get_unique_id()
-    neg_encode_node = create_node(assembler, "TextEncodeQwenImageEditPlus", "TextEncodeQwenImageEditPlus (Negative)")
-    neg_encode_node['inputs']['prompt'] = neg_text
-    if clip_connection:
-        neg_encode_node['inputs']['clip'] = clip_connection
-    if vae_connection:
-        neg_encode_node['inputs']['vae'] = vae_connection
-    for idx, s_id in enumerate(scaled_image_ids):
-        neg_encode_node['inputs'][f"image{idx+1}"] = [s_id, 0]
-    assembler.workflow[neg_encode_id] = neg_encode_node
-
-    pos_ref_id = assembler._get_unique_id()
-    pos_ref_node = create_node(assembler, "FluxKontextMultiReferenceLatentMethod", "Edit Model Reference Method")
-    pos_ref_node['inputs']['reference_latents_method'] = "index_timestep_zero"
-    pos_ref_node['inputs']['conditioning'] = [pos_encode_id, 0]
-    assembler.workflow[pos_ref_id] = pos_ref_node
-
-    neg_ref_id = assembler._get_unique_id()
-    neg_ref_node = create_node(assembler, "FluxKontextMultiReferenceLatentMethod", "Edit Model Reference Method")
-    neg_ref_node['inputs']['reference_latents_method'] = "index_timestep_zero"
-    neg_ref_node['inputs']['conditioning'] = [neg_encode_id, 0]
-    assembler.workflow[neg_ref_id] = neg_ref_node
-
-    existing_pos = assembler.workflow[ksampler_id]['inputs'].get('positive')
-    existing_neg = assembler.workflow[ksampler_id]['inputs'].get('negative')
-
-    has_krea2_edit = False
-    if existing_pos and isinstance(existing_pos, (list, tuple)) and len(existing_pos) > 0:
-        pos_node_id = existing_pos[0]
-        if pos_node_id in assembler.workflow:
-            pos_node = assembler.workflow[pos_node_id]
-            if isinstance(pos_node, dict) and pos_node.get('class_type') == 'Krea2EditGroundedEncode':
-                has_krea2_edit = True
-
-    if not has_krea2_edit:
-        for node in assembler.workflow.values():
-            if isinstance(node, dict) and node.get('class_type') in ['Krea2EditModelPatch', 'Krea2EditGroundedEncode']:
-                has_krea2_edit = True
-                break
-
-    if has_krea2_edit and existing_pos and existing_neg:
-        combine_pos_id = assembler._get_unique_id()
-        combine_pos_node = create_node(assembler, "ConditioningCombine", "Conditioning (Combine)")
-        combine_pos_node['inputs']['conditioning_1'] = existing_pos
-        combine_pos_node['inputs']['conditioning_2'] = [pos_ref_id, 0]
-        assembler.workflow[combine_pos_id] = combine_pos_node
-
-        combine_neg_id = assembler._get_unique_id()
-        combine_neg_node = create_node(assembler, "ConditioningCombine", "Conditioning (Combine)")
-        combine_neg_node['inputs']['conditioning_1'] = existing_neg
-        combine_neg_node['inputs']['conditioning_2'] = [neg_ref_id, 0]
-        assembler.workflow[combine_neg_id] = combine_neg_node
-
-        assembler.workflow[ksampler_id]['inputs']['positive'] = [combine_pos_id, 0]
-        assembler.workflow[ksampler_id]['inputs']['negative'] = [combine_neg_id, 0]
-    else:
-        assembler.workflow[ksampler_id]['inputs']['positive'] = [pos_ref_id, 0]
-        assembler.workflow[ksampler_id]['inputs']['negative'] = [neg_ref_id, 0]
-
-    if pos_prompt_id and pos_prompt_id in assembler.workflow:
-        del assembler.workflow[pos_prompt_id]
-
-    if neg_prompt_id and neg_prompt_id in assembler.workflow:
-        del assembler.workflow[neg_prompt_id]
-
-    print(f"Krea2 Style Reference Edit injector applied with {len(valid_images)} reference image(s). Original CLIPTextEncode nodes replaced.")
+import os
+from utils.app_utils import ensure_file_downloaded
+
+def create_node(assembler, class_type, title):
+    try:
+        node = assembler._get_node_template(class_type)
+    except Exception:
+        node = {
+            "inputs": {},
+            "class_type": class_type,
+            "_meta": {"title": title}
+        }
+    node['_meta']['title'] = title
+    return node
+
+def inject(assembler, chain_definition, chain_items):
+    if not chain_items:
+        return
+
+    valid_images = []
+    for item in chain_items:
+        if not item:
+            continue
+        img_path = item
+        if isinstance(item, dict):
+            img_path = item.get('image') or item.get('filename') or item.get('path')
+        if img_path:
+            valid_images.append(img_path)
+
+    if not valid_images:
+        return
+
+    valid_images = valid_images[:3]
+
+    lora_filename = "krea2_style_reference.safetensors"
+    try:
+        ensure_file_downloaded(lora_filename)
+    except Exception as e:
+        print(f"Warning: Failed to ensure '{lora_filename}' downloaded: {e}")
+
+    ksampler_name = chain_definition.get('ksampler_node', 'ksampler')
+    pos_prompt_name = chain_definition.get('pos_prompt_node', 'pos_prompt')
+    neg_prompt_name = chain_definition.get('neg_prompt_node', 'neg_prompt')
+    clip_loader_name = chain_definition.get('clip_loader_node', 'clip_loader')
+    vae_loader_name = chain_definition.get('vae_loader_node', 'vae_loader')
+
+    if ksampler_name not in assembler.node_map:
+        print(f"Warning: Target node '{ksampler_name}' for Krea2 Style Reference Edit chain not found. Skipping.")
+        return
+
+    ksampler_id = assembler.node_map[ksampler_name]
+
+    if 'model' not in assembler.workflow[ksampler_id]['inputs']:
+        print(f"Warning: KSampler node '{ksampler_name}' is missing 'model' input. Skipping.")
+        return
+
+    current_model_connection = assembler.workflow[ksampler_id]['inputs']['model']
+
+    vae_connection = None
+    if vae_loader_name in assembler.node_map:
+        vae_connection = [assembler.node_map[vae_loader_name], 0]
+    else:
+        for node_id, node in assembler.workflow.items():
+            if isinstance(node, dict) and node.get('class_type') == 'VAELoader':
+                vae_connection = [node_id, 0]
+                break
+
+    clip_connection = None
+    if clip_loader_name in assembler.node_map:
+        clip_connection = [assembler.node_map[clip_loader_name], 0]
+    elif pos_prompt_name in assembler.node_map:
+        pos_id = assembler.node_map[pos_prompt_name]
+        clip_connection = assembler.workflow[pos_id]['inputs'].get('clip')
+
+    scaled_image_ids = []
+    for i, img_filename in enumerate(valid_images):
+        load_id = assembler._get_unique_id()
+        load_node = create_node(assembler, "LoadImage", f"Load Reference Image {i+1}")
+        load_node['inputs']['image'] = img_filename
+        assembler.workflow[load_id] = load_node
+
+        scale_id = assembler._get_unique_id()
+        scale_node = create_node(assembler, "ImageScaleToTotalPixels", f"Scale Reference {i+1}")
+        scale_node['inputs']['upscale_method'] = "nearest-exact"
+        scale_node['inputs']['megapixels'] = 1
+        scale_node['inputs']['resolution_steps'] = 1
+        scale_node['inputs']['image'] = [load_id, 0]
+        assembler.workflow[scale_id] = scale_node
+        scaled_image_ids.append(scale_id)
+
+    lora_loader_id = assembler._get_unique_id()
+    lora_loader_node = create_node(assembler, "LoraLoaderModelOnly", "Load LoRA (Krea2 Style Reference)")
+    lora_loader_node['inputs']['lora_name'] = lora_filename
+    lora_loader_node['inputs']['strength_model'] = 1.0
+    lora_loader_node['inputs']['model'] = current_model_connection
+    assembler.workflow[lora_loader_id] = lora_loader_node
+
+    assembler.workflow[ksampler_id]['inputs']['model'] = [lora_loader_id, 0]
+
+    pos_prompt_id = assembler.node_map.get(pos_prompt_name)
+    neg_prompt_id = assembler.node_map.get(neg_prompt_name)
+
+    pos_text = ""
+    if pos_prompt_id and pos_prompt_id in assembler.workflow:
+        pos_text = assembler.workflow[pos_prompt_id]['inputs'].get('text', '')
+    elif hasattr(assembler, 'ui_values') and isinstance(assembler.ui_values, dict):
+        pos_text = assembler.ui_values.get('positive_prompt') or assembler.ui_values.get('prompt') or ''
+
+    if not pos_text:
+        for node_id, node in assembler.workflow.items():
+            if isinstance(node, dict):
+                cls = node.get('class_type', '')
+                if cls in ['Krea2EditGroundedEncode', 'TextEncodeQwenImageEditPlus', 'CLIPTextEncode']:
+                    t = node.get('inputs', {}).get('prompt') or node.get('inputs', {}).get('text')
+                    if t:
+                        pos_text = t
+                        break
+
+    neg_text = ""
+    if neg_prompt_id and neg_prompt_id in assembler.workflow:
+        neg_text = assembler.workflow[neg_prompt_id]['inputs'].get('text', '')
+    elif hasattr(assembler, 'ui_values') and isinstance(assembler.ui_values, dict):
+        neg_text = assembler.ui_values.get('negative_prompt') or assembler.ui_values.get('neg_prompt') or ''
+
+    pos_encode_id = assembler._get_unique_id()
+    pos_encode_node = create_node(assembler, "TextEncodeQwenImageEditPlus", "TextEncodeQwenImageEditPlus (Positive)")
+    pos_encode_node['inputs']['prompt'] = pos_text
+    if clip_connection:
+        pos_encode_node['inputs']['clip'] = clip_connection
+    if vae_connection:
+        pos_encode_node['inputs']['vae'] = vae_connection
+    for idx, s_id in enumerate(scaled_image_ids):
+        pos_encode_node['inputs'][f"image{idx+1}"] = [s_id, 0]
+    assembler.workflow[pos_encode_id] = pos_encode_node
+
+    neg_encode_id = assembler._get_unique_id()
+    neg_encode_node = create_node(assembler, "TextEncodeQwenImageEditPlus", "TextEncodeQwenImageEditPlus (Negative)")
+    neg_encode_node['inputs']['prompt'] = neg_text
+    if clip_connection:
+        neg_encode_node['inputs']['clip'] = clip_connection
+    if vae_connection:
+        neg_encode_node['inputs']['vae'] = vae_connection
+    for idx, s_id in enumerate(scaled_image_ids):
+        neg_encode_node['inputs'][f"image{idx+1}"] = [s_id, 0]
+    assembler.workflow[neg_encode_id] = neg_encode_node
+
+    pos_ref_id = assembler._get_unique_id()
+    pos_ref_node = create_node(assembler, "FluxKontextMultiReferenceLatentMethod", "Edit Model Reference Method")
+    pos_ref_node['inputs']['reference_latents_method'] = "index_timestep_zero"
+    pos_ref_node['inputs']['conditioning'] = [pos_encode_id, 0]
+    assembler.workflow[pos_ref_id] = pos_ref_node
+
+    neg_ref_id = assembler._get_unique_id()
+    neg_ref_node = create_node(assembler, "FluxKontextMultiReferenceLatentMethod", "Edit Model Reference Method")
+    neg_ref_node['inputs']['reference_latents_method'] = "index_timestep_zero"
+    neg_ref_node['inputs']['conditioning'] = [neg_encode_id, 0]
+    assembler.workflow[neg_ref_id] = neg_ref_node
+
+    assembler.workflow[ksampler_id]['inputs']['positive'] = [pos_ref_id, 0]
+    assembler.workflow[ksampler_id]['inputs']['negative'] = [neg_ref_id, 0]
+
+    if pos_prompt_id and pos_prompt_id in assembler.workflow:
+        del assembler.workflow[pos_prompt_id]
+
+    if neg_prompt_id and neg_prompt_id in assembler.workflow:
+        del assembler.workflow[neg_prompt_id]
+
+    print(f"Krea2 Style Reference Edit injector applied with {len(valid_images)} reference image(s). Original CLIPTextEncode nodes replaced.")

chain_injectors/qwen_image_edit_injector.py

@@ -1,113 +1,113 @@
-def create_node(assembler, class_type, title):
-    try:
-        node = assembler._get_node_template(class_type)
-    except Exception:
-        node = {
-            "inputs": {},
-            "class_type": class_type,
-            "_meta": {"title": title}
-        }
-    node['_meta']['title'] = title
-    return node
-
-def inject(assembler, chain_definition, chain_items):
-    if not chain_items:
-        return
-
-    valid_images = []
-    for item in chain_items:
-        if not item:
-            continue
-        img_path = item
-        if isinstance(item, dict):
-            img_path = item.get('image') or item.get('filename') or item.get('path')
-        if img_path:
-            valid_images.append(img_path)
-
-    if not valid_images:
-        return
-
-    valid_images = valid_images[:3]
-
-    ksampler_name = chain_definition.get('ksampler_node', 'ksampler')
-    pos_prompt_name = chain_definition.get('pos_prompt_node', 'pos_prompt')
-    neg_prompt_name = chain_definition.get('neg_prompt_node', 'neg_prompt')
-    vae_loader_name = chain_definition.get('vae_loader_node', 'vae_loader')
-    model_sampler_name = chain_definition.get('model_sampler_node', 'model_sampler')
-
-    if ksampler_name not in assembler.node_map:
-        print(f"Warning: Target node '{ksampler_name}' for Qwen-Image Edit chain not found. Skipping.")
-        return
-
-    ksampler_id = assembler.node_map[ksampler_name]
-    pos_prompt_id = assembler.node_map.get(pos_prompt_name)
-    neg_prompt_id = assembler.node_map.get(neg_prompt_name)
-
-    if not pos_prompt_id or not neg_prompt_id:
-        print("Warning: Positive or negative prompt node not found for Qwen-Image Edit chain. Skipping.")
-        return
-
-    vae_id = assembler.node_map.get(vae_loader_name)
-    if not vae_id:
-        for node_id, node in assembler.workflow.items():
-            if isinstance(node, dict) and node.get('class_type') == 'VAELoader':
-                vae_id = node_id
-                break
-
-    if vae_id:
-        assembler.workflow[pos_prompt_id]['inputs']['vae'] = [vae_id, 0]
-        assembler.workflow[neg_prompt_id]['inputs']['vae'] = [vae_id, 0]
-
-    for i, img_filename in enumerate(valid_images):
-        load_id = assembler._get_unique_id()
-        load_node = create_node(assembler, "LoadImage", f"Load Reference Image {i+1}")
-        load_node['inputs']['image'] = img_filename
-        assembler.workflow[load_id] = load_node
-
-        scale_id = assembler._get_unique_id()
-        scale_node = create_node(assembler, "ImageScaleToTotalPixels", f"Scale Reference {i+1}")
-        scale_node['inputs']['upscale_method'] = "lanczos"
-        scale_node['inputs']['megapixels'] = 1
-        scale_node['inputs']['resolution_steps'] = 1
-        scale_node['inputs']['image'] = [load_id, 0]
-        assembler.workflow[scale_id] = scale_node
-
-        image_key = f"image{i+1}"
-        assembler.workflow[pos_prompt_id]['inputs'][image_key] = [scale_id, 0]
-        assembler.workflow[neg_prompt_id]['inputs'][image_key] = [scale_id, 0]
-
-    pos_ref_id = assembler._get_unique_id()
-    pos_ref_node = create_node(assembler, "FluxKontextMultiReferenceLatentMethod", "Edit Model Reference Method")
-    pos_ref_node['inputs']['reference_latents_method'] = "index_timestep_zero"
-    pos_ref_node['inputs']['conditioning'] = [pos_prompt_id, 0]
-    assembler.workflow[pos_ref_id] = pos_ref_node
-
-    neg_ref_id = assembler._get_unique_id()
-    neg_ref_node = create_node(assembler, "FluxKontextMultiReferenceLatentMethod", "Edit Model Reference Method")
-    neg_ref_node['inputs']['reference_latents_method'] = "index_timestep_zero"
-    neg_ref_node['inputs']['conditioning'] = [neg_prompt_id, 0]
-    assembler.workflow[neg_ref_id] = neg_ref_node
-
-    assembler.workflow[ksampler_id]['inputs']['positive'] = [pos_ref_id, 0]
-    assembler.workflow[ksampler_id]['inputs']['negative'] = [neg_ref_id, 0]
-
-    model_sampler_id = assembler.node_map.get(model_sampler_name)
-    if not model_sampler_id:
-        for node_id, node in assembler.workflow.items():
-            if isinstance(node, dict) and node.get('class_type') == 'ModelSamplingAuraFlow':
-                model_sampler_id = node_id
-                break
-
-    if model_sampler_id and model_sampler_id in assembler.workflow:
-        assembler.workflow[model_sampler_id]['inputs']['shift'] = 3
-
-    current_model_connection = assembler.workflow[ksampler_id]['inputs']['model']
-    cfg_norm_id = assembler._get_unique_id()
-    cfg_norm_node = create_node(assembler, "CFGNorm", "CFGNorm")
-    cfg_norm_node['inputs']['strength'] = 1
-    cfg_norm_node['inputs']['pre_cfg'] = False
-    cfg_norm_node['inputs']['model'] = current_model_connection
-    assembler.workflow[cfg_norm_id] = cfg_norm_node
-    assembler.workflow[ksampler_id]['inputs']['model'] = [cfg_norm_id, 0]
-
-    print(f"Qwen-Image Edit injector applied with {len(valid_images)} reference image(s). Connected VAE dynamically.")
+def create_node(assembler, class_type, title):
+    try:
+        node = assembler._get_node_template(class_type)
+    except Exception:
+        node = {
+            "inputs": {},
+            "class_type": class_type,
+            "_meta": {"title": title}
+        }
+    node['_meta']['title'] = title
+    return node
+
+def inject(assembler, chain_definition, chain_items):
+    if not chain_items:
+        return
+
+    valid_images = []
+    for item in chain_items:
+        if not item:
+            continue
+        img_path = item
+        if isinstance(item, dict):
+            img_path = item.get('image') or item.get('filename') or item.get('path')
+        if img_path:
+            valid_images.append(img_path)
+
+    if not valid_images:
+        return
+
+    valid_images = valid_images[:3]
+
+    ksampler_name = chain_definition.get('ksampler_node', 'ksampler')
+    pos_prompt_name = chain_definition.get('pos_prompt_node', 'pos_prompt')
+    neg_prompt_name = chain_definition.get('neg_prompt_node', 'neg_prompt')
+    vae_loader_name = chain_definition.get('vae_loader_node', 'vae_loader')
+    model_sampler_name = chain_definition.get('model_sampler_node', 'model_sampler')
+
+    if ksampler_name not in assembler.node_map:
+        print(f"Warning: Target node '{ksampler_name}' for Qwen-Image Edit chain not found. Skipping.")
+        return
+
+    ksampler_id = assembler.node_map[ksampler_name]
+    pos_prompt_id = assembler.node_map.get(pos_prompt_name)
+    neg_prompt_id = assembler.node_map.get(neg_prompt_name)
+
+    if not pos_prompt_id or not neg_prompt_id:
+        print("Warning: Positive or negative prompt node not found for Qwen-Image Edit chain. Skipping.")
+        return
+
+    vae_id = assembler.node_map.get(vae_loader_name)
+    if not vae_id:
+        for node_id, node in assembler.workflow.items():
+            if isinstance(node, dict) and node.get('class_type') == 'VAELoader':
+                vae_id = node_id
+                break
+
+    if vae_id:
+        assembler.workflow[pos_prompt_id]['inputs']['vae'] = [vae_id, 0]
+        assembler.workflow[neg_prompt_id]['inputs']['vae'] = [vae_id, 0]
+
+    for i, img_filename in enumerate(valid_images):
+        load_id = assembler._get_unique_id()
+        load_node = create_node(assembler, "LoadImage", f"Load Reference Image {i+1}")
+        load_node['inputs']['image'] = img_filename
+        assembler.workflow[load_id] = load_node
+
+        scale_id = assembler._get_unique_id()
+        scale_node = create_node(assembler, "ImageScaleToTotalPixels", f"Scale Reference {i+1}")
+        scale_node['inputs']['upscale_method'] = "lanczos"
+        scale_node['inputs']['megapixels'] = 1
+        scale_node['inputs']['resolution_steps'] = 1
+        scale_node['inputs']['image'] = [load_id, 0]
+        assembler.workflow[scale_id] = scale_node
+
+        image_key = f"image{i+1}"
+        assembler.workflow[pos_prompt_id]['inputs'][image_key] = [scale_id, 0]
+        assembler.workflow[neg_prompt_id]['inputs'][image_key] = [scale_id, 0]
+
+    pos_ref_id = assembler._get_unique_id()
+    pos_ref_node = create_node(assembler, "FluxKontextMultiReferenceLatentMethod", "Edit Model Reference Method")
+    pos_ref_node['inputs']['reference_latents_method'] = "index_timestep_zero"
+    pos_ref_node['inputs']['conditioning'] = [pos_prompt_id, 0]
+    assembler.workflow[pos_ref_id] = pos_ref_node
+
+    neg_ref_id = assembler._get_unique_id()
+    neg_ref_node = create_node(assembler, "FluxKontextMultiReferenceLatentMethod", "Edit Model Reference Method")
+    neg_ref_node['inputs']['reference_latents_method'] = "index_timestep_zero"
+    neg_ref_node['inputs']['conditioning'] = [neg_prompt_id, 0]
+    assembler.workflow[neg_ref_id] = neg_ref_node
+
+    assembler.workflow[ksampler_id]['inputs']['positive'] = [pos_ref_id, 0]
+    assembler.workflow[ksampler_id]['inputs']['negative'] = [neg_ref_id, 0]
+
+    model_sampler_id = assembler.node_map.get(model_sampler_name)
+    if not model_sampler_id:
+        for node_id, node in assembler.workflow.items():
+            if isinstance(node, dict) and node.get('class_type') == 'ModelSamplingAuraFlow':
+                model_sampler_id = node_id
+                break
+
+    if model_sampler_id and model_sampler_id in assembler.workflow:
+        assembler.workflow[model_sampler_id]['inputs']['shift'] = 3
+
+    current_model_connection = assembler.workflow[ksampler_id]['inputs']['model']
+    cfg_norm_id = assembler._get_unique_id()
+    cfg_norm_node = create_node(assembler, "CFGNorm", "CFGNorm")
+    cfg_norm_node['inputs']['strength'] = 1
+    cfg_norm_node['inputs']['pre_cfg'] = False
+    cfg_norm_node['inputs']['model'] = current_model_connection
+    assembler.workflow[cfg_norm_id] = cfg_norm_node
+    assembler.workflow[ksampler_id]['inputs']['model'] = [cfg_norm_id, 0]
+
+    print(f"Qwen-Image Edit injector applied with {len(valid_images)} reference image(s). Connected VAE dynamically.")

chain_injectors/reference_image_injector.py

@@ -1,64 +1,64 @@
-import os
-
-def inject(assembler, chain_definition, chain_items):
-    if not chain_items:
-        return
-
-    valid_images = []
-    for item in chain_items:
-        if not item:
-            continue
-        img_path = item
-        if isinstance(item, dict):
-            img_path = item.get('image') or item.get('filename') or item.get('path')
-        if img_path:
-            valid_images.append(img_path)
-
-    if not valid_images:
-        return
-
-    text_encode_name = chain_definition.get('text_encode_node')
-    text_encode_id = None
-    if text_encode_name and text_encode_name in assembler.node_map:
-        text_encode_id = assembler.node_map[text_encode_name]
-    else:
-        for node_id, node in assembler.workflow.items():
-            if isinstance(node, dict) and node.get('class_type') == 'TextEncodeMageFlowEdit':
-                text_encode_id = node_id
-                break
-
-    if not text_encode_id or text_encode_id not in assembler.workflow:
-        print("Warning: TextEncodeMageFlowEdit node not found for Reference Image chain. Skipping.")
-        return
-
-    vae_node_name = chain_definition.get('vae_node', 'vae_loader')
-    vae_node_id = assembler.node_map.get(vae_node_name)
-    if not vae_node_id:
-        for node_id, node in assembler.workflow.items():
-            if isinstance(node, dict) and node.get('class_type') == 'VAELoader':
-                vae_node_id = node_id
-                break
-
-    if vae_node_id:
-        assembler.workflow[text_encode_id]['inputs']['vae'] = [vae_node_id, 0]
-
-    for i, img_filename in enumerate(valid_images):
-        load_id = assembler._get_unique_id()
-        load_node = assembler._get_node_template("LoadImage")
-        load_node['inputs']['image'] = img_filename
-        load_node['_meta']['title'] = f"Load Reference Image {i+1}"
-        assembler.workflow[load_id] = load_node
-
-        scale_id = assembler._get_unique_id()
-        scale_node = assembler._get_node_template("ImageScaleToTotalPixels")
-        scale_node['inputs']['megapixels'] = 1.0
-        scale_node['inputs']['upscale_method'] = "nearest-exact"
-        scale_node['inputs']['resolution_steps'] = 1
-        scale_node['inputs']['image'] = [load_id, 0]
-        scale_node['_meta']['title'] = f"Scale Reference {i+1}"
-        assembler.workflow[scale_id] = scale_node
-
-        input_key = f"images.image_{i+1}"
-        assembler.workflow[text_encode_id]['inputs'][input_key] = [scale_id, 0]
-
-    print(f"Reference Image injector applied. Injected {len(valid_images)} reference images to TextEncodeMageFlowEdit node '{text_encode_id}'.")
+import os
+
+def inject(assembler, chain_definition, chain_items):
+    if not chain_items:
+        return
+
+    valid_images = []
+    for item in chain_items:
+        if not item:
+            continue
+        img_path = item
+        if isinstance(item, dict):
+            img_path = item.get('image') or item.get('filename') or item.get('path')
+        if img_path:
+            valid_images.append(img_path)
+
+    if not valid_images:
+        return
+
+    text_encode_name = chain_definition.get('text_encode_node')
+    text_encode_id = None
+    if text_encode_name and text_encode_name in assembler.node_map:
+        text_encode_id = assembler.node_map[text_encode_name]
+    else:
+        for node_id, node in assembler.workflow.items():
+            if isinstance(node, dict) and node.get('class_type') == 'TextEncodeMageFlowEdit':
+                text_encode_id = node_id
+                break
+
+    if not text_encode_id or text_encode_id not in assembler.workflow:
+        print("Warning: TextEncodeMageFlowEdit node not found for Reference Image chain. Skipping.")
+        return
+
+    vae_node_name = chain_definition.get('vae_node', 'vae_loader')
+    vae_node_id = assembler.node_map.get(vae_node_name)
+    if not vae_node_id:
+        for node_id, node in assembler.workflow.items():
+            if isinstance(node, dict) and node.get('class_type') == 'VAELoader':
+                vae_node_id = node_id
+                break
+
+    if vae_node_id:
+        assembler.workflow[text_encode_id]['inputs']['vae'] = [vae_node_id, 0]
+
+    for i, img_filename in enumerate(valid_images):
+        load_id = assembler._get_unique_id()
+        load_node = assembler._get_node_template("LoadImage")
+        load_node['inputs']['image'] = img_filename
+        load_node['_meta']['title'] = f"Load Reference Image {i+1}"
+        assembler.workflow[load_id] = load_node
+
+        scale_id = assembler._get_unique_id()
+        scale_node = assembler._get_node_template("ImageScaleToTotalPixels")
+        scale_node['inputs']['megapixels'] = 1.0
+        scale_node['inputs']['upscale_method'] = "nearest-exact"
+        scale_node['inputs']['resolution_steps'] = 1
+        scale_node['inputs']['image'] = [load_id, 0]
+        scale_node['_meta']['title'] = f"Scale Reference {i+1}"
+        assembler.workflow[scale_id] = scale_node
+
+        input_key = f"images.image_{i+1}"
+        assembler.workflow[text_encode_id]['inputs'][input_key] = [scale_id, 0]
+
+    print(f"Reference Image injector applied. Injected {len(valid_images)} reference images to TextEncodeMageFlowEdit node '{text_encode_id}'.")

core/pipelines/workflow_executor.py

@@ -1,134 +1,134 @@
-import torch
-from collections import defaultdict, deque
-from typing import Dict, Any, List
-from comfy_integration.nodes import NODE_CLASS_MAPPINGS
-from utils.app_utils import get_value_at_index
-
-class WorkflowExecutor:
-    @staticmethod
-    def topological_sort(workflow: Dict[str, Any]) -> List[str]:
-        graph = defaultdict(list)
-        in_degree = {node_id: 0 for node_id in workflow}
-
-        for node_id, node_info in workflow.items():
-            for input_value in node_info.get('inputs', {}).values():
-                if isinstance(input_value, list) and len(input_value) == 2 and isinstance(input_value[0], str):
-                    source_node_id = input_value[0]
-                    if source_node_id in workflow:
-                        graph[source_node_id].append(node_id)
-                        in_degree[node_id] += 1
-
-        queue = deque([node_id for node_id, degree in in_degree.items() if degree == 0])
-        
-        sorted_nodes = []
-        while queue:
-            current_node_id = queue.popleft()
-            sorted_nodes.append(current_node_id)
-
-            for neighbor_node_id in graph[current_node_id]:
-                in_degree[neighbor_node_id] -= 1
-                if in_degree[neighbor_node_id] == 0:
-                    queue.append(neighbor_node_id)
-        
-        if len(sorted_nodes) != len(workflow):
-            raise RuntimeError("Workflow contains a cycle and cannot be executed.")
-            
-        return sorted_nodes
-
-    @staticmethod
-    def execute_workflow(workflow: Dict[str, Any], initial_objects: Dict[str, Any]):
-        with torch.no_grad():
-            computed_outputs = initial_objects
-            
-            try:
-                sorted_node_ids = WorkflowExecutor.topological_sort(workflow)
-                
-                final_node_id = None
-                for node_id in reversed(sorted_node_ids):
-                    if workflow[node_id].get('class_type') == 'SaveImage':
-                        final_node_id = node_id
-                        break
-                
-                if final_node_id:
-                    required_nodes = set()
-                    nodes_to_visit = [final_node_id]
-                    while nodes_to_visit:
-                        curr_id = nodes_to_visit.pop()
-                        if curr_id in required_nodes:
-                            continue
-                        required_nodes.add(curr_id)
-                        curr_info = workflow.get(curr_id, {})
-                        for input_val in curr_info.get('inputs', {}).values():
-                            if isinstance(input_val, list) and len(input_val) == 2 and isinstance(input_val[0], str):
-                                src_id = input_val[0]
-                                if src_id in workflow and src_id not in required_nodes:
-                                    nodes_to_visit.append(src_id)
-                    
-                    sorted_node_ids = [nid for nid in sorted_node_ids if nid in required_nodes]
-
-                print(f"--- [Workflow Executor] Execution order: {sorted_node_ids}")
-            except RuntimeError as e:
-                print("--- [Workflow Executor] ERROR: Failed to sort workflow. Dumping graph details. ---")
-                for node_id, node_info in workflow.items():
-                    print(f"  Node {node_id} ({node_info['class_type']}):")
-                    for input_name, input_value in node_info['inputs'].items():
-                         if isinstance(input_value, list) and len(input_value) == 2 and isinstance(input_value[0], str):
-                             print(f"    - {input_name} <- [{input_value[0]}, {input_value[1]}]")
-                raise e
-
-            for node_id in sorted_node_ids:
-                if node_id in computed_outputs:
-                    continue
-                    
-                node_info = workflow[node_id]
-                class_type = node_info['class_type']
-                
-                is_loader_with_filename = 'Loader' in class_type and any(key.endswith('_name') for key in node_info['inputs'])
-                if node_id in initial_objects and is_loader_with_filename:
-                    continue
-                
-                node_class = NODE_CLASS_MAPPINGS.get(class_type)
-                if node_class is None:
-                     raise RuntimeError(f"Could not find node class '{class_type}'. Is it imported in comfy_integration/nodes.py?")
-                
-                node_instance = node_class()
-                
-                kwargs = {}
-                for param_name, param_value in node_info['inputs'].items():
-                    if isinstance(param_value, list) and len(param_value) == 2 and isinstance(param_value[0], str):
-                        source_node_id, output_index = param_value
-                        if source_node_id not in computed_outputs:
-                            raise RuntimeError(f"Workflow integrity error: Output of node {source_node_id} needed for {node_id} but not yet computed.")
-                        
-                        source_output_tuple = computed_outputs[source_node_id]
-                        actual_value = get_value_at_index(source_output_tuple, output_index)
-                    else:
-                        actual_value = param_value
-
-                    if '.' in param_name:
-                        parent_key, child_key = param_name.split('.', 1)
-                        if parent_key not in kwargs or not isinstance(kwargs[parent_key], dict):
-                            kwargs[parent_key] = {}
-                        kwargs[parent_key][child_key] = actual_value
-                    else:
-                        kwargs[param_name] = actual_value
-
-                function_name = getattr(node_class, 'FUNCTION')
-                execution_method = getattr(node_instance, function_name)
-                
-                result = execution_method(**kwargs)
-                computed_outputs[node_id] = result
-            
-            final_node_id = None
-            for node_id in reversed(sorted_node_ids):
-                 if workflow[node_id]['class_type'] == 'SaveImage':
-                     final_node_id = node_id
-                     break
-            
-            if not final_node_id:
-                raise RuntimeError("Workflow does not contain a 'SaveImage' node as the output.")
-
-            save_image_inputs = workflow[final_node_id]['inputs']
-            image_source_node_id, image_source_index = save_image_inputs['images']
-            
-            return get_value_at_index(computed_outputs[image_source_node_id], image_source_index)
+import torch
+from collections import defaultdict, deque
+from typing import Dict, Any, List
+from comfy_integration.nodes import NODE_CLASS_MAPPINGS
+from utils.app_utils import get_value_at_index
+
+class WorkflowExecutor:
+    @staticmethod
+    def topological_sort(workflow: Dict[str, Any]) -> List[str]:
+        graph = defaultdict(list)
+        in_degree = {node_id: 0 for node_id in workflow}
+
+        for node_id, node_info in workflow.items():
+            for input_value in node_info.get('inputs', {}).values():
+                if isinstance(input_value, list) and len(input_value) == 2 and isinstance(input_value[0], str):
+                    source_node_id = input_value[0]
+                    if source_node_id in workflow:
+                        graph[source_node_id].append(node_id)
+                        in_degree[node_id] += 1
+
+        queue = deque([node_id for node_id, degree in in_degree.items() if degree == 0])
+        
+        sorted_nodes = []
+        while queue:
+            current_node_id = queue.popleft()
+            sorted_nodes.append(current_node_id)
+
+            for neighbor_node_id in graph[current_node_id]:
+                in_degree[neighbor_node_id] -= 1
+                if in_degree[neighbor_node_id] == 0:
+                    queue.append(neighbor_node_id)
+        
+        if len(sorted_nodes) != len(workflow):
+            raise RuntimeError("Workflow contains a cycle and cannot be executed.")
+            
+        return sorted_nodes
+
+    @staticmethod
+    def execute_workflow(workflow: Dict[str, Any], initial_objects: Dict[str, Any]):
+        with torch.no_grad():
+            computed_outputs = initial_objects
+            
+            try:
+                sorted_node_ids = WorkflowExecutor.topological_sort(workflow)
+                
+                final_node_id = None
+                for node_id in reversed(sorted_node_ids):
+                    if workflow[node_id].get('class_type') == 'SaveImage':
+                        final_node_id = node_id
+                        break
+                
+                if final_node_id:
+                    required_nodes = set()
+                    nodes_to_visit = [final_node_id]
+                    while nodes_to_visit:
+                        curr_id = nodes_to_visit.pop()
+                        if curr_id in required_nodes:
+                            continue
+                        required_nodes.add(curr_id)
+                        curr_info = workflow.get(curr_id, {})
+                        for input_val in curr_info.get('inputs', {}).values():
+                            if isinstance(input_val, list) and len(input_val) == 2 and isinstance(input_val[0], str):
+                                src_id = input_val[0]
+                                if src_id in workflow and src_id not in required_nodes:
+                                    nodes_to_visit.append(src_id)
+                    
+                    sorted_node_ids = [nid for nid in sorted_node_ids if nid in required_nodes]
+
+                print(f"--- [Workflow Executor] Execution order: {sorted_node_ids}")
+            except RuntimeError as e:
+                print("--- [Workflow Executor] ERROR: Failed to sort workflow. Dumping graph details. ---")
+                for node_id, node_info in workflow.items():
+                    print(f"  Node {node_id} ({node_info['class_type']}):")
+                    for input_name, input_value in node_info['inputs'].items():
+                         if isinstance(input_value, list) and len(input_value) == 2 and isinstance(input_value[0], str):
+                             print(f"    - {input_name} <- [{input_value[0]}, {input_value[1]}]")
+                raise e
+
+            for node_id in sorted_node_ids:
+                if node_id in computed_outputs:
+                    continue
+                    
+                node_info = workflow[node_id]
+                class_type = node_info['class_type']
+                
+                is_loader_with_filename = 'Loader' in class_type and any(key.endswith('_name') for key in node_info['inputs'])
+                if node_id in initial_objects and is_loader_with_filename:
+                    continue
+                
+                node_class = NODE_CLASS_MAPPINGS.get(class_type)
+                if node_class is None:
+                     raise RuntimeError(f"Could not find node class '{class_type}'. Is it imported in comfy_integration/nodes.py?")
+                
+                node_instance = node_class()
+                
+                kwargs = {}
+                for param_name, param_value in node_info['inputs'].items():
+                    if isinstance(param_value, list) and len(param_value) == 2 and isinstance(param_value[0], str):
+                        source_node_id, output_index = param_value
+                        if source_node_id not in computed_outputs:
+                            raise RuntimeError(f"Workflow integrity error: Output of node {source_node_id} needed for {node_id} but not yet computed.")
+                        
+                        source_output_tuple = computed_outputs[source_node_id]
+                        actual_value = get_value_at_index(source_output_tuple, output_index)
+                    else:
+                        actual_value = param_value
+
+                    if '.' in param_name:
+                        parent_key, child_key = param_name.split('.', 1)
+                        if parent_key not in kwargs or not isinstance(kwargs[parent_key], dict):
+                            kwargs[parent_key] = {}
+                        kwargs[parent_key][child_key] = actual_value
+                    else:
+                        kwargs[param_name] = actual_value
+
+                function_name = getattr(node_class, 'FUNCTION')
+                execution_method = getattr(node_instance, function_name)
+                
+                result = execution_method(**kwargs)
+                computed_outputs[node_id] = result
+            
+            final_node_id = None
+            for node_id in reversed(sorted_node_ids):
+                 if workflow[node_id]['class_type'] == 'SaveImage':
+                     final_node_id = node_id
+                     break
+            
+            if not final_node_id:
+                raise RuntimeError("Workflow does not contain a 'SaveImage' node as the output.")
+
+            save_image_inputs = workflow[final_node_id]['inputs']
+            image_source_node_id, image_source_index = save_image_inputs['images']
+            
+            return get_value_at_index(computed_outputs[image_source_node_id], image_source_index)

core/pipelines/workflow_recipes/_partials/_base_sampler.yaml

@@ -1,28 +1,28 @@
-nodes:
-  ksampler:
-    class_type: KSampler
-    title: "KSampler"
-    params:
-      denoise: 1.0
-  vae_decode:
-    class_type: VAEDecode
-    title: "VAE Decode"
-  save_image:
-    class_type: SaveImage
-    title: "Save Image"
-    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"
+nodes:
+  ksampler:
+    class_type: KSampler
+    title: "KSampler"
+    params:
+      denoise: 1.0
+  vae_decode:
+    class_type: VAEDecode
+    title: "VAE Decode"
+  save_image:
+    class_type: SaveImage
+    title: "Save Image"
+    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"
   filename_prefix: "save_image:filename_prefix"

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

@@ -1,67 +1,67 @@
-nodes:
-  boogu_prompt:
-    class_type: TextEncodeBooguEdit
-    title: "Text Encode Boogu Edit"
-  unet_loader:
-    class_type: UNETLoader
-    title: "Load Diffusion Model"
-    params:
-      weight_dtype: "default"
-  clip_loader:
-    class_type: CLIPLoader
-    title: "Load CLIP"
-    params:
-      type: "boogu"
-      device: "default"
-  vae_loader:
-    class_type: VAELoader
-    title: "Load VAE"
-
-connections:
-  - from: "unet_loader:0"
-    to: "ksampler:model"
-  - from: "clip_loader:0"
-    to: "boogu_prompt:clip"
-  - from: "boogu_prompt:0"
-    to: "ksampler:positive"
-  - from: "boogu_prompt:1"
-    to: "ksampler:negative"
-  - from: "vae_loader:0"
-    to: "vae_decode:vae"
-  - from: "vae_loader:0"
-    to: "vae_encode:vae"
-
-dynamic_lora_chains:
-  lora_chain:
-    template: "LoraLoader"
-    output_map:
-      "unet_loader:0": "model"
-      "clip_loader:0": "clip"
-    input_map:
-      "model": "model"
-      "clip": "clip"
-    end_input_map:
-      "model": ["ksampler:model"]
-      "clip": ["boogu_prompt:clip"]
-
-dynamic_conditioning_chains:
-  conditioning_chain:
-    ksampler_node: "ksampler"
-    clip_source: "clip_loader:0"
-
-dynamic_boogu_image_edit_chains:
-  boogu_image_edit_chain:
-    ksampler_node: "ksampler"
-    boogu_prompt_node: "boogu_prompt"
-    vae_loader_node: "vae_loader"
-
-dynamic_pid_chains:
-  pid_chain:
-    ksampler_node: "ksampler"
-
-ui_map:
-  positive_prompt: "boogu_prompt:prompt"
-  negative_prompt: "boogu_prompt:negative_prompt"
-  unet_name: "unet_loader:unet_name"
-  clip_name: "clip_loader:clip_name"
-  vae_name: "vae_loader:vae_name"
+nodes:
+  boogu_prompt:
+    class_type: TextEncodeBooguEdit
+    title: "Text Encode Boogu Edit"
+  unet_loader:
+    class_type: UNETLoader
+    title: "Load Diffusion Model"
+    params:
+      weight_dtype: "default"
+  clip_loader:
+    class_type: CLIPLoader
+    title: "Load CLIP"
+    params:
+      type: "boogu"
+      device: "default"
+  vae_loader:
+    class_type: VAELoader
+    title: "Load VAE"
+
+connections:
+  - from: "unet_loader:0"
+    to: "ksampler:model"
+  - from: "clip_loader:0"
+    to: "boogu_prompt:clip"
+  - from: "boogu_prompt:0"
+    to: "ksampler:positive"
+  - from: "boogu_prompt:1"
+    to: "ksampler:negative"
+  - from: "vae_loader:0"
+    to: "vae_decode:vae"
+  - from: "vae_loader:0"
+    to: "vae_encode:vae"
+
+dynamic_lora_chains:
+  lora_chain:
+    template: "LoraLoader"
+    output_map:
+      "unet_loader:0": "model"
+      "clip_loader:0": "clip"
+    input_map:
+      "model": "model"
+      "clip": "clip"
+    end_input_map:
+      "model": ["ksampler:model"]
+      "clip": ["boogu_prompt:clip"]
+
+dynamic_conditioning_chains:
+  conditioning_chain:
+    ksampler_node: "ksampler"
+    clip_source: "clip_loader:0"
+
+dynamic_boogu_image_edit_chains:
+  boogu_image_edit_chain:
+    ksampler_node: "ksampler"
+    boogu_prompt_node: "boogu_prompt"
+    vae_loader_node: "vae_loader"
+
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
+ui_map:
+  positive_prompt: "boogu_prompt:prompt"
+  negative_prompt: "boogu_prompt:negative_prompt"
+  unet_name: "unet_loader:unet_name"
+  clip_name: "clip_loader:clip_name"
+  vae_name: "vae_loader:vae_name"

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

@@ -1,91 +1,91 @@
-nodes:
-  pos_prompt:
-    class_type: TextEncodeQwenImageEditPlus
-    title: "Text Encode Qwen Image Edit Plus (Positive)"
-  neg_prompt:
-    class_type: TextEncodeQwenImageEditPlus
-    title: "Text Encode Qwen Image Edit Plus (Negative)"
-  unet_loader:
-    class_type: UNETLoader
-    title: "Load Qwen UNET"
-    params:
-      weight_dtype: "default"
-  vae_loader:
-    class_type: VAELoader
-    title: "Load Qwen VAE"
-  clip_loader:
-    class_type: CLIPLoader
-    title: "Load Qwen CLIP"
-    params:
-      type: "qwen_image"
-      device: "default"
-  model_sampler:
-    class_type: ModelSamplingAuraFlow
-    title: "ModelSamplingAuraFlow"
-    params:
-      shift: 3.1
-
-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: "vae_loader:0"
-    to: "vae_decode:vae"
-  - from: "vae_loader:0"
-    to: "vae_encode:vae"
-
-  - from: "pos_prompt:0"
-    to: "ksampler:positive"
-  - from: "neg_prompt:0"
-    to: "ksampler:negative"
-
-dynamic_lora_chains:
-  lora_chain:
-    template: "LoraLoader"
-    output_map:
-      "unet_loader:0": "model"
-      "clip_loader:0": "clip"
-    input_map:
-      "model": "model"
-      "clip": "clip"
-    end_input_map:
-      "model": ["model_sampler:model"]
-      "clip": ["pos_prompt:clip", "neg_prompt:clip"]
-
-dynamic_controlnet_chains:
-  controlnet_chain:
-    template: "ControlNetApplyAdvanced"
-    ksampler_node: "ksampler"
-    vae_source: "vae_loader:0"
-
-dynamic_conditioning_chains:
-  conditioning_chain:
-    ksampler_node: "ksampler"
-    clip_source: "clip_loader:0"
-
-dynamic_qwen_image_edit_chains:
-  qwen_image_edit_chain:
-    ksampler_node: "ksampler"
-    pos_prompt_node: "pos_prompt"
-    neg_prompt_node: "neg_prompt"
-    vae_loader_node: "vae_loader"
-    model_sampler_node: "model_sampler"
-
-dynamic_pid_chains:
-  pid_chain:
-    ksampler_node: "ksampler"
-
-ui_map:
-  positive_prompt: "pos_prompt:prompt"
-  negative_prompt: "neg_prompt:prompt"
-  unet_name: "unet_loader:unet_name"
-  vae_name: "vae_loader:vae_name"
+nodes:
+  pos_prompt:
+    class_type: TextEncodeQwenImageEditPlus
+    title: "Text Encode Qwen Image Edit Plus (Positive)"
+  neg_prompt:
+    class_type: TextEncodeQwenImageEditPlus
+    title: "Text Encode Qwen Image Edit Plus (Negative)"
+  unet_loader:
+    class_type: UNETLoader
+    title: "Load Qwen UNET"
+    params:
+      weight_dtype: "default"
+  vae_loader:
+    class_type: VAELoader
+    title: "Load Qwen VAE"
+  clip_loader:
+    class_type: CLIPLoader
+    title: "Load Qwen CLIP"
+    params:
+      type: "qwen_image"
+      device: "default"
+  model_sampler:
+    class_type: ModelSamplingAuraFlow
+    title: "ModelSamplingAuraFlow"
+    params:
+      shift: 3.1
+
+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: "vae_loader:0"
+    to: "vae_decode:vae"
+  - from: "vae_loader:0"
+    to: "vae_encode:vae"
+
+  - from: "pos_prompt:0"
+    to: "ksampler:positive"
+  - from: "neg_prompt:0"
+    to: "ksampler:negative"
+
+dynamic_lora_chains:
+  lora_chain:
+    template: "LoraLoader"
+    output_map:
+      "unet_loader:0": "model"
+      "clip_loader:0": "clip"
+    input_map:
+      "model": "model"
+      "clip": "clip"
+    end_input_map:
+      "model": ["model_sampler:model"]
+      "clip": ["pos_prompt:clip", "neg_prompt:clip"]
+
+dynamic_controlnet_chains:
+  controlnet_chain:
+    template: "ControlNetApplyAdvanced"
+    ksampler_node: "ksampler"
+    vae_source: "vae_loader:0"
+
+dynamic_conditioning_chains:
+  conditioning_chain:
+    ksampler_node: "ksampler"
+    clip_source: "clip_loader:0"
+
+dynamic_qwen_image_edit_chains:
+  qwen_image_edit_chain:
+    ksampler_node: "ksampler"
+    pos_prompt_node: "pos_prompt"
+    neg_prompt_node: "neg_prompt"
+    vae_loader_node: "vae_loader"
+    model_sampler_node: "model_sampler"
+
+dynamic_pid_chains:
+  pid_chain:
+    ksampler_node: "ksampler"
+
+ui_map:
+  positive_prompt: "pos_prompt:prompt"
+  negative_prompt: "neg_prompt:prompt"
+  unet_name: "unet_loader:unet_name"
+  vae_name: "vae_loader:vae_name"
   clip_name: "clip_loader:clip_name"

core/pipelines/workflow_recipes/unified_recipe.yaml

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

core/workflow_assembler.py

@@ -1,165 +1,165 @@
-import yaml
-import os
-import importlib
-from copy import deepcopy
-from comfy_integration.nodes import NODE_CLASS_MAPPINGS
-from chain_injectors import discover_injectors, get_registered_features
-from core.settings import FEATURES_CONFIG
-
-class WorkflowAssembler:
-    def __init__(self, recipe_path, dynamic_values=None):
-        self.base_path = os.path.dirname(recipe_path)
-        self.dynamic_values = dynamic_values or {}
-        self.node_counter = 0
-        self.workflow = {}
-        self.node_map = {}
-        
-        model_type = self.dynamic_values.get('model_type')
-        self._load_injector_config(model_type=model_type)
-        
-        self.recipe = self._load_and_merge_recipe(os.path.basename(recipe_path), self.dynamic_values)
-
-    def _load_injector_config(self, model_type=None):
-        self.global_injectors = discover_injectors()
-        registered_features = get_registered_features()
-
-        order = []
-        if model_type and model_type in FEATURES_CONFIG:
-            enabled_features = FEATURES_CONFIG[model_type].get('enabled_chains', [])
-            for feat in enabled_features:
-                if feat in registered_features:
-                    chain_key = registered_features[feat]['chain_type']
-                else:
-                    chain_key = f"dynamic_{feat}_chains"
-                if chain_key in self.global_injectors and chain_key not in order:
-                    order.append(chain_key)
-
-        for chain_key in self.global_injectors.keys():
-            if chain_key not in order:
-                order.append(chain_key)
-
-        self.injector_order = order
-
-    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):
-        self.ui_values = 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 ---")
-        
+import yaml
+import os
+import importlib
+from copy import deepcopy
+from comfy_integration.nodes import NODE_CLASS_MAPPINGS
+from chain_injectors import discover_injectors, get_registered_features
+from core.settings import FEATURES_CONFIG
+
+class WorkflowAssembler:
+    def __init__(self, recipe_path, dynamic_values=None):
+        self.base_path = os.path.dirname(recipe_path)
+        self.dynamic_values = dynamic_values or {}
+        self.node_counter = 0
+        self.workflow = {}
+        self.node_map = {}
+        
+        model_type = self.dynamic_values.get('model_type')
+        self._load_injector_config(model_type=model_type)
+        
+        self.recipe = self._load_and_merge_recipe(os.path.basename(recipe_path), self.dynamic_values)
+
+    def _load_injector_config(self, model_type=None):
+        self.global_injectors = discover_injectors()
+        registered_features = get_registered_features()
+
+        order = []
+        if model_type and model_type in FEATURES_CONFIG:
+            enabled_features = FEATURES_CONFIG[model_type].get('enabled_chains', [])
+            for feat in enabled_features:
+                if feat in registered_features:
+                    chain_key = registered_features[feat]['chain_type']
+                else:
+                    chain_key = f"dynamic_{feat}_chains"
+                if chain_key in self.global_injectors and chain_key not in order:
+                    order.append(chain_key)
+
+        for chain_key in self.global_injectors.keys():
+            if chain_key not in order:
+                order.append(chain_key)
+
+        self.injector_order = order
+
+    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):
+        self.ui_values = 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

mcp_tools/__init__.py

@@ -0,0 +1,50 @@
+def __getattr__(name):
+    if name in ("types", "server", "client", "shared"):
+        raise ImportError(f"No module named 'mcp.{name}' in local mcp package")
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+from .get_task_list import handle_get_task_list
+from .get_model_architecture_list import handle_get_model_architecture_list
+from .get_model_list import handle_get_model_list
+from .get_feature_list import handle_get_feature_list
+from .get_model_features import handle_get_model_features
+from .get_chain_schema import handle_get_chain_schema
+from .run_imagegen import handle_run_imagegen
+from .get_task_status import handle_get_task_status
+from .error_schema import make_error, make_validation_error, make_not_found_error
+from .mcp_gradio_integration import (
+    register_high_level_mcp_apis,
+    cleanup_dependencies_api_names,
+    patch_gradio_api_suppression,
+    HIGH_LEVEL_MCP_API_NAMES,
+)
+
+MCP_FUNCTIONS = [
+    handle_get_task_list,
+    handle_get_model_architecture_list,
+    handle_get_model_list,
+    handle_get_feature_list,
+    handle_get_model_features,
+    handle_run_imagegen,
+    handle_get_task_status,
+    handle_get_chain_schema,
+]
+
+__all__ = [
+    "handle_get_task_list",
+    "handle_get_model_architecture_list",
+    "handle_get_model_list",
+    "handle_get_feature_list",
+    "handle_get_model_features",
+    "handle_get_chain_schema",
+    "handle_run_imagegen",
+    "handle_get_task_status",
+    "make_error",
+    "make_validation_error",
+    "make_not_found_error",
+    "register_high_level_mcp_apis",
+    "cleanup_dependencies_api_names",
+    "patch_gradio_api_suppression",
+    "HIGH_LEVEL_MCP_API_NAMES",
+    "MCP_FUNCTIONS",
+]

mcp_tools/common.py

@@ -0,0 +1,335 @@
+"""
+MCP Common Utilities & Data Structures
+Contains YAML loading utilities, config file paths, task definitions, and async task database.
+"""
+
+import os
+import time
+import urllib.parse
+import urllib.request
+import base64
+import io
+import yaml
+from typing import Dict, Any
+from PIL import Image
+
+_PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+_YAML_DIR = os.path.join(_PROJECT_ROOT, "yaml")
+
+_MODEL_ARCHITECTURES_PATH = os.path.join(_YAML_DIR, "model_architectures.yaml")
+_MODEL_LIST_PATH = os.path.join(_YAML_DIR, "model_list.yaml")
+_MODEL_DEFAULTS_PATH = os.path.join(_YAML_DIR, "model_defaults.yaml")
+_IMAGE_GEN_FEATURES_PATH = os.path.join(_YAML_DIR, "image_gen_features.yaml")
+_CHAIN_FEATURES_PATH = os.path.join(_YAML_DIR, "chain_features.yaml")
+_CONSTANTS_PATH = os.path.join(_YAML_DIR, "constants.yaml")
+
+
+def _parse_image_param(image_param: Any) -> Any:
+    """Parse a Base64 Data URI, local file path, or PIL.Image into a PIL Image object. HTTP URLs are not supported."""
+    if isinstance(image_param, Image.Image):
+        return image_param
+
+    if not isinstance(image_param, str) or not image_param.strip():
+        return None
+
+    image_param = image_param.strip()
+
+    # Reject HTTP / HTTPS URL
+    if image_param.startswith("http://") or image_param.startswith("https://"):
+        raise ValueError(
+            "Image URLs are not supported. Please supply the image directly as a Base64 Data URI (e.g., 'data:image/png;base64,...')."
+        )
+
+    # Base64 Data URI (e.g. data:image/png;base64,...)
+    if image_param.startswith("data:image/"):
+        _, encoded = image_param.split(",", 1) if "," in image_param else ("", image_param)
+        data = base64.b64decode(encoded)
+        return Image.open(io.BytesIO(data))
+
+    # Base64 string without header
+    if len(image_param) > 100 and not os.path.exists(image_param):
+        try:
+            data = base64.b64decode(image_param)
+            return Image.open(io.BytesIO(data))
+        except Exception:
+            pass
+
+    # Local file path
+    if os.path.exists(image_param):
+        return Image.open(image_param)
+
+    raise ValueError(
+        "Invalid image parameter format. Expected a Base64 Data URI (e.g., 'data:image/png;base64,...') or local file path."
+    )
+
+
+def _load_yaml(filepath: str) -> dict:
+    """Safely load a YAML file, returning an empty dict if the file does not exist."""
+    if not os.path.exists(filepath):
+        print(f"Warning: YAML file not found: {filepath}")
+        return {}
+    with open(filepath, "r", encoding="utf-8") as f:
+        return yaml.safe_load(f) or {}
+
+
+_COMMON_OPTIONAL_INPUTS = [
+    "steps", "cfg", "sampler", "scheduler", "seed",
+    "negative_prompt", "batch_size", "chain", "async_execution",
+]
+
+_TASK_DEFINITIONS = [
+    {
+        "task_type": "txt2img",
+        "display_name": "Text-to-Image",
+        "description": "Generate images from text prompts. Canvas width and height must be specified.",
+        "required_inputs": ["prompt", "width", "height"],
+        "optional_inputs": _COMMON_OPTIONAL_INPUTS,
+    },
+    {
+        "task_type": "img2img",
+        "display_name": "Image-to-Image",
+        "description": "Perform global repaint and style transfer based on a source image. Denoise strength must be specified.",
+        "required_inputs": ["prompt", "image", "denoise"],
+        "optional_inputs": _COMMON_OPTIONAL_INPUTS,
+    },
+    {
+        "task_type": "inpaint",
+        "display_name": "Inpaint",
+        "description": "Repaint specified masked regions of the input image (with alpha mask/channel).",
+        "required_inputs": ["prompt", "image"],
+        "optional_inputs": ["denoise"] + _COMMON_OPTIONAL_INPUTS,
+    },
+    {
+        "task_type": "outpaint",
+        "display_name": "Outpaint",
+        "description": "Extend the canvas outward from the source image. Padding pixel values for top, bottom, left, and right must be specified.",
+        "required_inputs": ["prompt", "image", "pad_left", "pad_right", "pad_top", "pad_bottom"],
+        "optional_inputs": _COMMON_OPTIONAL_INPUTS,
+    },
+    {
+        "task_type": "hires_fix",
+        "display_name": "Hi-Res Fix / Upscale",
+        "description": "Enhance details and upscale an existing low-resolution image.",
+        "required_inputs": ["prompt", "image", "upscale_by"],
+        "optional_inputs": _COMMON_OPTIONAL_INPUTS,
+    },
+]
+
+_TASKS_DB: Dict[str, Dict[str, Any]] = {}
+
+
+class DummyProgress:
+    def __call__(self, progress=0.0, desc=None):
+        pass
+
+
+def _get_public_base_url() -> str:
+    """Auto-resolve the publicly accessible base URL (including protocol and port)."""
+    # 1. Explicit environment variable override
+    public_url = os.getenv("PUBLIC_URL") or os.getenv("BASE_URL")
+    if public_url:
+        return public_url.rstrip("/")
+
+    # 2. Hugging Face Space environment variable
+    space_host = os.getenv("SPACE_HOST")
+    if space_host:
+        if not space_host.startswith("http://") and not space_host.startswith("https://"):
+            return f"https://{space_host}"
+        return space_host.rstrip("/")
+
+    # 3. Local Gradio config fallback
+    try:
+        from core.settings import GRADIO_SERVER_NAME, SERVER_PORT
+    except ImportError:
+        GRADIO_SERVER_NAME = "127.0.0.1"
+        SERVER_PORT = 7860
+
+    server_name = os.getenv("GRADIO_SERVER_NAME", GRADIO_SERVER_NAME)
+    if server_name == "0.0.0.0":
+        server_name = "127.0.0.1"
+    port = os.getenv("GRADIO_SERVER_PORT", str(SERVER_PORT))
+
+    return f"http://{server_name}:{port}"
+
+
+def _execute_imagegen_pipeline(task_id: str, params: dict):
+    """Execute the image generation pipeline in the background and update _TASKS_DB."""
+    start_time = time.time()
+    try:
+        _TASKS_DB[task_id]["status"] = "processing"
+        _TASKS_DB[task_id]["progress"] = 10
+        _TASKS_DB[task_id]["updated_at"] = int(start_time)
+
+        from core.generation_logic import sd_image_pipeline
+
+        task_type = params["task_type"]
+        model = params["model"]
+        prompt = params["prompt"]
+
+        model_defaults = _load_yaml(_MODEL_DEFAULTS_PATH)
+        model_list = _load_yaml(_MODEL_LIST_PATH)
+        checkpoints = model_list.get("Checkpoint", {})
+        found_arch = None
+        for arch_name, arch_data in checkpoints.items():
+            if isinstance(arch_data, dict):
+                for m in arch_data.get("models", []):
+                    if m.get("display_name") == model:
+                        found_arch = arch_name
+                        break
+            if found_arch:
+                break
+
+        arch_defaults_section = model_defaults.get(found_arch, {}) if found_arch else {}
+        arch_level_defaults = arch_defaults_section.get("_defaults", {})
+        model_specific_defaults = arch_defaults_section.get(model, {})
+        global_defaults = model_defaults.get("Default", {})
+        merged_defaults = {**global_defaults, **arch_level_defaults, **model_specific_defaults}
+
+        steps = params.get("steps") if params.get("steps") is not None else merged_defaults.get("steps", 20)
+        cfg = params.get("cfg") if params.get("cfg") is not None else merged_defaults.get("cfg", 1.0)
+        sampler = params.get("sampler") or merged_defaults.get("sampler_name", "euler")
+        scheduler = params.get("scheduler") or merged_defaults.get("scheduler", "simple")
+
+        ui_inputs = {
+            "task_type": task_type,
+            "model_display_name": model,
+            "base_model_" + task_type: model,
+            "positive_prompt": prompt,
+            "negative_prompt": params.get("negative_prompt", merged_defaults.get("negative_prompt", "")),
+            "width": params.get("width", 1024),
+            "height": params.get("height", 1024),
+            "num_inference_steps": steps,
+            "guidance_scale": cfg,
+            "sampler": sampler,
+            "scheduler": scheduler,
+            "seed": params.get("seed", -1),
+            "batch_size": params.get("batch_size", 1),
+            "zero_gpu_duration": params.get("zero_gpu_duration"),
+            "denoise": params.get("denoise", 1.0),
+        }
+
+        if "image" in params and params["image"]:
+            pil_img = _parse_image_param(params["image"])
+            if pil_img:
+                if task_type == "img2img":
+                    ui_inputs["img2img_image"] = pil_img
+                    ui_inputs["img2img_denoise"] = params.get("denoise", 0.7)
+                elif task_type == "inpaint":
+                    ui_inputs["inpaint_image"] = pil_img
+                    ui_inputs["inpaint_denoise"] = params.get("denoise", 1.0)
+                elif task_type == "outpaint":
+                    ui_inputs["outpaint_image"] = pil_img
+                    ui_inputs["left"] = params.get("pad_left", 0)
+                    ui_inputs["right"] = params.get("pad_right", 0)
+                    ui_inputs["top"] = params.get("pad_top", 0)
+                    ui_inputs["bottom"] = params.get("pad_bottom", 0)
+                    ui_inputs["feathering"] = params.get("feathering", 10)
+                elif task_type == "hires_fix":
+                    ui_inputs["hires_image"] = pil_img
+                    ui_inputs["hires_upscaler"] = params.get("upscaler", "latent")
+                    ui_inputs["hires_scale_by"] = params.get("upscale_by", 2.0)
+                    ui_inputs["hires_denoise"] = params.get("denoise", 0.55)
+
+        chain = params.get("chain", [])
+        if chain:
+            lora_data = []
+            controlnet_data = []
+            ipadapter_data = []
+            style_data = []
+
+            for item in chain:
+                itype = item.get("injector_type")
+                if itype == "lora":
+                    lora_data.extend([
+                        item.get("lora_source", "Civitai"),
+                        item.get("lora_value", ""),
+                        item.get("scale", 1.0),
+                        None
+                    ])
+                elif itype in ("controlnet", "krea2_controlnet", "anima_controlnet_lllite"):
+                    controlnet_data.extend([
+                        item.get("control_net_name", ""),
+                        _parse_image_param(item.get("image")),
+                        item.get("strength", 1.0)
+                    ])
+                elif itype in ("ipadapter", "flux1_ipadapter", "sd3_ipadapter"):
+                    ipadapter_data.extend([
+                        item.get("preset", "STANDARD (medium strength)"),
+                        _parse_image_param(item.get("image")),
+                        item.get("weight", 1.0)
+                    ])
+                elif itype == "style":
+                    style_data.extend([
+                        _parse_image_param(item.get("image")),
+                        item.get("strength", 1.0)
+                    ])
+
+            if lora_data: ui_inputs["lora_data"] = lora_data
+            if controlnet_data: ui_inputs["controlnet_data"] = controlnet_data
+            if ipadapter_data: ui_inputs["ipadapter_data"] = ipadapter_data
+            if style_data: ui_inputs["style_data"] = style_data
+
+        _TASKS_DB[task_id]["progress"] = 50
+
+        # Execute Pipeline
+        output = sd_image_pipeline.run(ui_inputs=ui_inputs, progress=DummyProgress())
+
+        try:
+            from core.settings import OUTPUT_DIR
+        except ImportError:
+            OUTPUT_DIR = os.path.join(_PROJECT_ROOT, "output")
+
+        os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+        import tempfile
+        import gradio.processing_utils as pu
+
+        gradio_cache_dir = os.path.join(tempfile.gettempdir(), "gradio")
+        os.makedirs(gradio_cache_dir, exist_ok=True)
+
+        base_url = _get_public_base_url()
+        images = []
+        raw_list = output if isinstance(output, list) else ([output] if output else [])
+        for idx, item in enumerate(raw_list):
+            target_path = None
+            if hasattr(item, "save"):  # PIL Image
+                filename = f"mcp_{task_id}_{idx}.png"
+                filepath = os.path.join(OUTPUT_DIR, filename)
+                item.save(filepath)
+                target_path = filepath
+            elif isinstance(item, str) and os.path.exists(item):
+                target_path = item
+
+            if target_path:
+                try:
+                    cached_path = pu.save_file_to_cache(target_path, cache_dir=gradio_cache_dir)
+                    abs_path = os.path.abspath(cached_path).replace("\\", "/")
+                except Exception as e:
+                    print(f"Warning: Failed to cache image file to Gradio temp dir: {e}")
+                    abs_path = os.path.abspath(target_path).replace("\\", "/")
+
+                url = f"{base_url}/gradio_api/file={urllib.parse.quote(abs_path)}"
+                images.append(url)
+            elif item:
+                images.append(str(item))
+
+        execution_time = round(time.time() - start_time, 2)
+        _TASKS_DB[task_id]["status"] = "completed"
+        _TASKS_DB[task_id]["progress"] = 100
+        _TASKS_DB[task_id]["completed_at"] = int(time.time())
+        _TASKS_DB[task_id]["result"] = {
+            "images": images,
+            "seed": params.get("seed", -1),
+            "width": params.get("width", 1024),
+            "height": params.get("height", 1024),
+            "execution_time_seconds": execution_time,
+        }
+
+    except Exception as e:
+        _TASKS_DB[task_id]["status"] = "failed"
+        _TASKS_DB[task_id]["progress"] = 0
+        _TASKS_DB[task_id]["failed_at"] = int(time.time())
+        _TASKS_DB[task_id]["error"] = {
+            "code": "EXECUTION_ERROR",
+            "message": str(e),
+        }

mcp_tools/error_schema.py

@@ -0,0 +1,85 @@
+"""
+Unified MCP tool error response format.
+
+Error code enumeration:
+  - INVALID_PARAMS:         Parameter validation failed (missing required fields, type errors, value out of range)
+  - MODEL_NOT_FOUND:        The specified model name does not exist
+  - ARCHITECTURE_NOT_FOUND: The specified architecture name does not exist
+  - CHAIN_TYPE_NOT_FOUND:   The specified chain/injector type is invalid
+  - FEATURE_NOT_SUPPORTED:  The current model does not support the requested feature
+  - TASK_NOT_FOUND:         The async task ID does not exist
+  - MODEL_OOM:              GPU out of memory
+  - INTERNAL_ERROR:         Internal server error
+"""
+
+
+def make_error(code: str, message: str, details: dict = None) -> dict:
+    """
+    Construct a unified MCP tool error response.
+
+    Args:
+        code:    Error code (UPPER_SNAKE_CASE format)
+        message: Human-readable error description
+        details: Optional details dictionary
+
+    Returns:
+        Standardized error response dictionary
+    """
+    error = {
+        "error": {
+            "code": code,
+            "message": message,
+        }
+    }
+    if details:
+        error["error"]["details"] = details
+    return error
+
+
+def make_validation_error(
+    message: str = "Request validation failed.",
+    missing_fields: list = None,
+    invalid_fields: dict = None,
+) -> dict:
+    """
+    Construct a parameter validation failure error response.
+
+    Args:
+        message:        Error description
+        missing_fields: List of missing required field names
+        invalid_fields: Key-value pairs of invalid fields, key=field name, value=reason description
+
+    Returns:
+        Standardized INVALID_PARAMS error response
+    """
+    details = {}
+    if missing_fields:
+        details["missing_fields"] = missing_fields
+    if invalid_fields:
+        details["invalid_fields"] = invalid_fields
+    return make_error("INVALID_PARAMS", message, details if details else None)
+
+
+def make_not_found_error(resource_type: str, resource_id: str) -> dict:
+    """
+    Construct a resource-not-found error response.
+
+    Args:
+        resource_type: Resource type (e.g., "model", "architecture", "chain_type", "task")
+        resource_id:   Resource identifier
+
+    Returns:
+        Standardized *_NOT_FOUND error response
+    """
+    code_map = {
+        "model": "MODEL_NOT_FOUND",
+        "architecture": "ARCHITECTURE_NOT_FOUND",
+        "chain_type": "CHAIN_TYPE_NOT_FOUND",
+        "task": "TASK_NOT_FOUND",
+    }
+    code = code_map.get(resource_type, f"{resource_type.upper()}_NOT_FOUND")
+    return make_error(
+        code,
+        f"The specified {resource_type} '{resource_id}' was not found.",
+        {"resource_type": resource_type, "resource_id": resource_id},
+    )

mcp_tools/get_chain_schema.py

@@ -0,0 +1,32 @@
+"""
+MCP Tool: get_chain_schema
+Get the complete parameter schema and usage guide for a specified chain/injector type.
+"""
+
+from .common import _load_yaml, _CHAIN_FEATURES_PATH
+from .error_schema import make_validation_error, make_not_found_error
+
+
+def handle_get_chain_schema(chain_type: str) -> dict:
+    """Get the complete parameter schema and usage guide for a specified chain/injector type."""
+    if not chain_type:
+        return make_validation_error(
+            "Parameter 'chain_type' is required.",
+            missing_fields=["chain_type"],
+        )
+
+    chain_features = _load_yaml(_CHAIN_FEATURES_PATH)
+
+    if chain_type not in chain_features:
+        return make_not_found_error("chain_type", chain_type)
+
+    chain_data = chain_features[chain_type]
+    return {
+        "feature_name": chain_type,
+        "display_name": chain_data.get("display_name", chain_type),
+        "description": chain_data.get("description", ""),
+        "supported_tasks": chain_data.get("supported_tasks", []),
+        "max_count": chain_data.get("max_count", 1),
+        "usage_guideline": chain_data.get("usage_guideline", ""),
+        "parameters_schema": chain_data.get("parameters_schema", {}),
+    }

mcp_tools/get_feature_list.py

@@ -0,0 +1,29 @@
+"""
+MCP Tool: get_feature_list
+Get the list of supported advanced features along with their usage constraints and parameter schemas.
+"""
+
+from .common import _load_yaml, _CHAIN_FEATURES_PATH
+
+
+def handle_get_feature_list() -> list:
+    """Dynamically load the list of supported advanced features from chain_features.yaml."""
+    chain_features = _load_yaml(_CHAIN_FEATURES_PATH)
+    result = []
+
+    for chain_name, chain_data in chain_features.items():
+        if chain_data.get("visibility", "public") != "public":
+            continue
+
+        entry = {
+            "feature_name": chain_name,
+            "display_name": chain_data.get("display_name", chain_name),
+            "description": chain_data.get("description", ""),
+            "supported_tasks": chain_data.get("supported_tasks", []),
+            "max_count": chain_data.get("max_count", 1),
+            "usage_guideline": chain_data.get("usage_guideline", ""),
+            "parameters_schema": chain_data.get("parameters_schema", {}),
+        }
+        result.append(entry)
+
+    return result

mcp_tools/get_model_architecture_list.py

@@ -0,0 +1,36 @@
+"""
+MCP Tool: get_model_architecture_list
+Get all supported model architectures and their corresponding default resolutions.
+"""
+
+from .common import _load_yaml, _MODEL_ARCHITECTURES_PATH, _CONSTANTS_PATH
+
+
+def handle_get_model_architecture_list() -> list:
+    """Dynamically load all supported model architectures from model_architectures.yaml."""
+    arch_config = _load_yaml(_MODEL_ARCHITECTURES_PATH)
+    constants = _load_yaml(_CONSTANTS_PATH)
+    resolution_map = constants.get("RESOLUTION_MAP", {})
+    architectures = arch_config.get("architectures", {})
+    architecture_order = arch_config.get("architecture_order", list(architectures.keys()))
+
+    result = []
+    for arch_name in architecture_order:
+        if arch_name not in architectures:
+            continue
+        arch_data = architectures[arch_name]
+        model_type = arch_data.get("model_type", arch_name.lower())
+
+        default_res = [1024, 1024]
+        if model_type in resolution_map:
+            resolutions = resolution_map[model_type]
+            if resolutions:
+                first_key = next(iter(resolutions))
+                default_res = resolutions[first_key]
+
+        result.append({
+            "model_architecture": arch_name,
+            "default_resolution": default_res,
+        })
+
+    return result

mcp_tools/get_model_features.py

@@ -0,0 +1,111 @@
+"""
+MCP Tool: get_model_features
+Query metadata for a specified model, including supported task types, extended features, and default inference parameters.
+"""
+
+from .common import (
+    _load_yaml,
+    _MODEL_LIST_PATH,
+    _MODEL_DEFAULTS_PATH,
+    _IMAGE_GEN_FEATURES_PATH,
+    _MODEL_ARCHITECTURES_PATH,
+    _CHAIN_FEATURES_PATH,
+    _TASK_DEFINITIONS,
+)
+from .error_schema import make_validation_error, make_not_found_error
+
+
+def handle_get_model_features(model: str) -> dict:
+    """Query metadata for a specified model: supported task types, extended features, and default inference parameters."""
+    if not model:
+        return make_validation_error(
+            "Parameter 'model' is required.",
+            missing_fields=["model"],
+        )
+
+    model_list = _load_yaml(_MODEL_LIST_PATH)
+    model_defaults = _load_yaml(_MODEL_DEFAULTS_PATH)
+    features_config = _load_yaml(_IMAGE_GEN_FEATURES_PATH)
+    arch_config = _load_yaml(_MODEL_ARCHITECTURES_PATH)
+    chain_features = _load_yaml(_CHAIN_FEATURES_PATH)
+
+    found_arch = None
+    checkpoints = model_list.get("Checkpoint", {})
+    for arch_name, arch_data in checkpoints.items():
+        if not isinstance(arch_data, dict):
+            continue
+        for m in arch_data.get("models", []):
+            if m.get("display_name") == model:
+                found_arch = arch_name
+                break
+        if found_arch:
+            break
+
+    if not found_arch:
+        return make_not_found_error("model", model)
+
+    architectures = arch_config.get("architectures", {})
+    arch_info = architectures.get(found_arch, {})
+    model_type = arch_info.get("model_type", found_arch.lower())
+
+    arch_features = features_config.get(model_type, features_config.get("default", {}))
+    enabled_chains = arch_features.get("enabled_chains", [])
+
+    supported_features = []
+    for chain_name in enabled_chains:
+        if chain_name in chain_features:
+            chain_data = chain_features[chain_name]
+            visibility = chain_data.get("visibility", "public")
+            if visibility == "public":
+                supported_features.append(chain_name)
+            else:
+                generic_mapping = {
+                    "krea2_controlnet": "controlnet",
+                    "anima_controlnet_lllite": "controlnet",
+                    "controlnet_model_patch": "controlnet",
+                    "flux1_ipadapter": "ipadapter",
+                    "sd3_ipadapter": "ipadapter",
+                    "hidream_o1_reference": "reference_latent",
+                }
+                generic_name = generic_mapping.get(chain_name)
+                if generic_name and generic_name not in supported_features:
+                    supported_features.append(generic_name)
+
+    arch_defaults_section = model_defaults.get(found_arch, {})
+    arch_level_defaults = arch_defaults_section.get("_defaults", {})
+    model_specific_defaults = arch_defaults_section.get(model, {})
+    global_defaults = model_defaults.get("Default", {})
+
+    merged_defaults = {**global_defaults, **arch_level_defaults, **model_specific_defaults}
+
+    default_parameter = {
+        "sampler": merged_defaults.get("sampler_name", "euler"),
+        "scheduler": merged_defaults.get("scheduler", "simple"),
+        "steps": merged_defaults.get("steps", 20),
+        "cfg": merged_defaults.get("cfg", 1.0),
+    }
+
+    supported_tasks = [t["task_type"] for t in _TASK_DEFINITIONS]
+
+    result = {
+        "name": model,
+        "model_architecture": found_arch,
+        "supported_tasks": supported_tasks,
+        "supported_features": supported_features,
+        "default_parameter": default_parameter,
+    }
+
+    default_pos = model_specific_defaults.get(
+        "positive_prompt",
+        arch_level_defaults.get("positive_prompt", ""),
+    )
+    default_neg = model_specific_defaults.get(
+        "negative_prompt",
+        arch_level_defaults.get("negative_prompt", ""),
+    )
+    if default_pos:
+        result["default_positive_prompt"] = default_pos
+    if default_neg:
+        result["default_negative_prompt"] = default_neg
+
+    return result

mcp_tools/get_model_list.py

@@ -0,0 +1,64 @@
+"""
+MCP Tool: get_model_list
+Query the list of available image generation models, with optional filtering by model architecture.
+"""
+
+from .common import _load_yaml, _MODEL_LIST_PATH, _MODEL_DEFAULTS_PATH, _MODEL_ARCHITECTURES_PATH
+from .error_schema import make_not_found_error
+
+
+def handle_get_model_list(model_architecture: str = None) -> list | dict:
+    """Dynamically load the list of available image generation models from model_list.yaml."""
+    model_list = _load_yaml(_MODEL_LIST_PATH)
+    model_defaults = _load_yaml(_MODEL_DEFAULTS_PATH)
+    arch_config = _load_yaml(_MODEL_ARCHITECTURES_PATH)
+    valid_architectures = set(arch_config.get("architectures", {}).keys())
+
+    if model_architecture and model_architecture not in valid_architectures:
+        return make_not_found_error("architecture", model_architecture)
+
+    result = []
+    checkpoints = model_list.get("Checkpoint", {})
+
+    for arch_name, arch_data in checkpoints.items():
+        if model_architecture and arch_name != model_architecture:
+            continue
+        if not isinstance(arch_data, dict):
+            continue
+
+        models = arch_data.get("models", [])
+        if not isinstance(models, list):
+            continue
+
+        arch_defaults = model_defaults.get(arch_name, {})
+        arch_level_defaults = arch_defaults.get("_defaults", {})
+
+        for model in models:
+            display_name = model.get("display_name", "")
+            category = model.get("category", None)
+
+            model_specific_defaults = arch_defaults.get(display_name, {})
+
+            default_pos = model_specific_defaults.get(
+                "positive_prompt",
+                arch_level_defaults.get("positive_prompt", ""),
+            )
+            default_neg = model_specific_defaults.get(
+                "negative_prompt",
+                arch_level_defaults.get("negative_prompt", ""),
+            )
+
+            entry = {
+                "name": display_name,
+                "model_architecture": arch_name,
+            }
+            if category:
+                entry["category"] = category
+            if default_pos:
+                entry["default_positive_prompt"] = default_pos
+            if default_neg:
+                entry["default_negative_prompt"] = default_neg
+
+            result.append(entry)
+
+    return result

mcp_tools/get_task_list.py

@@ -0,0 +1,11 @@
+"""
+MCP Tool: get_task_list
+Get a list of all supported image generation task types along with their required/optional parameter lists.
+"""
+
+from .common import _TASK_DEFINITIONS
+
+
+def handle_get_task_list() -> list:
+    """Get a list of all supported image generation task types along with their required/optional parameter lists."""
+    return _TASK_DEFINITIONS

mcp_tools/get_task_status.py

@@ -0,0 +1,21 @@
+"""
+MCP Tool: get_task_status
+Query the processing progress and final results of an async image generation task.
+"""
+
+from .common import _TASKS_DB
+from .error_schema import make_validation_error, make_not_found_error
+
+
+def handle_get_task_status(task_id: str) -> dict:
+    """Query the processing progress and final results of an async image generation task."""
+    if not task_id:
+        return make_validation_error(
+            "Parameter 'task_id' is required.",
+            missing_fields=["task_id"],
+        )
+
+    if task_id not in _TASKS_DB:
+        return make_not_found_error("task", task_id)
+
+    return _TASKS_DB[task_id]

mcp_tools/mcp_gradio_integration.py

@@ -0,0 +1,137 @@
+"""
+MCP & Gradio Integration Module
+
+Provides:
+1. register_high_level_mcp_apis: Expose only 8 high-level abstract API/MCP endpoints (using gr.api without polluting the visual UI structure)
+2. cleanup_dependencies_api_names: Force cleanup of show_api attribute for non-high-level APIs in dependencies
+3. patch_gradio_api_suppression: No-op implementation retained for backward compatibility
+"""
+
+import json
+import gradio as gr
+
+from .get_task_list import handle_get_task_list
+from .get_model_architecture_list import handle_get_model_architecture_list
+from .get_model_list import handle_get_model_list
+from .get_feature_list import handle_get_feature_list
+from .get_model_features import handle_get_model_features
+from .get_chain_schema import handle_get_chain_schema
+from .run_imagegen import handle_run_imagegen
+from .get_task_status import handle_get_task_status
+
+HIGH_LEVEL_MCP_API_NAMES = {
+    "get_task_list",
+    "get_model_architecture_list",
+    "get_model_list",
+    "get_feature_list",
+    "get_model_features",
+    "run_imagegen",
+    "get_task_status",
+    "get_chain_schema",
+}
+
+
+def sanitize_keys(obj):
+    """Recursively ensure all dictionary keys are converted to str type to avoid Gradio 5 orjson TypeError: Dict key must be str."""
+    if isinstance(obj, dict):
+        return {str(k): sanitize_keys(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [sanitize_keys(x) for x in obj]
+    elif isinstance(obj, tuple):
+        return tuple(sanitize_keys(x) for x in obj)
+    return obj
+
+
+def patch_gradio_api_suppression():
+    """Retained for backward compatibility (no-op)."""
+    pass
+
+
+def cleanup_dependencies_api_names(demo):
+    """
+    Clean up residual auto-generated API names in demo.fns and demo.dependencies.
+    Force only the 8 high-level abstract MCP APIs to be exposed as public endpoints.
+    """
+    for fn in demo.fns.values():
+        api_name = getattr(fn, "api_name", None)
+        if api_name not in HIGH_LEVEL_MCP_API_NAMES:
+            fn.show_api = False
+
+    deps = getattr(demo, "dependencies", None)
+    if deps is None and hasattr(demo, "config") and isinstance(demo.config, dict):
+        deps = demo.config.get("dependencies", [])
+
+    if deps:
+        for dep in deps:
+            if isinstance(dep, dict):
+                api_name = dep.get("api_name")
+                if api_name not in HIGH_LEVEL_MCP_API_NAMES:
+                    dep["show_api"] = False
+
+    print("[MCP Protection] Cleaned up demo dependencies. Suppressed atomic API endpoints.")
+
+
+def register_high_level_mcp_apis(demo):
+    """
+    Explicitly register 8 high-level abstract MCP API endpoints on the Gradio demo using gr.api.
+    Using gr.api() never adds any visual UI components (such as Row, Textbox, Button, etc.), avoiding duplicate interface rendering.
+    """
+    def get_task_list() -> list:
+        """[Recommended Discovery Flow Step 1] Get a list of all supported image generation task types (txt2img, img2img, inpaint, outpaint, hires_fix) along with their required and optional parameter lists. Recommended flow: get_task_list -> get_model_architecture_list -> get_model_list -> [Path 1: Call run_imagegen directly (pass only required params) | Path 2: Call get_model_features to get official default hyperparams -> run_imagegen]."""
+        return sanitize_keys(handle_get_task_list())
+
+    def get_model_architecture_list() -> list:
+        """[Recommended Discovery Flow Step 2] Get a list of all supported model architectures (e.g., SD1.5, SDXL, FLUX, etc.) along with their default resolutions. It is recommended to call this tool before get_model_list to obtain valid model_architecture parameters for precise model filtering."""
+        return sanitize_keys(handle_get_model_architecture_list())
+
+    def get_model_list(model_architecture: str = "") -> list | dict:
+        """[Recommended Discovery Flow Step 3] Query the list of available image generation models. After obtaining models, choose one of two paths: 1. [Path 1 (Recommended - Minimal Mode)] Call run_imagegen directly with only required parameters. Do NOT guess steps/cfg/sampler/scheduler from experience; the server will automatically apply the model's optimal default hyperparameters. 2. [Path 2 (Explicit Alignment Mode)] First call get_model_features to query the model's officially recommended hyperparameters, then pass them to run_imagegen."""
+        arch = model_architecture.strip() if model_architecture else None
+        return sanitize_keys(handle_get_model_list(arch))
+
+    def get_feature_list() -> list:
+        """Get the list of supported advanced features along with their usage constraints and parameter schemas."""
+        return sanitize_keys(handle_get_feature_list())
+
+    def get_model_features(model: str = "") -> dict:
+        """Query metadata for the specified model, including supported task types, extended features, and official default inference parameters (steps, cfg, sampler, scheduler). This tool MUST be called when explicitly obtaining a model's optimal default hyperparameters (Path 2). Guessing or fabricating hyperparameters without querying is strictly prohibited."""
+        return sanitize_keys(handle_get_model_features(model.strip()))
+
+    def run_imagegen(json_params: str = "{}") -> dict:
+        """[Recommended Discovery Flow Step 4] Unified image generation task execution interface. Supports txt2img, img2img, and other tasks with chainable extended features. [IMPORTANT PARAMETER RULES] Do NOT guess or fabricate inference hyperparameters such as steps, cfg, sampler, scheduler! Path 1 (Recommended): Pass only required parameters (task_type, model, prompt, width, height), leave optional hyperparams empty (server uses optimal defaults). Path 2: If explicit hyperparams are needed, you MUST first call get_model_features to obtain official defaults before passing them."""
+        try:
+            if isinstance(json_params, dict):
+                params = json_params
+            else:
+                params = json.loads(json_params or "{}")
+        except Exception as e:
+            return {"error": {"code": "INVALID_JSON", "message": f"Failed to parse JSON params: {e}"}}
+        return sanitize_keys(handle_run_imagegen(params))
+
+    def get_task_status(task_id: str = "") -> dict:
+        """Query the progress, status, and final generated results of an async image generation task."""
+        return sanitize_keys(handle_get_task_status(task_id.strip()))
+
+    def get_chain_schema(chain_type: str = "") -> dict:
+        """Get the complete parameter schema and usage examples for a specified chain/injector type."""
+        return sanitize_keys(handle_get_chain_schema(chain_type.strip()))
+
+    funcs = [
+        get_task_list,
+        get_model_architecture_list,
+        get_model_list,
+        get_feature_list,
+        get_model_features,
+        run_imagegen,
+        get_task_status,
+        get_chain_schema,
+    ]
+
+    for func in funcs:
+        gr.api(func)
+
+    for fn in demo.fns.values():
+        if getattr(fn, "api_name", None) in HIGH_LEVEL_MCP_API_NAMES:
+            fn.show_api = True
+
+    print("[MCP Integration] Successfully registered 8 High-Level Abstract MCP APIs via gr.api().")

mcp_tools/run_imagegen.py

@@ -0,0 +1,76 @@
+"""
+MCP Tool: run_imagegen
+Unified image generation task submission and execution interface.
+"""
+
+import time
+import uuid
+import threading
+from .common import (
+    _load_yaml,
+    _MODEL_LIST_PATH,
+    _TASK_DEFINITIONS,
+    _TASKS_DB,
+    _execute_imagegen_pipeline,
+)
+from .error_schema import make_validation_error, make_not_found_error
+
+
+def handle_run_imagegen(params: dict) -> dict:
+    """Unified image generation task execution interface."""
+    if not isinstance(params, dict):
+        return make_validation_error("Request params must be an object.")
+
+    missing = []
+    for req_field in ["task_type", "model", "prompt"]:
+        if req_field not in params or not params[req_field]:
+            missing.append(req_field)
+    if missing:
+        return make_validation_error(
+            f"Missing required parameter(s): {', '.join(missing)}",
+            missing_fields=missing,
+        )
+
+    task_type = params["task_type"]
+    valid_tasks = [t["task_type"] for t in _TASK_DEFINITIONS]
+    if task_type not in valid_tasks:
+        return make_validation_error(
+            f"Invalid task_type '{task_type}'. Must be one of {valid_tasks}.",
+            invalid_fields={"task_type": f"Must be in {valid_tasks}"},
+        )
+
+    model_list = _load_yaml(_MODEL_LIST_PATH)
+    checkpoints = model_list.get("Checkpoint", {})
+    all_models = set()
+    for arch_name, arch_data in checkpoints.items():
+        if isinstance(arch_data, dict):
+            for m in arch_data.get("models", []):
+                all_models.add(m.get("display_name"))
+
+    if params["model"] not in all_models:
+        return make_not_found_error("model", params["model"])
+
+    task_id = f"img_task_{uuid.uuid4().hex[:10]}"
+    created_at = int(time.time())
+
+    _TASKS_DB[task_id] = {
+        "task_id": task_id,
+        "status": "queued",
+        "progress": 0,
+        "created_at": created_at,
+    }
+
+    async_exec = params.get("async_execution", False)
+
+    if async_exec:
+        t = threading.Thread(target=_execute_imagegen_pipeline, args=(task_id, params), daemon=True)
+        t.start()
+        return {
+            "status": "queued",
+            "task_id": task_id,
+            "poll_interval_ms": 2000,
+            "message": "Task queued successfully. Poll get_task_status for results.",
+        }
+    else:
+        _execute_imagegen_pipeline(task_id, params)
+        return _TASKS_DB[task_id]

mcp_tools/tool_handlers.py

@@ -0,0 +1,30 @@
+"""
+MCP Tool Handlers — Backward-compatible aggregation entry point.
+Core logic has been split into individual files (get_*.py and run_imagegen.py).
+"""
+
+from .get_task_list import handle_get_task_list
+from .get_model_architecture_list import handle_get_model_architecture_list
+from .get_model_list import handle_get_model_list
+from .get_feature_list import handle_get_feature_list
+from .get_model_features import handle_get_model_features
+from .get_chain_schema import handle_get_chain_schema
+from .run_imagegen import handle_run_imagegen
+from .get_task_status import handle_get_task_status
+from .common import (
+    _TASK_DEFINITIONS,
+    _TASKS_DB,
+    _load_yaml,
+    _execute_imagegen_pipeline,
+)
+
+__all__ = [
+    "handle_get_task_list",
+    "handle_get_model_architecture_list",
+    "handle_get_model_list",
+    "handle_get_feature_list",
+    "handle_get_model_features",
+    "handle_get_chain_schema",
+    "handle_run_imagegen",
+    "handle_get_task_status",
+]

requirements.txt

@@ -1,5 +1,5 @@
-comfyui-frontend-package==1.47.11
-comfyui-workflow-templates==0.11.23
+comfyui-frontend-package==1.47.12
+comfyui-workflow-templates==0.11.27
 comfyui-embedded-docs==0.5.9
 torch
 torchsde
@@ -23,7 +23,7 @@ SQLAlchemy>=2.0.0
 filelock
 av>=16.0.0
 comfy-kitchen==0.2.26
-comfy-aimdo==0.4.10
+comfy-aimdo==0.4.11
 requests
 simpleeval>=1.0.0
 blake3

ui/events/__init__.py

@@ -1,12 +1,12 @@
-from .main import attach_event_handlers
-from .config_loaders import (
-    load_controlnet_config,
-    get_cn_defaults,
-    load_anima_controlnet_lllite_config,
-    get_anima_cn_defaults,
-    load_diffsynth_controlnet_config,
-    get_diffsynth_cn_defaults,
-    load_krea2_controlnet_config,
-    get_krea2_cn_defaults,
-    load_ipadapter_config
+from .main import attach_event_handlers
+from .config_loaders import (
+    load_controlnet_config,
+    get_cn_defaults,
+    load_anima_controlnet_lllite_config,
+    get_anima_cn_defaults,
+    load_diffsynth_controlnet_config,
+    get_diffsynth_cn_defaults,
+    load_krea2_controlnet_config,
+    get_krea2_cn_defaults,
+    load_ipadapter_config
 )

ui/events/config_loaders.py

@@ -1,168 +1,168 @@
-import os
-import yaml
-from functools import lru_cache
-
-@lru_cache(maxsize=1)
-def load_controlnet_config():
-    _PROJECT_ROOT = os.path.dirname(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", {})
-    except Exception as e:
-        print(f"Error loading controlnet_models.yaml: {e}")
-        return {}
-
-
-def get_cn_defaults(arch_val):
-    cn_full_config = load_controlnet_config()
-    cn_config = cn_full_config.get(arch_val, [])
-    
-    if not cn_config:
-        return [], None, [], None, "None"
-        
-    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
-                
-    return all_types, default_type, series_choices, default_series, filepath
-
-
-@lru_cache(maxsize=1)
-def load_anima_controlnet_lllite_config():
-    _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-    _CN_MODEL_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'anima_controlnet_lllite_models.yaml')
-    try:
-        print("--- Loading anima_controlnet_lllite_models.yaml ---")
-        with open(_CN_MODEL_LIST_PATH, 'r', encoding='utf-8') as f:
-            config = yaml.safe_load(f)
-        print("--- ✅ anima_controlnet_lllite_models.yaml loaded successfully ---")
-        return config.get("Anima_ControlNet_Lllite", [])
-    except Exception as e:
-        print(f"Error loading anima_controlnet_lllite_models.yaml: {e}")
-        return []
-
-
-def get_anima_cn_defaults():
-    cn_config = load_anima_controlnet_lllite_config()
-    if not cn_config:
-        return [], None, [], None, "None"
-    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
-    return all_types, default_type, series_choices, default_series, filepath
-
-
-@lru_cache(maxsize=1)
-def load_diffsynth_controlnet_config():
-    _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-    _CN_MODEL_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'diffsynth_controlnet_models.yaml')
-    try:
-        print("--- Loading diffsynth_controlnet_models.yaml ---")
-        with open(_CN_MODEL_LIST_PATH, 'r', encoding='utf-8') as f:
-            config = yaml.safe_load(f)
-        print("--- ✅ diffsynth_controlnet_models.yaml loaded successfully ---")
-        return config.get("DiffSynth_ControlNet", {})
-    except Exception as e:
-        print(f"Error loading diffsynth_controlnet_models.yaml: {e}")
-        return {}
-
-
-def get_diffsynth_cn_defaults(arch_val):
-    cn_full_config = load_diffsynth_controlnet_config()
-    cn_config = cn_full_config.get(arch_val, [])
-    
-    if not cn_config:
-        return [], None, [], None, "None"
-        
-    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
-                
-    return all_types, default_type, series_choices, default_series, filepath
-
-
-@lru_cache(maxsize=1)
-def load_krea2_controlnet_config():
-    _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-    _CN_MODEL_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'krea2_controlnet_models.yaml')
-    try:
-        print("--- Loading krea2_controlnet_models.yaml ---")
-        with open(_CN_MODEL_LIST_PATH, 'r', encoding='utf-8') as f:
-            config = yaml.safe_load(f)
-        print("--- ✅ krea2_controlnet_models.yaml loaded successfully ---")
-        return config.get("Krea2_ControlNet", [])
-    except Exception as e:
-        print(f"Error loading krea2_controlnet_models.yaml: {e}")
-        return []
-
-def get_krea2_cn_defaults():
-    cn_config = load_krea2_controlnet_config()
-    if not cn_config:
-        return [], None, [], None, "None"
-        
-    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
-                
-    return all_types, default_type, series_choices, default_series, filepath
-
-
-@lru_cache(maxsize=1)
-def load_ipadapter_config():
-    _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-    _IPA_MODEL_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'ipadapter.yaml')
-    try:
-        print("--- Loading ipadapter.yaml ---")
-        with open(_IPA_MODEL_LIST_PATH, 'r', encoding='utf-8') as f:
-            config = yaml.safe_load(f)
-        print("--- ✅ ipadapter.yaml loaded successfully ---")
-        return config
-    except Exception as e:
-        print(f"Error loading ipadapter.yaml: {e}")
+import os
+import yaml
+from functools import lru_cache
+
+@lru_cache(maxsize=1)
+def load_controlnet_config():
+    _PROJECT_ROOT = os.path.dirname(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", {})
+    except Exception as e:
+        print(f"Error loading controlnet_models.yaml: {e}")
+        return {}
+
+
+def get_cn_defaults(arch_val):
+    cn_full_config = load_controlnet_config()
+    cn_config = cn_full_config.get(arch_val, [])
+    
+    if not cn_config:
+        return [], None, [], None, "None"
+        
+    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
+                
+    return all_types, default_type, series_choices, default_series, filepath
+
+
+@lru_cache(maxsize=1)
+def load_anima_controlnet_lllite_config():
+    _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    _CN_MODEL_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'anima_controlnet_lllite_models.yaml')
+    try:
+        print("--- Loading anima_controlnet_lllite_models.yaml ---")
+        with open(_CN_MODEL_LIST_PATH, 'r', encoding='utf-8') as f:
+            config = yaml.safe_load(f)
+        print("--- ✅ anima_controlnet_lllite_models.yaml loaded successfully ---")
+        return config.get("Anima_ControlNet_Lllite", [])
+    except Exception as e:
+        print(f"Error loading anima_controlnet_lllite_models.yaml: {e}")
+        return []
+
+
+def get_anima_cn_defaults():
+    cn_config = load_anima_controlnet_lllite_config()
+    if not cn_config:
+        return [], None, [], None, "None"
+    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
+    return all_types, default_type, series_choices, default_series, filepath
+
+
+@lru_cache(maxsize=1)
+def load_diffsynth_controlnet_config():
+    _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    _CN_MODEL_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'diffsynth_controlnet_models.yaml')
+    try:
+        print("--- Loading diffsynth_controlnet_models.yaml ---")
+        with open(_CN_MODEL_LIST_PATH, 'r', encoding='utf-8') as f:
+            config = yaml.safe_load(f)
+        print("--- ✅ diffsynth_controlnet_models.yaml loaded successfully ---")
+        return config.get("DiffSynth_ControlNet", {})
+    except Exception as e:
+        print(f"Error loading diffsynth_controlnet_models.yaml: {e}")
+        return {}
+
+
+def get_diffsynth_cn_defaults(arch_val):
+    cn_full_config = load_diffsynth_controlnet_config()
+    cn_config = cn_full_config.get(arch_val, [])
+    
+    if not cn_config:
+        return [], None, [], None, "None"
+        
+    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
+                
+    return all_types, default_type, series_choices, default_series, filepath
+
+
+@lru_cache(maxsize=1)
+def load_krea2_controlnet_config():
+    _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    _CN_MODEL_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'krea2_controlnet_models.yaml')
+    try:
+        print("--- Loading krea2_controlnet_models.yaml ---")
+        with open(_CN_MODEL_LIST_PATH, 'r', encoding='utf-8') as f:
+            config = yaml.safe_load(f)
+        print("--- ✅ krea2_controlnet_models.yaml loaded successfully ---")
+        return config.get("Krea2_ControlNet", [])
+    except Exception as e:
+        print(f"Error loading krea2_controlnet_models.yaml: {e}")
+        return []
+
+def get_krea2_cn_defaults():
+    cn_config = load_krea2_controlnet_config()
+    if not cn_config:
+        return [], None, [], None, "None"
+        
+    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
+                
+    return all_types, default_type, series_choices, default_series, filepath
+
+
+@lru_cache(maxsize=1)
+def load_ipadapter_config():
+    _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    _IPA_MODEL_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'ipadapter.yaml')
+    try:
+        print("--- Loading ipadapter.yaml ---")
+        with open(_IPA_MODEL_LIST_PATH, 'r', encoding='utf-8') as f:
+            config = yaml.safe_load(f)
+        print("--- ✅ ipadapter.yaml loaded successfully ---")
+        return config
+    except Exception as e:
+        print(f"Error loading ipadapter.yaml: {e}")
         return {}

ui/layout.py

@@ -37,4 +37,22 @@ def build_ui(event_handler_function):
         
         event_handler_function(ui_components, demo)
 
+        try:
+            import mcp_tools as mcp
+            if hasattr(mcp, "register_high_level_mcp_apis"):
+                mcp.register_high_level_mcp_apis(demo)
+                mcp.cleanup_dependencies_api_names(demo)
+            elif hasattr(mcp, 'MCP_FUNCTIONS') and isinstance(mcp.MCP_FUNCTIONS, list):
+                for func in mcp.MCP_FUNCTIONS:
+                    gr.api(func)
+                    print(f"✅ Registered MCP API endpoint: '{func.__name__}'")
+        except Exception as e:
+            print(f"⚠️ Warning registering MCP functions: {e}")
+
+        # Disable API exposure for all atomic UI event handlers
+        high_level_names = getattr(mcp, "HIGH_LEVEL_MCP_API_NAMES", set())
+        for fn in demo.fns.values():
+            if getattr(fn, "api_name", None) not in high_level_names:
+                fn.show_api = False
+
     return demo

yaml/image_gen_features.yaml

@@ -1,200 +1,200 @@
-# Feature names in enabled_chains correspond 1-to-1 with chain_injectors/<name>_injector.py
-krea-2:
-  enabled_chains:
-  - lora
-  - krea2_controlnet
-  - krea2_identity_edit
-  - krea2_style_reference
-  - pid
-
-mage-flow:
-  enabled_chains:
-  - reference_image
-
-joyai-image:
-  enabled_chains:
-  - joyai_image
-
-boogu-image:
-  enabled_chains:
-  - lora
-  - boogu_image_edit
-  - pid
-
-pixeldit:
-  enabled_chains:
-  - conditioning
-
-ideogram-4:
-  enabled_chains:
-  - vae
-  - pid
-
-lens:
-  enabled_chains:
-  - conditioning
-  - pid
-
-flux2-kv:
-  enabled_chains:
-  - lora
-  - reference_latent
-  - conditioning
-  - vae
-  - pid
-
-flux2:
-  enabled_chains:
-  - lora
-  - reference_latent
-  - conditioning
-  - vae
-  - pid
-
-ernie-image:
-  enabled_chains:
-  - conditioning
-  - pid
-
-z-image:
-  enabled_chains:
-  - lora
-  - diffsynth_controlnet
-  - conditioning
-  - vae
-  - pid
-
-qwen-image:
-  enabled_chains:
-  - lora
-  - conditioning
-  - controlnet
-  - qwen_image_edit
-  - vae
-  - pid
-
-longcat-image:
-  enabled_chains:
-  - lora
-  - conditioning
-  - pid
-
-cosmos-predict2:
-  enabled_chains:
-  - conditioning
-  - vae
-
-anima:
-  enabled_chains:
-  - lora
-  - conditioning
-  - anima_controlnet_lllite
-  - vae
-  - pid
-
-newbie-image:
-  enabled_chains:
-  - lora
-  - embedding
-  - conditioning
-  - vae
-  - pid
-
-kandinsky-5:
-  enabled_chains:
-  - conditioning
-  - vae
-  - pid
-
-ovis-image:
-  enabled_chains:
-  - conditioning
-  - vae
-  - pid
-
-hunyuanimage:
-  enabled_chains:
-  - conditioning
-  - vae
-
-chroma1-radiance:
-  enabled_chains:
-  - conditioning
-
-chroma1:
-  enabled_chains:
-  - conditioning
-  - vae
-  - pid
-
-omnigen2:
-  enabled_chains:
-  - reference_latent
-  - conditioning
-  - pid
-
-lumina:
-  enabled_chains:
-  - lora
-  - embedding
-  - conditioning
-  - vae
-  - pid
-
-hidream-o1:
-  enabled_chains:
-  - lora
-  - conditioning
-  - hidream_o1_smoothing
-  - hidream_o1_reference
-
-hidream-i1:
-  enabled_chains:
-  - lora
-  - conditioning
-  - pid
-
-flux1:
-  enabled_chains:
-  - lora
-  - flux1_ipadapter
-  - conditioning
-  - style
-  - controlnet
-  - vae
-  - pid
-
-auraflow:
-  enabled_chains:
-  - lora
-  - conditioning
-  - vae
-
-sd35:
-  enabled_chains:
-  - lora
-  - sd3_ipadapter
-  - embedding
-  - conditioning
-  - controlnet
-  - vae
-  - pid
-
-sdxl:
-  enabled_chains:
-  - lora
-  - ipadapter
-  - embedding
-  - conditioning
-  - controlnet
-  - vae
-  - pid
-
-sd15:
-  enabled_chains:
-  - lora
-  - ipadapter
-  - embedding
-  - conditioning
-  - controlnet
+# Feature names in enabled_chains correspond 1-to-1 with chain_injectors/<name>_injector.py
+krea-2:
+  enabled_chains:
+  - lora
+  - krea2_controlnet
+  - krea2_identity_edit
+  - krea2_style_reference
+  - pid
+
+mage-flow:
+  enabled_chains:
+  - reference_image
+
+joyai-image:
+  enabled_chains:
+  - joyai_image
+
+boogu-image:
+  enabled_chains:
+  - lora
+  - boogu_image_edit
+  - pid
+
+pixeldit:
+  enabled_chains:
+  - conditioning
+
+ideogram-4:
+  enabled_chains:
+  - vae
+  - pid
+
+lens:
+  enabled_chains:
+  - conditioning
+  - pid
+
+flux2-kv:
+  enabled_chains:
+  - lora
+  - reference_latent
+  - conditioning
+  - vae
+  - pid
+
+flux2:
+  enabled_chains:
+  - lora
+  - reference_latent
+  - conditioning
+  - vae
+  - pid
+
+ernie-image:
+  enabled_chains:
+  - conditioning
+  - pid
+
+z-image:
+  enabled_chains:
+  - lora
+  - diffsynth_controlnet
+  - conditioning
+  - vae
+  - pid
+
+qwen-image:
+  enabled_chains:
+  - lora
+  - conditioning
+  - controlnet
+  - qwen_image_edit
+  - vae
+  - pid
+
+longcat-image:
+  enabled_chains:
+  - lora
+  - conditioning
+  - pid
+
+cosmos-predict2:
+  enabled_chains:
+  - conditioning
+  - vae
+
+anima:
+  enabled_chains:
+  - lora
+  - conditioning
+  - anima_controlnet_lllite
+  - vae
+  - pid
+
+newbie-image:
+  enabled_chains:
+  - lora
+  - embedding
+  - conditioning
+  - vae
+  - pid
+
+kandinsky-5:
+  enabled_chains:
+  - conditioning
+  - vae
+  - pid
+
+ovis-image:
+  enabled_chains:
+  - conditioning
+  - vae
+  - pid
+
+hunyuanimage:
+  enabled_chains:
+  - conditioning
+  - vae
+
+chroma1-radiance:
+  enabled_chains:
+  - conditioning
+
+chroma1:
+  enabled_chains:
+  - conditioning
+  - vae
+  - pid
+
+omnigen2:
+  enabled_chains:
+  - reference_latent
+  - conditioning
+  - pid
+
+lumina:
+  enabled_chains:
+  - lora
+  - embedding
+  - conditioning
+  - vae
+  - pid
+
+hidream-o1:
+  enabled_chains:
+  - lora
+  - conditioning
+  - hidream_o1_smoothing
+  - hidream_o1_reference
+
+hidream-i1:
+  enabled_chains:
+  - lora
+  - conditioning
+  - pid
+
+flux1:
+  enabled_chains:
+  - lora
+  - flux1_ipadapter
+  - conditioning
+  - style
+  - controlnet
+  - vae
+  - pid
+
+auraflow:
+  enabled_chains:
+  - lora
+  - conditioning
+  - vae
+
+sd35:
+  enabled_chains:
+  - lora
+  - sd3_ipadapter
+  - embedding
+  - conditioning
+  - controlnet
+  - vae
+  - pid
+
+sdxl:
+  enabled_chains:
+  - lora
+  - ipadapter
+  - embedding
+  - conditioning
+  - controlnet
+  - vae
+  - pid
+
+sd15:
+  enabled_chains:
+  - lora
+  - ipadapter
+  - embedding
+  - conditioning
+  - controlnet
   - vae