From 287299e0e72c3ac9c029cb922072ab56d5c86767 Mon Sep 17 00:00:00 2001 From: F1LIP3 Date: Sat, 15 Aug 2026 14:27:34 -0300 Subject: [PATCH] Use comfy.quant_ops NVFP4 layout and exact MM3_NVFP4=1 gate - Import TensorCoreNVFP4Layout from comfy.quant_ops so ComfyUI's quantize override and Params contract apply; comfy_kitchen stays lazy for the dequantize helpers. - Gate the conversion on the exact documented value MM3_NVFP4=1 in both load paths. --- comfy/sd.py | 2 +- comfy/text_encoders/minimax_music.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/comfy/sd.py b/comfy/sd.py index 4d3323992..068bf7df3 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -2267,7 +2267,7 @@ def load_diffusion_model_state_dict(sd, model_options={}, metadata=None, disable parameters = comfy.utils.calculate_parameters(sd) weight_dtype = comfy.utils.weight_dtype(sd) - if os.environ.get("MM3_NVFP4") and metadata is not None and metadata.get("comfy_model") == "minimax_music3_dit": + if os.environ.get("MM3_NVFP4") == "1" and metadata is not None and metadata.get("comfy_model") == "minimax_music3_dit": comfy.text_encoders.minimax_music._to_nvfp4(sd) load_device = model_options.get("load_device", model_management.get_torch_device()) diff --git a/comfy/text_encoders/minimax_music.py b/comfy/text_encoders/minimax_music.py index b6a7dae2d..b03a08795 100644 --- a/comfy/text_encoders/minimax_music.py +++ b/comfy/text_encoders/minimax_music.py @@ -16,7 +16,7 @@ def _to_nvfp4(sd): bytes, so the AR model's resident set fits VRAM and aimdo stops paging. """ import comfy_kitchen as ck - from comfy_kitchen.tensor.nvfp4 import TensorCoreNVFP4Layout + from comfy.quant_ops import TensorCoreNVFP4Layout dtype_code = ck.DTYPE_TO_CODE[torch.float32] for key in list(sd): @@ -132,7 +132,7 @@ class MiniMaxMusic3TEModel(MiniMaxMusic3AR): return hidden.unsqueeze(0), None, {} def load_state_dict(self, state_dict, strict=True, assign=False): - if os.environ.get("MM3_NVFP4"): + if os.environ.get("MM3_NVFP4") == "1": _to_nvfp4(state_dict) if self.model.pruned_embedding is None: self.model.pruned_embedding = "model.embed_tokens_prefill.weight" in state_dict