From acba342a630f6053d890359aa91b7b54d40a38dc Mon Sep 17 00:00:00 2001 From: Artiprocher Date: Wed, 14 Jan 2026 16:29:43 +0800 Subject: [PATCH] fix RMSNorm precision --- diffsynth/models/z_image_dit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diffsynth/models/z_image_dit.py b/diffsynth/models/z_image_dit.py index f157f38..6e8866a 100644 --- a/diffsynth/models/z_image_dit.py +++ b/diffsynth/models/z_image_dit.py @@ -6,7 +6,7 @@ import torch.nn as nn import torch.nn.functional as F from torch.nn.utils.rnn import pad_sequence -from torch.nn import RMSNorm +from .general_modules import RMSNorm from ..core.attention import attention_forward from ..core.device.npu_compatible_device import IS_NPU_AVAILABLE from ..core.gradient import gradient_checkpoint_forward