diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 9b65cd97bd3e..f929a9840afb 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -377,6 +377,8 @@ title: LuminaNextDiT2DModel - local: api/models/minimax_h3_transformer3d title: MiniMaxH3Transformer3DModel + - local: api/models/minimax_music3_transformer + title: MiniMaxMusic3Transformer1DModel - local: api/models/mochi_transformer3d title: MochiTransformer3DModel - local: api/models/motif_video_transformer_3d @@ -701,6 +703,8 @@ title: LTX-2 - local: api/pipelines/ltx_video title: LTXVideo + - local: api/pipelines/minimax_music3 + title: MiniMax Music 3 - local: api/pipelines/minimax_h3 title: MiniMax-H3 - local: api/pipelines/mochi diff --git a/docs/source/en/api/models/minimax_music3_transformer.md b/docs/source/en/api/models/minimax_music3_transformer.md new file mode 100644 index 000000000000..3f62199490ab --- /dev/null +++ b/docs/source/en/api/models/minimax_music3_transformer.md @@ -0,0 +1,22 @@ + + +# MiniMaxMusic3Transformer1DModel + +The 2.4B flow-matching Diffusion Transformer of [MiniMax Music 3](https://huggingface.co/MiniMaxAI/MiniMax-Music3). It +denoises 128-channel Flow-VAE audio latents conditioned on the per-frame hidden states of the model's autoregressive +language-model stage, prepending the flow-matching timestep as an extra sequence token (a Stable-Audio-lineage +continuous transformer with partial rotary attention and GLU feedforwards). + +## MiniMaxMusic3Transformer1DModel + +[[autodoc]] MiniMaxMusic3Transformer1DModel diff --git a/docs/source/en/api/pipelines/minimax_music3.md b/docs/source/en/api/pipelines/minimax_music3.md new file mode 100644 index 000000000000..7645d71a8861 --- /dev/null +++ b/docs/source/en/api/pipelines/minimax_music3.md @@ -0,0 +1,116 @@ + + +# MiniMax Music 3 + +[MiniMax Music 3](https://huggingface.co/MiniMaxAI/MiniMax-Music3) is a music generation model that produces complete +songs up to five minutes long from lyrics and a music description, with expressive vocals and long-range structure. + +The model is a hybrid of an autoregressive and a diffusion stage: an 8B Qwen3-based global language model predicts one +semantic audio token per frame while a small depth decoder fills in seven residual RVQ codebooks, and their fused +hidden states condition a 2.4B flow-matching transformer that produces Flow-VAE latents in overlapping chunks. A +DAC-style decoder turns the latents into 44.1 kHz stereo audio. + +## Usage + +MiniMax Music 3 is available as a modular pipeline. + +```py +import soundfile as sf +import torch +from diffusers import ModularPipeline + +pipe = ModularPipeline.from_pretrained("MiniMaxAI/MiniMax-Music3") +pipe.load_components(dtype=torch.bfloat16) +pipe.to("cuda") + +lyrics = """[verse] +Morning light filtering through the pine +Every quiet street is yours and mine +[chorus] +Softly the world begins to breathe""" + +prompt = ( + "Genre: acoustic pop. BPM: 96. Key: C major. Warm and intimate, building gently into the chorus. " + "Vocals: soft female lead, close and breathy, light stacked harmonies in the chorus. " + "Arrangement: fingerpicked guitar and soft piano; brushed drums and upright bass enter in the chorus." +) + +audio = pipe( + prompt=prompt, + lyrics=lyrics, + audio_duration=60.0, + generator=torch.Generator("cuda").manual_seed(7), + output="audios", +)[0] + +sf.write("minimax_music3.wav", audio.T, pipe.sampling_rate) +``` + +## Reduce memory usage + +Refer to the [Reduce memory usage](../../optimization/memory) guide for more details about the various memory saving +techniques. + +The full pipeline needs ~23 GB of VRAM in bfloat16. With automatic CPU offloading a generation runs in ~22 GB of free +VRAM, and additionally group-offloading the language model fits in 8 GB. + +```py +import torch +from diffusers import ComponentsManager, ModularPipeline +from diffusers.hooks.group_offloading import apply_group_offloading + +manager = ComponentsManager() +manager.enable_auto_cpu_offload(device="cuda") +pipe = ModularPipeline.from_pretrained("MiniMaxAI/MiniMax-Music3", components_manager=manager) +pipe.load_components(dtype=torch.bfloat16) + +# Only needed below ~22 GB of free VRAM — slower, but fits in 8 GB. +apply_group_offloading( + pipe.language_model, onload_device=torch.device("cuda"), offload_type="leaf_level", use_stream=True +) +``` + +## Tips + +- Structure tags such as `[intro]`, `[verse]`, `[pre-chorus]`, `[chorus]`, `[bridge]`, `[instrumental]`, `[solo]`, and + `[outro]` must each be on their own line in `lyrics`. Text on the same line as a leading tag is dropped by the + model's input contract. +- The music description controls the vocals: describe the vocal gender and timbre explicitly (e.g. "warm female + vocal") or the model may drift instrumental. For fine-grained control, structure the description into global + metadata (genre, BPM, key, emotional progression), vocal details, and arrangement. +- `audio_duration` is an upper bound — the language model may end the song earlier with a stop token. The + autoregressive stage generates 25 frames per second of audio and dominates the runtime. +- The classifier-free guidance scale of the flow-matching stage is a guider setting (the reference inference value is + 1.7): swap it with `pipe.update_components(guider=ClassifierFreeGuidance(guidance_scale=...))`. +- The pipeline returns the vocoder's native 44.1 kHz stereo output. The reference server additionally resamples to 32 + kHz; apply your own resampling if you need that exact rate. + +## MiniMaxMusic3ModularPipeline + +[[autodoc]] MiniMaxMusic3ModularPipeline + +## MiniMaxMusic3Blocks + +[[autodoc]] MiniMaxMusic3Blocks + +## MiniMaxMusic3ConditionEncoder + +[[autodoc]] MiniMaxMusic3ConditionEncoder + +## MiniMaxMusic3RVQDepthDecoder + +[[autodoc]] MiniMaxMusic3RVQDepthDecoder + +## MiniMaxMusic3Vocoder + +[[autodoc]] MiniMaxMusic3Vocoder diff --git a/scripts/convert_minimax_music3_to_diffusers.py b/scripts/convert_minimax_music3_to_diffusers.py new file mode 100644 index 000000000000..c6cef069c27e --- /dev/null +++ b/scripts/convert_minimax_music3_to_diffusers.py @@ -0,0 +1,269 @@ +# Conversion script for MiniMax Music 3 (https://huggingface.co/MiniMaxAI/MiniMax-Music3). +# +# Original checkpoint layout: +# flowmatching_vae.pth flow-matching DiT + condition projection +# dav.pth Flow-VAE (DAC-style) decoder +# qwen_7B/qwen_7B/ Qwen3 backbone + audio embedding + RVQ depth decoder (sharded safetensors) +# qwen_7B/qwen3-8B-tokenizer-music/ music tokenizer +# +# Usage: +# python scripts/convert_minimax_music3_to_diffusers.py \ +# --checkpoint_dir MiniMaxAI/MiniMax-Music3 --output_path ./minimax-music3-diffusers + +import argparse +import json +import os + +import torch +from safetensors.torch import load_file + +from diffusers import ( + FlowMatchEulerDiscreteScheduler, + MiniMaxMusic3ConditionEncoder, + MiniMaxMusic3RVQDepthDecoder, + MiniMaxMusic3Transformer1DModel, + MiniMaxMusic3Vocoder, +) + + +def load_dit_state_dict(checkpoint_dir: str) -> dict: + return torch.load(os.path.join(checkpoint_dir, "flowmatching_vae.pth"), map_location="cpu", weights_only=True) + + +def load_dav_state_dict(checkpoint_dir: str) -> dict: + return torch.load(os.path.join(checkpoint_dir, "dav.pth"), map_location="cpu", weights_only=True) + + +def load_qwen_state_dict(checkpoint_dir: str) -> dict: + qwen_dir = os.path.join(checkpoint_dir, "qwen_7B", "qwen_7B") + with open(os.path.join(qwen_dir, "model.safetensors.index.json")) as f: + index = json.load(f) + state_dict = {} + for filename in sorted(set(index["weight_map"].values())): + state_dict.update(load_file(os.path.join(qwen_dir, filename), device="cpu")) + return state_dict + + +def convert_transformer(dit_state_dict: dict) -> MiniMaxMusic3Transformer1DModel: + prefix = "diffusion_transformer." + converted = { + "time_proj.weight": dit_state_dict[prefix + "timestep_features.weight"], + "time_embed.linear_1.weight": dit_state_dict[prefix + "to_timestep_embed.0.weight"], + "time_embed.linear_1.bias": dit_state_dict[prefix + "to_timestep_embed.0.bias"], + "time_embed.linear_2.weight": dit_state_dict[prefix + "to_timestep_embed.2.weight"], + "time_embed.linear_2.bias": dit_state_dict[prefix + "to_timestep_embed.2.bias"], + "preprocess_conv.weight": dit_state_dict[prefix + "preprocess_conv.weight"], + "postprocess_conv.weight": dit_state_dict[prefix + "postprocess_conv.weight"], + "proj_in.weight": dit_state_dict[prefix + "transformer.project_in.weight"], + "proj_out.weight": dit_state_dict[prefix + "transformer.project_out.weight"], + } + num_layers = 0 + while prefix + f"transformer.layers.{num_layers}.pre_norm.gamma" in dit_state_dict: + num_layers += 1 + for i in range(num_layers): + original = prefix + f"transformer.layers.{i}." + target = f"transformer_blocks.{i}." + converted[target + "norm1.weight"] = dit_state_dict[original + "pre_norm.gamma"] + converted[target + "norm1.bias"] = dit_state_dict[original + "pre_norm.beta"] + query, key, value = dit_state_dict[original + "self_attn.to_qkv.weight"].chunk(3, dim=0) + converted[target + "attn.to_q.weight"] = query + converted[target + "attn.to_k.weight"] = key + converted[target + "attn.to_v.weight"] = value + converted[target + "attn.to_out.0.weight"] = dit_state_dict[original + "self_attn.to_out.weight"] + converted[target + "norm2.weight"] = dit_state_dict[original + "ff_norm.gamma"] + converted[target + "norm2.bias"] = dit_state_dict[original + "ff_norm.beta"] + converted[target + "ff_in.weight"] = dit_state_dict[original + "ff.ff.0.proj.weight"] + converted[target + "ff_in.bias"] = dit_state_dict[original + "ff.ff.0.proj.bias"] + converted[target + "ff_out.weight"] = dit_state_dict[original + "ff.ff.2.weight"] + converted[target + "ff_out.bias"] = dit_state_dict[original + "ff.ff.2.bias"] + + with torch.device("meta"): + transformer = MiniMaxMusic3Transformer1DModel(num_layers=num_layers) + transformer.load_state_dict(converted, strict=True, assign=True) + return transformer + + +def convert_condition_encoder(dit_state_dict: dict) -> MiniMaxMusic3ConditionEncoder: + converted = { + "layer_weight_logits": dit_state_dict["cond_layer_logits"], + "layer_scale": dit_state_dict["cond_layer_scale"], + "proj.weight": dit_state_dict["latent_conditioners.0.weight"], + "proj.bias": dit_state_dict["latent_conditioners.0.bias"], + } + with torch.device("meta"): + condition_encoder = MiniMaxMusic3ConditionEncoder() + condition_encoder.load_state_dict(converted, strict=True, assign=True) + return condition_encoder + + +def convert_vocoder(dav_state_dict: dict) -> MiniMaxMusic3Vocoder: + converted = { + "dec_in_proj.weight": dav_state_dict["dec_in_proj.weight"], + "dec_in_proj.bias": dav_state_dict["dec_in_proj.bias"], + } + # The reference decoder is one nn.Sequential: [conv_in, block*4, snake, conv_out, tanh]. + for suffix in ("weight_g", "weight_v", "bias"): + converted[f"conv_in.{suffix}"] = dav_state_dict[f"decoder.model.0.{suffix}"] + converted[f"conv_out.{suffix}"] = dav_state_dict[f"decoder.model.6.{suffix}"] + converted["snake_out.alpha"] = dav_state_dict["decoder.model.5.alpha"] + for block_index in range(4): + original = f"decoder.model.{block_index + 1}.block." + target = f"blocks.{block_index}." + converted[target + "snake1.alpha"] = dav_state_dict[original + "0.alpha"] + for suffix in ("weight_g", "weight_v", "bias"): + converted[target + f"conv_t1.{suffix}"] = dav_state_dict[original + f"1.{suffix}"] + for unit_index, unit_name in ((2, "res_unit1"), (3, "res_unit2"), (4, "res_unit3")): + converted[target + f"{unit_name}.snake1.alpha"] = dav_state_dict[original + f"{unit_index}.block.0.alpha"] + converted[target + f"{unit_name}.snake2.alpha"] = dav_state_dict[original + f"{unit_index}.block.2.alpha"] + for suffix in ("weight_g", "weight_v", "bias"): + converted[target + f"{unit_name}.conv1.{suffix}"] = dav_state_dict[ + original + f"{unit_index}.block.1.{suffix}" + ] + converted[target + f"{unit_name}.conv2.{suffix}"] = dav_state_dict[ + original + f"{unit_index}.block.3.{suffix}" + ] + + vocoder = MiniMaxMusic3Vocoder() + vocoder.load_state_dict(converted, strict=True) + return vocoder + + +def convert_rvq_depth_decoder(qwen_state_dict: dict, model_config: dict) -> MiniMaxMusic3RVQDepthDecoder: + prefix = "model.audio_decoder." + converted = { + "audio_embeddings.weight": qwen_state_dict["model.audio_extra_embedding.weight"], + "projection.weight": qwen_state_dict[prefix + "projection.weight"], + "pos_embedding.weight": qwen_state_dict[prefix + "pos_embedding.weight"], + "norm.weight": qwen_state_dict[prefix + "norm.weight"], + } + num_codebooks = int(model_config["audio_num_codebooks"]) + for i in range(num_codebooks - 1): + converted[f"audio_heads.{i}.weight"] = qwen_state_dict[prefix + f"audio_heads.{i}.weight"] + num_layers = int(model_config["decoder_num_layers"]) + for i in range(num_layers): + original = prefix + f"layers.{i}." + target = f"layers.{i}." + converted[target + "input_layernorm.weight"] = qwen_state_dict[original + "input_layernorm.weight"] + converted[target + "post_attention_layernorm.weight"] = qwen_state_dict[ + original + "post_attention_layernorm.weight" + ] + converted[target + "attn.to_q.weight"] = qwen_state_dict[original + "self_attn.q_proj.weight"] + converted[target + "attn.to_k.weight"] = qwen_state_dict[original + "self_attn.k_proj.weight"] + converted[target + "attn.to_v.weight"] = qwen_state_dict[original + "self_attn.v_proj.weight"] + converted[target + "attn.to_out.weight"] = qwen_state_dict[original + "self_attn.o_proj.weight"] + for proj in ("gate_proj", "up_proj", "down_proj"): + converted[target + proj + ".weight"] = qwen_state_dict[original + f"mlp.{proj}.weight"] + + with torch.device("meta"): + rvq_depth_decoder = MiniMaxMusic3RVQDepthDecoder( + hidden_size=int(model_config["hidden_size"]), + num_layers=num_layers, + num_attention_heads=int(model_config["decoder_num_heads"]), + intermediate_size=int(model_config["decoder_intermediate_size"]), + audio_vocab_size=int(model_config["audio_vocab_size"]), + num_codebooks=num_codebooks, + ) + rvq_depth_decoder.load_state_dict(converted, strict=True, assign=True) + return rvq_depth_decoder + + +def convert_language_model(qwen_state_dict: dict, model_config: dict): + from transformers import Qwen3Config, Qwen3ForCausalLM + + config = Qwen3Config( + vocab_size=model_config["vocab_size"], + hidden_size=model_config["hidden_size"], + intermediate_size=model_config["intermediate_size"], + num_hidden_layers=model_config["num_hidden_layers"], + num_attention_heads=model_config["num_attention_heads"], + num_key_value_heads=model_config["num_key_value_heads"], + head_dim=model_config["head_dim"], + max_position_embeddings=model_config.get("max_position_embeddings", 10240), + rope_theta=model_config.get("rope_theta", 1000000), + rms_norm_eps=model_config.get("rms_norm_eps", 1e-6), + tie_word_embeddings=model_config.get("tie_word_embeddings", False), + ) + backbone_state_dict = { + key: value + for key, value in qwen_state_dict.items() + if not key.startswith(("model.audio_extra_embedding", "model.audio_decoder.")) + } + with torch.device("meta"): + language_model = Qwen3ForCausalLM(config) + language_model.load_state_dict(backbone_state_dict, strict=True, assign=True) + return language_model + + +def main(args): + checkpoint_dir = args.checkpoint_dir + if not os.path.isdir(checkpoint_dir): + from huggingface_hub import snapshot_download + + checkpoint_dir = snapshot_download(checkpoint_dir) + + with open(os.path.join(checkpoint_dir, "qwen_7B", "qwen_7B", "config.json")) as f: + model_config = json.load(f) + + dit_state_dict = load_dit_state_dict(checkpoint_dir) + transformer = convert_transformer(dit_state_dict).to(args.dtype) + condition_encoder = convert_condition_encoder(dit_state_dict).to(args.dtype) + del dit_state_dict + vocoder = convert_vocoder(load_dav_state_dict(checkpoint_dir)).to(args.dtype) + + qwen_state_dict = load_qwen_state_dict(checkpoint_dir) + rvq_depth_decoder = convert_rvq_depth_decoder(qwen_state_dict, model_config).to(torch.bfloat16) + language_model = convert_language_model(qwen_state_dict, model_config) + del qwen_state_dict + + from transformers import AutoTokenizer + + from diffusers import MiniMaxMusic3Blocks + + tokenizer = AutoTokenizer.from_pretrained(os.path.join(checkpoint_dir, "qwen_7B", "qwen3-8B-tokenizer-music")) + # num_train_timesteps=1 keeps `scheduler.timesteps` equal to the flow-matching time in [0, 1] that the + # transformer's Fourier embedding expects. + scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1, shift=1.0, invert_sigmas=True) + + pipeline = MiniMaxMusic3Blocks().init_pipeline() + pipeline.update_components( + language_model=language_model, + rvq_depth_decoder=rvq_depth_decoder, + condition_encoder=condition_encoder, + transformer=transformer, + vocoder=vocoder, + tokenizer=tokenizer, + scheduler=scheduler, + ) + pipeline.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + + # save_pretrained bakes the local output path into the modular index's loading specs; point them at the + # Hub repo the components will be uploaded to instead. + index_path = os.path.join(args.output_path, "modular_model_index.json") + with open(index_path) as f: + index = json.load(f) + for entry in index.values(): + if isinstance(entry, list) and len(entry) == 3 and isinstance(entry[2], dict): + if entry[2].get("pretrained_model_name_or_path") == args.output_path: + entry[2]["pretrained_model_name_or_path"] = args.repo_id + with open(index_path, "w") as f: + json.dump(index, f, indent=2, sort_keys=True) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--checkpoint_dir", + type=str, + default="MiniMaxAI/MiniMax-Music3", + help="Local directory or Hugging Face Hub repo id of the original checkpoint.", + ) + parser.add_argument("--output_path", type=str, required=True) + parser.add_argument( + "--repo_id", + type=str, + default="MiniMaxAI/MiniMax-Music3", + help="Hub repo id the converted components will live in (written into the modular index loading specs).", + ) + parser.add_argument("--dtype", type=lambda name: getattr(torch, name), default="float32") + args = parser.parse_args() + main(args) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 66cfa442908b..4bec0f5bd7ff 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -313,6 +313,10 @@ "Lumina2Transformer2DModel", "LuminaNextDiT2DModel", "MiniMaxH3Transformer3DModel", + "MiniMaxMusic3ConditionEncoder", + "MiniMaxMusic3RVQDepthDecoder", + "MiniMaxMusic3Transformer1DModel", + "MiniMaxMusic3Vocoder", "MochiTransformer3DModel", "ModelMixin", "MotifVideoTransformer3DModel", @@ -543,6 +547,8 @@ "LTXModularPipeline", "MiniMaxH3Blocks", "MiniMaxH3ModularPipeline", + "MiniMaxMusic3Blocks", + "MiniMaxMusic3ModularPipeline", "QwenImageAutoBlocks", "QwenImageEditAutoBlocks", "QwenImageEditModularPipeline", @@ -1171,6 +1177,10 @@ Lumina2Transformer2DModel, LuminaNextDiT2DModel, MiniMaxH3Transformer3DModel, + MiniMaxMusic3ConditionEncoder, + MiniMaxMusic3RVQDepthDecoder, + MiniMaxMusic3Transformer1DModel, + MiniMaxMusic3Vocoder, MochiTransformer3DModel, ModelMixin, MotifVideoTransformer3DModel, @@ -1380,6 +1390,8 @@ LTXModularPipeline, MiniMaxH3Blocks, MiniMaxH3ModularPipeline, + MiniMaxMusic3Blocks, + MiniMaxMusic3ModularPipeline, QwenImageAutoBlocks, QwenImageEditAutoBlocks, QwenImageEditModularPipeline, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index f9f40899f4a9..d7f2c75c68e7 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -59,9 +59,11 @@ _import_structure["autoencoders.autoencoder_vidtok"] = ["AutoencoderVidTok"] _import_structure["autoencoders.consistency_decoder_vae"] = ["ConsistencyDecoderVAE"] _import_structure["autoencoders.ltx2_diffusion_decoder"] = ["LTX2VideoDiffusionDecoderModel"] + _import_structure["autoencoders.minimax_music3_vocoder"] = ["MiniMaxMusic3Vocoder"] _import_structure["autoencoders.vq_model"] = ["VQModel"] _import_structure["cache_utils"] = ["CacheMixin"] _import_structure["condition_embedders.condition_embedder_anima"] = ["AnimaTextConditioner"] + _import_structure["condition_embedders.condition_embedder_minimax_music3"] = ["MiniMaxMusic3ConditionEncoder"] _import_structure["controlnets.controlnet"] = ["ControlNetModel"] _import_structure["controlnets.controlnet_cosmos"] = ["CosmosControlNetModel"] _import_structure["controlnets.controlnet_flux"] = ["FluxControlNetModel", "FluxMultiControlNetModel"] @@ -92,6 +94,7 @@ _import_structure["transformers.hunyuan_transformer_2d"] = ["HunyuanDiT2DModel"] _import_structure["transformers.latte_transformer_3d"] = ["LatteTransformer3DModel"] _import_structure["transformers.lumina_nextdit2d"] = ["LuminaNextDiT2DModel"] + _import_structure["transformers.minimax_music3_rvq_depth_decoder"] = ["MiniMaxMusic3RVQDepthDecoder"] _import_structure["transformers.pixart_transformer_2d"] = ["PixArtTransformer2DModel"] _import_structure["transformers.prior_transformer"] = ["PriorTransformer"] _import_structure["transformers.sana_transformer"] = ["SanaTransformer2DModel"] @@ -132,6 +135,7 @@ _import_structure["transformers.transformer_ltx2"] = ["LTX2VideoTransformer3DModel"] _import_structure["transformers.transformer_lumina2"] = ["Lumina2Transformer2DModel"] _import_structure["transformers.transformer_minimax_h3"] = ["MiniMaxH3Transformer3DModel"] + _import_structure["transformers.transformer_minimax_music3"] = ["MiniMaxMusic3Transformer1DModel"] _import_structure["transformers.transformer_mochi"] = ["MochiTransformer3DModel"] _import_structure["transformers.transformer_motif_video"] = ["MotifVideoTransformer3DModel"] _import_structure["transformers.transformer_nucleusmoe_image"] = ["NucleusMoEImageTransformer2DModel"] @@ -199,10 +203,11 @@ Cosmos3AVAEAudioTokenizer, LongCatAudioDiTVae, LTX2VideoDiffusionDecoderModel, + MiniMaxMusic3Vocoder, VQModel, ) from .cache_utils import CacheMixin - from .condition_embedders import AnimaTextConditioner + from .condition_embedders import AnimaTextConditioner, MiniMaxMusic3ConditionEncoder from .controlnets import ( ControlNetModel, ControlNetUnionModel, @@ -269,6 +274,8 @@ Lumina2Transformer2DModel, LuminaNextDiT2DModel, MiniMaxH3Transformer3DModel, + MiniMaxMusic3RVQDepthDecoder, + MiniMaxMusic3Transformer1DModel, MochiTransformer3DModel, MotifVideoTransformer3DModel, NucleusMoEImageTransformer2DModel, diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index 2ba58f588825..7d24611c825e 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/src/diffusers/models/autoencoders/__init__.py @@ -29,4 +29,5 @@ from .autoencoder_vidtok import AutoencoderVidTok from .consistency_decoder_vae import ConsistencyDecoderVAE from .ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel +from .minimax_music3_vocoder import MiniMaxMusic3Vocoder from .vq_model import VQModel diff --git a/src/diffusers/models/autoencoders/minimax_music3_vocoder.py b/src/diffusers/models/autoencoders/minimax_music3_vocoder.py new file mode 100644 index 000000000000..bf325634e451 --- /dev/null +++ b/src/diffusers/models/autoencoders/minimax_music3_vocoder.py @@ -0,0 +1,115 @@ +# Copyright 2026 The MiniMax Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math + +import torch +import torch.nn as nn +from torch.nn.utils import weight_norm + +from ...configuration_utils import ConfigMixin, register_to_config +from ..modeling_utils import ModelMixin + + +class MiniMaxMusic3Snake1d(nn.Module): + def __init__(self, channels: int): + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + shape = hidden_states.shape + hidden_states = hidden_states.reshape(shape[0], shape[1], -1) + hidden_states = hidden_states + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * hidden_states).pow(2) + return hidden_states.reshape(shape) + + +class MiniMaxMusic3VocoderResidualUnit(nn.Module): + def __init__(self, dim: int, dilation: int): + super().__init__() + pad = (7 - 1) * dilation // 2 + self.snake1 = MiniMaxMusic3Snake1d(dim) + self.conv1 = weight_norm(nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad)) + self.snake2 = MiniMaxMusic3Snake1d(dim) + self.conv2 = weight_norm(nn.Conv1d(dim, dim, kernel_size=1)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + residual = self.conv2(self.snake2(self.conv1(self.snake1(hidden_states)))) + return hidden_states + residual + + +class MiniMaxMusic3VocoderBlock(nn.Module): + def __init__(self, input_dim: int, output_dim: int, stride: int): + super().__init__() + self.snake1 = MiniMaxMusic3Snake1d(input_dim) + self.conv_t1 = weight_norm( + nn.ConvTranspose1d( + input_dim, output_dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2) + ) + ) + self.res_unit1 = MiniMaxMusic3VocoderResidualUnit(output_dim, dilation=1) + self.res_unit2 = MiniMaxMusic3VocoderResidualUnit(output_dim, dilation=3) + self.res_unit3 = MiniMaxMusic3VocoderResidualUnit(output_dim, dilation=9) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv_t1(self.snake1(hidden_states)) + hidden_states = self.res_unit1(hidden_states) + hidden_states = self.res_unit2(hidden_states) + return self.res_unit3(hidden_states) + + +class MiniMaxMusic3Vocoder(ModelMixin, ConfigMixin): + r""" + The Flow-VAE waveform decoder of MiniMax Music 3 (a DAC-style decoder). It decodes flow-matched latents of shape + `(batch, latent_channels, length)` into stereo waveforms at `sampling_rate`; the two audio channels are decoded as + two folded `latent_channels // 2` streams. + """ + + @register_to_config + def __init__( + self, + latent_channels: int = 128, + decoder_input_dim: int = 1024, + decoder_hidden_dim: int = 1536, + upsampling_ratios: tuple = (8, 8, 4, 2), + sampling_rate: int = 44100, + ): + super().__init__() + self.dec_in_proj = nn.Conv1d(latent_channels // 2, decoder_input_dim, kernel_size=1) + self.conv_in = weight_norm(nn.Conv1d(decoder_input_dim, decoder_hidden_dim, kernel_size=7, padding=3)) + blocks = [] + output_dim = decoder_hidden_dim + for index, stride in enumerate(upsampling_ratios): + input_dim = decoder_hidden_dim // (2**index) + output_dim = decoder_hidden_dim // (2 ** (index + 1)) + blocks.append(MiniMaxMusic3VocoderBlock(input_dim, output_dim, stride)) + self.blocks = nn.ModuleList(blocks) + self.snake_out = MiniMaxMusic3Snake1d(output_dim) + self.conv_out = weight_norm(nn.Conv1d(output_dim, 1, kernel_size=7, padding=3)) + + def forward(self, latents: torch.Tensor) -> torch.Tensor: + r""" + Args: + latents (`torch.Tensor` of shape `(batch, latent_channels, length)`): + Flow-matched Flow-VAE latents. + + Returns: + `torch.Tensor` of shape `(batch, 2, samples)`: the stereo waveform in `[-1, 1]`. + """ + batch_size, _, length = latents.shape + hidden_states = latents.reshape(batch_size * 2, self.config.latent_channels // 2, length) + hidden_states = self.conv_in(self.dec_in_proj(hidden_states)) + for block in self.blocks: + hidden_states = block(hidden_states) + waveform = torch.tanh(self.conv_out(self.snake_out(hidden_states))) + return waveform.reshape(batch_size, 2, -1) diff --git a/src/diffusers/models/condition_embedders/__init__.py b/src/diffusers/models/condition_embedders/__init__.py index 3a92469a13ce..2c54cbcde7bc 100644 --- a/src/diffusers/models/condition_embedders/__init__.py +++ b/src/diffusers/models/condition_embedders/__init__.py @@ -3,3 +3,4 @@ if is_torch_available(): from .condition_embedder_anima import AnimaTextConditioner + from .condition_embedder_minimax_music3 import MiniMaxMusic3ConditionEncoder diff --git a/src/diffusers/models/condition_embedders/condition_embedder_minimax_music3.py b/src/diffusers/models/condition_embedders/condition_embedder_minimax_music3.py new file mode 100644 index 000000000000..2ee3b51fa0e7 --- /dev/null +++ b/src/diffusers/models/condition_embedders/condition_embedder_minimax_music3.py @@ -0,0 +1,76 @@ +# Copyright 2026 The MiniMax Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ..modeling_utils import ModelMixin + + +class MiniMaxMusic3ConditionEncoder(ModelMixin, ConfigMixin): + r""" + Projects the per-frame hidden states of the autoregressive stage onto the Flow-VAE latent timeline. + + Each generated frame carries `num_condition_layers` hidden states of size `condition_hidden_dim` (one from the + language model and one per residual codebook step). They are mixed with learned softmax weights, projected, and + resampled from the language-model frame rate to the latent frame rate with nearest-neighbor interpolation. + """ + + @register_to_config + def __init__( + self, + condition_hidden_dim: int = 4096, + num_condition_layers: int = 8, + out_dim: int = 2048, + input_sampling_rate: int = 24000, + input_hop_length: int = 960, + output_sampling_rate: int = 44100, + output_hop_length: int = 512, + ): + super().__init__() + self.layer_weight_logits = nn.Parameter(torch.zeros(num_condition_layers)) + self.layer_scale = nn.Parameter(torch.ones(1)) + self.proj = nn.Conv1d(condition_hidden_dim, out_dim, kernel_size=3, padding=1) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + r""" + Args: + hidden_states (`torch.Tensor` of shape `(batch, frames, num_condition_layers * condition_hidden_dim)`): + Concatenated per-frame hidden states from the autoregressive stage. + + Returns: + `torch.Tensor` of shape `(batch, latent_length, out_dim)`: the latent-aligned conditioning sequence. + """ + batch_size, num_frames, _ = hidden_states.shape + num_layers = self.config.num_condition_layers + hidden_states = hidden_states.transpose(1, 2) + hidden_states = hidden_states.reshape(batch_size, num_layers, self.config.condition_hidden_dim, num_frames) + layer_weights = torch.softmax(self.layer_weight_logits, dim=0).to(hidden_states.dtype) + hidden_states = torch.einsum("blht,l->bht", hidden_states, layer_weights) + hidden_states = self.layer_scale.to(hidden_states.dtype) * hidden_states + hidden_states = self.proj(hidden_states) + latent_length = max( + 1, + int( + num_frames + * self.config.output_sampling_rate + / self.config.input_sampling_rate + * self.config.input_hop_length + / self.config.output_hop_length + ), + ) + hidden_states = F.interpolate(hidden_states, size=latent_length, mode="nearest") + return hidden_states.transpose(1, 2) diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 2333acc06762..e9fcfcf320dc 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -11,6 +11,7 @@ from .hunyuan_transformer_2d import HunyuanDiT2DModel from .latte_transformer_3d import LatteTransformer3DModel from .lumina_nextdit2d import LuminaNextDiT2DModel + from .minimax_music3_rvq_depth_decoder import MiniMaxMusic3RVQDepthDecoder from .pixart_transformer_2d import PixArtTransformer2DModel from .prior_transformer import PriorTransformer from .sana_transformer import SanaTransformer2DModel @@ -51,6 +52,7 @@ from .transformer_ltx2 import LTX2VideoTransformer3DModel from .transformer_lumina2 import Lumina2Transformer2DModel from .transformer_minimax_h3 import MiniMaxH3Transformer3DModel + from .transformer_minimax_music3 import MiniMaxMusic3Transformer1DModel from .transformer_mochi import MochiTransformer3DModel from .transformer_motif_video import MotifVideoTransformer3DModel from .transformer_nucleusmoe_image import NucleusMoEImageTransformer2DModel diff --git a/src/diffusers/models/transformers/minimax_music3_rvq_depth_decoder.py b/src/diffusers/models/transformers/minimax_music3_rvq_depth_decoder.py new file mode 100644 index 000000000000..97bee84933c5 --- /dev/null +++ b/src/diffusers/models/transformers/minimax_music3_rvq_depth_decoder.py @@ -0,0 +1,142 @@ +# Copyright 2026 The MiniMax Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ..attention import AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn +from ..modeling_utils import ModelMixin +from ..normalization import RMSNorm + + +class MiniMaxMusic3DepthAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__(self, attn: "MiniMaxMusic3DepthAttention", hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, seq_len, _ = hidden_states.shape + + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + query = query.view(batch_size, seq_len, attn.heads, attn.head_dim) + key = key.view(batch_size, seq_len, attn.heads, attn.head_dim) + value = value.view(batch_size, seq_len, attn.heads, attn.head_dim) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + is_causal=True, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.flatten(2, 3).to(query.dtype) + return attn.to_out(hidden_states) + + +class MiniMaxMusic3DepthAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = MiniMaxMusic3DepthAttnProcessor + _available_processors = [MiniMaxMusic3DepthAttnProcessor] + + def __init__(self, dim: int, heads: int, processor: Optional[MiniMaxMusic3DepthAttnProcessor] = None): + super().__init__() + self.heads = heads + self.head_dim = dim // heads + self.to_q = nn.Linear(dim, dim, bias=False) + self.to_k = nn.Linear(dim, dim, bias=False) + self.to_v = nn.Linear(dim, dim, bias=False) + self.to_out = nn.Linear(dim, dim, bias=False) + if processor is None: + processor = self._default_processor_cls() + self.set_processor(processor) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.processor(self, hidden_states) + + +class MiniMaxMusic3DepthDecoderBlock(nn.Module): + def __init__(self, dim: int, heads: int, intermediate_size: int): + super().__init__() + self.input_layernorm = RMSNorm(dim, eps=1e-6, elementwise_affine=True) + self.attn = MiniMaxMusic3DepthAttention(dim, heads) + self.post_attention_layernorm = RMSNorm(dim, eps=1e-6, elementwise_affine=True) + self.gate_proj = nn.Linear(dim, intermediate_size, bias=False) + self.up_proj = nn.Linear(dim, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, dim, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states + self.attn(self.input_layernorm(hidden_states)) + norm_states = self.post_attention_layernorm(hidden_states) + return hidden_states + self.down_proj(F.silu(self.gate_proj(norm_states)) * self.up_proj(norm_states)) + + +class MiniMaxMusic3RVQDepthDecoder(ModelMixin, ConfigMixin): + r""" + The local language model of MiniMax Music 3. Within each audio frame it autoregressively predicts the seven + residual RVQ codebooks (c1..c7) from the global language model's hidden state and the frame's semantic code, and + exposes the per-step hidden states that condition the flow-matching transformer. + + It also owns the embedding table for the residual codebooks, which the pipeline uses to embed complete frames for + the global language model's feedback loop. + """ + + @register_to_config + def __init__( + self, + hidden_size: int = 4096, + num_layers: int = 4, + num_attention_heads: int = 16, + intermediate_size: int = 6144, + audio_vocab_size: int = 1024, + num_codebooks: int = 8, + max_position_embeddings: int = 16, + ): + super().__init__() + self.audio_embeddings = nn.Embedding(audio_vocab_size * (num_codebooks - 1), hidden_size) + self.projection = nn.Linear(hidden_size, hidden_size, bias=False) + self.pos_embedding = nn.Embedding(max_position_embeddings, hidden_size) + self.layers = nn.ModuleList( + [ + MiniMaxMusic3DepthDecoderBlock(hidden_size, num_attention_heads, intermediate_size) + for _ in range(num_layers) + ] + ) + self.norm = RMSNorm(hidden_size, eps=1e-6, elementwise_affine=True) + self.audio_heads = nn.ModuleList( + [nn.Linear(hidden_size, audio_vocab_size, bias=False) for _ in range(num_codebooks - 1)] + ) + + def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor: + r""" + Args: + inputs_embeds (`torch.Tensor` of shape `(batch, steps, hidden_size)`): + Projected depth-sequence embeddings: the global hidden state followed by the embedded codes sampled so + far, each passed through `projection`. + + Returns: + `torch.Tensor` of shape `(batch, steps, hidden_size)`: normalized hidden states; the last step feeds the + next codebook head. + """ + positions = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + hidden_states = inputs_embeds + self.pos_embedding(positions).unsqueeze(0) + for layer in self.layers: + hidden_states = layer(hidden_states) + return self.norm(hidden_states) diff --git a/src/diffusers/models/transformers/transformer_minimax_music3.py b/src/diffusers/models/transformers/transformer_minimax_music3.py new file mode 100644 index 000000000000..6ff174fe5e87 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_minimax_music3.py @@ -0,0 +1,242 @@ +# Copyright 2026 The MiniMax Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from typing import Optional, Tuple + +import torch +import torch.nn as nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils.torch_utils import lru_cache_unless_export +from ..attention import AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn +from ..embeddings import TimestepEmbedding +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin + + +class MiniMaxMusic3FourierEmbedding(nn.Module): + """Random Fourier features over the flow-matching time in `[0, 1]`. The projection is a trained checkpoint weight.""" + + def __init__(self, embedding_dim: int): + super().__init__() + self.weight = nn.Parameter(torch.randn(embedding_dim // 2, 1)) + + def forward(self, timestep: torch.Tensor) -> torch.Tensor: + angles = 2.0 * math.pi * timestep.unsqueeze(-1) @ self.weight.T + return torch.cat((angles.cos(), angles.sin()), dim=-1) + + +class MiniMaxMusic3RotaryEmbedding(nn.Module): + """Partial rotary embedding: only the first `rotary_dim` dimensions of each head rotate.""" + + def __init__(self, rotary_dim: int, theta: float = 10000.0): + super().__init__() + self.rotary_dim = rotary_dim + self.theta = theta + + @lru_cache_unless_export(maxsize=32) + def forward(self, seq_len: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]: + inv_freq = 1.0 / (self.theta ** (torch.arange(0, self.rotary_dim, 2, device=device).float() / self.rotary_dim)) + steps = torch.arange(seq_len, device=device, dtype=torch.float32) + freqs = torch.outer(steps, inv_freq) + freqs = torch.cat((freqs, freqs), dim=-1) + return freqs.cos().contiguous(), freqs.sin().contiguous() + + +def _apply_partial_rotary_emb( + hidden_states: torch.Tensor, rotary_emb: Tuple[torch.Tensor, torch.Tensor] +) -> torch.Tensor: + # hidden_states: [batch, seq, heads, head_dim]; only the leading rotary dims rotate. + cos, sin = rotary_emb + rotary_dim = cos.shape[-1] + cos = cos[:, None, :].to(hidden_states.dtype) + sin = sin[:, None, :].to(hidden_states.dtype) + rotated = hidden_states[..., :rotary_dim] + half_first, half_second = rotated.chunk(2, dim=-1) + rotate_half = torch.cat((-half_second, half_first), dim=-1) + rotated = rotated * cos + rotate_half * sin + return torch.cat((rotated, hidden_states[..., rotary_dim:]), dim=-1) + + +class MiniMaxMusic3AttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: "MiniMaxMusic3Attention", + hidden_states: torch.Tensor, + rotary_emb: Tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + batch_size, seq_len, _ = hidden_states.shape + + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + query = query.view(batch_size, seq_len, attn.heads, attn.head_dim) + key = key.view(batch_size, seq_len, attn.heads, attn.head_dim) + value = value.view(batch_size, seq_len, attn.heads, attn.head_dim) + + query = _apply_partial_rotary_emb(query, rotary_emb) + key = _apply_partial_rotary_emb(key, rotary_emb) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.flatten(2, 3).to(query.dtype) + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +class MiniMaxMusic3Attention(nn.Module, AttentionModuleMixin): + _default_processor_cls = MiniMaxMusic3AttnProcessor + _available_processors = [MiniMaxMusic3AttnProcessor] + + def __init__(self, dim: int, heads: int, head_dim: int, processor: Optional[MiniMaxMusic3AttnProcessor] = None): + super().__init__() + self.heads = heads + self.head_dim = head_dim + self.inner_dim = heads * head_dim + self.to_q = nn.Linear(dim, self.inner_dim, bias=False) + self.to_k = nn.Linear(dim, self.inner_dim, bias=False) + self.to_v = nn.Linear(dim, self.inner_dim, bias=False) + self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, dim, bias=False), nn.Dropout(0.0)]) + if processor is None: + processor = self._default_processor_cls() + self.set_processor(processor) + + def forward(self, hidden_states: torch.Tensor, rotary_emb: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: + return self.processor(self, hidden_states, rotary_emb=rotary_emb) + + +class MiniMaxMusic3TransformerBlock(nn.Module): + def __init__(self, dim: int, heads: int, head_dim: int, ff_inner_dim: int): + super().__init__() + self.norm1 = nn.LayerNorm(dim) + self.attn = MiniMaxMusic3Attention(dim, heads, head_dim) + self.norm2 = nn.LayerNorm(dim) + self.ff_in = nn.Linear(dim, ff_inner_dim * 2) + self.ff_out = nn.Linear(ff_inner_dim, dim) + + def forward(self, hidden_states: torch.Tensor, rotary_emb: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: + hidden_states = hidden_states + self.attn(self.norm1(hidden_states), rotary_emb) + gate_states, gate = self.ff_in(self.norm2(hidden_states)).chunk(2, dim=-1) + hidden_states = hidden_states + self.ff_out(gate_states * torch.nn.functional.silu(gate)) + return hidden_states + + +class MiniMaxMusic3Transformer1DModel(ModelMixin, ConfigMixin): + r""" + The flow-matching diffusion transformer of MiniMax Music 3. It denoises Flow-VAE audio latents conditioned on + per-frame hidden states produced by the autoregressive language-model stage. + + Inputs are 1D latent sequences of shape `(batch, in_channels, length)`. The conditioning signal + (`encoder_hidden_states`, shape `(batch, length, condition_dim)`) must already be aligned to the latent timeline — + see `MiniMaxMusic3ConditionEncoder`. The flow-matching `timestep` runs from 0 (noise) to 1 (data). + """ + + _supports_gradient_checkpointing = True + _no_split_modules = ["MiniMaxMusic3TransformerBlock"] + _repeated_blocks = ["MiniMaxMusic3TransformerBlock"] + _skip_layerwise_casting_patterns = ["time_proj", "norm"] + + @register_to_config + def __init__( + self, + in_channels: int = 128, + condition_dim: int = 2048, + num_layers: int = 36, + num_attention_heads: int = 32, + attention_head_dim: int = 64, + ff_inner_dim: int = 8192, + rotary_dim: int = 32, + fourier_embedding_dim: int = 256, + ): + super().__init__() + inner_dim = num_attention_heads * attention_head_dim + # The transformer input concatenates [latent, zeros(in_channels), condition] along channels. + concat_channels = 2 * in_channels + condition_dim + + self.time_proj = MiniMaxMusic3FourierEmbedding(fourier_embedding_dim) + self.time_embed = TimestepEmbedding(fourier_embedding_dim, inner_dim) + + self.preprocess_conv = nn.Conv1d(concat_channels, concat_channels, 1, bias=False) + self.proj_in = nn.Linear(concat_channels, inner_dim, bias=False) + self.rotary_emb = MiniMaxMusic3RotaryEmbedding(rotary_dim) + self.transformer_blocks = nn.ModuleList( + [ + MiniMaxMusic3TransformerBlock(inner_dim, num_attention_heads, attention_head_dim, ff_inner_dim) + for _ in range(num_layers) + ] + ) + self.proj_out = nn.Linear(inner_dim, in_channels, bias=False) + self.postprocess_conv = nn.Conv1d(in_channels, in_channels, 1, bias=False) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + return_dict: bool = True, + ) -> Tuple[torch.Tensor] | Transformer2DModelOutput: + r""" + Args: + hidden_states (`torch.Tensor` of shape `(batch, in_channels, length)`): + Noisy Flow-VAE latents. + timestep (`torch.Tensor` of shape `(batch,)`): + Flow-matching time in `[0, 1]`, where 0 is pure noise and 1 is data. + encoder_hidden_states (`torch.Tensor` of shape `(batch, length, condition_dim)`): + Frame-aligned conditioning from `MiniMaxMusic3ConditionEncoder`. Pass zeros for the unconditional + branch of classifier-free guidance. + return_dict (`bool`, defaults to `True`): + Whether to return a [`~models.modeling_outputs.Transformer2DModelOutput`] instead of a plain tuple. + + Returns: + The predicted flow-matching velocity with the same shape as `hidden_states`. + """ + zeros = torch.zeros_like(hidden_states) + hidden_states = torch.cat((hidden_states, zeros, encoder_hidden_states.transpose(1, 2)), dim=1) + hidden_states = self.preprocess_conv(hidden_states) + hidden_states + hidden_states = hidden_states.transpose(1, 2) + + temb = self.time_embed(self.time_proj(timestep)) + + hidden_states = self.proj_in(hidden_states) + # The timestep embedding is prepended as one extra token and removed after the blocks. + hidden_states = torch.cat((temb.unsqueeze(1), hidden_states), dim=1) + rotary_emb = self.rotary_emb(hidden_states.shape[1], hidden_states.device) + + for block in self.transformer_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(block, hidden_states, rotary_emb) + else: + hidden_states = block(hidden_states, rotary_emb) + + hidden_states = self.proj_out(hidden_states[:, 1:]) + hidden_states = hidden_states.transpose(1, 2) + hidden_states = self.postprocess_conv(hidden_states) + hidden_states + + if not return_dict: + return (hidden_states,) + return Transformer2DModelOutput(sample=hidden_states) diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 0dc5e0b17740..a8a23a39d517 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -137,6 +137,10 @@ "MiniMaxH3Blocks", "MiniMaxH3ModularPipeline", ] + _import_structure["minimax_music3"] = [ + "MiniMaxMusic3Blocks", + "MiniMaxMusic3ModularPipeline", + ] _import_structure["z_image"] = [ "ZImageAutoBlocks", "ZImageModularPipeline", @@ -196,6 +200,10 @@ MiniMaxH3Blocks, MiniMaxH3ModularPipeline, ) + from .minimax_music3 import ( + MiniMaxMusic3Blocks, + MiniMaxMusic3ModularPipeline, + ) from .modular_pipeline import ( AutoPipelineBlocks, BlockState, diff --git a/src/diffusers/modular_pipelines/minimax_music3/__init__.py b/src/diffusers/modular_pipelines/minimax_music3/__init__.py new file mode 100644 index 000000000000..bff3e432117b --- /dev/null +++ b/src/diffusers/modular_pipelines/minimax_music3/__init__.py @@ -0,0 +1,47 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["modular_blocks_minimax_music3"] = ["MiniMaxMusic3Blocks"] + _import_structure["modular_pipeline"] = ["MiniMaxMusic3ModularPipeline"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 + else: + from .modular_blocks_minimax_music3 import MiniMaxMusic3Blocks + from .modular_pipeline import MiniMaxMusic3ModularPipeline +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/modular_pipelines/minimax_music3/before_denoise.py b/src/diffusers/modular_pipelines/minimax_music3/before_denoise.py new file mode 100644 index 000000000000..be58527d681f --- /dev/null +++ b/src/diffusers/modular_pipelines/minimax_music3/before_denoise.py @@ -0,0 +1,73 @@ +# Copyright 2026 The MiniMax Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import InputParam, OutputParam +from .modular_pipeline import MiniMaxMusic3ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +# Hidden-state chunking: the autoregressive frames are decoded in 200-frame windows with a 100-frame hop; neighboring +# windows overlap by ~344 latent frames (~3.445 latents per frame), of which the trailing 86 latent frames +# (86 * 512 samples) are kept from the previous window when cropping the decoded waveform. +_CHUNK_FRAMES = 200 +_CHUNK_HOP = 100 + + +class MiniMaxMusic3PrepareChunksStep(ModularPipelineBlocks): + model_name = "minimax-music3" + + @property + def description(self) -> str: + return ( + "Chunk bookkeeping step that splits the autoregressive frames into 200-frame windows with a 100-frame " + "hop; each window is flow-matched with the previous window's trailing latents as an overlap prompt." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "frame_hiddens", + required=True, + type_hint=torch.Tensor, + description="Per-frame hidden states generated by the autoregressive step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "chunk_starts", + type_hint=list, + description="Frame index at which each 200-frame denoising window starts.", + ), + ] + + @torch.no_grad() + def __call__(self, components: MiniMaxMusic3ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + num_frames = block_state.frame_hiddens.shape[1] + block_state.chunk_starts = ( + [0] if num_frames <= _CHUNK_FRAMES else list(range(0, num_frames - _CHUNK_HOP, _CHUNK_HOP)) + ) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/minimax_music3/decoders.py b/src/diffusers/modular_pipelines/minimax_music3/decoders.py new file mode 100644 index 000000000000..2472542f1563 --- /dev/null +++ b/src/diffusers/modular_pipelines/minimax_music3/decoders.py @@ -0,0 +1,95 @@ +# Copyright 2026 The MiniMax Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import torch + +from ...models import MiniMaxMusic3Vocoder +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import MiniMaxMusic3ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +# Neighboring windows overlap by ~344 latent frames: when stitching the decoded waveforms, every window after the first +# drops its leading 86 latent frames and every window before the last drops its trailing 344 - 86 latent frames, so +# the kept spans tile the full song. +_CROP_LEFT_LATENT = 86 +_CROP_RIGHT_LATENT = 344 - 86 + + +class MiniMaxMusic3VocoderDecodeStep(ModularPipelineBlocks): + model_name = "minimax-music3" + + @property + def description(self) -> str: + return ( + "Decode step that vocodes each window's Flow-VAE latents into a waveform, crops the overlapping spans, " + "and stitches the windows into the final stereo waveform at 44.1 kHz." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("vocoder", MiniMaxMusic3Vocoder)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "latent_chunks", + required=True, + type_hint=list, + description="List of per-window denoised latent tensors (uncropped). Can be generated in denoise step.", + ), + InputParam("output_type", default="np", type_hint=str, description="Output format: 'np' or 'pt'."), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "audios", + type_hint=torch.Tensor | np.ndarray, + description="The generated stereo waveform of shape `(batch, channels, samples)` in `[-1, 1]`.", + ), + ] + + @staticmethod + def check_inputs(block_state): + if block_state.output_type not in ["np", "pt"]: + raise ValueError(f"Invalid output_type: {block_state.output_type}") + + @torch.no_grad() + def __call__(self, components: MiniMaxMusic3ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(block_state) + + hop_length = components.latent_hop_length + num_chunks = len(block_state.latent_chunks) + waveform_chunks = [] + for chunk_index, latents in enumerate(block_state.latent_chunks): + waveform = components.vocoder(latents.to(components.vocoder.dtype)) + left = 0 if chunk_index == 0 else _CROP_LEFT_LATENT * hop_length + right = 0 if chunk_index == num_chunks - 1 else _CROP_RIGHT_LATENT * hop_length + waveform_chunks.append(waveform[..., left : waveform.shape[-1] - right]) + + audios = torch.cat(waveform_chunks, dim=-1).float().clamp(-1.0, 1.0) + if block_state.output_type == "np": + audios = audios.cpu().numpy() + block_state.audios = audios + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/minimax_music3/denoise.py b/src/diffusers/modular_pipelines/minimax_music3/denoise.py new file mode 100644 index 000000000000..d65155af07c4 --- /dev/null +++ b/src/diffusers/modular_pipelines/minimax_music3/denoise.py @@ -0,0 +1,328 @@ +# Copyright 2026 The MiniMax Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import torch + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...models import MiniMaxMusic3ConditionEncoder, MiniMaxMusic3Transformer1DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import ( + BlockState, + LoopSequentialPipelineBlocks, + ModularPipelineBlocks, + PipelineState, +) +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .before_denoise import _CHUNK_FRAMES +from .modular_pipeline import MiniMaxMusic3ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +# Neighboring windows overlap by ~344 latent frames (100-frame hop at ~3.445 latents per frame); only the first 172 +# of them are spliced and blended, from the previous window's carry spanning latent frames [L - 344, L - 172). +_OVERLAP_LATENT_LENGTH = 172 + + +class MiniMaxMusic3ChunkConditionStep(ModularPipelineBlocks): + model_name = "minimax-music3" + + @property + def description(self) -> str: + return ( + "Chunk conditioning step that projects the window's per-frame hidden states onto the Flow-VAE latent " + "timeline and splices in the previous window's conditioning over the overlapping latent frames." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("condition_encoder", MiniMaxMusic3ConditionEncoder), + ComponentSpec("transformer", MiniMaxMusic3Transformer1DModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "frame_hiddens", + required=True, + type_hint=torch.Tensor, + description="Per-frame hidden states generated by the autoregressive step.", + ), + InputParam( + "chunk_starts", + required=True, + type_hint=list, + description="Frame index at which each 200-frame denoising window starts.", + ), + ] + + @torch.no_grad() + def __call__(self, components: MiniMaxMusic3ModularPipeline, block_state: BlockState, k: int): + device = components._execution_device + + chunk_start = block_state.chunk_starts[k] + chunk_end = min(chunk_start + _CHUNK_FRAMES, block_state.frame_hiddens.shape[1]) + condition = components.condition_encoder(block_state.frame_hiddens[:, chunk_start:chunk_end].to(device)) + condition = condition.to(components.transformer.dtype) + + overlap = 0 + if block_state.previous_latent is not None: + overlap = min(block_state.previous_latent.shape[-1], condition.shape[1]) + condition[:, :overlap] = block_state.previous_condition[:, :overlap] + + block_state.condition = condition + block_state.overlap = overlap + return components, block_state + + +class MiniMaxMusic3ChunkPrepareLatentsStep(ModularPipelineBlocks): + model_name = "minimax-music3" + + @property + def description(self) -> str: + return ( + "Chunk latent preparation step that draws the window's initial noise and snapshots the noise over the " + "overlapping latent frames as the blending prompt." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", MiniMaxMusic3Transformer1DModel)] + + @property + def inputs(self) -> list[InputParam]: + return [InputParam.template("generator")] + + @torch.no_grad() + def __call__(self, components: MiniMaxMusic3ModularPipeline, block_state: BlockState, k: int): + device = components._execution_device + + latents = randn_tensor( + (1, components.num_channels_latents, block_state.condition.shape[1]), + generator=block_state.generator, + device=device, + dtype=block_state.condition.dtype, + ) + block_state.noise_prompt = latents[..., : block_state.overlap].clone() if block_state.overlap > 0 else None + block_state.latents = latents + return components, block_state + + +class MiniMaxMusic3ChunkSetTimestepsStep(ModularPipelineBlocks): + model_name = "minimax-music3" + + @property + def description(self) -> str: + return "Chunk scheduler step that resets the flow-matching sigma schedule for the current window." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "num_inference_steps", + default=30, + type_hint=int, + description="Number of flow-matching Euler steps per chunk.", + ), + ] + + @torch.no_grad() + def __call__(self, components: MiniMaxMusic3ModularPipeline, block_state: BlockState, k: int): + device = components._execution_device + + sigmas = np.linspace(1.0, 1.0 / block_state.num_inference_steps, block_state.num_inference_steps) + components.scheduler.set_timesteps(sigmas=sigmas, device=device) + block_state.timesteps = components.scheduler.timesteps + return components, block_state + + +class MiniMaxMusic3ChunkDenoiseInner(ModularPipelineBlocks): + model_name = "minimax-music3" + + @property + def description(self) -> str: + return ( + "Inner denoising loop that flow-matches one window's latents over the scheduler timesteps. The guider " + "manages the conditional/unconditional transformer passes (the unconditional conditioning is all zeros), " + "and the overlapping latent frames are blended toward the previous window's trailing latents at every " + "step so neighboring windows share their boundary." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", MiniMaxMusic3Transformer1DModel), + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 1.7}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "num_inference_steps", + default=30, + type_hint=int, + description="Number of flow-matching Euler steps per chunk.", + ), + ] + + @torch.no_grad() + def __call__(self, components: MiniMaxMusic3ModularPipeline, block_state: BlockState, k: int): + latents = block_state.latents + timesteps = block_state.timesteps + overlap = block_state.overlap + + # The unconditional branch conditions on zeros, not on a re-encoded empty prompt. + guider_inputs = { + "encoder_hidden_states": (block_state.condition, torch.zeros_like(block_state.condition)), + } + + for i, t in enumerate(timesteps): + if overlap > 0: + time_value = t.to(latents.dtype) + latents[..., :overlap] = (1.0 - (1.0 - 1e-6) * time_value) * block_state.noise_prompt + ( + time_value * block_state.previous_latent[..., :overlap] + ) + # The transformer consumes the scheduler timestep directly: flow-matching time in [0, 1], 0 = noise. + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + guider_state = components.guider.prepare_inputs(guider_inputs) + + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + cond_kwargs = {key: getattr(guider_state_batch, key) for key in guider_inputs} + guider_state_batch.noise_pred = components.transformer( + hidden_states=latents, + timestep=timestep, + return_dict=False, + **cond_kwargs, + )[0] + components.guider.cleanup_models(components.transformer) + + velocity = components.guider(guider_state)[0] + latents = components.scheduler.step(velocity, t, latents, return_dict=False)[0] + block_state.progress_bar.update() + + block_state.latents = latents + return components, block_state + + +class MiniMaxMusic3ChunkUpdateStep(ModularPipelineBlocks): + model_name = "minimax-music3" + + @property + def description(self) -> str: + return ( + "Post-denoising update step that restores the previous window's latents over the overlap, appends the " + "window's latents to the chunk list, and carries the trailing latents and conditioning to the next window." + ) + + @torch.no_grad() + def __call__(self, components: MiniMaxMusic3ModularPipeline, block_state: BlockState, k: int): + latents = block_state.latents + if block_state.overlap > 0: + latents[..., : block_state.overlap] = block_state.previous_latent[..., : block_state.overlap] + + overlap_start = max(0, latents.shape[-1] - 2 * _OVERLAP_LATENT_LENGTH) + overlap_end = max(overlap_start, latents.shape[-1] - _OVERLAP_LATENT_LENGTH) + block_state.previous_latent = latents[..., overlap_start:overlap_end] + block_state.previous_condition = block_state.condition[:, overlap_start:overlap_end] + + block_state.latent_chunks.append(latents) + return components, block_state + + +class MiniMaxMusic3ChunkLoopWrapper(LoopSequentialPipelineBlocks): + model_name = "minimax-music3" + + @property + def description(self) -> str: + return ( + "Pipeline block that iterates over the 200-frame denoising windows. At each window it runs sub-blocks " + "for conditioning, latent preparation, scheduler reset, denoising, and the overlap carry update." + ) + + @property + def loop_inputs(self) -> list[InputParam]: + return [ + InputParam( + "chunk_starts", + required=True, + type_hint=list, + description="Frame index at which each 200-frame denoising window starts.", + ), + ] + + @property + def loop_intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "latent_chunks", + type_hint=list, + description="List of per-window denoised latent tensors (uncropped).", + ), + ] + + @torch.no_grad() + def __call__(self, components: MiniMaxMusic3ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + block_state.latent_chunks = [] + block_state.previous_latent = None + block_state.previous_condition = None + + num_chunks = len(block_state.chunk_starts) + with self.progress_bar(total=num_chunks * block_state.num_inference_steps) as progress_bar: + block_state.progress_bar = progress_bar + for k in range(num_chunks): + components, block_state = self.loop_step(components, block_state, k=k) + block_state.progress_bar = None + + self.set_block_state(state, block_state) + return components, state + + +class MiniMaxMusic3ChunkDenoiseStep(MiniMaxMusic3ChunkLoopWrapper): + block_classes = [ + MiniMaxMusic3ChunkConditionStep, + MiniMaxMusic3ChunkPrepareLatentsStep, + MiniMaxMusic3ChunkSetTimestepsStep, + MiniMaxMusic3ChunkDenoiseInner, + MiniMaxMusic3ChunkUpdateStep, + ] + block_names = ["condition", "prepare_latents", "set_timesteps", "denoise_inner", "update_chunk"] + + @property + def description(self) -> str: + return ( + "Chunk denoise step that iterates over the 200-frame denoising windows.\n" + "At each window: condition -> prepare_latents -> set_timesteps -> denoise_inner -> update_chunk." + ) diff --git a/src/diffusers/modular_pipelines/minimax_music3/encoders.py b/src/diffusers/modular_pipelines/minimax_music3/encoders.py new file mode 100644 index 000000000000..6fae35be3ffe --- /dev/null +++ b/src/diffusers/modular_pipelines/minimax_music3/encoders.py @@ -0,0 +1,361 @@ +# Copyright 2026 The MiniMax Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re +from typing import Optional + +import torch +import torch.nn.functional as F +from transformers import Qwen2Tokenizer, Qwen3ForCausalLM + +from ...hooks.group_offloading import _is_group_offload_enabled +from ...models import MiniMaxMusic3RVQDepthDecoder +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import MiniMaxMusic3ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +# The prompt template and its token ids are part of the checkpoint contract: even whitespace-level changes to the +# assembled prompt change the generated audio. +_IM_START, _IM_END = "<|im_start|>", "<|im_end|>" +_CAPTION_START, _CAPTION_END = "<|caption_start|>", "<|caption_end|>" +_LYRICS_START, _LYRICS_END = "<|lyrics_start|>", "<|lyrics_end|>" +_AUDIO_START = "<|audio_start|>" +_AUDIO_END_TOKEN_ID = 151670 +_AUDIO_CFG_TOKEN_ID = 151654 +_AUDIO_CODE_OFFSET = 151675 +_SEMANTIC_VOCAB_SIZE = 16384 +_MAX_PROMPT_TOKENS = 5_000 +_MAX_AUDIO_FRAMES = 9_000 + +# The autoregressive stage's sampling parameters are fixed by the reference inference recipe. +_AR_CFG_SCALE = 1.5 +_AR_CFG_TOP_K = 50 +_AR_SAMPLING_TOP_K = 50 + +_SPECIAL_TAG_RE = re.compile(r"<\|([^|]*)\|>") +_LEADING_TAGS_RE = re.compile(r"^[ \t]*((?:\[[^\]]+\][ \t]*)+)") + + +def _clean_caption(caption: str) -> str: + def _rewrite_special_tag(match: re.Match) -> str: + inner = match.group(1).strip() + parts = inner.split(None, 1) + return f"{parts[0]} is {parts[1]}" if len(parts) == 2 else inner + + text = _SPECIAL_TAG_RE.sub(_rewrite_special_tag, caption) + # Strip the markdown forms accepted by the model's input contract. + lines_out = [] + for line in text.splitlines(): + line = re.sub(r"^\s{0,3}#{1,6}\s+", "", line) + line = re.sub(r"^\s*[*+-]\s+", "", line) + line = re.sub(r"^\s*\*\s+", "", line) + while "**" in line: + updated = re.sub(r"\*\*([^*]+)\*\*", r"\1", line) + if updated == line: + break + line = updated + line = re.sub(r"(? str: + # Keep only consecutive structural tags (e.g. "[verse]") at the start of a line; text on a tag line is dropped. + output = [] + for line in lyrics.split("\n"): + match = _LEADING_TAGS_RE.match(line) + output.append(match.group(1).strip() if match else line) + text = "\n".join(output) + text = text.replace("] ", "]\n") + text = text.replace(" [", "\n[") + text = text.replace(" ^ ", "\n") + text = re.sub(r"\[([^\]]+)\]", lambda match: f"[{match.group(1).lower()}]", text) + return f"[start]\n{text}" + + +def _sample_top_k(logits: torch.Tensor, generator: Optional[torch.Generator]) -> torch.Tensor: + values = torch.nan_to_num(logits.float(), nan=-1e9, posinf=1e9, neginf=-1e9) + top_k = min(_AR_SAMPLING_TOP_K, values.shape[-1]) + threshold = torch.topk(values, top_k, dim=-1).values[..., -1, None] + values = values.masked_fill(values < threshold, -float("inf")) + probs = torch.nan_to_num(F.softmax(values, dim=-1), nan=0.0) + probs = probs / probs.sum(dim=-1, keepdim=True).clamp_min(1e-12) + # Sample on the generator's device so a CPU generator gives device-independent results (the diffusers convention). + sample_device = generator.device if generator is not None else probs.device + return torch.multinomial(probs.to(sample_device), 1, generator=generator).squeeze(-1).to(probs.device) + + +def _embed_audio_frame( + language_model: Qwen3ForCausalLM, + rvq_depth_decoder: MiniMaxMusic3RVQDepthDecoder, + frame_codes: torch.Tensor, +) -> torch.Tensor: + # frame_codes: [2, num_codebooks]. Sum the semantic-code embedding with the residual-code embeddings. + num_codebooks = rvq_depth_decoder.config.num_codebooks + embeds = language_model.model.embed_tokens(frame_codes[:, :1] + _AUDIO_CODE_OFFSET) + offsets = ( + torch.arange(num_codebooks - 1, device=frame_codes.device) * rvq_depth_decoder.config.audio_vocab_size + ).unsqueeze(0) + extra = rvq_depth_decoder.audio_embeddings(frame_codes[:, 1:] + offsets).sum(dim=1, keepdim=True) + embeds = embeds + extra.to(embeds.dtype) + return embeds * num_codebooks**-0.5 + + +def _generate_depth_codes( + language_model: Qwen3ForCausalLM, + rvq_depth_decoder: MiniMaxMusic3RVQDepthDecoder, + last_hidden: torch.Tensor, + semantic_code: torch.Tensor, + generator: Optional[torch.Generator], +): + # Autoregressively sample the residual codes c1..c7 for one frame and collect their hidden states. + num_codebooks = rvq_depth_decoder.config.num_codebooks + sequence = [rvq_depth_decoder.projection(last_hidden).unsqueeze(1)] + code_embed = language_model.model.embed_tokens(semantic_code + _AUDIO_CODE_OFFSET) + sequence.append(rvq_depth_decoder.projection(code_embed).unsqueeze(1)) + codes = [semantic_code] + hidden_parts = [] + for index in range(1, num_codebooks): + hidden = rvq_depth_decoder(torch.cat(sequence, dim=1))[:, -1] + hidden_parts.append(hidden[:1]) + logits = rvq_depth_decoder.audio_heads[index - 1](hidden) + conditional, unconditional = logits[:1].float(), logits[1:2].float() + logits = unconditional + (conditional - unconditional) * _AR_CFG_SCALE + # The sampled code is repeated so the language-model feedback keeps the [conditional, unconditional] rows. + code = _sample_top_k(logits, generator).repeat(2) + codes.append(code) + if index < num_codebooks - 1: + embed = rvq_depth_decoder.audio_embeddings(code + (index - 1) * rvq_depth_decoder.config.audio_vocab_size) + sequence.append(rvq_depth_decoder.projection(embed).unsqueeze(1)) + return torch.stack(codes, dim=1), torch.cat(hidden_parts, dim=-1) + + +class MiniMaxMusic3TokenizeStep(ModularPipelineBlocks): + model_name = "minimax-music3" + + @property + def description(self) -> str: + return ( + "Tokenize step that assembles the checkpoint's special-token prompt from the music description and the " + "lyrics and tokenizes it into the conditional/unconditional token id pair." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("tokenizer", Qwen2Tokenizer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "prompt", + required=True, + type_hint=str, + description="The music description (genre, mood, vocals, instrumentation, arrangement).", + ), + InputParam( + "lyrics", + required=True, + type_hint=str, + description=( + "The lyrics to sing. Structure tags such as `[verse]` or `[chorus]` must each be on their own " + "line; text on the same line as a leading tag is dropped by the checkpoint's input contract." + ), + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "text_ids", + type_hint=torch.Tensor, + description=( + "Token ids of shape `[2, sequence_length]` holding the conditional prompt and its classifier-free " + "counterpart (every token except the first and the two trailing structure tokens replaced by the " + "audio-CFG token)." + ), + ), + ] + + @staticmethod + def check_inputs(block_state): + if not isinstance(block_state.prompt, str) or not block_state.prompt.strip(): + raise ValueError( + f"`prompt` (the music description) must be a non-empty string, got {block_state.prompt!r}" + ) + if not isinstance(block_state.lyrics, str) or not block_state.lyrics.strip(): + raise ValueError(f"`lyrics` must be a non-empty string, got {block_state.lyrics!r}") + + @torch.no_grad() + def __call__(self, components: MiniMaxMusic3ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(block_state) + + text = ( + f"{_IM_START}{_CAPTION_START}{_clean_caption(block_state.prompt)}{_CAPTION_END}" + f"{_LYRICS_START}{_normalize_lyrics(block_state.lyrics)}{_LYRICS_END}{_IM_END}{_AUDIO_START}" + ) + input_ids = components.tokenizer(text, return_tensors="pt")["input_ids"] + if input_ids.shape[1] > _MAX_PROMPT_TOKENS: + raise ValueError( + f"The assembled prompt has {input_ids.shape[1]} tokens; the maximum is {_MAX_PROMPT_TOKENS}" + ) + unconditional_ids = input_ids.clone() + unconditional_ids[:, 1:-2] = _AUDIO_CFG_TOKEN_ID + block_state.text_ids = torch.cat((input_ids, unconditional_ids), dim=0).to(components._execution_device) + + self.set_block_state(state, block_state) + return components, state + + +class MiniMaxMusic3AutoregressiveStep(ModularPipelineBlocks): + model_name = "minimax-music3" + + @property + def description(self) -> str: + return ( + "Autoregressive generation step: frame by frame, the global language model samples a semantic code with " + "classifier-free guidance and the depth decoder samples the residual codes; the concatenated per-frame " + "hidden states condition the flow-matching stage." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("language_model", Qwen3ForCausalLM), + ComponentSpec("rvq_depth_decoder", MiniMaxMusic3RVQDepthDecoder), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "text_ids", + required=True, + type_hint=torch.Tensor, + description="Tokenized conditional/unconditional prompt pair generated by the text encoder step.", + ), + InputParam( + "audio_duration", + default=60.0, + type_hint=float, + description=( + "Upper bound on the generated audio length in seconds. The language model may stop earlier. " + "Capped at 9000 frames (six minutes)." + ), + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "frame_hiddens", + type_hint=torch.Tensor, + description=( + "Concatenated per-frame hidden states of shape `[1, frames, num_codebooks * hidden_size]` that " + "condition the flow-matching stage." + ), + ), + ] + + @staticmethod + def check_inputs(block_state): + if block_state.audio_duration <= 0: + raise ValueError(f"`audio_duration` must be positive, got {block_state.audio_duration}") + + @torch.no_grad() + def __call__(self, components: MiniMaxMusic3ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(block_state) + + text_ids = block_state.text_ids + max_frames = min(int(block_state.audio_duration * components.frame_rate), _MAX_AUDIO_FRAMES) + if max_frames == 0: + raise ValueError( + f"`audio_duration` {block_state.audio_duration} is shorter than one audio frame " + f"(1 / {components.frame_rate} s)" + ) + generator = block_state.generator + + language_model = components.language_model + rvq_depth_decoder = components.rvq_depth_decoder + # Trigger CPU-offload hooks by hand (same workaround as minimax_h3): the autoregressive loop calls + # submodules (`embed_tokens`, `lm_head`, depth-decoder heads) while the hook wraps only the top-level + # `forward`. The language model goes first — placing it can evict other models but never the reverse, + # and both models are used on every frame, so a placement must not evict the other. + hooked = [ + model for model in (language_model, rvq_depth_decoder) if getattr(model, "_hf_hook", None) is not None + ] + for model in hooked: + model._hf_hook.pre_forward(model) + resident = [model for model in hooked if not _is_group_offload_enabled(model)] + if len(resident) == 2 and resident[0].device != resident[1].device: + raise RuntimeError( + "The language model and the RVQ depth decoder must fit on the device together for autoregressive " + "generation; there is not enough free device memory under CPU offloading." + ) + text_embeds = language_model.model.embed_tokens(text_ids) + output = language_model.model(inputs_embeds=text_embeds, use_cache=True) + past_key_values = output.past_key_values + last_hidden = output.last_hidden_state[:, -1] + + vocab_mask = torch.ones(language_model.config.vocab_size, dtype=torch.bool, device=text_ids.device) + vocab_mask[_AUDIO_CODE_OFFSET : _AUDIO_CODE_OFFSET + _SEMANTIC_VOCAB_SIZE] = False + vocab_mask[_AUDIO_END_TOKEN_ID] = False + + frame_hiddens = [] + # The first decode step only advances the state past `<|audio_start|>` and is not an emitted frame. + for frame_index in range(max_frames + 1): + logits = language_model.lm_head(last_hidden).float() + logits = logits.masked_fill(vocab_mask, -float("inf")) + conditional, unconditional = logits[0:1], logits[1:2] + guided = unconditional + (conditional - unconditional) * _AR_CFG_SCALE + # Restrict the guided distribution to the conditional branch's top candidates, then re-mask: guidance on + # two `-inf` logits produces NaN on masked positions. + threshold = torch.topk(conditional, _AR_CFG_TOP_K, dim=-1).values[..., -1, None] + guided = guided.masked_fill(conditional < threshold, -float("inf")) + guided = guided.masked_fill(vocab_mask.unsqueeze(0), -float("inf")) + sampled = _sample_top_k(guided, generator) + if int(sampled.item()) == _AUDIO_END_TOKEN_ID: + break + + semantic_code = sampled - _AUDIO_CODE_OFFSET + frame_codes, depth_hidden = _generate_depth_codes( + language_model, rvq_depth_decoder, last_hidden, semantic_code.repeat(2), generator + ) + if frame_index > 0: + frame_hiddens.append(torch.cat((last_hidden[:1], depth_hidden), dim=-1)) + if len(frame_hiddens) >= max_frames: + break + feedback = _embed_audio_frame(language_model, rvq_depth_decoder, frame_codes) + output = language_model.model(inputs_embeds=feedback, past_key_values=past_key_values, use_cache=True) + past_key_values = output.past_key_values + last_hidden = output.last_hidden_state[:, -1] + + if not frame_hiddens: + raise ValueError("MiniMax Music 3 generated zero audio frames; the prompt ended generation immediately") + block_state.frame_hiddens = torch.stack(frame_hiddens, dim=1) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/minimax_music3/modular_blocks_minimax_music3.py b/src/diffusers/modular_pipelines/minimax_music3/modular_blocks_minimax_music3.py new file mode 100644 index 000000000000..2fee8c484ec6 --- /dev/null +++ b/src/diffusers/modular_pipelines/minimax_music3/modular_blocks_minimax_music3.py @@ -0,0 +1,188 @@ +# Copyright 2026 The MiniMax Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import torch + +from ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import MiniMaxMusic3PrepareChunksStep +from .decoders import MiniMaxMusic3VocoderDecodeStep +from .denoise import MiniMaxMusic3ChunkDenoiseStep +from .encoders import MiniMaxMusic3AutoregressiveStep, MiniMaxMusic3TokenizeStep + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +MiniMaxMusic3SemanticGeneratorBlocks = InsertableDict( + [ + ("tokenize", MiniMaxMusic3TokenizeStep()), + ("generate", MiniMaxMusic3AutoregressiveStep()), + ] +) + + +# auto_docstring +class MiniMaxMusic3SemanticGenerationStep(SequentialPipelineBlocks): + """ + Semantic generation step that assembles and tokenizes the checkpoint's special-token prompt from the music + description and the lyrics, then runs the autoregressive stage over it to produce the per-frame hidden states that + condition the flow-matching stage. + + Components: + tokenizer (`Qwen2Tokenizer`) language_model (`Qwen3ForCausalLM`) rvq_depth_decoder + (`MiniMaxMusic3RVQDepthDecoder`) + + Inputs: + prompt (`str`): + The music description (genre, mood, vocals, instrumentation, arrangement). + lyrics (`str`): + The lyrics to sing. Structure tags such as `[verse]` or `[chorus]` must each be on their own line; text + on the same line as a leading tag is dropped by the checkpoint's input contract. + audio_duration (`float`, *optional*, defaults to 60.0): + Upper bound on the generated audio length in seconds. The language model may stop earlier. Capped at 9000 + frames (six minutes). + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + text_ids (`Tensor`): + Token ids of shape `[2, sequence_length]` holding the conditional prompt and its classifier-free + counterpart (every token except the first and the two trailing structure tokens replaced by the audio-CFG + token). + frame_hiddens (`Tensor`): + Concatenated per-frame hidden states of shape `[1, frames, num_codebooks * hidden_size]` that condition + the flow-matching stage. + """ + + model_name = "minimax-music3" + block_classes = MiniMaxMusic3SemanticGeneratorBlocks.values() + block_names = MiniMaxMusic3SemanticGeneratorBlocks.keys() + + @property + def description(self) -> str: + return ( + "Semantic generation step that assembles and tokenizes the checkpoint's special-token prompt from the " + "music description and the lyrics, then runs the autoregressive stage over it to produce the per-frame " + "hidden states that condition the flow-matching stage." + ) + + +MiniMaxMusic3CoreDenoiseBlocks = InsertableDict( + [ + ("prepare_chunks", MiniMaxMusic3PrepareChunksStep()), + ("denoise", MiniMaxMusic3ChunkDenoiseStep()), + ] +) + + +# auto_docstring +class MiniMaxMusic3CoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that splits the autoregressive frames into 200-frame windows and flow-matches each window's + Flow-VAE latents from noise, blending every window into the previous one over their overlap. + + Components: + condition_encoder (`MiniMaxMusic3ConditionEncoder`) transformer (`MiniMaxMusic3Transformer1DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) + + Inputs: + frame_hiddens (`Tensor`): + Per-frame hidden states generated by the autoregressive step. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 30): + Number of flow-matching Euler steps per chunk. + + Outputs: + chunk_starts (`list`): + Frame index at which each 200-frame denoising window starts. + latent_chunks (`list`): + List of per-window denoised latent tensors (uncropped). + """ + + model_name = "minimax-music3" + block_classes = MiniMaxMusic3CoreDenoiseBlocks.values() + block_names = MiniMaxMusic3CoreDenoiseBlocks.keys() + + @property + def description(self) -> str: + return ( + "Core denoise step that splits the autoregressive frames into 200-frame windows and flow-matches each " + "window's Flow-VAE latents from noise, blending every window into the previous one over their overlap." + ) + + +# auto_docstring +class MiniMaxMusic3Blocks(SequentialPipelineBlocks): + """ + Modular pipeline for lyrics- and caption-conditioned music generation using MiniMax Music 3. An autoregressive + Qwen3 language model generates per-frame semantic codes and hidden states from the lyrics and the music + description; a flow-matching transformer turns the hidden states into Flow-VAE latents chunk by chunk; and a + DAC-style vocoder decodes them into a stereo waveform at 44.1 kHz. + + Components: + tokenizer (`Qwen2Tokenizer`) language_model (`Qwen3ForCausalLM`) rvq_depth_decoder + (`MiniMaxMusic3RVQDepthDecoder`) condition_encoder (`MiniMaxMusic3ConditionEncoder`) transformer + (`MiniMaxMusic3Transformer1DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) vocoder (`MiniMaxMusic3Vocoder`) + + Inputs: + prompt (`str`): + The music description (genre, mood, vocals, instrumentation, arrangement). + lyrics (`str`): + The lyrics to sing. Structure tags such as `[verse]` or `[chorus]` must each be on their own line; text + on the same line as a leading tag is dropped by the checkpoint's input contract. + audio_duration (`float`, *optional*, defaults to 60.0): + Upper bound on the generated audio length in seconds. The language model may stop earlier. Capped at 9000 + frames (six minutes). + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 30): + Number of flow-matching Euler steps per chunk. + output_type (`str`, *optional*, defaults to np): + Output format: 'np' or 'pt'. + + Outputs: + audios (`Tensor | ndarray`): + The generated stereo waveform of shape `(batch, channels, samples)` in `[-1, 1]`. + """ + + block_classes = [ + MiniMaxMusic3SemanticGenerationStep, + MiniMaxMusic3CoreDenoiseStep, + MiniMaxMusic3VocoderDecodeStep, + ] + block_names = ["semantic_generator", "denoise", "decode"] + + @property + def description(self) -> str: + return ( + "Modular pipeline for lyrics- and caption-conditioned music generation using MiniMax Music 3. " + "An autoregressive Qwen3 language model generates per-frame semantic codes and hidden states from the " + "lyrics and the music description; a flow-matching transformer turns the hidden states into Flow-VAE " + "latents chunk by chunk; and a DAC-style vocoder decodes them into a stereo waveform at 44.1 kHz." + ) + + @property + def outputs(self): + return [ + OutputParam( + "audios", + type_hint=torch.Tensor | np.ndarray, + description="The generated stereo waveform of shape `(batch, channels, samples)` in `[-1, 1]`.", + ), + ] diff --git a/src/diffusers/modular_pipelines/minimax_music3/modular_pipeline.py b/src/diffusers/modular_pipelines/minimax_music3/modular_pipeline.py new file mode 100644 index 000000000000..5afc03fb8959 --- /dev/null +++ b/src/diffusers/modular_pipelines/minimax_music3/modular_pipeline.py @@ -0,0 +1,60 @@ +# Copyright 2026 The MiniMax Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...utils import logging +from ..modular_pipeline import ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class MiniMaxMusic3ModularPipeline(ModularPipeline): + """ + A ModularPipeline for lyrics- and caption-conditioned music generation with MiniMax Music 3. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "MiniMaxMusic3Blocks" + + @property + def sampling_rate(self): + sampling_rate = 44100 + if hasattr(self, "vocoder") and self.vocoder is not None: + sampling_rate = int(self.vocoder.config.sampling_rate) + return sampling_rate + + @property + def frame_rate(self): + # Frames per second of the autoregressive stage (25 Hz for the released checkpoint). + frame_rate = 25.0 + if hasattr(self, "condition_encoder") and self.condition_encoder is not None: + config = self.condition_encoder.config + frame_rate = config.input_sampling_rate / config.input_hop_length + return frame_rate + + @property + def latent_hop_length(self): + # Waveform samples per Flow-VAE latent frame. + latent_hop_length = 512 + if hasattr(self, "condition_encoder") and self.condition_encoder is not None: + latent_hop_length = int(self.condition_encoder.config.output_hop_length) + return latent_hop_length + + @property + def num_channels_latents(self): + num_channels_latents = 128 + if hasattr(self, "transformer") and self.transformer is not None: + num_channels_latents = self.transformer.config.in_channels + return num_channels_latents diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 88cea1b78b1d..1ff8968abc9b 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -155,6 +155,7 @@ def _helios_pyramid_map_fn(config_dict=None): ("ltx2", _create_default_map_fn("LTX2ModularPipeline")), ("ltx2.5", _create_default_map_fn("LTX25ModularPipeline")), ("minimax-h3", _create_default_map_fn("MiniMaxH3ModularPipeline")), + ("minimax-music3", _create_default_map_fn("MiniMaxMusic3ModularPipeline")), ("ernie-image", _create_default_map_fn("ErnieImageModularPipeline")), ] ) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index b8c370bfd1e6..04bb0ceb7143 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -1740,6 +1740,66 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class MiniMaxMusic3ConditionEncoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class MiniMaxMusic3RVQDepthDecoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class MiniMaxMusic3Transformer1DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class MiniMaxMusic3Vocoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class MochiTransformer3DModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 11e69fcdccf1..ebe36d242253 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -602,6 +602,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class MiniMaxMusic3Blocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class MiniMaxMusic3ModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class QwenImageAutoBlocks(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/models/transformers/test_models_transformer_minimax_music3.py b/tests/models/transformers/test_models_transformer_minimax_music3.py new file mode 100644 index 000000000000..261df4351da8 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_minimax_music3.py @@ -0,0 +1,93 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from diffusers import MiniMaxMusic3Transformer1DModel +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, +) + + +enable_full_determinism() + + +class MiniMaxMusic3Transformer1DTesterConfig(BaseModelTesterConfig): + @property + def main_input_name(self) -> str: + return "hidden_states" + + @property + def model_class(self): + return MiniMaxMusic3Transformer1DModel + + @property + def output_shape(self) -> tuple[int, ...]: + return (8, 24) + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict[str, int]: + return { + "in_channels": 8, + "condition_dim": 16, + "num_layers": 2, + "num_attention_heads": 2, + "attention_head_dim": 8, + "ff_inner_dim": 32, + "rotary_dim": 4, + "fourier_embedding_dim": 8, + } + + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + batch_size = 1 + latent_length = 24 + in_channels = 8 + condition_dim = 16 + + return { + "hidden_states": randn_tensor( + (batch_size, in_channels, latent_length), generator=self.generator, device=torch_device + ), + "encoder_hidden_states": randn_tensor( + (batch_size, latent_length, condition_dim), generator=self.generator, device=torch_device + ), + "timestep": torch.full((batch_size,), 0.5, device=torch_device), + } + + +class TestMiniMaxMusic3Transformer1DModel(MiniMaxMusic3Transformer1DTesterConfig, ModelTesterMixin): + pass + + +class TestMiniMaxMusic3Transformer1DMemory(MiniMaxMusic3Transformer1DTesterConfig, MemoryTesterMixin): + pass + + +class TestMiniMaxMusic3Transformer1DTorchCompile(MiniMaxMusic3Transformer1DTesterConfig, TorchCompileTesterMixin): + pass + + +class TestMiniMaxMusic3Transformer1DAttention(MiniMaxMusic3Transformer1DTesterConfig, AttentionTesterMixin): + pass diff --git a/tests/modular_pipelines/minimax_music3/__init__.py b/tests/modular_pipelines/minimax_music3/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/modular_pipelines/minimax_music3/test_modular_pipeline_minimax_music3.py b/tests/modular_pipelines/minimax_music3/test_modular_pipeline_minimax_music3.py new file mode 100644 index 000000000000..0197a6d34809 --- /dev/null +++ b/tests/modular_pipelines/minimax_music3/test_modular_pipeline_minimax_music3.py @@ -0,0 +1,142 @@ +# Copyright 2026 The HuggingFace Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import pytest +import torch + +from diffusers import ( + MiniMaxMusic3Blocks, + MiniMaxMusic3ConditionEncoder, + MiniMaxMusic3ModularPipeline, +) + +from ...testing_utils import enable_full_determinism, torch_device +from ..test_modular_pipelines_common import ModularPipelineTesterMixin + + +enable_full_determinism() + + +class MiniMaxMusic3ConditionEncoderFastTests(unittest.TestCase): + def test_condition_encoder_output_shape(self): + condition_encoder = MiniMaxMusic3ConditionEncoder( + condition_hidden_dim=16, + num_condition_layers=8, + out_dim=16, + input_sampling_rate=24000, + input_hop_length=960, + output_sampling_rate=44100, + output_hop_length=512, + ) + hidden_states = torch.randn(1, 5, 8 * 16) + + output = condition_encoder(hidden_states) + + # 5 frames at 25 Hz resampled to the ~86.13 Hz latent rate. + expected_length = int(5 * 44100 / 24000 * 960 / 512) + self.assertEqual(output.shape, (1, expected_length, 16)) + + +class TestMiniMaxMusic3ModularPipelineFast(ModularPipelineTesterMixin): + _SAMPLING_SKIP = ( + "MiniMax Music 3's autoregressive stage samples discrete audio tokens; any numeric perturbation " + "(dtype, offload-induced kernel changes) legitimately flips sampled tokens, so value-equality " + "comparisons cannot hold. Shape and determinism-at-fixed-seed are covered by the other tests." + ) + + def test_float16_inference(self): + pytest.skip(self._SAMPLING_SKIP) + + def test_components_auto_cpu_offload_inference_consistent(self): + pytest.skip(self._SAMPLING_SKIP) + + def test_unload_components_auto_cpu_offload(self): + pytest.skip(self._SAMPLING_SKIP) + + def test_num_images_per_prompt(self): + pytest.skip("The pipeline generates a single waveform per call; there is no num_images_per_prompt input.") + + def test_save_from_pretrained(self, tmp_path): + # the common implementation indexes 4-D image outputs; compare the audio waveform directly + pipe = self.get_pipeline().to(torch_device) + pipe.save_pretrained(str(tmp_path)) + from diffusers import ModularPipeline + + reloaded = ModularPipeline.from_pretrained(str(tmp_path)) + reloaded.load_components() + reloaded.to(torch_device) + inputs = self.get_dummy_inputs() + audio = pipe(**inputs, output="audios") + inputs = self.get_dummy_inputs() + audio_reloaded = reloaded(**inputs, output="audios") + assert audio.shape == audio_reloaded.shape + assert torch.allclose(audio, audio_reloaded, atol=1e-4) + + pipeline_class = MiniMaxMusic3ModularPipeline + pipeline_blocks_class = MiniMaxMusic3Blocks + pretrained_model_name_or_path = "hf-internal-testing/tiny-minimax-music3-modular-pipe" + params = frozenset(["prompt", "lyrics", "audio_duration"]) + # The pipeline generates a single waveform per call; `prompt` and `lyrics` are single strings. + batch_params = frozenset() + optional_params = frozenset(["num_inference_steps", "output_type"]) + output_name = "audios" + expected_workflow_defaults = { + None: { + "components": { + "tokenizer": "Qwen2Tokenizer", + "language_model": "Qwen3ForCausalLM", + "rvq_depth_decoder": "MiniMaxMusic3RVQDepthDecoder", + "condition_encoder": "MiniMaxMusic3ConditionEncoder", + "transformer": "MiniMaxMusic3Transformer1DModel", + "scheduler": "FlowMatchEulerDiscreteScheduler", + "guider": "ClassifierFreeGuidance", + "vocoder": "MiniMaxMusic3Vocoder", + }, + "required_inputs": ["prompt", "lyrics"], + "inputs": { + "audio_duration": 60.0, + "generator": None, + "num_inference_steps": 30, + "output_type": "np", + }, + "component_configs": {"guider": {"guidance_scale": 1.7}}, + }, + } + + def get_dummy_inputs(self, seed=0): + return { + "prompt": "a bright synth pop song with warm female vocals", + "lyrics": "[verse]\nhello world\n[chorus]\nsing with me", + "audio_duration": 0.2, + "num_inference_steps": 2, + "generator": self.get_generator(seed), + "output_type": "pt", + } + + def test_inference_batch_consistent(self): + pytest.skip("MiniMax Music 3 generates a single waveform per call; `prompt` and `lyrics` are single strings.") + + def test_inference_batch_single_identical(self): + pytest.skip("MiniMax Music 3 generates a single waveform per call; `prompt` and `lyrics` are single strings.") + + def test_output_is_stereo_waveform(self): + pipe = self.get_pipeline() + + audio = pipe(**self.get_dummy_inputs(), output="audios") + + assert audio.shape[0] == 1 + assert audio.shape[1] == 2 + assert audio.abs().max() <= 1.0