|
|
|
@ -178,7 +178,7 @@ class DoubleStreamBlock(nn.Module):
|
|
|
|
|
txt += txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
|
|
|
|
|
|
|
|
|
|
if txt.dtype == torch.float16:
|
|
|
|
|
txt = txt.clip(-65504, 65504)
|
|
|
|
|
txt = torch.nan_to_num(txt, nan=0.0, posinf=65504, neginf=-65504)
|
|
|
|
|
|
|
|
|
|
return img, txt
|
|
|
|
|
|
|
|
|
@ -233,7 +233,7 @@ class SingleStreamBlock(nn.Module):
|
|
|
|
|
output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
|
|
|
|
|
x += mod.gate * output
|
|
|
|
|
if x.dtype == torch.float16:
|
|
|
|
|
x = x.clip(-65504, 65504)
|
|
|
|
|
x = torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504)
|
|
|
|
|
return x
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|