From c19130957de9841da0192c9e86d34cdd6d04b2c1 Mon Sep 17 00:00:00 2001
From: ussoewwin <136552381+ussoewwin@users.noreply.github.com>
Date: Tue, 7 Jul 2026 16:38:55 +0900
Subject: [PATCH] \ fix: isolate controlnet routing by model family\

\Add strict route classification for ZI,Nunchaku QI,standard QI,and Krea2. Prevent fallback mixing by raising on unknown routes.\
---
 nodes/controlnet.py | 87 ++++++++++++++++++++++++++-------------------
 1 file changed, 50 insertions(+), 37 deletions(-)

diff --git a/nodes/controlnet.py b/nodes/controlnet.py
index a271594..6849c2c 100644
--- a/nodes/controlnet.py
+++ b/nodes/controlnet.py
@@ -8,6 +8,43 @@ import logging
 logger = logging.getLogger(__name__)
 
 
+def _classify_controlnet_target(model, model_patch):
+    """
+    Classify ControlNet target route strictly to avoid mixing model families.
+
+    Returns one of:
+      - "zimage"
+      - "nunchaku_qwenimage"
+      - "qwenimage_standard"
+      - "krea2"
+      - "unknown"
+    """
+    if isinstance(model_patch.model, comfy.ldm.lumina.controlnet.ZImage_Control):
+        return "zimage"
+
+    if not hasattr(model, "model"):
+        return "unknown"
+
+    model_base_name = model.model.__class__.__name__
+    if model_base_name == "NunchakuQwenImage":
+        return "nunchaku_qwenimage"
+
+    if not hasattr(model.model, "diffusion_model"):
+        return "unknown"
+
+    dm = model.model.diffusion_model
+    dm_name = dm.__class__.__name__
+    dm_module = getattr(dm.__class__, "__module__", "")
+
+    if dm_name in ("ComfyQwenImageWrapper", "NunchakuQwenImageTransformer2DModel"):
+        return "qwenimage_standard"
+
+    if dm_name == "SingleStreamDiT" or "comfy.ldm.krea2" in dm_module:
+        return "krea2"
+
+    return "unknown"
+
+
 class DiffSynthCnetPatch:
     def __init__(self, model_patch, vae, image, strength, mask=None):
         self.model_patch = model_patch
@@ -232,45 +269,16 @@ class NunchakuQwenImageDiffsynthControlnet:
                 mask = mask.unsqueeze(2)
             mask = 1.0 - mask
 
-        # Detect ControlNet model type
-        is_zimage_control = isinstance(model_patch.model, comfy.ldm.lumina.controlnet.ZImage_Control)
+        route = _classify_controlnet_target(model, model_patch)
+        logger.info(f"[ControlNet] route={route}")
 
-        # Detect Nunchaku Qwen Image model
-        # Nunchaku's _forward() processes patches_replace but ignores patches["double_block"],
-        # so we need to use set_model_patch_replace() instead of set_model_double_block_patch().
-        is_nunchaku_qwenimage = False
-        is_krea2_model = False
-        if not is_zimage_control and hasattr(model, 'model'):
-            # Check model base class name (NunchakuQwenImage from ComfyUI-nunchaku model_base)
-            model_base_name = model.model.__class__.__name__
-            if model_base_name == 'NunchakuQwenImage':
-                is_nunchaku_qwenimage = True
-            # Check diffusion model class name
-            elif hasattr(model.model, 'diffusion_model'):
-                dm_name = model.model.diffusion_model.__class__.__name__
-                if dm_name in ('ComfyQwenImageWrapper', 'NunchakuQwenImageTransformer2DModel'):
-                    is_nunchaku_qwenimage = True
-                # Krea2 backbone in ComfyUI (comfy.ldm.krea2.model.SingleStreamDiT)
-                elif dm_name == 'SingleStreamDiT':
-                    is_krea2_model = True
-                else:
-                    dm_module = getattr(model.model.diffusion_model.__class__, "__module__", "")
-                    if "comfy.ldm.krea2" in dm_module:
-                        is_krea2_model = True
-
-        logger.info(
-            f"[ControlNet] is_zimage={is_zimage_control}, "
-            f"is_nunchaku_qwenimage={is_nunchaku_qwenimage}, "
-            f"is_krea2={is_krea2_model}"
-        )
-
-        if is_zimage_control:
+        if route == "zimage":
             # ZImage ControlNet (works for both standard and Nunchaku Z-Image)
             logger.info("[ControlNet] Using ZImageControlPatch")
             patch = ZImageControlPatch(model_patch, vae, image, strength, mask=mask)
             model_patched.set_model_noise_refiner_patch(patch)
             model_patched.set_model_double_block_patch(patch)
-        elif is_nunchaku_qwenimage:
+        elif route == "nunchaku_qwenimage":
             # Nunchaku Qwen Image: use patches_replace since _forward ignores patches["double_block"]
             logger.info("[ControlNet] Using DiffSynthCnetBlockReplace for Nunchaku Qwen Image")
             cnet_patch = DiffSynthCnetPatch(model_patch, vae, image, strength, mask)
@@ -290,13 +298,18 @@ class NunchakuQwenImageDiffsynthControlnet:
                     "dit", "double_block", i
                 )
             logger.info(f"[ControlNet] Registered {num_blocks} patches_replace entries for Nunchaku Qwen Image")
-        elif is_krea2_model:
-            # Krea2 model path: keep control processing in this patcher (not in LoRA loader).
+        elif route == "krea2":
+            # Krea2 path is isolated here by explicit routing.
             logger.info("[ControlNet] Using DiffSynthCnetPatch (krea2)")
             model_patched.set_model_double_block_patch(DiffSynthCnetPatch(model_patch, vae, image, strength, mask))
-        else:
+        elif route == "qwenimage_standard":
             # Standard Qwen Image: use patches["double_block"] (processed by ComfyUI's _forward)
-            logger.info("[ControlNet] Using DiffSynthCnetPatch (standard)")
+            logger.info("[ControlNet] Using DiffSynthCnetPatch (qwenimage_standard)")
             model_patched.set_model_double_block_patch(DiffSynthCnetPatch(model_patch, vae, image, strength, mask))
+        else:
+            raise RuntimeError(
+                "Unsupported model/controlnet route. "
+                "Control routing is strict to avoid mixing QI/ZI/Nunchaku/Krea2 branches."
+            )
 
         return (model_patched,)