diff --git a/comfy/ldm/minimax_music/ar.py b/comfy/ldm/minimax_music/ar.py index 28215cb58..2a8935318 100644 --- a/comfy/ldm/minimax_music/ar.py +++ b/comfy/ldm/minimax_music/ar.py @@ -43,15 +43,17 @@ def sample_topk(logits, top_k, generator): class RVQAttention(nn.Module): - def __init__(self, hidden_size, num_heads, dtype, device, operations): + def __init__(self, hidden_size, num_heads, merged_qkv, dtype, device, operations): super().__init__() self.num_heads = num_heads self.head_dim = hidden_size // num_heads - self.merged_qkv = None - self.qkv_proj = operations.Linear(hidden_size, hidden_size * 3, bias=False, dtype=dtype, device=device) - self.q_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device) - self.k_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device) - self.v_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device) + self.merged_qkv = merged_qkv + if merged_qkv: + self.qkv_proj = operations.Linear(hidden_size, hidden_size * 3, bias=False, dtype=dtype, device=device) + else: + self.q_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device) + self.k_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device) + self.v_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device) self.o_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device) def forward(self, x): @@ -81,12 +83,14 @@ class RVQRMSNorm(nn.Module): class RVQMLP(nn.Module): - def __init__(self, hidden_size, intermediate_size, dtype, device, operations): + def __init__(self, hidden_size, intermediate_size, merged_mlp, dtype, device, operations): super().__init__() - self.merged_mlp = None - self.gate_up_proj = operations.Linear(hidden_size, intermediate_size * 2, bias=False, dtype=dtype, device=device) - self.gate_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device) - self.up_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device) + self.merged_mlp = merged_mlp + if merged_mlp: + self.gate_up_proj = operations.Linear(hidden_size, intermediate_size * 2, bias=False, dtype=dtype, device=device) + else: + self.gate_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device) + self.up_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device) self.down_proj = operations.Linear(intermediate_size, hidden_size, bias=False, dtype=dtype, device=device) def forward(self, x): @@ -96,12 +100,12 @@ class RVQMLP(nn.Module): class RVQDecoderBlock(nn.Module): - def __init__(self, hidden_size, num_heads, intermediate_size, dtype, device, operations): + def __init__(self, hidden_size, num_heads, intermediate_size, merged_qkv, merged_mlp, dtype, device, operations): super().__init__() self.input_layernorm = RVQRMSNorm(hidden_size, dtype, device) - self.self_attn = RVQAttention(hidden_size, num_heads, dtype, device, operations) + self.self_attn = RVQAttention(hidden_size, num_heads, merged_qkv, dtype, device, operations) self.post_attention_layernorm = RVQRMSNorm(hidden_size, dtype, device) - self.mlp = RVQMLP(hidden_size, intermediate_size, dtype, device, operations) + self.mlp = RVQMLP(hidden_size, intermediate_size, merged_mlp, dtype, device, operations) def forward(self, x): x = x + self.self_attn(self.input_layernorm(x)) @@ -113,6 +117,8 @@ class RVQDepthDecoder(nn.Module): super().__init__() hidden_size = int(config["hidden_size"]) audio_vocab_size = int(config["audio_vocab_size"]) + merged_qkv = config.get("decoder_merged_qkv", False) + merged_mlp = config.get("decoder_merged_mlp", False) num_codebooks = int(config["audio_num_codebooks"]) self.projection = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device) self.pos_embedding = operations.Embedding(16, hidden_size, dtype=dtype, device=device) @@ -125,6 +131,8 @@ class RVQDepthDecoder(nn.Module): hidden_size, int(config["decoder_num_heads"]), int(config["decoder_intermediate_size"]), + merged_qkv, + merged_mlp, dtype, device, operations, @@ -148,8 +156,6 @@ class MiniMaxMusic3AR(nn.Module): qwen_config = Qwen3_8BConfig(**{key: value for key, value in config.items() if key in config_fields}) qwen_config.lm_head = False qwen_config.fixed_kv = True - qwen_config.merged_qkv = None - qwen_config.merged_mlp = None self.model = Llama2_(qwen_config, device=device, dtype=dtype, ops=operations) self.model.prefetch_dynamic_vbars = True self.model.graph_dynamic_vbar_blocks = True diff --git a/comfy/sd.py b/comfy/sd.py index 94f4f284f..4bdaa978c 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -1716,6 +1716,7 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip if quant is not None: model_options = model_options.copy() model_options["quantization_metadata"] = quant + clip_target.params["projection_config"] = comfy.text_encoders.minimax_music.detect_merged_config(clip_data[0]) clip_target.clip = comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel clip_target.tokenizer = comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer elif te_model == TEModel.CLIP_G: diff --git a/comfy/supported_models.py b/comfy/supported_models.py index d6d3c857f..33c378435 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -2218,7 +2218,10 @@ class MiniMaxMusic3(supported_models_base.BASE): return model_base.ModelType.FLOW def clip_target(self, state_dict={}): - return supported_models_base.ClipTarget(comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer, comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel) + detect = comfy.text_encoders.minimax_music.detect_merged_config(state_dict, self.text_encoder_key_prefix[0]) + target = supported_models_base.ClipTarget(comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer, comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel) + target.params["projection_config"] = detect + return target class LongCatImage(supported_models_base.BASE): diff --git a/comfy/text_encoders/llama.py b/comfy/text_encoders/llama.py index ddac547e0..f182e5147 100644 --- a/comfy/text_encoders/llama.py +++ b/comfy/text_encoders/llama.py @@ -17,13 +17,6 @@ import comfy.clip_model from . import qwen_vl -def detect_merged_config(state_dict, prefix="", layer_prefix="model.layers.0."): - return { - "merged_qkv": "{}{}self_attn.qkv_proj.weight".format(prefix, layer_prefix) in state_dict, - "merged_mlp": "{}{}mlp.gate_up_proj.weight".format(prefix, layer_prefix) in state_dict, - } - - @dataclass class FixedKV: key: torch.Tensor @@ -528,9 +521,9 @@ class Attention(nn.Module): ops = ops or nn self.kv_size = self.num_kv_heads * self.head_dim self.merged_qkv = getattr(config, "merged_qkv", False) - if self.merged_qkv is not False: + if self.merged_qkv: self.qkv_proj = ops.Linear(config.hidden_size, self.inner_size + self.kv_size * 2, bias=config.qkv_bias, device=device, dtype=dtype) - if self.merged_qkv is not True: + else: self.q_proj = ops.Linear(config.hidden_size, self.inner_size, bias=config.qkv_bias, device=device, dtype=dtype) self.k_proj = ops.Linear(config.hidden_size, self.kv_size, bias=config.qkv_bias, device=device, dtype=dtype) self.v_proj = ops.Linear(config.hidden_size, self.kv_size, bias=config.qkv_bias, device=device, dtype=dtype) @@ -627,9 +620,9 @@ class MLP(nn.Module): super().__init__() intermediate_size = intermediate_size or config.intermediate_size self.merged_mlp = getattr(config, "merged_mlp", False) - if self.merged_mlp is not False: + if self.merged_mlp: self.gate_up_proj = ops.Linear(config.hidden_size, intermediate_size * 2, bias=False, device=device, dtype=dtype) - if self.merged_mlp is not True: + else: self.gate_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype) self.up_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype) self.down_proj = ops.Linear(intermediate_size, config.hidden_size, bias=False, device=device, dtype=dtype) diff --git a/comfy/text_encoders/minimax_music.py b/comfy/text_encoders/minimax_music.py index 37d072664..c88d463cc 100644 --- a/comfy/text_encoders/minimax_music.py +++ b/comfy/text_encoders/minimax_music.py @@ -2,7 +2,6 @@ import torch from tokenizers import Tokenizer import comfy.ops -import comfy.text_encoders.llama from comfy.ldm.minimax_music.ar import CFG_SCALE, CFG_TOP_K, MAX_AUDIO_FRAMES, MiniMaxMusic3AR from comfy.ldm.minimax_music.prompt import SPECIAL_TOKEN_IDS, build_prompt @@ -26,6 +25,15 @@ MODEL_CONFIG = { } +def detect_merged_config(state_dict, prefix=""): + return { + "merged_qkv": "{}model.layers.0.self_attn.qkv_proj.weight".format(prefix) in state_dict, + "merged_mlp": "{}model.layers.0.mlp.gate_up_proj.weight".format(prefix) in state_dict, + "decoder_merged_qkv": "{}model.audio_decoder.layers.0.self_attn.qkv_proj.weight".format(prefix) in state_dict, + "decoder_merged_mlp": "{}model.audio_decoder.layers.0.mlp.gate_up_proj.weight".format(prefix) in state_dict, + } + + class MiniMaxMusic3Tokenizer: def __init__(self, embedding_directory=None, tokenizer_data={}): tokenizer_json = tokenizer_data.get("tokenizer_json") @@ -58,13 +66,13 @@ class MiniMaxMusic3Tokenizer: class MiniMaxMusic3TEModel(MiniMaxMusic3AR): - def __init__(self, device="cpu", dtype=None, model_options={}): + def __init__(self, device="cpu", dtype=None, model_options={}, projection_config=None): dtype = torch.bfloat16 quant_config = model_options.get("quantization_metadata", None) operations = model_options.get("custom_operations", None) if operations is None: operations = comfy.ops.mixed_precision_ops(quant_config, dtype) if quant_config is not None else comfy.ops.manual_cast - super().__init__(MODEL_CONFIG, dtype, device, operations) + super().__init__({**MODEL_CONFIG, **(projection_config or {})}, dtype, device, operations) self.dtypes = {dtype} self.execution_device = device @@ -91,26 +99,6 @@ class MiniMaxMusic3TEModel(MiniMaxMusic3AR): return hidden.unsqueeze(0), None, {} def load_state_dict(self, state_dict, strict=True, assign=False): - def select_projections(layers, config): - for layer in layers: - if layer.self_attn.merged_qkv is None: - if config["merged_qkv"]: - del layer.self_attn.q_proj, layer.self_attn.k_proj, layer.self_attn.v_proj - else: - del layer.self_attn.qkv_proj - layer.self_attn.merged_qkv = config["merged_qkv"] - if layer.mlp.merged_mlp is None: - if config["merged_mlp"]: - del layer.mlp.gate_proj, layer.mlp.up_proj - else: - del layer.mlp.gate_up_proj - layer.mlp.merged_mlp = config["merged_mlp"] - - select_projections(self.model.layers, comfy.text_encoders.llama.detect_merged_config(state_dict)) - select_projections( - self.model.audio_decoder.layers, - comfy.text_encoders.llama.detect_merged_config(state_dict, layer_prefix="model.audio_decoder.layers.0."), - ) if self.model.pruned_embedding is None: self.model.pruned_embedding = "model.embed_tokens_prefill.weight" in state_dict if self.model.pruned_embedding: