ussoewwin-ComfyUI-QwenImageLoraLoader/md/ZIMAGETURBO_CONTROLNET_FIX.md

Detailed Explanation of Z-Image-Turbo ControlNet Support Fix in ComfyUI-QwenImageLoraLoader

Reason for Additional Fix

The unofficial Z-Image-Turbo loader (ComfyUI-nunchaku-unofficial-z-image-turbo-loader) monkey-patched NunchakuZImageTransformer2DModel.forward to accept transformer_options and control parameters.

However, the ComfyZImageTurboWrapper class in ComfyUI-QwenImageLoraLoader was not correctly passing these parameters to the underlying NunchakuZImageTransformer2DModel.forward.

ComfyZImageTurboWrapper is a wrapper that separates LoRA composition and forward pass, passing arguments received in the forward method to NunchakuZImageTransformer2DModel.forward via the _execute_model method. We needed to add transformer_options and control to this transmission path.

Modified File

ComfyUI-QwenImageLoraLoader/wrappers/zimageturbo.py

Fix Details

Fix 1: forward Method Signature (Lines 76-85)

Before

def forward(
        self,
        x,
        timestep,
        context=None,
        y=None,
        guidance=None,
        **kwargs,
):

After

def forward(
        self,
        x,
        timestep,
        context=None,
        y=None,
        guidance=None,
        control=None,
        transformer_options={},
        **kwargs,
):

Explanation

Added control and transformer_options to the forward method signature to receive them from ComfyUI's ModelPatcher.

Fix 2: transformer_options Merge Processing in forward Method (Lines 93-102)

Fix Content

# Remove guidance, transformer_options, and attention_mask from kwargs
if "guidance" in kwargs:
    kwargs.pop("guidance")
if "transformer_options" in kwargs:
    if isinstance(transformer_options, dict) and isinstance(kwargs["transformer_options"], dict):
        transformer_options = {**transformer_options, **kwargs.pop("transformer_options")}
    else:
        kwargs.pop("transformer_options")
if "attention_mask" in kwargs:
    kwargs.pop("attention_mask")

Explanation

ComfyUI's ModelPatcher can pass parameters either as explicit arguments or via **kwargs. To handle both cases correctly:

1. If transformer_options is passed both as an explicit parameter and in kwargs, merge them with the explicit parameter taking precedence (dictionary spread operator {**transformer_options, **kwargs["transformer_options"]} ensures explicit parameter values override kwargs values)
2. Remove guidance and attention_mask from kwargs to prevent duplicate parameter errors, as these are handled separately
3. This merging strategy ensures compatibility with different ComfyUI versions that may pass parameters in different ways

Fix 3: Argument Passing to _execute_model Method (Lines 219, 221)

Fix Content

Modified calls to _execute_model within the forward method to pass control and transformer_options:

with cache_context(self._cache_context):
    out = self._execute_model(x, timestep, context, y, guidance, control, transformer_options, **kwargs)
else:
    out = self._execute_model(x, timestep, context, y, guidance, control, transformer_options, **kwargs)

Explanation

By passing control and transformer_options to _execute_model, they can be transmitted to the underlying NunchakuZImageTransformer2DModel.forward.

Fix 4: _execute_model Method Signature and Parameter Cleanup (Line 239-247)

Before

def _execute_model(self, x, timestep, context, y, guidance, **kwargs):

After

def _execute_model(self, x, timestep, context, y, guidance, control, transformer_options, **kwargs):
    """Helper function to run the Z-Image-Turbo model's forward pass."""
    # Get ref_latents from kwargs before removing it
    ref_latents_value = kwargs.pop("ref_latents", None)
    
    # Remove guidance, transformer_options, and attention_mask from kwargs
    kwargs.pop("guidance", None)
    kwargs.pop("transformer_options", None)
    kwargs.pop("attention_mask", None)
    ...

Explanation

1. Added control and transformer_options as explicit parameters to the _execute_model signature to match the forward method's parameter list
2. Extract ref_latents from kwargs early (line 242) before cleaning kwargs, as ref_latents is needed later in both customized_forward and direct forward paths
3. Remove guidance, transformer_options, and attention_mask from kwargs to prevent duplicate argument errors when passing **kwargs to downstream methods
4. This cleanup ensures that these parameters are passed explicitly where needed, rather than through **kwargs, providing better control and preventing conflicts

Fix 5: Transmission When Using customized_forward (Lines 282-300)

Fix Content

Processing when customized_forward is set (when using the custom forward from the unofficial ZIT loader):

if self.customized_forward:
    with torch.inference_mode():
        forward_kwargs_cleaned = {k: v for k, v in self.forward_kwargs.items() if k not in ("guidance", "ref_latents", "transformer_options", "attention_mask")}
        transformer_options_cleaned = dict(transformer_options) if transformer_options else {}
        transformer_options_cleaned.pop("guidance", None)
        transformer_options_cleaned.pop("ref_latents", None)
        
        return self.customized_forward(
            self.model,
            hidden_states=x,
            encoder_hidden_states=context,
            timestep=timestep,
            guidance=guidance if self.config.get("guidance_embed", False) else None,
            ref_latents=ref_latents_value,
            control=control,
            transformer_options=transformer_options_cleaned,
            **forward_kwargs_cleaned,
            **kwargs,
        )

Explanation

Pass control and transformer_options when calling customized_forward. Remove guidance and ref_latents from transformer_options (to avoid duplication).

Additional Details:

1. forward_kwargs cleaning (line 134): The self.forward_kwargs dictionary may contain parameters that conflict with explicit arguments passed to customized_forward. We filter out guidance, ref_latents, transformer_options, and attention_mask to prevent duplicate argument errors.

2. transformer_options cleaning (lines 135-137): Create a copy of transformer_options to avoid modifying the original dictionary, then remove guidance and ref_latents if they exist. This prevents duplicate parameter passing, as guidance is passed explicitly based on self.config.get("guidance_embed", False), and ref_latents is extracted from kwargs earlier and passed as ref_latents_value.

3. ref_latents handling: The ref_latents_value was extracted from kwargs in _execute_model (line 242) before cleaning kwargs. It is passed explicitly to customized_forward rather than through transformer_options, providing clearer separation of concerns.

4. guidance conditional passing: guidance is only passed to customized_forward if self.config.get("guidance_embed", False) is True, allowing for conditional guidance embedding support.

Fix 6: Transmission When Directly Calling forward (Lines 389-409)

Fix Content

Processing when customized_forward is not set (when directly calling NunchakuZImageTransformer2DModel.forward):

# Call Z-Image-Turbo forward with correct signature
# Pass control and transformer_options to allow Model Patcher (double_block patches) to work
# Check if the model's forward method accepts control parameter
import inspect
forward_sig = inspect.signature(self.model.forward)
forward_params = set(forward_sig.parameters.keys())

zimage_kwargs_clean = {k: v for k, v in zimage_kwargs.items() if k not in ('control', 'transformer_options')}

# Build kwargs based on what the forward method accepts
forward_kwargs = zimage_kwargs_clean.copy()
if 'control' in forward_params:
    forward_kwargs['control'] = control
if 'transformer_options' in forward_params:
    forward_kwargs['transformer_options'] = transformer_options

model_output = self.model(
    x_list,
    t_zimage,
    cap_feats=cap_feats,
    **forward_kwargs
)

Explanation

• Check the signature of self.model.forward using inspect.signature
• Add control and transformer_options to forward_kwargs only if they are accepted
• This allows the _patched_forward (monkey-patched forward method) from the unofficial ZIT loader to correctly receive transformer_options and control

Additional Details:

1. zimage_kwargs preparation (lines 365-371, not shown in snippet): Before the dynamic parameter check, zimage_kwargs is prepared with Z-Image-Turbo specific parameters like patch_size, f_patch_size, and return_dict=False. These are extracted from kwargs and prepared separately.

2. zimage_kwargs_clean (line 171): Remove control and transformer_options from zimage_kwargs if they were accidentally included, ensuring they are only passed through the dynamic parameter checking mechanism.

3. Dynamic parameter inspection (lines 391-393):

   • inspect.signature(self.model.forward) retrieves the signature of the model's forward method
   • forward_params = set(forward_sig.parameters.keys()) creates a set of parameter names for fast lookup
   • This checks the actual signature at runtime, not the static type definition

4. Conditional parameter addition (lines 399-402): Only add control and transformer_options to forward_kwargs if they are accepted by the forward method. This is crucial because:

   • The original NunchakuZImageTransformer2DModel.forward does not accept these parameters
   • The unofficial ZIT loader monkey-patches the forward method to add these parameters
   • This check ensures compatibility whether the monkey-patch is applied or not

5. Model call structure (lines 404-409): The model is called with positional arguments (x_list, t_zimage, cap_feats) followed by **forward_kwargs, which includes control and transformer_options if accepted. This matches the expected signature pattern for Z-Image-Turbo models.

6. Why inspect.signature is necessary: Without dynamic parameter checking, passing control or transformer_options to a forward method that doesn't accept them would raise a TypeError: forward() got an unexpected keyword argument. Using inspect.signature allows the code to adapt to the actual method signature at runtime, supporting both patched and unpatched versions of the model.

Key Design Decisions

Dynamic Parameter Checking with inspect.signature

Using inspect.signature to dynamically check whether the forward method accepts control and transformer_options. This ensures:

1. Backward compatibility is maintained (no errors in older versions that don't accept these parameters)
2. Works appropriately regardless of whether the monkey patch from the unofficial ZIT loader is applied
3. Easy to adapt to future extensions

Integration with Unofficial ZIT Loader

This fix enables the following integration flow:

1. ComfyUI's ModelPatcher passes transformer_options and control to ComfyZImageTurboWrapper.forward
2. ComfyZImageTurboWrapper.forward passes these to _execute_model
3. _execute_model checks with inspect.signature before passing to NunchakuZImageTransformer2DModel.forward
4. The _patched_forward (monkey-patched forward) from the unofficial ZIT loader receives transformer_options and control and applies the double_block patch

Summary

Modified the ComfyZImageTurboWrapper class in ComfyUI-QwenImageLoraLoader to correctly pass control and transformer_options parameters.

Main changes:

1. Added control and transformer_options to the forward method signature
2. Added control and transformer_options to the _execute_model method signature
3. Implemented using inspect.signature to pass these parameters only if the underlying forward method accepts them

This enables the ControlNet support monkey patch added in the unofficial ZIT loader to work correctly, achieving compatibility between ControlNet and LoRA.

Technical Flow Diagram

The complete parameter transmission flow:

ComfyUI ModelPatcher
    ↓ (passes control, transformer_options)
ComfyZImageTurboWrapper.forward()
    ↓ (merges transformer_options if in kwargs, removes from kwargs)
    ↓ (passes control, transformer_options as explicit parameters)
ComfyZImageTurboWrapper._execute_model()
    ↓ (extracts ref_latents from kwargs, cleans kwargs)
    ↓
    ├─→ customized_forward path (if self.customized_forward is set)
    │       ↓ (cleans forward_kwargs and transformer_options)
    │       ↓ (passes control, transformer_options explicitly)
    │   Custom Forward Function (from unofficial ZIT loader)
    │       ↓ (receives control, transformer_options)
    │   NunchakuZImageTransformer2DModel.forward() (patched)
    │       ↓ (applies double_block patches using transformer_options and control)
    │
    └─→ Direct forward path (if customized_forward is None)
            ↓ (prepares zimage_kwargs, uses inspect.signature)
            ↓ (conditionally adds control, transformer_options to forward_kwargs)
            NunchakuZImageTransformer2DModel.forward() (patched or unpatched)
                ↓ (if patched: receives control, transformer_options, applies double_block patches)
                ↓ (if unpatched: ignores extra parameters, works normally)

Parameter Cleanup Strategy

The fix implements a comprehensive parameter cleanup strategy to prevent duplicate argument errors:

1. forward method level (lines 93-102): Merge transformer_options from kwargs with explicit parameter, remove conflicting parameters from kwargs
2. _execute_model level (lines 241-247): Extract ref_latents early, remove guidance, transformer_options, and attention_mask from kwargs
3. customized_forward path (lines 284-287): Clean forward_kwargs and transformer_options to remove duplicates before passing to customized forward
4. Direct forward path (lines 395-402): Use dynamic parameter checking to only pass control and transformer_options if accepted

This multi-level cleanup ensures that parameters are passed explicitly where needed and not duplicated through **kwargs, preventing runtime errors and ensuring correct parameter precedence.