Skip to content

v1.3.1 - ZIT Diffsynth Union ControlNet Support

Choose a tag to compare

View all tags
@ussoewwin ussoewwin released this 19 Jan 02:47
· 271 commits to main since this release
1.3.1
294c266

Release Notes v1.3.1 - ZIT Diffsynth Union ControlNet Support

Overview

This release adds full Diffsynth Union ControlNet support for both standard Z-Image Turbo (ZIT) and Nunchaku ZIT models. This implementation replaces ComfyUI's nodes_model_patch.py design directly, integrating it into the Forge environment while maintaining compatibility.

Table of Contents

  1. Overview
  2. Files Added/Modified
  3. ComfyUI Design Replacement Structure
  4. Key Code Implementation
  5. RecursionError Fix
  6. Important Design Points

Implementation Overview

ZIT ControlNet operates with a different mechanism than standard ControlNet:

  • Standard ControlNet: Uses get_control() method to generate control signals and injects them into the model's forward pass
  • ZIT Diffsynth ControlNet: Works as a "model patch", directly applying patches to the model's internal blocks through transformer_options

This implementation enables Union ControlNet functionality for both standard ZIT and Nunchaku ZIT models, allowing multiple ControlNet models to be used simultaneously.

Files Added/Modified

1. New Files

modules_forge/supported_controlnet_zit.py

Overview: Complete ZIT ControlNet implementation file. This is a direct copy of ComfyUI's comfy_extras/nodes_model_patch.py, adapted minimally for the Forge environment.

Key Classes:

  • ZImageControlPatch: Class for applying ZIT ControlNet patches
  • ZImageVAEWrapper: Wrapper to adapt Forge VAE to ComfyUI ControlNet interface
  • ZImageControlNetPatcher: Patcher for integrating into Forge's ControlNet system
  • ModelPatchLoader: ComfyUI-compatible model patch loader

Important Design Decisions:

  • ComfyUI's nodes_model_patch.py was copied completely, with no speculative modifications
  • ZImageControlPatch.__call__ method is identical to ComfyUI's implementation
  • Maintains the implementation of directly calling ZImage_Control model via self.model_patch.model(...)

2. Modified Files

backend/modules/k_model.py

Changes:

  1. Added import comfy (lines 8-10)

    import comfy.patcher_extension
import comfy.ldm.lumina.model  # For NextDiT check in apply_model
    • Required to use comfy.patcher_extension.WrapperExecutor
    • Prevents UnboundLocalError

  2. Added ZIT Model Detection Logic (lines 236-250)

    # CRITICAL: Check if this is a ZIT model (NextDiT) BEFORE applying ZIT-specific logic
is_zit_model = False
try:
    NextDiT = comfy.ldm.lumina.model.NextDiT
    if isinstance(self.diffusion_model, NextDiT):
        is_zit_model = True
    else:
        # Fallback: check by type name
        model_type_name = type(self.diffusion_model).__name__
        if model_type_name == "NextDiT":
            is_zit_model = True
except (ImportError, AttributeError, TypeError):
    is_zit_model = False

  3. ZIT-Specific extra_conds Processing (lines 252-297)

    • Added processing for attention_mask and num_tokens
    • Fully replicates ComfyUI's Lumina2.extra_conds logic
    • Adds transformer_options to kwargs (required for NextDiT.forward() to get wrappers)

  4. RecursionError Fix (lines 145-208)

    • Added ZIT-specific patch management logic to apply_model method
    • Prevents double patch application
    • Removes stale patches

backend/patcher/base.py

Changes:

  1. Added set_model_double_block_patch Method (lines 2914-2944)

    def set_model_double_block_patch(self, patch):
    """
    Set double_block patch for ZIT ControlNet.
    This method is ZIT models (NextDiT) ONLY.
    Other models (SD1.5, SDXL, Flux1, Qwen Image) will ignore this patch.
    """
    # CRITICAL: Verify this is a ZIT model (NextDiT) before applying patch
    is_zit_model = False
    try:
        import comfy.ldm.lumina.model
        NextDiT = comfy.ldm.lumina.model.NextDiT
        if hasattr(self.model, 'diffusion_model'):
            if isinstance(self.model.diffusion_model, NextDiT):
                is_zit_model = True
            else:
                model_type_name = type(self.model.diffusion_model).__name__
                if model_type_name == "NextDiT":
                    is_zit_model = True
    except (ImportError, AttributeError, TypeError):
        is_zit_model = False
    
    if not is_zit_model:
        return
    
    self.set_model_patch(patch, "double_block")

  2. Added set_model_noise_refiner_patch Method (lines 2946-2976)

    • Sets noise_refiner patch with similar logic to double_block

Design Intent:

  • Ignore patches for non-ZIT models (no error raised)
  • Safety checks to prevent ZIT-specific patches from being applied to other model types

backend/loader.py

Changes:

  1. NextDiT Initialization operations Parameter Setting (lines 278-294) 
    # NextDiT requires operations parameter
# For standard ZIT: use ComfyUI's manual_cast operations (same as ComfyUI's Lumina2)
# For Nunchaku ZIT: use disable_weight_init (torch.nn wrapper) since using_forge_operations(operations=False) uses torch.nn
if _nz:
    # Nunchaku ZIT: operations=False (will use torch.nn directly in using_forge_operations context)
    model_loader = lambda c: NextDiT(**c)
else:
    # Standard ZIT: use ComfyUI's manual_cast operations (same as ComfyUI BaseModel)
    import comfy.ops
    model_loader = lambda c: NextDiT(**c, operations=comfy.ops.manual_cast)

Important Points:

  • Standard ZIT: Uses comfy.ops.manual_cast (same as ComfyUI)
  • Nunchaku ZIT: Uses operations=False with using_forge_operations context to directly use torch.nn

modules_forge/supported_controlnet.py

Changes:

  1. Added ZIT ControlNet Detection Logic (lines 106-117)

    # ZIT ControlNet (from model_patches folder)
# Detected by 'control_all_x_embedder.2-1.weight' key
if 'control_all_x_embedder.2-1.weight' in controlnet_data:
    try:
        from modules_forge.supported_controlnet_zit import ZImageControlNetPatcher
        result = ZImageControlNetPatcher.try_build_from_state_dict(controlnet_data, ckpt_path)
        if result is not None:
            print(f"[ControlNet] Loaded ZIT ControlNet (model patch): {ckpt_path}")
        return result
    except Exception as e:
        print(f"Failed to load ZIT ControlNet: {e}")
        return None

  2. ZIT ControlNet Patcher Registration (lines 258-262)

    # Register ZIT ControlNet patcher
try:
    from modules_forge.supported_controlnet_zit import ZImageControlNetPatcher
    add_supported_control_model(ZImageControlNetPatcher)
except ImportError:
    pass  # ZIT ControlNet support may not be available

ComfyUI Design Replacement Structure

ComfyUI's Original Design Flow

  1. Node Definition (comfy_extras/nodes_model_patch.py)

    • ZImageFunControlnet node: Calls diffsynth_controlnet function
    • diffsynth_controlnet function: Creates ZImageControlPatch and registers it with ModelPatcher

  2. Patch Registration (comfy/model_patcher.py)

    • set_model_noise_refiner_patch(): Registers noise_refiner patch in model_options["transformer_options"]["patches"]["noise_refiner"]
    • set_model_double_block_patch(): Registers double_block patch in model_options["transformer_options"]["patches"]["double_block"]

  3. Patch Application (comfy/ldm/lumina/model.py - NextDiT.forward())

    • Retrieves patches from transformer_options
    • Applies patches in noise_refiner and double_block blocks
    • Calls patch's __call__ method to generate control signals

Forge Replacement Flow

  1. ControlNet Patcher (modules_forge/supported_controlnet_zit.py)

    • ZImageControlNetPatcher.process_before_every_sampling(): Called before Forge sampling
    • Creates ZImageControlPatch and registers it with UnetPatcher

  2. Patch Registration (backend/patcher/base.py)

    • set_model_noise_refiner_patch(): Registers noise_refiner patch after ZIT model check
    • set_model_double_block_patch(): Registers double_block patch after ZIT model check
    • Patches registered in UnetPatcher.model_options["transformer_options"]["patches"]

  3. Patch Application (backend/modules/k_model.py → NextDiT.forward())

    • KModel.apply_model(): Merges transformer_options and passes to NextDiT.forward()
    • NextDiT.forward(): Applies patches with the same logic as ComfyUI

Important Design Differences

Item ComfyUI Forge
Patch Registration Location ModelPatcher.model_options UnetPatcher.model_options (via KModel)
Model Structure BaseModel → NextDiT KModel → NextDiT
Patch Registration Timing Node execution (diffsynth_controlnet) Before sampling (process_before_every_sampling)
VAE Processing ComfyUI VAE (direct use) Forge VAE (wrapped with ZImageVAEWrapper)
ControlNet Integration diffsynth_controlnet node ZImageControlNetPatcher.process_before_every_sampling

Complete ComfyUI Code Copy

modules_forge/supported_controlnet_zit.py directly copies the following classes and functions from ComfyUI's comfy_extras/nodes_model_patch.py:

  • ZImageControlPatch class: Identical implementation
    • __init__ method: Image encoding processing
    • encode_latent_cond method: Inpainting-compatible latent encoding
    • __call__ method: Patch application logic (completely identical to ComfyUI, lines 402-423)
  • z_image_convert function: Model state dictionary conversion
  • ModelPatchLoader class: Model patch loader (ComfyUI node compatible)

Important Point: These implementations are direct copies of ComfyUI's code with no speculative modifications.

Key Code Implementation

1. ZImageControlPatch Class (modules_forge/supported_controlnet_zit.py)

__init__ Method (lines 325-355)

def __init__(self, model_patch, vae, image, strength, inpaint_image=None, mask=None):
    self.model_patch = model_patch
    self.vae = vae
    self.image = image
    self.inpaint_image = inpaint_image
    self.mask = mask
    self.strength = strength
    self.is_inpaint = self.model_patch.model.additional_in_dim > 0
    
    # Image encoding processing
    self.encoded_image = self.encode_latent_cond(self.image, self.inpaint_image)
    self.encoded_image_size = (self.image.shape[1], self.image.shape[2])
    self.temp_data = None

Functionality:

  • Encodes images for ControlNet
  • Supports inpainting (when additional_in_dim > 0)
  • Stores encoded image size

__call__ Method (lines 384-451)

def __call__(self, kwargs):
    x = kwargs.get("x")
    img = kwargs.get("img")
    img_input = kwargs.get("img_input")
    txt = kwargs.get("txt")
    pe = kwargs.get("pe")
    vec = kwargs.get("vec")
    block_index = kwargs.get("block_index")
    block_type = kwargs.get("block_type", "")
    
    # Check encoded image size and re-encode if needed
    if self.encoded_image is None or self.encoded_image_size != (x.shape[-2] * spacial_compression, x.shape[-1] * spacial_compression):
        # Re-encoding processing
        ...
    
    # Calculate ControlNet block index
    cnet_blocks = self.model_patch.model.n_control_layers
    div = round(30 / cnet_blocks)
    cnet_index = (block_index // div)
    cnet_index_float = (block_index / div)
    
    # Direct call to ZImage_Control model (same implementation as ComfyUI)
    if self.temp_data is None or self.temp_data[0] > cnet_index:
        if block_type == "noise_refiner":
            self.temp_data = (-3, (None, self.model_patch.model(txt, self.encoded_image.to(img.dtype), pe, vec)))
        else:
            self.temp_data = (-1, (None, self.model_patch.model(txt, self.encoded_image.to(img.dtype), pe, vec)))
    
    # Processing based on block type
    if block_type == "noise_refiner":
        next_layer = self.temp_data[0] + 1
        self.temp_data = (next_layer, self.model_patch.model.forward_noise_refiner_block(block_index, self.temp_data[1][1], img_input[:, :self.temp_data[1][1].shape[1]], None, pe, vec))
        if self.temp_data[1][0] is not None:
            img[:, :self.temp_data[1][0].shape[1]] += (self.temp_data[1][0] * self.strength)
    else:
        # double_block processing
        while self.temp_data[0] < cnet_index and (self.temp_data[0] + 1) < cnet_blocks:
            next_layer = self.temp_data[0] + 1
            self.temp_data = (next_layer, self.model_patch.model.forward_control_block(next_layer, self.temp_data[1][1], img_input[:, :self.temp_data[1][1].shape[1]], None, pe, vec))
        
        if cnet_index_float == self.temp_data[0]:
            img[:, :self.temp_data[1][0].shape[1]] += (self.temp_data[1][0] * self.strength)
            if cnet_blocks == self.temp_data[0] + 1:
                self.temp_data = None
    
    kwargs['img'] = img
    return kwargs

Important Points:

  • Directly calls ZImage_Control model via self.model_patch.model(...) (same as ComfyUI)
  • Caches intermediate results with temp_data (avoids recalculation for same ControlNet block)
  • Branches processing based on block_type ("noise_refiner" or "double_block")

2. ZImageControlNetPatcher Class (modules_forge/supported_controlnet_zit.py)

try_build_from_state_dict Method (lines 589-631)

@staticmethod
def try_build_from_state_dict(controlnet_data, ckpt_path):
    if 'control_all_x_embedder.2-1.weight' not in controlnet_data:
        return None
    
    try:
        sd = z_image_convert(controlnet_data)
        dtype = comfy.utils.weight_dtype(sd)
        
        # Auto-detect model configuration
        config = {}
        if 'control_layers.4.adaLN_modulation.0.weight' not in sd:
            config['n_control_layers'] = 3
            config['additional_in_dim'] = 17
            config['refiner_control'] = True
        if 'control_layers.14.adaLN_modulation.0.weight' in sd:
            config['n_control_layers'] = 15
            config['additional_in_dim'] = 17
            config['refiner_control'] = True
        
        # Create ZImage_Control model
        model = comfy.ldm.lumina.controlnet.ZImage_Control(
            device=comfy.model_management.unet_offload_device(),
            dtype=dtype,
            operations=comfy.ops.manual_cast,
            **config
        )
        
        model.load_state_dict(sd)
        
        # Wrap with ModelPatcher
        model_patcher = comfy.model_patcher.ModelPatcher(
            model,
            load_device=comfy.model_management.get_torch_device(),
            offload_device=comfy.model_management.unet_offload_device()
        )
        
        return ZImageControlNetPatcher(model_patcher)
    except Exception as e:
        traceback.print_exc()
        return None

Functionality:

  • Detects ZIT ControlNet model (determined by control_all_x_embedder.2-1.weight key)
  • Auto-detects model configuration (n_control_layers, additional_in_dim, etc.)
  • Wraps with ComfyUI's ModelPatcher and returns

process_before_every_sampling Method (lines 638-672)

def process_before_every_sampling(self, process, cond, mask, *args, **kwargs):
    unet = process.sd_model.forge_objects.unet
    
    # Create VAE wrapper
    vae_raw = process.sd_model.forge_objects.vae
    if not hasattr(vae_raw, "spacial_compression_encode"):
        vae = ZImageVAEWrapper(vae_raw)
    else:
        vae = vae_raw
    
    # Preprocess image and mask
    image = cond
    if image is not None:
        if image.ndim == 4 and image.shape[1] <= 4:
            image = image.movedim(1, -1)  # BCHW → BHWC
        image = image[:, :, :, :3]
    
    if mask is not None:
        if mask.ndim == 3:
            mask = mask.unsqueeze(1)
        if mask.ndim == 4:
            mask = mask.unsqueeze(2)
        mask = 1.0 - mask  # Invert mask
    
    # Create and register ZImageControlPatch
    if isinstance(self.model_patcher.model, comfy.ldm.lumina.controlnet.ZImage_Control):
        patch = ZImageControlPatch(self.model_patcher, vae, image, self.strength, inpaint_image=None, mask=mask)
        unet.set_model_noise_refiner_patch(patch)
        unet.set_model_double_block_patch(patch)

Functionality:

  • Called before sampling
  • Creates VAE wrapper (adapts Forge VAE to ComfyUI interface)
  • Preprocesses image and mask
  • Creates ZImageControlPatch and registers it with UnetPatcher

3. KModel.apply_model Method (backend/modules/k_model.py)

ZIT Model Detection and Processing (lines 236-297)

# CRITICAL: Check if this is a ZIT model (NextDiT) BEFORE applying ZIT-specific logic
is_zit_model = False
try:
    NextDiT = comfy.ldm.lumina.model.NextDiT
    if isinstance(self.diffusion_model, NextDiT):
        is_zit_model = True
    else:
        model_type_name = type(self.diffusion_model).__name__
        if model_type_name == "NextDiT":
            is_zit_model = True
except (ImportError, AttributeError, TypeError):
    is_zit_model = False

# ZIT models (NextDiT) ONLY: Apply ComfyUI Lumina2.extra_conds logic
if is_zit_model:
    attention_mask = extra_conds_clean.pop("attention_mask", None)
    num_tokens = None
    
    if attention_mask is not None:
        if torch.numel(attention_mask) != attention_mask.sum():
            extra_conds_clean["attention_mask"] = attention_mask
        num_tokens = max(1, torch.sum(attention_mask).item())
    
    if num_tokens is None and context is not None:
        num_tokens = context.shape[1]
    
    # CRITICAL: transformer_options must be in kwargs for NextDiT.forward()
    extra_conds_clean["transformer_options"] = transformer_options
    
    # Pass transformer_options to NextDiT.forward()
    if num_tokens is not None:
        model_output = self.diffusion_model(xc, t, context=context, num_tokens=num_tokens, attention_mask=attention_mask, control=control, **extra_conds_clean).float()
    else:
        model_output = self.diffusion_model(xc, t, context=context, control=control, **extra_conds_clean).float()
else:
    # Other models: Normal forward without ZIT-specific parameters
    model_output = self.diffusion_model(xc, t, context=context, control=control, **extra_conds_clean).float()

Functionality:

  • Detects ZIT model (NextDiT)
  • Processes ZIT-specific parameters (attention_mask, num_tokens)
  • Adds transformer_options to kwargs (required for NextDiT.forward() to get wrappers)

RecursionError Fix

Root Cause

RecursionError had two root causes:

  1. Double Patching

    • Forge has two locations for storing patches:
      • transformer_options: Temporary options (only for that generation)
      • self.model_options: Persistent options (maintained by model object)
    • Simple merging of these causes the same patch to be applied twice

  2. Stale Patches

    • When switching models, old patches remain in self.model_options
    • Even when new ControlNet patches are provided, old patches also remain
    • System attempts to merge these, causing infinite loops

Fix Implementation (backend/modules/k_model.py lines 139-208)

Actual implementation code:

def apply_model(self, x, t, c_concat=None, c_crossattn=None, control=None, transformer_options=None, **kwargs):
    # Use WrapperExecutor same as ComfyUI BaseModel.apply_model
    if transformer_options is None:
        transformer_options = {}
    
    # ZIT-ONLY FIX: Handle ZIT patches to prevent double application and stale patches
    # 
    # Problem 1: Double application - ZIT patches exist in both provided_patches AND self_patches
    # Solution: Skip self.model_options merge when ZIT patches are already in provided_patches
    #
    # Problem 2: Stale patches - After model switch, ZIT patches remain in self_patches
    #            but no new ZIT patches are in provided_patches, causing RecursionError
    # Solution: Clear self.model_options ZIT patches when detected as stale
    #
    # This fix ONLY affects NextDiT models - other models use the original merge logic.
    skip_model_options_merge = False
    provided_patches = transformer_options.get("patches", {})
    zit_patches_in_provided = "noise_refiner" in provided_patches or "double_block" in provided_patches
    
    # Get current model type and self patches
    model_type_name = type(self.diffusion_model).__name__
    self_patches = self.model_options.get("transformer_options", {}).get("patches", {})
    zit_patches_in_self = "noise_refiner" in self_patches or "double_block" in self_patches
    
    # Case 1: ZIT patches in provided - skip merge to prevent duplication
    if zit_patches_in_provided and model_type_name == "NextDiT":
        skip_model_options_merge = True
    
    # Case 2: Stale ZIT patches detected - self has ZIT patches but provided doesn't
    # This happens after model switch: old patches remain in self.model_options
    elif zit_patches_in_self and not zit_patches_in_provided:
        # Clear the stale transformer_options to prevent applying old patches
        self.model_options["transformer_options"] = {}
        skip_model_options_merge = True  # Nothing to merge now
    
    # Merge transformer_options from model_options before passing to WrapperExecutor
    final_transformer_options = {}
    
    # ZIT-ONLY: Skip self.model_options merge if ZIT patches already in transformer_options
    if not skip_model_options_merge and "transformer_options" in self.model_options:
        import copy
        final_transformer_options = copy.deepcopy(self.model_options["transformer_options"])
    
    # Merge with provided transformer_options (from sampling_function_inner)
    if transformer_options:
        # Merge patches (same logic as sampling_function.py lines 250-259)
        if "patches" in transformer_options:
            if "patches" not in final_transformer_options:
                final_transformer_options["patches"] = {}
            cur_patches = final_transformer_options["patches"].copy()
            for patch_name, patches in transformer_options["patches"].items():
                if patch_name in cur_patches:
                    cur_patches[patch_name] = cur_patches[patch_name] + patches
                else:
                    cur_patches[patch_name] = patches
            final_transformer_options["patches"] = cur_patches
        # Merge other transformer_options (override with provided values)
        for key, value in transformer_options.items():
            if key != "patches":
                final_transformer_options[key] = value
    
    transformer_options = final_transformer_options
    
    return comfy.patcher_extension.WrapperExecutor.new_class_executor(
        self._apply_model,
        self,
        comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.APPLY_MODEL, transformer_options)
    ).execute(x, t, c_concat, c_crossattn, control, transformer_options, **kwargs)

Fix Points:

  1. Prevent Double Application (Case 1):

    • When provided_patches contains ZIT patches ("noise_refiner" or "double_block")
    • And model type is NextDiT
    • Sets skip_model_options_merge = True to skip merge with self.model_options
    • Starts with empty final_transformer_options = {}, using only transformer_options

  2. Remove Stale Patches (Case 2):

    • When self_patches has ZIT patches but provided_patches doesn't
    • Determines model switch occurred (old patches remain from previous model)
    • Clears self.model_options["transformer_options"] = {} to remove stale patches
    • Sets skip_model_options_merge = True

  3. Normal Merge Processing:

    • When skip_model_options_merge = False, executes normal merge processing
    • Copies self.model_options["transformer_options"] with copy.deepcopy
    • Merges with transformer_options (combines patches as arrays)

  4. WrapperExecutor Usage:

    • Uses WrapperExecutor same as ComfyUI to execute APPLY_MODEL level wrappers
    • Passes transformer_options to _apply_model method

Fix Effects:

  • Prevents double patch application
  • Removes stale patches during model switching
  • Resolves RecursionError
  • ZIT ControlNet operates stably

Important Design Points

1. Complete ComfyUI Replication

  • modules_forge/supported_controlnet_zit.py is a complete copy of ComfyUI's comfy_extras/nodes_model_patch.py
  • No speculative modifications added
  • ZImageControlPatch.__call__ method is identical to ComfyUI's implementation (lines 402-423)
  • Maintains implementation of directly calling ZImage_Control model via self.model_patch.model(...)

2. ZIT-Specific Patch Methods

  • set_model_double_block_patch(): Applies patch to ZIT model's double_block
  • set_model_noise_refiner_patch(): Applies patch to ZIT model's noise_refiner
  • Ignores patches for non-ZIT models (no error raised)

3. VAE Wrapper

  • Forge VAE may not have spacial_compression_encode() method
  • Wraps with ZImageVAEWrapper to adapt to ComfyUI ControlNet interface

4. Differences Between Standard ZIT and Nunchaku ZIT

NextDiT Initialization operations Setting (backend/loader.py lines 278-294)

Standard ZIT:

# Standard ZIT: use ComfyUI's manual_cast operations (same as ComfyUI BaseModel)
guess.unet_config["operations"] = comfy.ops.manual_cast
model_loader = lambda c: NextDiT(**c)
  • Uses comfy.ops.manual_cast (same as ComfyUI)
  • Same logic as ComfyUI's BaseModel.__init__

Nunchaku ZIT:

# Nunchaku ZIT: use disable_weight_init (torch.nn wrapper) since operations=False uses torch.nn directly
guess.unet_config["operations"] = comfy.ops.disable_weight_init
model_loader = lambda c: NextDiT(**c)
  • Uses comfy.ops.disable_weight_init (torch.nn wrapper)
  • Uses torch.nn directly within using_forge_operations(operations=False) context

Important Points:

  • NextDiT.__init__ requires operations parameter (cannot be None)
  • using_forge_operations(operations=False) only affects context, not NextDiT.__init__
  • Standard ZIT and Nunchaku ZIT use different operations, but both use the same ControlNet implementation (ZImageControlNetPatcher)

ControlNet Implementation Commonality

The difference between standard ZIT and Nunchaku ZIT is only the operations parameter during NextDiT initialization. The ControlNet implementation (ZImageControlNetPatcher) is shared:

  • Uses same ZImageControlPatch class
  • Uses same ZImageControlNetPatcher class
  • Uses same patch registration methods (set_model_noise_refiner_patch, set_model_double_block_patch)

5. Union ControlNet Support

  • Multiple ZIT ControlNet models can be used simultaneously
  • Each ControlNet patch operates independently
  • Caches intermediate results for each ControlNet with temp_data
  • Patches are registered as arrays and applied sequentially

6. Complete Replacement Structure of ComfyUI Design

ComfyUI Design Flow

1. ZImageFunControlnet Node
   ↓
2. diffsynth_controlnet() Function
   - Creates ZImageControlPatch
   - Calls ModelPatcher.set_model_noise_refiner_patch()
   - Calls ModelPatcher.set_model_double_block_patch()
   ↓
3. ModelPatcher.model_options["transformer_options"]["patches"]
   - "noise_refiner": [ZImageControlPatch, ...]
   - "double_block": [ZImageControlPatch, ...]
   ↓
4. BaseModel.apply_model()
   - Passes transformer_options to NextDiT.forward()
   ↓
5. NextDiT.forward()
   - Retrieves patches from transformer_options["patches"]
   - Applies patches in noise_refiner and double_block blocks
   - Calls patch's __call__() method

Forge Replacement Flow

1. ZImageControlNetPatcher.process_before_every_sampling()
   - Called before sampling
   - Creates ZImageControlPatch
   - Calls UnetPatcher.set_model_noise_refiner_patch()
   - Calls UnetPatcher.set_model_double_block_patch()
   ↓
2. UnetPatcher.model_options["transformer_options"]["patches"]
   - "noise_refiner": [ZImageControlPatch, ...]
   - "double_block": [ZImageControlPatch, ...]
   ↓
3. KModel.apply_model()
   - Merges transformer_options (applies RecursionError fix logic)
   - Uses WrapperExecutor to call _apply_model()
   - Passes transformer_options to NextDiT.forward()
   ↓
4. NextDiT.forward() (same as ComfyUI)
   - Retrieves patches from transformer_options["patches"]
   - Applies patches in noise_refiner and double_block blocks
   - Calls patch's __call__() method

Replacement Points

  1. Node-based to Function-based: Changed from ComfyUI's node execution to Forge's process_before_every_sampling method call

  2. ModelPatcher to UnetPatcher: Changed from ComfyUI's ModelPatcher to Forge's UnetPatcher (via KModel)

  3. BaseModel to KModel: Changed from ComfyUI's BaseModel.apply_model to Forge's KModel.apply_model

  4. Patch Application Logic is Identical: Patch application logic in NextDiT.forward() is completely identical to ComfyUI

  5. RecursionError Fix: Fixes Forge-specific issues (double patch application and stale patches) in KModel.apply_model

Summary

This implementation enables Diffsynth Union ControlNet for both standard ZIT and Nunchaku ZIT models. By directly replacing ComfyUI's design, we've maintained compatibility while integrating it into the Forge environment.

Particularly important:

  1. Complete copy of ComfyUI's code, avoiding speculative modifications
  2. Added ZIT-specific patch methods without affecting other model types
  3. Fixed RecursionError root causes (double patch application and stale patches)

With these implementations, ZIT ControlNet operates stably and Union ControlNet functionality works correctly.

Assets 2
Loading