From a8ea6953ec15c3ba43e114710d56779cb51cc4ee Mon Sep 17 00:00:00 2001 From: vanDuven Date: Mon, 8 Dec 2025 05:43:00 -0500 Subject: [PATCH 1/2] Fix Z Image FP16 overflow via downscaling --- comfy/ldm/lumina/model.py | 12 ++++++++++-- comfy/supported_models.py | 2 +- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/comfy/ldm/lumina/model.py b/comfy/ldm/lumina/model.py index c47df49ca..389d479fb 100644 --- a/comfy/ldm/lumina/model.py +++ b/comfy/ldm/lumina/model.py @@ -119,6 +119,9 @@ class JointAttention(nn.Module): xv = xv.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3) output = optimized_attention_masked(xq.movedim(1, 2), xk.movedim(1, 2), xv.movedim(1, 2), self.n_local_heads, x_mask, skip_reshape=True, transformer_options=transformer_options) + if output.dtype == torch.float16: + output.div_(4) + return self.out(output) @@ -175,8 +178,12 @@ class FeedForward(nn.Module): def _forward_silu_gating(self, x1, x3): return clamp_fp16(F.silu(x1) * x3) - def forward(self, x): - return self.w2(self._forward_silu_gating(self.w1(x), self.w3(x))) + def forward(self, x, apply_fp16_downscale=False): + x3 = self.w3(x) + if x.dtype == torch.float16 and apply_fp16_downscale: + x3.div_(32) + + return self.w2(self._forward_silu_gating(self.w1(x), x3)) class JointTransformerBlock(nn.Module): @@ -287,6 +294,7 @@ class JointTransformerBlock(nn.Module): x = x + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2( clamp_fp16(self.feed_forward( modulate(self.ffn_norm1(x), scale_mlp), + apply_fp16_downscale=True, )) ) else: diff --git a/comfy/supported_models.py b/comfy/supported_models.py index 383c82c3e..fa15fac3c 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -1001,7 +1001,7 @@ class Lumina2(supported_models_base.BASE): unet_extra_config = {} latent_format = latent_formats.Flux - supported_inference_dtypes = [torch.bfloat16, torch.float32] + supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32] vae_key_prefix = ["vae."] text_encoder_key_prefix = ["text_encoders."] From 5905513e3207fb54df4507f9ce23f866b3a791f9 Mon Sep 17 00:00:00 2001 From: vanDuven Date: Thu, 11 Dec 2025 01:06:26 -0500 Subject: [PATCH 2/2] revert fp16 support for Lumina2 --- comfy/supported_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/supported_models.py b/comfy/supported_models.py index fa15fac3c..383c82c3e 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -1001,7 +1001,7 @@ class Lumina2(supported_models_base.BASE): unet_extra_config = {} latent_format = latent_formats.Flux - supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32] + supported_inference_dtypes = [torch.bfloat16, torch.float32] vae_key_prefix = ["vae."] text_encoder_key_prefix = ["text_encoders."]