Merge pull request #4753 from resonantsky/dev

FP16 Clamp fix for Z-Image-Turbo, enable FP16 on ZiT.

pipelines/model_z_image.py

@@ -3,6 +3,7 @@ import diffusers
 from modules import shared, devices, sd_models, model_quant, sd_hijack_te
 from modules.logger import log
 from pipelines import generic
+from pipelines.z_image.patch_fp16 import apply_patches
 
 
 def load_nunchaku():
@@ -22,6 +23,7 @@ def load_nunchaku():
 
 
 def load_z_image(checkpoint_info, diffusers_load_config=None):
+    apply_patches()
     if diffusers_load_config is None:
         diffusers_load_config = {}
     repo_id = sd_models.path_to_repo(checkpoint_info)

pipelines/z_image/patch_fp16.py

@@ -0,0 +1,109 @@
+import torch
+import torch.nn.functional as F
+
+
+def _clamp_fp16(x: torch.Tensor) -> torch.Tensor:
+    """Replace NaN/inf in FP16 tensors to prevent overflow propagation (matches ComfyUI behaviour)."""
+    if x.dtype == torch.float16:
+        return torch.nan_to_num(x, nan=0.0, posinf=65504.0, neginf=-65504.0)
+    return x
+
+
+_applied = False
+
+
+def apply_patches():
+    """Replace FP16-unsafe method bodies in the diffusers ZImage transformer.
+
+    Patches:
+      - FeedForward._forward_silu_gating   : wraps SiLU gate output in _clamp_fp16
+      - ZImageTransformerBlock.forward     : wraps attn/FFN outputs in _clamp_fp16
+      - FinalLayer.forward                 : upstream-equivalent (adaLN calls unchanged)
+
+    _clamp_fp16 is a no-op for bf16/fp32 — zero behavioural change for non-fp16 users.
+
+    All replacements are closures; select_per_token is captured from the diffusers
+    module at call time. No names are injected into the diffusers module namespace.
+
+    Uses modules.patches.patch() — idempotent, reversible via modules.patches.undo().
+    """
+    global _applied
+    if _applied:
+        return
+    _applied = True
+    from modules import patches as sdnext_patches
+    import diffusers.models.transformers.transformer_z_image as m
+
+    _select_per_token = m.select_per_token  # upstream helper, captured by closure
+
+    # ------------------------------------------------------------------
+    # FeedForward._forward_silu_gating
+    # ------------------------------------------------------------------
+    def _patched_forward_silu_gating(self, x1, x3):
+        return _clamp_fp16(F.silu(x1) * x3)
+
+    # ------------------------------------------------------------------
+    # ZImageTransformerBlock.forward
+    # ------------------------------------------------------------------
+    def _patched_zimage_block_forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        adaln_input: torch.Tensor | None = None,
+        noise_mask: torch.Tensor | None = None,
+        adaln_noisy: torch.Tensor | None = None,
+        adaln_clean: torch.Tensor | None = None,
+    ):
+        if self.modulation:
+            seq_len = x.shape[1]
+            if noise_mask is not None:
+                mod_noisy = self.adaLN_modulation(adaln_noisy)
+                mod_clean = self.adaLN_modulation(adaln_clean)
+                scale_msa_noisy, gate_msa_noisy, scale_mlp_noisy, gate_mlp_noisy = mod_noisy.chunk(4, dim=1)
+                scale_msa_clean, gate_msa_clean, scale_mlp_clean, gate_mlp_clean = mod_clean.chunk(4, dim=1)
+                gate_msa_noisy, gate_mlp_noisy = gate_msa_noisy.tanh(), gate_mlp_noisy.tanh()
+                gate_msa_clean, gate_mlp_clean = gate_msa_clean.tanh(), gate_mlp_clean.tanh()
+                scale_msa_noisy, scale_mlp_noisy = 1.0 + scale_msa_noisy, 1.0 + scale_mlp_noisy
+                scale_msa_clean, scale_mlp_clean = 1.0 + scale_msa_clean, 1.0 + scale_mlp_clean
+                scale_msa = _select_per_token(scale_msa_noisy, scale_msa_clean, noise_mask, seq_len)
+                scale_mlp = _select_per_token(scale_mlp_noisy, scale_mlp_clean, noise_mask, seq_len)
+                gate_msa  = _select_per_token(gate_msa_noisy,  gate_msa_clean,  noise_mask, seq_len)
+                gate_mlp  = _select_per_token(gate_mlp_noisy,  gate_mlp_clean,  noise_mask, seq_len)
+            else:
+                mod = self.adaLN_modulation(adaln_input)
+                scale_msa, gate_msa, scale_mlp, gate_mlp = mod.unsqueeze(1).chunk(4, dim=2)
+                gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
+                scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
+            attn_out = self.attention(
+                self.attention_norm1(x) * scale_msa, attention_mask=attn_mask, freqs_cis=freqs_cis
+            )
+            x = x + gate_msa * self.attention_norm2(_clamp_fp16(attn_out))
+            x = x + gate_mlp * self.ffn_norm2(_clamp_fp16(self.feed_forward(self.ffn_norm1(x) * scale_mlp)))
+        else:
+            attn_out = self.attention(self.attention_norm1(x), attention_mask=attn_mask, freqs_cis=freqs_cis)
+            x = x + self.attention_norm2(_clamp_fp16(attn_out))
+            x = x + self.ffn_norm2(_clamp_fp16(self.feed_forward(self.ffn_norm1(x))))
+        return x
+
+    # ------------------------------------------------------------------
+    # FinalLayer.forward
+    # ------------------------------------------------------------------
+    def _patched_final_layer_forward(self, x, c=None, noise_mask=None, c_noisy=None, c_clean=None):
+        seq_len = x.shape[1]
+        if noise_mask is not None:
+            scale_noisy = 1.0 + self.adaLN_modulation(c_noisy)
+            scale_clean = 1.0 + self.adaLN_modulation(c_clean)
+            scale = _select_per_token(scale_noisy, scale_clean, noise_mask, seq_len)
+        else:
+            assert c is not None, "Either c or (c_noisy, c_clean) must be provided"
+            scale = 1.0 + self.adaLN_modulation(c)
+            scale = scale.unsqueeze(1)
+        x = self.norm_final(x) * scale
+        x = self.linear(x)
+        return x
+
+    sdnext_patches.patch(__name__, m.FeedForward,            '_forward_silu_gating', _patched_forward_silu_gating)
+    sdnext_patches.patch(__name__, m.ZImageTransformerBlock, 'forward',              _patched_zimage_block_forward)
+    sdnext_patches.patch(__name__, m.FinalLayer,             'forward',              _patched_final_layer_forward)
+