From 407c54dc47353a3e309e8c54865b65f37ac25034 Mon Sep 17 00:00:00 2001 From: resonantsky Date: Mon, 13 Apr 2026 10:36:57 +0200 Subject: [PATCH 1/2] FP16 Clamp fix for Z-Image-Turbo, enable FP16 on ZiT. --- pipelines/model_z_image.py | 2 + pipelines/model_z_image_clampFP16.py | 109 +++++++++++++++++++++++++++ 2 files changed, 111 insertions(+) create mode 100644 pipelines/model_z_image_clampFP16.py diff --git a/pipelines/model_z_image.py b/pipelines/model_z_image.py index 6f9b8cf76..7e4487c82 100644 --- a/pipelines/model_z_image.py +++ b/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.model_z_image_clampFP16 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) diff --git a/pipelines/model_z_image_clampFP16.py b/pipelines/model_z_image_clampFP16.py new file mode 100644 index 000000000..4318a3549 --- /dev/null +++ b/pipelines/model_z_image_clampFP16.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) + From 0c9e86564129bef54a073c39df7934fdbae4de3d Mon Sep 17 00:00:00 2001 From: resonantsky Date: Mon, 13 Apr 2026 11:18:56 +0200 Subject: [PATCH 2/2] Restructure FP16 clamp patch into pipelines/z_image/patch_fp16.py subpackage --- pipelines/model_z_image.py | 2 +- pipelines/z_image/__init__.py | 0 pipelines/{model_z_image_clampFP16.py => z_image/patch_fp16.py} | 0 3 files changed, 1 insertion(+), 1 deletion(-) create mode 100644 pipelines/z_image/__init__.py rename pipelines/{model_z_image_clampFP16.py => z_image/patch_fp16.py} (100%) diff --git a/pipelines/model_z_image.py b/pipelines/model_z_image.py index 7e4487c82..048810776 100644 --- a/pipelines/model_z_image.py +++ b/pipelines/model_z_image.py @@ -3,7 +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.model_z_image_clampFP16 import apply_patches +from pipelines.z_image.patch_fp16 import apply_patches def load_nunchaku(): diff --git a/pipelines/z_image/__init__.py b/pipelines/z_image/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/pipelines/model_z_image_clampFP16.py b/pipelines/z_image/patch_fp16.py similarity index 100% rename from pipelines/model_z_image_clampFP16.py rename to pipelines/z_image/patch_fp16.py