Add ControlNet, DiffSynth ControlNet & VAE injectors

chain_injectors/controlnet_injector.py

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

chain_injectors/diffsynth_controlnet_injector.py

@@ -0,0 +1,75 @@
+def inject(assembler, chain_definition, chain_items):
+    if not chain_items:
+        return
+
+    model_sampler_name = chain_definition.get('model_sampler_node')
+    ksampler_name = chain_definition.get('ksampler_node', 'ksampler')
+    
+    target_node_id = None
+    target_input_name = 'model'
+    
+    if model_sampler_name and model_sampler_name in assembler.node_map:
+        model_sampler_id = assembler.node_map[model_sampler_name]
+        if target_input_name in assembler.workflow[model_sampler_id]['inputs']:
+            target_node_id = model_sampler_id
+            print(f"ControlNet Model Patch injector targeting ModelSamplingAuraFlow node '{model_sampler_name}'.")
+
+    if not target_node_id:
+        if ksampler_name in assembler.node_map:
+            ksampler_id = assembler.node_map[ksampler_name]
+            if target_input_name in assembler.workflow[ksampler_id]['inputs']:
+                target_node_id = ksampler_id
+                print(f"ControlNet Model Patch injector targeting KSampler node '{ksampler_name}'.")
+        else:
+            print(f"Warning: Neither ModelSamplingAuraFlow node '{model_sampler_name}' nor KSampler node '{ksampler_name}' found for ControlNet patch chain. Skipping.")
+            return
+            
+    if not target_node_id:
+        print(f"Warning: Could not find a valid 'model' input on target nodes. Skipping ControlNet patch chain.")
+        return
+
+    current_model_connection = assembler.workflow[target_node_id]['inputs'][target_input_name]
+    
+    vae_source_str = chain_definition.get('vae_source')
+    vae_connection = None
+    if vae_source_str:
+        try:
+            vae_node_name, vae_idx_str = vae_source_str.split(':')
+            if vae_node_name in assembler.node_map:
+                vae_connection = [assembler.node_map[vae_node_name], int(vae_idx_str)]
+            else:
+                print(f"Warning: VAE source node '{vae_node_name}' not found for ControlNet patch chain. VAE will not be connected.")
+        except ValueError:
+            print(f"Warning: Invalid 'vae_source' format '{vae_source_str}' for ControlNet patch chain. Expected 'node_name:index'. VAE will not be connected.")
+    else:
+        print(f"Warning: 'vae_source' not defined for ControlNet patch chain definition. VAE may not be connected.")
+
+    for item_data in chain_items:
+        patch_loader_id = assembler._get_unique_id()
+        patch_loader_node = assembler._get_node_template("ModelPatchLoader")
+        patch_loader_node['inputs']['name'] = item_data['control_net_name']
+        assembler.workflow[patch_loader_id] = patch_loader_node
+
+        image_loader_id = assembler._get_unique_id()
+        image_loader_node = assembler._get_node_template("LoadImage")
+        image_loader_node['inputs']['image'] = item_data['image']
+        assembler.workflow[image_loader_id] = image_loader_node
+
+        apply_cn_id = assembler._get_unique_id()
+        apply_cn_node = assembler._get_node_template(chain_definition['template'])
+        
+        apply_cn_node['inputs']['strength'] = item_data.get('strength', 1.0)
+        apply_cn_node['inputs']['model'] = current_model_connection
+        apply_cn_node['inputs']['model_patch'] = [patch_loader_id, 0]
+        apply_cn_node['inputs']['image'] = [image_loader_id, 0]
+        
+        if 'vae' in apply_cn_node['inputs'] and vae_connection:
+            apply_cn_node['inputs']['vae'] = vae_connection
+        
+        assembler.workflow[apply_cn_id] = apply_cn_node
+
+        current_model_connection = [apply_cn_id, 0]
+
+    assembler.workflow[target_node_id]['inputs'][target_input_name] = current_model_connection
+    
+    print(f"ControlNet Model Patch injector applied. Target 'model' input re-routed through {len(chain_items)} patch(es).")

chain_injectors/vae_injector.py

@@ -0,0 +1,30 @@
+def inject(assembler, chain_definition, chain_items):
+    if not chain_items:
+        return
+
+    vae_name = chain_items[0] if isinstance(chain_items, list) else chain_items
+    if not vae_name or vae_name == "None":
+        return
+
+    targets = chain_definition.get('targets', [])
+    if not targets:
+        return
+
+    vae_loader_id = assembler._get_unique_id()
+    vae_loader_node = assembler._get_node_template("VAELoader")
+    vae_loader_node['inputs']['vae_name'] = vae_name
+    assembler.workflow[vae_loader_id] = vae_loader_node
+
+    injected_count = 0
+    for target_str in targets:
+        try:
+            node_name, input_name = target_str.split(':')
+            if node_name in assembler.node_map:
+                node_id = assembler.node_map[node_name]
+                assembler.workflow[node_id]['inputs'][input_name] = [vae_loader_id, 0]
+                injected_count += 1
+        except ValueError:
+            print(f"Warning: Invalid VAE injector target format '{target_str}'. Expected 'node_name:input_name'.")
+    
+    if injected_count > 0:
+        print(f"VAE injector applied. Rerouted {injected_count} connection(s) to new VAELoader ({vae_name}).")

yaml/controlnet_models.yaml

@@ -0,0 +1,177 @@
+ControlNet:
+  Qwen-Image:
+    - Filepath: "Qwen-Image-InstantX-ControlNet-Union.safetensors"
+      Series: "InstantX Union"
+      Type: ["Canny", "Soft Edge", "Depth", "Pose"]
+  FLUX.1:
+    - Filepath: "FLUX.1-dev-ControlNet-Union-Pro-2.0.safetensors"
+      Series: "Shakker Labs Union Pro"
+      Type: ["Canny", "Tile", "Depth", "Blur", "Pose", "Gray", "Low Quality"]
+    - Filepath: "flux-canny-controlnet-v3.safetensors"
+      Series: "XLabs-AI"
+      Type: ["Canny"]
+    - Filepath: "flux-depth-controlnet-v3.safetensors"
+      Series: "XLabs-AI"
+      Type: ["Depth (Midas)"]
+    - Filepath: "flux-hed-controlnet-v3.safetensors"
+      Series: "XLabs-AI"
+      Type: ["HED"]
+  SD3.5:
+    - Filepath: "sd3.5_large_controlnet_blur.safetensors"
+      Series: "[Large] StabilityAI"
+      Type: ["Blur"]
+    - Filepath: "sd3.5_large_controlnet_canny.safetensors"
+      Series: "[Large] StabilityAI"
+      Type: ["Canny"]
+    - Filepath: "sd3.5_large_controlnet_depth.safetensors"
+      Series: "[Large] StabilityAI"
+      Type: ["Depth"]
+  SDXL:
+    - Filepath: "controlnet-union-sdxl-1.0_promax.safetensors"
+      Series: "xinsir Union"
+      Type: ["Tile Deblur", "Tile variation", "Tile Super Resolution", "Image Inpainting", "Image Outpainting", "OpenPose", "Depth", "Canny", "Lineart", "Anime Lineart", "Mlsd", "Scribble", "Hed", "Pidi(Softedge)", "Teed", "Segment", "Normal"]
+    - Filepath: "controlnet-tile-sdxl-1.0.safetensors"
+      Series: "xinsir"
+      Type: ["Tile"]
+    - Filepath: "controlnet-canny-sdxl-1.0_V2.safetensors"
+      Series: "xinsir"
+      Type: ["Canny"]
+    - Filepath: "controlnet-openpose-sdxl-1.0.safetensors"
+      Series: "xinsir"
+      Type: ["OpenPose"]
+    - Filepath: "controlnet-openpose-sdxl-1.0.safetensors"
+      Series: "xinsir"
+      Type: ["OpenPose(Twins)"]
+    - Filepath: "controlnet-depth-sdxl-1.0"
+      Series: "xinsir"
+      Type: ["Depth"]
+    - Filepath: "controlnet-scribble-sdxl-1.0.safetensors"
+      Series: "xinsir"
+      Type: ["Scribble"]
+    - Filepath: "anime-painter.safetensors"
+      Series: "xinsir"
+      Type: ["Anime Painter"]
+      # SDXL-NoobAI
+    - Filepath: "noob_sdxl_controlnet_canny.fp16.safetensors"
+      Series: "NoobAI"
+      Type: ["Canny"]
+    - Filepath: "noob_sdxl_controlnet_depth.fp16.safetensors"
+      Series: "NoobAI"
+      Type: ["Depth"]
+    - Filepath: "noob-sdxl-controlnet-lineart_anime.fp16.safetensors"
+      Series: "NoobAI"
+      Type: ["Anime Lineart"]
+    - Filepath: "noob-sdxl-controlnet-lineart_realistic.fp16.safetensors"
+      Series: "NoobAI"
+      Type: ["Realistic Lineart"]
+    - Filepath: "noob-sdxl-controlnet-manga_line.fp16.safetensors"
+      Series: "NoobAI"
+      Type: ["Manga Lineart"]
+    - Filepath: "noob-sdxl-controlnet-normal.fp16.safetensors"
+      Series: "NoobAI"
+      Type: ["Normal"]
+    - Filepath: "noob-sdxl-controlnet-softedge_hed.fp16.safetensors"
+      Series: "NoobAI"
+      Type: ["SoftEdge (HED)"]
+    - Filepath: "noob-sdxl-controlnet-tile.fp16.safetensors"
+      Series: "NoobAI"
+      Type: ["Tile"]
+    - Filepath: "NoobAI_Inpainting_ControlNet.safetensors"
+      Series: "NoobAI"
+      Type: ["Inpainting"]
+    - Filepath: "noob_openpose_pre.safetensors"
+      Series: "NoobAI"
+      Type: ["OpenPose"]
+    - Filepath: "noobaiXLControlnet_epsDepthMidasV11.safetensors"
+      Series: "NoobAI"
+      Type: ["Depth (Midas)"]
+    - Filepath: "noobaiXLControlnet_epsNormalMidas.safetensors"
+      Series: "NoobAI"
+      Type: ["Normal (Midas)"]
+    - Filepath: "noobaiXLControlnet_epsScribbleHed.safetensors"
+      Series: "NoobAI"
+      Type: ["Scribble (HED)"]
+    - Filepath: "noobaiXLControlnet_epsScribblePidinet.safetensors"
+      Series: "NoobAI"
+      Type: ["Scribble (PiDiNet)"]
+    - Filepath: "noobaiXLControlnet_epsDepthMidas.safetensors"
+      Series: "NoobAI"
+      Type: ["Depth (Midas)"]
+    - Filepath: "noobaiXLControlnet_openposeModel.safetensors"
+      Series: "NoobAI"
+      Type: ["OpenPose"]
+      # SDXL - Illustrious-XL-v0.1
+    - Filepath: "illustriousXLv0.1_Canny_fp16.safetensors"
+      Series: "MIC-Lab/illustriousXLv0.1_controlnet"
+      Type: ["Canny"]
+    - Filepath: "illustriousXLv0.1_Lineart_fp16.safetensors"
+      Series: "MIC-Lab/illustriousXLv0.1_controlnet"
+      Type: ["Lineart"]
+    - Filepath: "illustriousXLv0.1_Softedge_fp16.safetensors"
+      Series: "MIC-Lab/illustriousXLv0.1_controlnet"
+      Type: ["SoftEdge"]
+    - Filepath: "illustriousXLv0.1_Tile_fp16.safetensors"
+      Series: "MIC-Lab/illustriousXLv0.1_controlnet"
+      Type: ["Tile"]
+    - Filepath: "illustriousXLv0.1_depth_midas_fp16.safetensors"
+      Series: "MIC-Lab/illustriousXLv0.1_controlnet"
+      Type: ["Depth (Midas)"]
+    - Filepath: "illustriousXLv0.1_inpainting_fp16.safetensors"
+      Series: "MIC-Lab/illustriousXLv0.1_controlnet"
+      Type: ["Inpainting"]
+      # SDXL - Illustrious-XL-v1.1
+    - Filepath: "illustriousXLv1.1_canny_fp16.safetensors"
+      Series: "MIC-Lab/illustriousXLv1.1_controlnet"
+      Type: ["Canny"]
+    - Filepath: "illustriousXLv1.1_depth_midas_fp16.safetensors"
+      Series: "MIC-Lab/illustriousXLv1.1_controlnet"
+      Type: ["Depth (Midas)"]
+    - Filepath: "illustriousXLv1.1_inpainting_fp16.safetensors"
+      Series: "MIC-Lab/illustriousXLv1.1_controlnet"
+      Type: ["Inpainting"]
+    - Filepath: "illustriousXLv1.1_tile_fp16.safetensors"
+      Series: "MIC-Lab/illustriousXLv1.1_controlnet"
+      Type: ["Tile"]
+  SD1.5:
+    - Filepath: "control_v11p_sd15_canny_fp16.safetensors"
+      Series: "Standard"
+      Type: ["Canny"]
+    - Filepath: "control_v11f1p_sd15_depth_fp16.safetensors"
+      Series: "Standard"
+      Type: ["Depth"]
+    - Filepath: "control_v11p_sd15_openpose_fp16.safetensors"
+      Series: "Standard"
+      Type: ["OpenPose"]
+    - Filepath: "control_v11p_sd15_lineart_fp16.safetensors"
+      Series: "Standard"
+      Type: ["Lineart"]
+    - Filepath: "control_v11p_sd15_softedge_fp16.safetensors"
+      Series: "Standard"
+      Type: ["SoftEdge"]
+    - Filepath: "control_v11p_sd15_scribble_fp16.safetensors"
+      Series: "Standard"
+      Type: ["Scribble"]
+    - Filepath: "control_v11p_sd15_seg_fp16.safetensors"
+      Series: "Standard"
+      Type: ["Segmentation"]
+    - Filepath: "control_v11p_sd15_normalbae_fp16.safetensors"
+      Series: "Standard"
+      Type: ["Normal BAE"]
+    - Filepath: "control_v11p_sd15_mlsd_fp16.safetensors"
+      Series: "Standard"
+      Type: ["MLSD"]
+    - Filepath: "control_v11p_sd15_inpaint_fp16.safetensors"
+      Series: "Standard"
+      Type: ["Inpaint"]
+    - Filepath: "control_v11f1e_sd15_tile_fp16.safetensors"
+      Series: "Standard"
+      Type: ["Tile"]
+    - Filepath: "control_v11e_sd15_shuffle_fp16.safetensors"
+      Series: "Standard"
+      Type: ["Shuffle"]
+    - Filepath: "control_v11e_sd15_ip2p_fp16.safetensors"
+      Series: "Standard"
+      Type: ["Instruct P2P"]
+    - Filepath: "control_v11p_sd15s2_lineart_anime_fp16.safetensors"
+      Series: "Standard"
+      Type: ["Anime Lineart"]

yaml/diffsynth_controlnet_models.yaml

@@ -0,0 +1,8 @@
+DiffSynth_ControlNet:
+  Z-Image:
+    - Filepath: "Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.safetensors"
+      Series: "alibaba-pai Controlnet Union 2.1 8steps"
+      Type: ["Pose", "Canny", "HED", "Depth", "MLSD"]
+    - Filepath: "Z-Image-Turbo-Fun-Controlnet-Tile-2.1-8steps.safetensors"
+      Series: "alibaba-pai Controlnet Union 2.1 8steps"
+      Type: ["Tile"]