From 8a4ff747bd919afffe23b68c2e3480295b91998b Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Wed, 11 Jun 2025 12:13:29 -0700 Subject: [PATCH] Fix mistake in last commit. (#8496) * Move to right place. --- comfy/ldm/flux/model.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/comfy/ldm/flux/model.py b/comfy/ldm/flux/model.py index 09dd2482..846703d5 100644 --- a/comfy/ldm/flux/model.py +++ b/comfy/ldm/flux/model.py @@ -159,6 +159,9 @@ class Flux(nn.Module): if add is not None: 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) for i, block in enumerate(self.single_blocks): @@ -187,9 +190,6 @@ class Flux(nn.Module): if add is not None: 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 = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)