ussoewwin-ComfyUI-QwenImageLoraLoader/ZIMAGETURBO_CONTROLNET_FIX.md

ComfyUI-QwenImageLoraLoader側のZ-Image-Turbo ControlNet対応修正の詳細解説

追加修正が必要になった原因

非公式Z-Image-Turboローダー側（ComfyUI-nunchaku-unofficial-z-image-turbo-loader）でNunchakuZImageTransformer2DModel.forwardをモンキーパッチし、transformer_optionsとcontrolパラメータを受け付けるようにした。

しかし、ComfyUI-QwenImageLoraLoader側のComfyZImageTurboWrapperクラスが、これらのパラメータを下位のNunchakuZImageTransformer2DModel.forwardに正しく伝達していなかった。

ComfyZImageTurboWrapperはLoRA合成とフォワードパスを分離するラッパーで、forwardメソッドで受け取った引数を_execute_modelメソッド経由でNunchakuZImageTransformer2DModel.forwardに渡す。この伝達経路にtransformer_optionsとcontrolを追加する必要があった。

修正したファイル

ComfyUI-QwenImageLoraLoader/wrappers/zimageturbo.py

修正内容

修正1: forwardメソッドのシグネチャ（76行目から85行目）

修正前

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

修正後

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

意味の説明

forwardメソッドのシグネチャにcontrolとtransformer_optionsを追加。ComfyUIのModelPatcherからこれらを受け取れるようにした。

修正2: forwardメソッド内のtransformer_optionsマージ処理（93行目から100行目）

修正内容

# 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")

意味の説明

kwargsにtransformer_optionsが含まれている場合、明示的なパラメータとマージ。明示的なパラメータが優先。

修正3: _execute_modelメソッドへの引数伝達（219行目、221行目）

修正内容

forwardメソッド内で_execute_modelを呼び出す際に、controlと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)

意味の説明

_execute_modelにcontrolとtransformer_optionsを渡すことで、下位のNunchakuZImageTransformer2DModel.forwardまで伝達できるようにした。

修正4: _execute_modelメソッドのシグネチャ（239行目）

修正前

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

修正後

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

意味の説明

_execute_modelのシグネチャにcontrolとtransformer_optionsを追加。

修正5: customized_forward使用時の伝達（282行目から300行目）

修正内容

customized_forwardが設定されている場合（非公式ZITローダー側のカスタムフォワードを使用する場合）の処理：

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,
        )

意味の説明

customized_forward呼び出し時にcontrolとtransformer_optionsを渡す。transformer_optionsからguidanceとref_latentsを除去（重複を避けるため）。

修正6: 直接forward呼び出し時の伝達（389行目から409行目）

修正内容

customized_forwardが設定されていない場合（直接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
)

意味の説明

• inspect.signatureでself.model.forwardのシグネチャを確認
• controlとtransformer_optionsが受け付けられる場合のみ、それらをforward_kwargsに追加
• これにより、非公式ZITローダー側の_patched_forward（モンキーパッチされたforwardメソッド）が正しくtransformer_optionsとcontrolを受け取れる

重要な設計判断

inspect.signatureによる動的パラメータチェック

inspect.signatureを使って、forwardメソッドがcontrolとtransformer_optionsを受け付けるかどうかを動的に確認する。これにより：

1. 後方互換性が保たれる（これらのパラメータを受け付けない旧バージョンでもエラーにならない）
2. 非公式ZITローダー側のモンキーパッチが適用されているかどうかに関係なく、適切に動作する
3. 将来的な拡張に対応しやすい

非公式ZITローダー側との連携

この修正により、以下の連携フローが実現される：

1. ComfyUIのModelPatcherがtransformer_optionsとcontrolをComfyZImageTurboWrapper.forwardに渡す
2. ComfyZImageTurboWrapper.forwardがこれらを_execute_modelに渡す
3. _execute_modelがinspect.signatureで確認してからNunchakuZImageTransformer2DModel.forwardに渡す
4. 非公式ZITローダー側の_patched_forward（モンキーパッチされたforward）がtransformer_optionsとcontrolを受け取り、double_blockパッチを適用する

まとめ

ComfyUI-QwenImageLoraLoader側のComfyZImageTurboWrapperクラスを修正し、controlとtransformer_optionsパラメータを正しく伝達するようにした。

主な変更点は：

1. forwardメソッドのシグネチャにcontrolとtransformer_optionsを追加
2. _execute_modelメソッドのシグネチャにcontrolとtransformer_optionsを追加
3. inspect.signatureを使って、下位のforwardメソッドがこれらのパラメータを受け付ける場合のみ渡すように実装

これにより、非公式ZITローダー側で追加したControlNet対応のモンキーパッチが正しく動作し、ControlNetとLoRAの両立が実現される。