fix: patch for multiple controlnets (#78)

--------- Co-authored-by: ryanontheinstide <[email protected]> Co-authored-by: John Mull <[email protected]>
yondonfu · Feb 25, 2025 · bd51d4e · bd51d4e
1 parent 020ded5
commit bd51d4e
Show file tree

Hide file tree

Showing 7 changed files with 2,322 additions and 1 deletion.
diff --git a/configs/models.yaml b/configs/models.yaml
@@ -37,4 +37,23 @@ models:
     name: "ControlNet Depth"
     url: "https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors/resolve/main/control_v11f1p_sd15_depth_fp16.safetensors"
     path: "controlnet/control_v11f1p_sd15_depth_fp16.safetensors"
-    type: "controlnet" 
+    type: "controlnet" 
+
+  controlnet-mediapipe-face:
+    name: "ControlNet MediaPipe Face"
+    url: "https://huggingface.co/CrucibleAI/ControlNetMediaPipeFace/resolve/main/control_v2p_sd15_mediapipe_face.safetensors"
+    path: "controlnet/control_v2p_sd15_mediapipe_face.safetensors"
+    type: "controlnet"
+
+  # Lora models
+  ral-polygon-sd15:
+    name: "ral-polygon-sd15"
+    url: "https://huggingface.co/Livepeer-Studio/comfystream_loras/resolve/main/ral-polygon-sd15.safetensors"
+    path: "loras/SD1.5/ral-polygon-sd15.safetensors"
+    type: "lora"
+
+  ral-chrome-sdxl:
+    name: "ral-chrome-sd15"
+    url: "https://civitai.com/api/download/models/276570?type=Model&format=SafeTensor"
+    path: "loras/SD1.5/ral-chrome-sd15.safetensors"
+    type: "lora"
diff --git a/configs/nodes.yaml b/configs/nodes.yaml
@@ -61,3 +61,10 @@ nodes:
     name: "ComfyUI Background Edit"
     url: "https://github.com/yondonfu/ComfyUI-Background-Edit"
     type: "utility"
+
+  comfyui_controlnet_aux:
+    name: "ComfyUI ControlNet Auxiliary"
+    url: "https://github.com/Fannovel16/comfyui_controlnet_aux"
+    type: "utility"
+    dependencies: 
+      - "mediapipe==0.10.8"
diff --git a/nodes/tensor_utils/prestartup_script.py b/nodes/tensor_utils/prestartup_script.py
@@ -0,0 +1,53 @@
+import torch
+
+print("Loading pre-startup script for controlnet torch.compile() compatibility...")
+
+def patch_controlnet_for_stream():
+    """Patch ControlNet for better compatibility with torch.compile()"""
+    try:
+        from comfy.controlnet import ControlBase
+        original_control_merge = ControlBase.control_merge
+
+        def wrapped_control_merge(self, control, control_prev, output_dtype):
+            # Mark CUDA graph step at start
+            if torch.cuda.is_available() and hasattr(torch.compiler, 'cudagraph_mark_step_begin'):
+                torch.compiler.cudagraph_mark_step_begin()
+                torch.cuda.synchronize()
+
+            # Deep clone control outputs to prevent CUDA graph overwrites
+            control = {
+                k: [t.clone() if t is not None else None for t in v]
+                for k, v in control.items()
+            }
+
+            # Deep clone previous control if it exists
+            if control_prev is not None:
+                control_prev = {
+                    k: [t.clone() if t is not None else None for t in v]
+                    for k, v in control_prev.items()
+                }
+
+            # Get result from original merge function
+            result = original_control_merge(self, control, control_prev, output_dtype)
+
+            # Clone all output tensors
+            result = {
+                k: [t.clone() if t is not None else None for t in v]
+                for k, v in result.items()
+            }
+
+            # Mark CUDA graph step at end
+            if torch.cuda.is_available() and hasattr(torch.compiler, 'cudagraph_mark_step_begin'):
+                torch.compiler.cudagraph_mark_step_begin()
+                torch.cuda.synchronize()
+
+            return result
+
+        # Apply the patch
+        ControlBase.control_merge = wrapped_control_merge
+        print("Successfully patched ControlNet for torch.compile() compatibility")
+    except Exception as e:
+        print(f"Warning: Failed to patch ControlNet: {str(e)}")
+
+# Apply patch when module is imported
+patch_controlnet_for_stream() 
diff --git a/workflows/sd15_multi_cnet_depthanything_face.json b/workflows/sd15_multi_cnet_depthanything_face.json
@@ -0,0 +1,295 @@
+{
+  "2": {
+    "inputs": {
+      "engine": "depth_anything_v2_vitl_fp16.engine",
+      "images": [
+        "26",
+        0
+      ]
+    },
+    "class_type": "DepthAnythingTensorrt",
+    "_meta": {
+      "title": "Depth Anything Tensorrt"
+    }
+  },
+  "3": {
+    "inputs": {
+      "unet_name": "ComfyUI_STAT_SD15_$stat-b-1-h-512-w-512_00001_.engine",
+      "model_type": "SD15"
+    },
+    "class_type": "TensorRTLoader",
+    "_meta": {
+      "title": "TensorRT Loader"
+    }
+  },
+  "5": {
+    "inputs": {
+      "text": "beautiful pencil sketch, masterpiece ",
+      "clip": [
+        "18",
+        0
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Prompt)"
+    }
+  },
+  "6": {
+    "inputs": {
+      "text": "",
+      "clip": [
+        "18",
+        0
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Prompt)"
+    }
+  },
+  "7": {
+    "inputs": {
+      "seed": 59962529669489,
+      "steps": 2,
+      "cfg": 1,
+      "sampler_name": "lcm",
+      "scheduler": "normal",
+      "denoise": 1,
+      "model": [
+        "3",
+        0
+      ],
+      "positive": [
+        "19",
+        0
+      ],
+      "negative": [
+        "19",
+        1
+      ],
+      "latent_image": [
+        "16",
+        0
+      ]
+    },
+    "class_type": "KSampler",
+    "_meta": {
+      "title": "KSampler"
+    }
+  },
+  "8": {
+    "inputs": {
+      "control_net_name": "1.5/control_v11f1p_sd15_depth_fp16.safetensors"
+    },
+    "class_type": "ControlNetLoader",
+    "_meta": {
+      "title": "Load ControlNet Model"
+    }
+  },
+  "9": {
+    "inputs": {
+      "strength": 0.9,
+      "start_percent": 0,
+      "end_percent": 1,
+      "positive": [
+        "5",
+        0
+      ],
+      "negative": [
+        "6",
+        0
+      ],
+      "control_net": [
+        "10",
+        0
+      ],
+      "image": [
+        "2",
+        0
+      ]
+    },
+    "class_type": "ControlNetApplyAdvanced",
+    "_meta": {
+      "title": "Apply ControlNet"
+    }
+  },
+  "10": {
+    "inputs": {
+      "backend": "inductor",
+      "fullgraph": false,
+      "mode": "reduce-overhead",
+      "controlnet": [
+        "8",
+        0
+      ]
+    },
+    "class_type": "TorchCompileLoadControlNet",
+    "_meta": {
+      "title": "TorchCompileLoadControlNet"
+    }
+  },
+  "11": {
+    "inputs": {
+      "vae_name": "taesd"
+    },
+    "class_type": "VAELoader",
+    "_meta": {
+      "title": "Load VAE"
+    }
+  },
+  "13": {
+    "inputs": {
+      "backend": "inductor",
+      "fullgraph": true,
+      "mode": "reduce-overhead",
+      "compile_encoder": true,
+      "compile_decoder": true,
+      "vae": [
+        "11",
+        0
+      ]
+    },
+    "class_type": "TorchCompileLoadVAE",
+    "_meta": {
+      "title": "TorchCompileLoadVAE"
+    }
+  },
+  "14": {
+    "inputs": {
+      "samples": [
+        "7",
+        0
+      ],
+      "vae": [
+        "13",
+        0
+      ]
+    },
+    "class_type": "VAEDecode",
+    "_meta": {
+      "title": "VAE Decode"
+    }
+  },
+  "15": {
+    "inputs": {
+      "images": [
+        "14",
+        0
+      ]
+    },
+    "class_type": "PreviewImage",
+    "_meta": {
+      "title": "Preview Image"
+    }
+  },
+  "16": {
+    "inputs": {
+      "width": 512,
+      "height": 512,
+      "batch_size": 1
+    },
+    "class_type": "EmptyLatentImage",
+    "_meta": {
+      "title": "Empty Latent Image"
+    }
+  },
+  "18": {
+    "inputs": {
+      "stop_at_clip_layer": -2,
+      "clip": [
+        "27",
+        1
+      ]
+    },
+    "class_type": "CLIPSetLastLayer",
+    "_meta": {
+      "title": "CLIP Set Last Layer"
+    }
+  },
+  "19": {
+    "inputs": {
+      "strength": 0.9,
+      "start_percent": 0,
+      "end_percent": 1,
+      "positive": [
+        "9",
+        0
+      ],
+      "negative": [
+        "9",
+        1
+      ],
+      "control_net": [
+        "21",
+        0
+      ],
+      "image": [
+        "22",
+        0
+      ]
+    },
+    "class_type": "ControlNetApplyAdvanced",
+    "_meta": {
+      "title": "Apply ControlNet"
+    }
+  },
+  "20": {
+    "inputs": {
+      "control_net_name": "control_v2p_sd15_mediapipe_face.safetensors"
+    },
+    "class_type": "ControlNetLoader",
+    "_meta": {
+      "title": "Load ControlNet Model"
+    }
+  },
+  "21": {
+    "inputs": {
+      "backend": "inductor",
+      "fullgraph": false,
+      "mode": "reduce-overhead",
+      "controlnet": [
+        "20",
+        0
+      ]
+    },
+    "class_type": "TorchCompileLoadControlNet",
+    "_meta": {
+      "title": "TorchCompileLoadControlNet"
+    }
+  },
+  "22": {
+    "inputs": {
+      "max_faces": 1,
+      "min_confidence": 0.5,
+      "resolution": 512,
+      "image": [
+        "26",
+        0
+      ]
+    },
+    "class_type": "MediaPipe-FaceMeshPreprocessor",
+    "_meta": {
+      "title": "MediaPipe Face Mesh"
+    }
+  },
+  "26": {
+    "inputs": {
+      "image": "pasted/image (2).png",
+      "upload": "image"
+    },
+    "class_type": "LoadImage",
+    "_meta": {
+      "title": "Load Image"
+    }
+  },
+  "27": {
+    "inputs": {
+      "ckpt_name": "dreamshaper_8.safetensors"
+    },
+    "class_type": "CheckpointLoaderSimple",
+    "_meta": {
+      "title": "Load Checkpoint"
+    }
+  }
+}