wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-12-17 01:52:59 +08:00
Code Issues Actions 13 Packages Projects Releases Wiki Activity

Merge 5905513e32 into da2bfb5b0a

Browse Source
This commit is contained in:
vanDuven 2025-12-13 07:31:59 +00:00 committed by GitHub
commit d197be933d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 10 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
comfy/ldm/lumina/model.py
View File

@ -119,6 +119,9 @@ class JointAttention(nn.Module):
xv = xv.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3) 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) 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) return self.out(output)
@ -175,8 +178,12 @@ class FeedForward(nn.Module):
def _forward_silu_gating(self, x1, x3): def _forward_silu_gating(self, x1, x3):
return clamp_fp16(F.silu(x1) * x3) return clamp_fp16(F.silu(x1) * x3)
def forward(self, x): def forward(self, x, apply_fp16_downscale=False):
return self.w2(self._forward_silu_gating(self.w1(x), self.w3(x))) 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): class JointTransformerBlock(nn.Module):
@ -287,6 +294,7 @@ class JointTransformerBlock(nn.Module):
x = x + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2( x = x + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2(
clamp_fp16(self.feed_forward( clamp_fp16(self.feed_forward(
modulate(self.ffn_norm1(x), scale_mlp), modulate(self.ffn_norm1(x), scale_mlp),
apply_fp16_downscale=True,
)) ))
) )
else: else: