From c27ebeb1c2b7f163973f1dc9c57ee035214d6737 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 6 Sep 2024 03:21:52 -0400 Subject: [PATCH] Fix onnx export not working on flux. --- comfy/ldm/common_dit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/ldm/common_dit.py b/comfy/ldm/common_dit.py index 9016abc4..5aebaf9e 100644 --- a/comfy/ldm/common_dit.py +++ b/comfy/ldm/common_dit.py @@ -14,7 +14,7 @@ def pad_to_patch_size(img, patch_size=(2, 2), padding_mode="circular"): rms_norm_torch = None def rms_norm(x, weight, eps=1e-6): - if rms_norm_torch is not None: + if rms_norm_torch is not None and not (torch.jit.is_tracing() or torch.jit.is_scripting()): return rms_norm_torch(x, weight.shape, weight=comfy.ops.cast_to(weight, dtype=x.dtype, device=x.device), eps=eps) else: rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + eps)