298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-06-02 01:22:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

Revert "This doesn't seem to be needed on chroma. (#8209)" (#8210)

Browse Source 
This reverts commit 7e84bf53737879ace37a68dc93e0df7704a53514.
This commit is contained in:
comfyanonymous 2025-05-20 02:39:55 -07:00 committed by GitHub
parent 7e84bf5373
commit 87f9130778
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
comfy/ldm/chroma/layers.py
View File

@ -109,6 +109,9 @@ class DoubleStreamBlock(nn.Module):
txt += txt_mod1.gate * self.txt_attn.proj(txt_attn)
txt += txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
if txt.dtype == torch.float16:
txt = torch.nan_to_num(txt, nan=0.0, posinf=65504, neginf=-65504)
return img, txt
@ -160,6 +163,8 @@ class SingleStreamBlock(nn.Module):
# compute activation in mlp stream, cat again and run second linear layer
output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
x += mod.gate * output
if x.dtype == torch.float16:
x = torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504)
return x