import dataclasses import hashlib import torch from torch import nn import comfy.model_management import comfy.model_prefetch import comfy.ops import comfy.utils from comfy.ldm.modules.attention import optimized_attention_for_device from comfy.text_encoders.llama import Llama2_, Qwen3_8BConfig from .prompt import AUDIO_CODE_OFFSET, SPECIAL_TOKEN_IDS CFG_SCALE = 1.5 CFG_TOP_K = 50 C0_VOCAB_SIZE = 16384 MAX_PROMPT_TOKENS = 5000 MAX_AUDIO_FRAMES = 9000 AUDIO_FRAMES_PER_SECOND = 25 def derive_seed(seed, *parts): digest = hashlib.blake2b(digest_size=8, person=b"minimax-ttm") digest.update(int(seed).to_bytes(8, "little", signed=False)) for part in parts: value = str(part).encode("utf-8") digest.update(len(value).to_bytes(4, "little")) digest.update(value) return int.from_bytes(digest.digest(), "little") & ((1 << 63) - 1) def sample_topk(logits, top_k, generator): values = torch.nan_to_num(logits.float(), nan=-1e9, posinf=1e9, neginf=-1e9) top_k = min(top_k, values.shape[-1]) threshold = torch.topk(values, top_k, dim=-1).values[..., -1, None] values = values.masked_fill(values < threshold, -float("inf")) probabilities = torch.nan_to_num(torch.softmax(values, dim=-1), nan=0.0) probabilities = probabilities / probabilities.sum(dim=-1, keepdim=True).clamp_min(1e-12) return torch.multinomial(probabilities, 1, generator=generator).squeeze(-1) class RVQAttention(nn.Module): 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 = 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): batch, length, hidden_size = x.shape if self.merged_qkv: q, k, v = self.qkv_proj(x).chunk(3, dim=-1) else: q = self.q_proj(x) k = self.k_proj(x) v = self.v_proj(x) q = q.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2) k = k.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2) v = v.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2) mask = torch.full((length, length), torch.finfo(q.dtype).min, device=q.device, dtype=q.dtype).triu_(1) attention = optimized_attention_for_device(q.device, mask=True, small_input=True) out = attention(q, k, v, self.num_heads, mask=mask, skip_reshape=True) return self.o_proj(out) class RVQRMSNorm(nn.Module): def __init__(self, hidden_size, dtype, device): super().__init__() self.weight = nn.Parameter(torch.empty(hidden_size, dtype=dtype, device=device)) def forward(self, x): return torch.nn.functional.rms_norm(x, (x.shape[-1],), comfy.ops.cast_to_input(self.weight, x), 1e-6) class RVQMLP(nn.Module): def __init__(self, hidden_size, intermediate_size, merged_mlp, dtype, device, operations): super().__init__() 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): if self.merged_mlp: return comfy.ops.linear_input_act(self.down_proj, self.gate_up_proj(x), "swiglu") return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x)) class RVQDecoderBlock(nn.Module): 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, merged_qkv, dtype, device, operations) self.post_attention_layernorm = RVQRMSNorm(hidden_size, dtype, device) 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)) return x + self.mlp(self.post_attention_layernorm(x)) class RVQDepthDecoder(nn.Module): def __init__(self, config, dtype, device, operations): 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) self.audio_heads = nn.ModuleList([ operations.Linear(hidden_size, audio_vocab_size, bias=False, dtype=dtype, device=device) for _ in range(num_codebooks - 1) ]) self.layers = nn.ModuleList([ RVQDecoderBlock( hidden_size, int(config["decoder_num_heads"]), int(config["decoder_intermediate_size"]), merged_qkv, merged_mlp, dtype, device, operations, ) for _ in range(int(config["decoder_num_layers"])) ]) self.norm = RVQRMSNorm(hidden_size, dtype, device) def forward(self, sequence): positions = torch.arange(sequence.shape[1], device=sequence.device) x = sequence + self.pos_embedding(positions, out_dtype=sequence.dtype).unsqueeze(0) for layer in self.layers: x = layer(x) return self.norm(x) class MiniMaxMusic3AR(nn.Module): def __init__(self, config, dtype, device, operations): super().__init__() config_fields = {field.name for field in dataclasses.fields(Qwen3_8BConfig)} 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 self.model = Llama2_(qwen_config, device=device, dtype=dtype, ops=operations) self.model.prefetch_dynamic_vbars = True self.model.graph_dynamic_vbar_blocks = True self.model.lm_head = operations.Linear(qwen_config.hidden_size, qwen_config.vocab_size, bias=False, dtype=dtype, device=device) self.model.lm_head_pruned = operations.Linear(qwen_config.hidden_size, C0_VOCAB_SIZE + 1, bias=False, dtype=dtype, device=device) self.model.embed_tokens_prefill = operations.Embedding(AUDIO_CODE_OFFSET, qwen_config.hidden_size, dtype=dtype, device=device) self.model.embed_tokens_audio = operations.Embedding(C0_VOCAB_SIZE, qwen_config.hidden_size, dtype=dtype, device=device) self.model.pruned_lm_head = None self.model.pruned_embedding = None self.model.audio_extra_embedding = operations.Embedding( int(config["audio_vocab_size"]) * (int(config["audio_num_codebooks"]) - 1), qwen_config.hidden_size, dtype=dtype, device=device, ) self.model.audio_decoder = RVQDepthDecoder(config, dtype, device, operations) self.audio_vocab_size = int(config["audio_vocab_size"]) self.num_codebooks = int(config["audio_num_codebooks"]) self.embedding_scale = self.num_codebooks ** -0.5 def _guided_c0(self, logits, cfg_scale, top_k): conditioned = logits[0:1].float() unconditioned = logits[1:2].float() guided = unconditioned + (conditioned - unconditioned) * cfg_scale threshold = torch.topk(conditioned, top_k, dim=-1).values[..., -1, None] return guided.masked_fill(conditioned < threshold, -float("inf")) def _depth_codes(self, hidden, c0, c0_embed, generator, execution_dtype, cfg_scale, top_k): decoder = self.model.audio_decoder sequence = [decoder.projection(hidden).unsqueeze(1)] sequence.append(decoder.projection(c0_embed).unsqueeze(1)) codes = [c0] hidden_parts = [] for index in range(1, self.num_codebooks): out = decoder(torch.cat(sequence, dim=1))[:, -1] hidden_parts.append(out[:1].detach()) logits = decoder.audio_heads[index - 1](out) conditioned = logits[:1].float() unconditioned = logits[1:2].float() code = sample_topk(unconditioned + (conditioned - unconditioned) * cfg_scale, top_k, generator).repeat(2) codes.append(code) if index < self.num_codebooks - 1: embedding = self.model.audio_extra_embedding( code + (index - 1) * self.audio_vocab_size, out_dtype=execution_dtype, ) sequence.append(decoder.projection(embedding).unsqueeze(1)) return torch.stack(codes, dim=1), torch.cat(hidden_parts, dim=-1) def _embed_c0(self, codes, execution_dtype): if self.model.pruned_embedding: return self.model.embed_tokens_audio(codes, out_dtype=execution_dtype) return self.model.embed_tokens(codes + AUDIO_CODE_OFFSET, out_dtype=execution_dtype) def _embed_audio_frame(self, codes, execution_dtype): c0 = self._embed_c0(codes[:, 0], execution_dtype) offsets = torch.arange(self.num_codebooks - 1, device=codes.device) * self.audio_vocab_size extra = self.model.audio_extra_embedding(codes[:, 1:] + offsets.unsqueeze(0), out_dtype=execution_dtype).sum(dim=1) return ((c0 + extra) * self.embedding_scale).unsqueeze(1) def _sample_c0(self, hidden, cfg_scale, top_k, generator, vocab_mask): if self.model.pruned_lm_head: guided = self._guided_c0(self.model.lm_head_pruned(hidden).float(), cfg_scale, top_k) code = sample_topk(guided, top_k, generator) stop_token = 0 offset = 1 else: logits = self.model.lm_head(hidden).float() stop_token = SPECIAL_TOKEN_IDS["<|audio_end|>"] logits = logits.masked_fill(vocab_mask, -float("inf")) guided = self._guided_c0(logits, cfg_scale, top_k).masked_fill(vocab_mask, -float("inf")) code = sample_topk(guided, top_k, generator) offset = AUDIO_CODE_OFFSET return torch.where(code == stop_token, 0, code - offset), code, stop_token def generate(self, input_ids, seed, max_audio_frames, device, cfg_scale=CFG_SCALE, top_k=CFG_TOP_K): prompt_tokens = int(input_ids.shape[1]) if prompt_tokens > MAX_PROMPT_TOKENS: raise ValueError(f"MiniMax Music3 prompt has {prompt_tokens} tokens; maximum is {MAX_PROMPT_TOKENS}") input_ids = input_ids.to(device) if comfy.model_management.should_use_bf16(device): execution_dtype = torch.bfloat16 else: execution_dtype = torch.float32 unconditioned = input_ids.clone() unconditioned[:, 1:-2] = SPECIAL_TOKEN_IDS["<|audio_cfg|>"] text_ids = torch.cat((input_ids, unconditioned), dim=0) if self.model.pruned_embedding: text_embeds = self.model.embed_tokens_prefill(text_ids, out_dtype=execution_dtype) else: text_embeds = self.model.embed_tokens(text_ids, out_dtype=execution_dtype) decode_limit = min(int(max_audio_frames), MAX_AUDIO_FRAMES) past = self.model.init_kv_cache(2, prompt_tokens + decode_limit + 1, device, execution_dtype) output = self.model(None, embeds=text_embeds, past_key_values=past, dtype=execution_dtype) last_hidden = output[0][:, -1] past = output[2] generator = torch.Generator(device=device).manual_seed(derive_seed(seed, "ar")) decoder = self.model.audio_decoder depth_io = { "hidden": torch.empty_like(last_hidden), "c0": torch.empty((last_hidden.shape[0],), dtype=torch.long, device=device), "c0_embed": torch.empty_like(last_hidden), "codes": torch.empty((last_hidden.shape[0], self.num_codebooks), dtype=torch.long, device=device), "depth_hidden": torch.empty((1, last_hidden.shape[-1] * (self.num_codebooks - 1)), dtype=execution_dtype, device=device), } decoder._comfy_cross_step_state = depth_io comfy.model_management._register_cross_step(decoder) hidden_frames = [] pending_code = None stop_token = None pending_event = None pending_hidden = None progress = comfy.utils.ProgressBar(decode_limit) cuda_device = torch.device(device).type == "cuda" vocab_mask = None if not self.model.pruned_lm_head: vocab_mask = torch.ones(self.model.vocab_size, dtype=torch.bool, device=device) vocab_mask[AUDIO_CODE_OFFSET:AUDIO_CODE_OFFSET + C0_VOCAB_SIZE] = False vocab_mask[SPECIAL_TOKEN_IDS["<|audio_end|>"]] = False for frame_index in comfy.utils.model_trange(decode_limit + 1, desc="AR sampling"): comfy.model_management.throw_exception_if_processing_interrupted() if pending_code is not None: if pending_event is not None: pending_event.synchronize() if int(pending_code.item()) == stop_token: pending_hidden = None break if pending_hidden is not None: hidden_frames.append(pending_hidden) progress.update_absolute(len(hidden_frames)) if len(hidden_frames) >= decode_limit: break c0, code_or_stop, stop_token = self._sample_c0(last_hidden, cfg_scale, top_k, generator, vocab_mask) if pending_code is None: pending_code = torch.empty_like(code_or_stop, device="cpu", pin_memory=cuda_device) if cuda_device: pending_event = torch.cuda.Event() pending_code.copy_(code_or_stop, non_blocking=cuda_device) if pending_event is not None: pending_event.record() c0 = c0.repeat(2) c0_embed = self._embed_c0(c0, execution_dtype) depth_io["hidden"].copy_(last_hidden) depth_io["c0"].copy_(c0) depth_io["c0_embed"].copy_(c0_embed) def depth_core(): codes, depth_hidden = self._depth_codes( depth_io["hidden"], depth_io["c0"], depth_io["c0_embed"], generator, execution_dtype, cfg_scale, top_k ) depth_io["codes"].copy_(codes) depth_io["depth_hidden"].copy_(depth_hidden) depth_queue = comfy.model_prefetch.make_prefetch_queue( [[decoder, self.model.audio_extra_embedding]], device, {"prefetch_dynamic_vbars": True} ) comfy.model_prefetch.prefetch_queue_pop( depth_queue, device, decoder, execution_dtype, core=depth_core, enable_graph=True, generator=generator ) comfy.model_prefetch.prefetch_queue_pop(depth_queue, device, None) feedback_codes = depth_io["codes"] depth_hidden = depth_io["depth_hidden"] frame_hidden = torch.cat((last_hidden[:1].detach(), depth_hidden), dim=-1) if frame_index > 0: pending_hidden = frame_hidden[0].clone() feedback = self._embed_audio_frame(feedback_codes, execution_dtype) output = self.model(None, embeds=feedback, past_key_values=past, dtype=execution_dtype) last_hidden = output[0][:, -1] past = output[2] if pending_hidden is not None and len(hidden_frames) < decode_limit: if pending_event is not None: pending_event.synchronize() if int(pending_code.item()) != stop_token: hidden_frames.append(pending_hidden) if not hidden_frames: raise ValueError("MiniMax Music3 generated zero audio frames") return torch.stack(hidden_frames).to(device="cpu")