From 25f810d004a7a349524dfab2b95d5fdf00bb9c52 Mon Sep 17 00:00:00 2001 From: Haoming Date: Sat, 3 Jan 2026 14:46:47 +0800 Subject: [PATCH] nvfp4 --- backend/nn/svdq.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/backend/nn/svdq.py b/backend/nn/svdq.py index cf6e842d..7adfa20d 100644 --- a/backend/nn/svdq.py +++ b/backend/nn/svdq.py @@ -910,4 +910,20 @@ class NunchakuZImageModel(NextDiT): def load_state_dict(self, sd, *args, **kwargs): sd = patch_z_image_state_dict(sd) + return self._load_state_dict(sd, *args, **kwargs) + + def _load_state_dict(self, sd, *args, **kwargs): + state_dict = self.state_dict() + for k in state_dict.keys(): + if k not in sd: + if "dummy" in k: + continue + if ".wcscales" not in k: + raise ValueError(f"Key {k} not found in state_dict") + sd[k] = torch.ones_like(state_dict[k]) + for n, m in self.named_modules(): + if isinstance(m, SVDQW4A4Linear): + if m.wtscale is not None: + m.wtscale = sd.pop(f"{n}.wtscale", 1.0) + return super().load_state_dict(sd, *args, **kwargs)