From 12134fb9a39309421995768acb073979b88753a3 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Wed, 11 Jun 2025 15:12:09 -0400 Subject: [PATCH] Move to right place. --- comfy/ldm/flux/model.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/comfy/ldm/flux/model.py b/comfy/ldm/flux/model.py index 09dd2482c..db2ef3d6e 100644 --- a/comfy/ldm/flux/model.py +++ b/comfy/ldm/flux/model.py @@ -159,6 +159,10 @@ 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 +191,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)