Fix mistake in last commit. (#8496)

* Move to right place.
This commit is contained in:
comfyanonymous 2025-06-11 12:13:29 -07:00 committed by GitHub
parent af1eb58be8
commit 8a4ff747bd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -159,6 +159,9 @@ class Flux(nn.Module):
if add is not None: if add is not None:
img += add img += add
if img.dtype == torch.float16:
img = torch.nan_to_num(img, nan=0.0, posinf=65504, neginf=-65504)
img = torch.cat((txt, img), 1) img = torch.cat((txt, img), 1)
for i, block in enumerate(self.single_blocks): for i, block in enumerate(self.single_blocks):
@ -187,9 +190,6 @@ class Flux(nn.Module):
if add is not None: if add is not None:
img[:, txt.shape[1] :, ...] += add img[:, txt.shape[1] :, ...] += add
if img.dtype == torch.float16:
img = torch.nan_to_num(img, nan=0.0, posinf=65504, neginf=-65504)
img = img[:, txt.shape[1] :, ...] img = img[:, txt.shape[1] :, ...]
img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels) img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)