Add v1-v4 RVQ encoder collection and reference-audio integrations

LICENSE-CODE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, Any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

README.md

@@ -0,0 +1,220 @@
+---
+license: other
+library_name: pytorch
+datasets:
+- bghira/minimax-music3-rvq-reverse-distillation
+tags:
+- audio
+- music
+- minimax-music-3
+- rvq
+- reference-audio
+- comfyui
+- diffusers
+---
+
+# Open RVQ Encoders for MiniMax Music 3
+
+## What this release does
+
+MiniMax Music 3 normally creates its internal music codes from text and lyrics. The original audio-to-code encoder was not released. These models add an approximate path from a reference recording back to those codes.
+
+The useful result is **0.8748 mean condition-replay cosine** on 130 held-out generated tracks. In plain English: after a reference track is reduced to eight integers per 25 Hz frame and replayed through the official language-model path, v4 recovers most of the conditioning signal that the diffusion model would have received from the original codes.
+
+This is not 87% exact-token accuracy. It is not an audio-quality score. The RVQ code space is redundant: different code tuples can lead to nearly the same downstream condition. Exact token labels therefore understate functional accuracy.
+
+The project started with a 41M-parameter, single-GPU community proof at **0.6633** replay cosine. The SimpleTuner experiments moved through width and music-feature alignment, then reached **0.8748** by making the acoustic codebooks autoregressive across RVQ depth.
+
+Use **v4** unless reproducing an experiment.
+
+## Files
+
+| File | Parameters | Experiment | Replay cosine |
+|---|---:|---|---:|
+| `minimax_music3_rvq_encoder_v1_41m_independent_heads.safetensors` | 40,978,944 | Baseline; eight independent heads | 0.7624* |
+| `minimax_music3_rvq_encoder_v2_155m_wide_independent_heads.safetensors` | 154,736,064 | Wider shared encoder | 0.7698 |
+| `minimax_music3_rvq_encoder_v3_155m_mert_aligned_independent_heads.safetensors` | 154,736,064 | v2 plus training-only MERT alignment | 0.7703 |
+| `minimax_music3_rvq_encoder_v4_169m_autoregressive_depth_recommended.safetensors` | 169,008,576 | Causal acoustic decoder across codebook depth | **0.8748** |
+
+Each weight file has a same-named `.json` configuration file in [`encoders/`](encoders/).
+
+\* The v1 replay evaluation used the final checkpoint. This collection ships the v1 card-recommended step-17,500 checkpoint. V2-v4 replay scores and shipped checkpoints are final checkpoints.
+
+![Condition replay comparison](assets/condition-embedding-replay-comparison.png)
+
+Metric: per-track mean cosine between replayed and stored condition embeddings. Evaluation set: 130 exact-alignment records from the held-out split. True sampled codes score 0.9999 as the pipeline control. Release provenance is in [`experiment-summary.json`](experiment-summary.json). Full summary statistics and quantiles are in [`evaluation/condition-replay-aggregate.json`](evaluation/condition-replay-aggregate.json).
+
+## Experiment arc
+
+### Community proof
+
+`Serveurperso` trained a 41M encoder on one GPU and demonstrated that exact token agreement was not required. Their checkpoint reached 0.6633 condition-replay cosine. Their implementation independently identified the 345-latent stitched hop, the extra warm-up code row, and the final partial-chunk interpolation rule. The SimpleTuner dataset and trainer work independently derived the same alignment contract.
+
+That result changed the question from "can this work?" to "which errors matter downstream?"
+
+### v1: working baseline
+
+- 41M parameters.
+- DAV latents as input.
+- One semantic head and seven independent acoustic heads.
+- Teacher top-50 KL plus hard cross-entropy.
+- μP initialization and μTransfer shape metadata.
+- DDP training on the reverse-distillation corpus.
+- 0.7624 replay cosine at the evaluated final checkpoint.
+
+Detailed card: [`SimpleTuner/open-rvq-encoder-minimax-music3-41m-v1`](https://huggingface.co/SimpleTuner/open-rvq-encoder-minimax-music3-41m-v1).
+
+### v2: width
+
+- Width increased from 512 to 1,088.
+- Parameters increased from 41M to 155M.
+- Polynomial decay with warm-up.
+- Replay improved from 0.7624 to 0.7698.
+
+Width helped token metrics. It barely moved the downstream replay metric. Capacity was not the main limit.
+
+Detailed card: [`SimpleTuner/open-rvq-encoder-minimax-music3-155m-v2`](https://huggingface.co/SimpleTuner/open-rvq-encoder-minimax-music3-155m-v2).
+
+### v3: MERT alignment
+
+- Same exported architecture as v2.
+- Frozen MERT features used only as a training target.
+- Mid-layer cosine alignment loss annealed to zero.
+- No MERT weights or projection are present in the exported encoder.
+- Replay improved from 0.7698 to 0.7703.
+
+The measured gain was too small to justify an inference dependency or a larger MERT experiment on this corpus.
+
+Detailed card: [`SimpleTuner/open-rvq-encoder-minimax-music3-155m-v3`](https://huggingface.co/SimpleTuner/open-rvq-encoder-minimax-music3-155m-v3).
+
+### v4: autoregressive depth
+
+- Shared encoder kept at v2 width.
+- The seven acoustic heads were replaced by a causal decoder.
+- Codebook `k` is conditioned on the semantic code and acoustic codebooks `< k`.
+- Parameters increased to 169M.
+- Replay improved from 0.7703 to 0.8748.
+
+Free-running exact acoustic top-1 decreased slightly. Teacher-forced acoustic top-1 increased to 18.42%. Replay improved by 0.1045. This is the central result: code-tuple compatibility matters more than matching one sampled token tuple.
+
+Detailed card: [`SimpleTuner/open-rvq-encoder-minimax-music3-169m-v4`](https://huggingface.co/SimpleTuner/open-rvq-encoder-minimax-music3-169m-v4).
+
+## ComfyUI
+
+Requirements:
+
+- current ComfyUI with native MiniMax Music 3 support;
+- official ComfyUI MiniMax diffusion model, text encoder, and DAV decoder;
+- original `dav.pth`, because the Comfy DAV file is decoder-only;
+- this repository cloned with Git LFS.
+
+```bash
+git lfs install
+git clone https://huggingface.co/SimpleTuner/open-rvq-encoder-minimax-music3
+
+cd ComfyUI
+python3 -m venv .venv
+.venv/bin/pip install -r requirements.txt huggingface_hub
+
+.venv/bin/hf download Comfy-Org/MiniMax-Music-3 \
+  diffusion_models/minimax_music3_dit_fp16.safetensors \
+  text_encoders/minimax_music3_text_encoder_pruned_int8_convrot.safetensors \
+  vae/minimax_music3_dav.safetensors \
+  --local-dir models
+
+.venv/bin/hf download MiniMaxAI/MiniMax-Music3 dav.pth \
+  --local-dir models/vae
+
+ln -s /path/to/open-rvq-encoder-minimax-music3/comfyui_open_rvq \
+  custom_nodes/comfyui_open_rvq
+```
+
+Restart ComfyUI. Load [`comfyui_workflow_example.json`](comfyui_workflow_example.json). Upload a reference audio file. Select v4 in **MiniMax Music3 RVQ Reference Encoder Loader**.
+
+The node package reads the encoder files directly from this clone. They can instead be placed in:
+
+```text
+ComfyUI/models/minimax_music3_rvq_encoders/
+```
+
+The tested graph used ComfyUI 0.33.0, one NVIDIA L40S, the pruned int8 text encoder, the fp16 diffusion model, a 32.28-second held-out reference, five Euler steps, and the v4 encoder. It completed and produced a full-length non-silent stereo FLAC. Use 30 steps for normal output.
+
+## Diffusers
+
+The MiniMax Music 3 Diffusers integration is currently installed from its integration commit:
+
+```bash
+python3 -m venv .venv
+source .venv/bin/activate
+pip install \
+  git+https://github.com/huggingface/diffusers@dafe3733fcfdbf3c48915fe77be3aef65b5d6a2d \
+  transformers accelerate safetensors huggingface_hub soundfile torchaudio
+```
+
+Download [`minimax_music3_reference_adapter.py`](minimax_music3_reference_adapter.py) into the project, then:
+
+```python
+import soundfile as sf
+import torch
+from diffusers import ModularPipeline
+
+from minimax_music3_reference_adapter import (
+    MiniMaxMusic3ReferenceAdapter,
+    install_diffusers_reference_adapter,
+)
+
+install_diffusers_reference_adapter()
+
+pipe = ModularPipeline.from_pretrained(
+    "MiniMaxAI/MiniMax-Music3",
+    torch_dtype=torch.bfloat16,
+).to("cuda")
+adapter = MiniMaxMusic3ReferenceAdapter.from_pretrained()
+
+audio, sample_rate = sf.read("reference.flac", always_2d=True)
+waveform = torch.from_numpy(audio.T).float()
+generator = torch.Generator(device="cpu").manual_seed(42)
+
+frame_hiddens, predicted_codes = adapter.encode_reference(
+    pipe,
+    waveform,
+    sample_rate,
+    prompt="rock",
+    lyrics="[instrumental]",
+    generator=generator,
+    device="cuda",
+)
+
+result = pipe(
+    frame_hiddens=frame_hiddens,
+    generator=torch.Generator(device="cpu").manual_seed(42),
+    num_inference_steps=30,
+    output_type="pt",
+)
+```
+
+The patch only adds a precomputed-`frame_hiddens` bypass to the modular pipeline. It does not replace MiniMax model code.
+
+## Limits
+
+- These are approximate encoders, not the original MiniMax encoder.
+- No original encoder weights or source were used.
+- Training data is synthetic MiniMax Music 3 output, not MiniMax's training set.
+- Real-audio generalization is not established.
+- Context is 128 frames, or 5.12 seconds. There is no cross-window encoder state.
+- Reference replay still needs the official MiniMax language model and RVQ depth decoder.
+- v4 uses greedy code selection. Other decoding strategies remain untested.
+
+## Credits
+
+- [`bghira`](https://huggingface.co/bghira): reverse-distillation dataset, independent alignment derivation, SimpleTuner trainer, compute, v1-v4 runs, and replay experiments.
+- [`marduk191`](https://huggingface.co/marduk191): WAV/code samples, early mel-encoder proof, corpus-scale observations, and encoder experiments.
+- [`scragnog`](https://huggingface.co/scragnog): SimpleTuner calibration, relative-weight analysis, caption-cache and rollout-seam findings, and GGML interoperability tests.
+- [`Serveurperso`](https://huggingface.co/Serveurperso): independent 41M encoder, single-GPU proof, replay evaluation, and independent stitched-timeline findings.
+- [`dernet`](https://huggingface.co/dernet): RVQ supervision, internal-alignment, and tokenizer reverse-engineering analysis.
+
+## License
+
+`minimax_music3_reference_adapter.py` and `comfyui_open_rvq/` are Apache-2.0 licensed. See [`LICENSE-CODE`](LICENSE-CODE).
+
+The encoder weights were trained on outputs from MiniMax Music 3. Use of the weights, official model components, and generated output remains subject to the applicable MiniMax Music 3 terms. The MERT license also applies to the v3 training provenance; no MERT weights are distributed here.

comfyui_open_rvq/__init__.py

@@ -0,0 +1,6 @@
+# Copyright 2026 SimpleTuner contributors
+# Licensed under the Apache License, Version 2.0
+
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

comfyui_open_rvq/nodes.py

@@ -0,0 +1,248 @@
+# Copyright 2026 SimpleTuner contributors
+#
+# 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
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import torch
+
+import comfy.model_management
+import comfy.model_prefetch
+import comfy.ops
+import folder_paths
+from comfy.ldm.minimax_music.ar import AUDIO_FRAMES_PER_SECOND, CFG_SCALE, CFG_TOP_K, MiniMaxMusic3AR
+from comfy.text_encoders.minimax_music import MiniMaxMusic3TEModel
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from minimax_music3_reference_adapter import MiniMaxMusic3ReferenceAdapter  # noqa: E402
+
+
+ENCODER_FOLDER = "minimax_music3_rvq_encoders"
+folder_paths.add_model_folder_path(ENCODER_FOLDER, str(Path(folder_paths.models_dir) / ENCODER_FOLDER), is_default=True)
+bundled_encoders = REPO_ROOT / "encoders"
+if bundled_encoders.is_dir():
+    folder_paths.add_model_folder_path(ENCODER_FOLDER, str(bundled_encoders))
+
+
+def _teacher_depth_hidden(model: MiniMaxMusic3AR, hidden, codes, execution_dtype):
+    decoder = model.model.audio_decoder
+    sequence = [decoder.projection(hidden).unsqueeze(1)]
+    semantic = model._embed_c0(codes[:, 0], execution_dtype)
+    sequence.append(decoder.projection(semantic).unsqueeze(1))
+    hidden_parts = []
+    for index in range(1, model.num_codebooks):
+        depth_hidden = decoder(torch.cat(sequence, dim=1))[:, -1]
+        hidden_parts.append(depth_hidden[:1])
+        if index < model.num_codebooks - 1:
+            embedding = model.model.audio_extra_embedding(
+                codes[:, index] + (index - 1) * model.audio_vocab_size,
+                out_dtype=execution_dtype,
+            )
+            sequence.append(decoder.projection(embedding).unsqueeze(1))
+    return torch.cat(hidden_parts, dim=-1)
+
+
+def _run_depth(model: MiniMaxMusic3AR, device, execution_dtype, core):
+    decoder = model.model.audio_decoder
+    queue = comfy.model_prefetch.make_prefetch_queue(
+        [[decoder, model.model.audio_extra_embedding]],
+        device,
+        {"prefetch_dynamic_vbars": True},
+    )
+    comfy.model_prefetch.prefetch_queue_pop(
+        queue,
+        device,
+        decoder,
+        execution_dtype,
+        core=core,
+        enable_graph=False,
+    )
+    comfy.model_prefetch.prefetch_queue_pop(queue, device, None)
+
+
+@torch.inference_mode()
+def replay_reference_codes(model: MiniMaxMusic3AR, input_ids, codes, seed, device, cfg_scale, top_k):
+    if codes.ndim != 2 or codes.shape[1] != model.num_codebooks or codes.shape[0] == 0:
+        raise ValueError(f"reference codes must have shape [frames, {model.num_codebooks}]")
+    prompt_tokens = int(input_ids.shape[1])
+    input_ids = input_ids.to(device)
+    codes = codes.to(device=device, dtype=torch.long)
+    execution_dtype = torch.bfloat16 if comfy.model_management.should_use_bf16(device) else torch.float32
+    unconditioned = input_ids.clone()
+    from comfy.ldm.minimax_music.prompt import AUDIO_CODE_OFFSET, SPECIAL_TOKEN_IDS
+
+    unconditioned[:, 1:-2] = SPECIAL_TOKEN_IDS["<|audio_cfg|>"]
+    text_ids = torch.cat((input_ids, unconditioned), dim=0)
+    if model.model.pruned_embedding:
+        text_embeds = model.model.embed_tokens_prefill(text_ids, out_dtype=execution_dtype)
+    else:
+        text_embeds = model.model.embed_tokens(text_ids, out_dtype=execution_dtype)
+    past = model.model.init_kv_cache(2, prompt_tokens + codes.shape[0] + 2, device, execution_dtype)
+    output = model.model(None, embeds=text_embeds, past_key_values=past, dtype=execution_dtype)
+    last_hidden, past = output[0][:, -1], output[2]
+    from comfy.ldm.minimax_music.ar import derive_seed
+
+    generator = torch.Generator(device=device).manual_seed(derive_seed(seed, "ar"))
+    vocab_mask = None
+    if not model.model.pruned_lm_head:
+        vocab_mask = torch.ones(model.model.vocab_size, dtype=torch.bool, device=device)
+        vocab_mask[AUDIO_CODE_OFFSET : AUDIO_CODE_OFFSET + 16384] = False
+        vocab_mask[SPECIAL_TOKEN_IDS["<|audio_end|>"]] = False
+
+    warmup_c0, warmup_or_stop, stop_token = model._sample_c0(last_hidden, cfg_scale, top_k, generator, vocab_mask)
+    if int(warmup_or_stop.item()) == stop_token:
+        raise ValueError("the selected seed ended during the required AR warm-up frame")
+    warmup_c0 = warmup_c0.repeat(2)
+    warmup_embed = model._embed_c0(warmup_c0, execution_dtype)
+    warmup: dict[str, torch.Tensor] = {}
+
+    def warmup_core():
+        warmup["codes"], _ = model._depth_codes(
+            last_hidden,
+            warmup_c0,
+            warmup_embed,
+            generator,
+            execution_dtype,
+            cfg_scale,
+            top_k,
+        )
+
+    _run_depth(model, device, execution_dtype, warmup_core)
+    feedback = model._embed_audio_frame(warmup["codes"], execution_dtype)
+    output = model.model(None, embeds=feedback, past_key_values=past, dtype=execution_dtype)
+    last_hidden, past = output[0][:, -1], output[2]
+
+    frames = []
+    for frame_index in range(codes.shape[0]):
+        comfy.model_management.throw_exception_if_processing_interrupted()
+        frame_codes = codes[frame_index].unsqueeze(0).repeat(2, 1)
+        result: dict[str, torch.Tensor] = {}
+
+        def teacher_core():
+            result["hidden"] = _teacher_depth_hidden(model, last_hidden, frame_codes, execution_dtype)
+
+        _run_depth(model, device, execution_dtype, teacher_core)
+        frames.append(torch.cat((last_hidden[:1], result["hidden"]), dim=-1)[0].cpu())
+        if frame_index + 1 < codes.shape[0]:
+            feedback = model._embed_audio_frame(frame_codes, execution_dtype)
+            output = model.model(None, embeds=feedback, past_key_values=past, dtype=execution_dtype)
+            last_hidden, past = output[0][:, -1], output[2]
+    return torch.stack(frames)
+
+
+_ORIGINAL_ENCODE_TOKEN_WEIGHTS = MiniMaxMusic3TEModel.encode_token_weights
+
+
+def _encode_token_weights_with_reference(self, token_weight_pairs):
+    codes = token_weight_pairs.get("minimax_reference_codes")
+    if codes is None:
+        return _ORIGINAL_ENCODE_TOKEN_WEIGHTS(self, token_weight_pairs)
+    token_ids = [token for token, _ in token_weight_pairs["minimax_music3"][0]]
+    input_ids = torch.tensor([token_ids], dtype=torch.long)
+    hidden = replay_reference_codes(
+        self,
+        input_ids,
+        codes,
+        int(token_weight_pairs["seed"]),
+        self.execution_device,
+        float(token_weight_pairs["cfg_scale"]),
+        int(token_weight_pairs["top_k"]),
+    )
+    return hidden.unsqueeze(0), None, {}
+
+
+if not getattr(MiniMaxMusic3TEModel, "_simpletuner_reference_patch", False):
+    MiniMaxMusic3TEModel.encode_token_weights = _encode_token_weights_with_reference
+    MiniMaxMusic3TEModel._simpletuner_reference_patch = True
+
+
+class MiniMaxMusic3RVQReferenceEncoderLoader:
+    @classmethod
+    def INPUT_TYPES(cls):
+        encoders = [name for name in folder_paths.get_filename_list(ENCODER_FOLDER) if name.lower().endswith(".safetensors")]
+        dav_files = [name for name in folder_paths.get_filename_list("vae") if name.lower().endswith((".pth", ".pt"))]
+        return {"required": {"encoder": (encoders,), "dav": (dav_files,)}}
+
+    RETURN_TYPES = ("MINIMAX_MUSIC3_RVQ_REFERENCE_ENCODER",)
+    RETURN_NAMES = ("reference_encoder",)
+    FUNCTION = "load"
+    CATEGORY = "loaders/minimax music"
+
+    def load(self, encoder, dav):
+        encoder_path = Path(folder_paths.get_full_path_or_raise(ENCODER_FOLDER, encoder))
+        config_path = encoder_path.with_suffix(".json")
+        if not config_path.is_file():
+            raise FileNotFoundError(f"Missing RVQ encoder config next to weights: {config_path.name}")
+        dav_path = folder_paths.get_full_path_or_raise("vae", dav)
+        return (MiniMaxMusic3ReferenceAdapter.from_files(encoder_path, config_path, dav_path),)
+
+
+class MiniMaxMusic3ReferenceAudioEncode:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "clip": ("CLIP",),
+                "reference_encoder": ("MINIMAX_MUSIC3_RVQ_REFERENCE_ENCODER",),
+                "audio": ("AUDIO",),
+                "caption": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+                "lyrics": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
+            },
+            "optional": {
+                "cfg_scale": ("FLOAT", {"default": CFG_SCALE, "min": 0.0, "max": 100.0, "step": 0.1}),
+                "top_k": ("INT", {"default": CFG_TOP_K, "min": 1, "max": 16384}),
+            },
+        }
+
+    RETURN_TYPES = ("CONDITIONING", "FLOAT")
+    RETURN_NAMES = ("conditioning", "seconds")
+    FUNCTION = "encode"
+    CATEGORY = "conditioning/minimax music"
+
+    def encode(self, clip, reference_encoder, audio, caption, lyrics, seed, cfg_scale=CFG_SCALE, top_k=CFG_TOP_K):
+        codes = reference_encoder.predict_codes(
+            audio["waveform"],
+            int(audio["sample_rate"]),
+            device=comfy.model_management.get_torch_device(),
+        )
+        tokens = clip.tokenize(
+            caption,
+            lyrics=lyrics,
+            seed=seed,
+            max_audio_frames=codes.shape[0],
+            cfg_scale=cfg_scale,
+            top_k=top_k,
+        )
+        tokens["minimax_reference_codes"] = codes
+        conditioning = clip.encode_from_tokens_scheduled(tokens)
+        for cond in conditioning:
+            hidden = cond[0]
+            cond[1]["conditioning_scale"] = torch.ones(
+                (hidden.shape[0], 1, 1),
+                device=hidden.device,
+                dtype=hidden.dtype,
+            )
+        return conditioning, codes.shape[0] / AUDIO_FRAMES_PER_SECOND
+
+
+NODE_CLASS_MAPPINGS = {
+    "MiniMaxMusic3RVQReferenceEncoderLoader": MiniMaxMusic3RVQReferenceEncoderLoader,
+    "MiniMaxMusic3ReferenceAudioEncode": MiniMaxMusic3ReferenceAudioEncode,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "MiniMaxMusic3RVQReferenceEncoderLoader": "MiniMax Music3 RVQ Reference Encoder Loader",
+    "MiniMaxMusic3ReferenceAudioEncode": "MiniMax Music3 Reference Audio Encode",
+}

comfyui_workflow_example.json

@@ -0,0 +1,700 @@
+{
+  "id": "f9df4058-b0bd-409c-a2e7-c41c8d68d56f",
+  "revision": 0,
+  "last_node_id": 46,
+  "last_link_id": 12,
+  "nodes": [
+    {
+      "id": 6,
+      "type": "UNETLoader",
+      "pos": [
+        40,
+        80
+      ],
+      "size": [
+        480,
+        90
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "showAdvanced": false,
+      "inputs": [
+        {
+          "localized_name": "unet_name",
+          "name": "unet_name",
+          "type": "COMBO",
+          "widget": {
+            "name": "unet_name"
+          }
+        }
+      ],
+      "outputs": [
+        {
+          "localized_name": "MODEL",
+          "name": "MODEL",
+          "type": "MODEL",
+          "links": [
+            9
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.31.0",
+        "Node name for S&R": "UNETLoader",
+        "models": [
+          {
+            "name": "minimax_music3_dit_fp16.safetensors",
+            "url": "https://huggingface.co/Comfy-Org/MiniMax-Music-3/resolve/main/diffusion_models/minimax_music3_dit_fp16.safetensors",
+            "directory": "diffusion_models"
+          }
+        ]
+      },
+      "widgets_values": [
+        "minimax_music3_dit_fp16.safetensors",
+        "default"
+      ]
+    },
+    {
+      "id": 3,
+      "type": "CLIPLoader",
+      "pos": [
+        40,
+        250
+      ],
+      "size": [
+        480,
+        120
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        {
+          "localized_name": "clip_name",
+          "name": "clip_name",
+          "type": "COMBO",
+          "widget": {
+            "name": "clip_name"
+          }
+        }
+      ],
+      "outputs": [
+        {
+          "localized_name": "CLIP",
+          "name": "CLIP",
+          "type": "CLIP",
+          "links": [
+            1
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.31.0",
+        "Node name for S&R": "CLIPLoader",
+        "models": [
+          {
+            "name": "minimax_music3_text_encoder_pruned_int8_convrot.safetensors",
+            "url": "https://huggingface.co/Comfy-Org/MiniMax-Music-3/resolve/main/text_encoders/minimax_music3_text_encoder_pruned_int8_convrot.safetensors",
+            "directory": "text_encoders"
+          }
+        ]
+      },
+      "widgets_values": [
+        "minimax_music3_text_encoder_pruned_int8_convrot.safetensors",
+        "minimax",
+        "default"
+      ]
+    },
+    {
+      "id": 7,
+      "type": "VAELoader",
+      "pos": [
+        40,
+        740
+      ],
+      "size": [
+        480,
+        70
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "localized_name": "vae_name",
+          "name": "vae_name",
+          "type": "COMBO",
+          "widget": {
+            "name": "vae_name"
+          }
+        }
+      ],
+      "outputs": [
+        {
+          "localized_name": "VAE",
+          "name": "VAE",
+          "type": "VAE",
+          "links": [
+            11
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.31.0",
+        "Node name for S&R": "VAELoader",
+        "models": [
+          {
+            "name": "minimax_music3_dav.safetensors",
+            "url": "https://huggingface.co/Comfy-Org/MiniMax-Music-3/resolve/main/vae/minimax_music3_dav.safetensors",
+            "directory": "vae"
+          }
+        ]
+      },
+      "widgets_values": [
+        "minimax_music3_dav.safetensors"
+      ]
+    },
+    {
+      "id": 10,
+      "type": "ConditioningZeroOut",
+      "pos": [
+        1390,
+        440
+      ],
+      "size": [
+        230,
+        40
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "localized_name": "conditioning",
+          "name": "conditioning",
+          "type": "CONDITIONING",
+          "link": 5
+        }
+      ],
+      "outputs": [
+        {
+          "localized_name": "CONDITIONING",
+          "name": "CONDITIONING",
+          "type": "CONDITIONING",
+          "links": [
+            7
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.31.0",
+        "Node name for S&R": "ConditioningZeroOut"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 15,
+      "type": "EmptyMiniMaxMusic3LatentAudio",
+      "pos": [
+        1390,
+        220
+      ],
+      "size": [
+        350,
+        90
+      ],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "localized_name": "seconds",
+          "name": "seconds",
+          "type": "FLOAT",
+          "widget": {
+            "name": "seconds"
+          },
+          "link": 6
+        }
+      ],
+      "outputs": [
+        {
+          "localized_name": "LATENT",
+          "name": "LATENT",
+          "type": "LATENT",
+          "links": [
+            8
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.31.0",
+        "Node name for S&R": "EmptyMiniMaxMusic3LatentAudio"
+      },
+      "widgets_values": [
+        120,
+        1
+      ]
+    },
+    {
+      "id": 9,
+      "type": "KSampler",
+      "pos": [
+        1740,
+        220
+      ],
+      "size": [
+        370,
+        270
+      ],
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "localized_name": "model",
+          "name": "model",
+          "type": "MODEL",
+          "link": 9
+        },
+        {
+          "localized_name": "positive",
+          "name": "positive",
+          "type": "CONDITIONING",
+          "link": 4
+        },
+        {
+          "localized_name": "negative",
+          "name": "negative",
+          "type": "CONDITIONING",
+          "link": 7
+        },
+        {
+          "localized_name": "latent_image",
+          "name": "latent_image",
+          "type": "LATENT",
+          "link": 8
+        },
+        {
+          "localized_name": "seed",
+          "name": "seed",
+          "type": "INT",
+          "widget": {
+            "name": "seed"
+          }
+        }
+      ],
+      "outputs": [
+        {
+          "localized_name": "LATENT",
+          "name": "LATENT",
+          "type": "LATENT",
+          "links": [
+            10
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.31.0",
+        "Node name for S&R": "KSampler"
+      },
+      "widgets_values": [
+        42,
+        "fixed",
+        30,
+        1.0,
+        "euler",
+        "simple",
+        1
+      ]
+    },
+    {
+      "id": 12,
+      "type": "VAEDecodeAudio",
+      "pos": [
+        2160,
+        220
+      ],
+      "size": [
+        230,
+        60
+      ],
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "localized_name": "samples",
+          "name": "samples",
+          "type": "LATENT",
+          "link": 10
+        },
+        {
+          "localized_name": "vae",
+          "name": "vae",
+          "type": "VAE",
+          "link": 11
+        }
+      ],
+      "outputs": [
+        {
+          "localized_name": "AUDIO",
+          "name": "AUDIO",
+          "type": "AUDIO",
+          "links": [
+            12
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.31.0",
+        "Node name for S&R": "VAEDecodeAudio"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 13,
+      "type": "MiniMaxMusic3ReferenceAudioEncode",
+      "pos": [
+        810,
+        220
+      ],
+      "size": [
+        520,
+        640
+      ],
+      "flags": {},
+      "order": 8,
+      "mode": 0,
+      "showAdvanced": true,
+      "inputs": [
+        {
+          "localized_name": "clip",
+          "name": "clip",
+          "type": "CLIP",
+          "link": 1
+        },
+        {
+          "localized_name": "reference_encoder",
+          "name": "reference_encoder",
+          "type": "MINIMAX_MUSIC3_RVQ_REFERENCE_ENCODER",
+          "link": 2
+        },
+        {
+          "localized_name": "audio",
+          "name": "audio",
+          "type": "AUDIO",
+          "link": 3
+        },
+        {
+          "localized_name": "caption",
+          "name": "caption",
+          "type": "STRING",
+          "widget": {
+            "name": "caption"
+          }
+        },
+        {
+          "localized_name": "lyrics",
+          "name": "lyrics",
+          "type": "STRING",
+          "widget": {
+            "name": "lyrics"
+          }
+        },
+        {
+          "localized_name": "seed",
+          "name": "seed",
+          "type": "INT",
+          "widget": {
+            "name": "seed"
+          }
+        },
+        {
+          "localized_name": "cfg_scale",
+          "name": "cfg_scale",
+          "type": "FLOAT",
+          "widget": {
+            "name": "cfg_scale"
+          }
+        },
+        {
+          "localized_name": "top_k",
+          "name": "top_k",
+          "type": "INT",
+          "widget": {
+            "name": "top_k"
+          }
+        }
+      ],
+      "outputs": [
+        {
+          "localized_name": "conditioning",
+          "name": "conditioning",
+          "type": "CONDITIONING",
+          "links": [
+            4,
+            5
+          ]
+        },
+        {
+          "localized_name": "seconds",
+          "name": "seconds",
+          "type": "FLOAT",
+          "links": [
+            6
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfyui_open_rvq",
+        "Node name for S&R": "MiniMaxMusic3ReferenceAudioEncode"
+      },
+      "widgets_values": [
+        "rock",
+        "[instrumental]",
+        299871780,
+        1.5,
+        50
+      ]
+    },
+    {
+      "id": 44,
+      "type": "MiniMaxMusic3RVQReferenceEncoderLoader",
+      "pos": [
+        40,
+        420
+      ],
+      "size": [
+        650,
+        100
+      ],
+      "flags": {},
+      "order": 9,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "localized_name": "reference_encoder",
+          "name": "reference_encoder",
+          "type": "MINIMAX_MUSIC3_RVQ_REFERENCE_ENCODER",
+          "links": [
+            2
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfyui_open_rvq",
+        "Node name for S&R": "MiniMaxMusic3RVQReferenceEncoderLoader",
+        "models": [
+          {
+            "name": "minimax_music3_rvq_encoder_v4_169m_autoregressive_depth_recommended.safetensors",
+            "url": "https://huggingface.co/SimpleTuner/open-rvq-encoder-minimax-music3/resolve/main/encoders/minimax_music3_rvq_encoder_v4_169m_autoregressive_depth_recommended.safetensors",
+            "directory": "minimax_music3_rvq_encoders"
+          },
+          {
+            "name": "dav.pth",
+            "url": "https://huggingface.co/MiniMaxAI/MiniMax-Music3/resolve/main/dav.pth",
+            "directory": "vae"
+          }
+        ]
+      },
+      "widgets_values": [
+        "minimax_music3_rvq_encoder_v4_169m_autoregressive_depth_recommended.safetensors",
+        "dav.pth"
+      ]
+    },
+    {
+      "id": 45,
+      "type": "LoadAudio",
+      "pos": [
+        40,
+        570
+      ],
+      "size": [
+        650,
+        100
+      ],
+      "flags": {},
+      "order": 10,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "localized_name": "AUDIO",
+          "name": "AUDIO",
+          "type": "AUDIO",
+          "links": [
+            3
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "Node name for S&R": "LoadAudio"
+      },
+      "widgets_values": [
+        "reference_audio.flac"
+      ]
+    },
+    {
+      "id": 46,
+      "type": "SaveAudioAdvanced",
+      "pos": [
+        2480,
+        220
+      ],
+      "size": [
+        750,
+        160
+      ],
+      "flags": {},
+      "order": 11,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": 12
+        }
+      ],
+      "outputs": [
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": null
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.31.0"
+      },
+      "widgets_values": [
+        "audio/minimax_music3_rvq_reference",
+        "flac"
+      ]
+    }
+  ],
+  "links": [
+    [
+      1,
+      3,
+      0,
+      13,
+      0,
+      "CLIP"
+    ],
+    [
+      2,
+      44,
+      0,
+      13,
+      1,
+      "MINIMAX_MUSIC3_RVQ_REFERENCE_ENCODER"
+    ],
+    [
+      3,
+      45,
+      0,
+      13,
+      2,
+      "AUDIO"
+    ],
+    [
+      4,
+      13,
+      0,
+      9,
+      1,
+      "CONDITIONING"
+    ],
+    [
+      5,
+      13,
+      0,
+      10,
+      0,
+      "CONDITIONING"
+    ],
+    [
+      6,
+      13,
+      1,
+      15,
+      0,
+      "FLOAT"
+    ],
+    [
+      7,
+      10,
+      0,
+      9,
+      2,
+      "CONDITIONING"
+    ],
+    [
+      8,
+      15,
+      0,
+      9,
+      3,
+      "LATENT"
+    ],
+    [
+      9,
+      6,
+      0,
+      9,
+      0,
+      "MODEL"
+    ],
+    [
+      10,
+      9,
+      0,
+      12,
+      0,
+      "LATENT"
+    ],
+    [
+      11,
+      7,
+      0,
+      12,
+      1,
+      "VAE"
+    ],
+    [
+      12,
+      12,
+      0,
+      46,
+      0,
+      "AUDIO"
+    ]
+  ],
+  "groups": [],
+  "definitions": {
+    "subgraphs": []
+  },
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.72,
+      "offset": [
+        150,
+        100
+      ]
+    },
+    "frontendVersion": "1.49.6"
+  },
+  "version": 0.4
+}

encoders/minimax_music3_rvq_encoder_v1_41m_independent_heads.json

@@ -0,0 +1,28 @@
+{
+  "codebook_vocab_sizes": [
+    16384,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024
+  ],
+  "conv_dilations": [
+    1,
+    3,
+    9
+  ],
+  "d_model": 512,
+  "dropout": 0.1,
+  "ff_mult": 4,
+  "latent_channels": 128,
+  "max_position_embeddings": 128,
+  "mup": true,
+  "mup_attention_multiplier": 8.0,
+  "mup_output_mult": 1.0,
+  "mup_readout_zero_init": true,
+  "num_heads": 8,
+  "num_layers": 8
+}

encoders/minimax_music3_rvq_encoder_v1_41m_independent_heads.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9683ba0a719fc8f462fb0013a413e22e59e3b29f50d90f34b25658fb827adb6e
+size 163931632

encoders/minimax_music3_rvq_encoder_v2_155m_wide_independent_heads.json

@@ -0,0 +1,28 @@
+{
+  "codebook_vocab_sizes": [
+    16384,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024
+  ],
+  "conv_dilations": [
+    1,
+    3,
+    9
+  ],
+  "d_model": 1088,
+  "dropout": 0.1,
+  "ff_mult": 4,
+  "latent_channels": 128,
+  "max_position_embeddings": 128,
+  "mup": true,
+  "mup_attention_multiplier": 8.0,
+  "mup_output_mult": 1.0,
+  "mup_readout_zero_init": true,
+  "num_heads": 17,
+  "num_layers": 8
+}

encoders/minimax_music3_rvq_encoder_v2_155m_wide_independent_heads.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:47dfffb7a76d95589752df0c59ff6eaf3496f6fcf5408d324218887e471023c7
+size 618960456

encoders/minimax_music3_rvq_encoder_v3_155m_mert_aligned_independent_heads.json

@@ -0,0 +1,28 @@
+{
+  "codebook_vocab_sizes": [
+    16384,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024
+  ],
+  "conv_dilations": [
+    1,
+    3,
+    9
+  ],
+  "d_model": 1088,
+  "dropout": 0.1,
+  "ff_mult": 4,
+  "latent_channels": 128,
+  "max_position_embeddings": 128,
+  "mup": true,
+  "mup_attention_multiplier": 8.0,
+  "mup_output_mult": 1.0,
+  "mup_readout_zero_init": true,
+  "num_heads": 17,
+  "num_layers": 8
+}

encoders/minimax_music3_rvq_encoder_v3_155m_mert_aligned_independent_heads.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:356e97fea65c486a28814dc778b52645807ab8f784c6021d2403f688a2bfe788
+size 618960456

encoders/minimax_music3_rvq_encoder_v4_169m_autoregressive_depth_recommended.json

@@ -0,0 +1,34 @@
+{
+  "codebook_vocab_sizes": [
+    16384,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024
+  ],
+  "conv_dilations": [
+    1,
+    3,
+    9
+  ],
+  "d_model": 1088,
+  "depth_decoder": true,
+  "depth_decoder_dim": 512,
+  "depth_decoder_dropout": 0.1,
+  "depth_decoder_ff_mult": 4,
+  "depth_decoder_heads": 8,
+  "depth_decoder_layers": 2,
+  "dropout": 0.1,
+  "ff_mult": 4,
+  "latent_channels": 128,
+  "max_position_embeddings": 128,
+  "mup": true,
+  "mup_attention_multiplier": 8.0,
+  "mup_output_mult": 1.0,
+  "mup_readout_zero_init": true,
+  "num_heads": 17,
+  "num_layers": 8
+}

encoders/minimax_music3_rvq_encoder_v4_169m_autoregressive_depth_recommended.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8fa93a7db2e5a090442d9492ea3f9f07664ce5bdb2ea7890559a4408f049744
+size 676055232

evaluation/condition-replay-aggregate.json

@@ -0,0 +1,225 @@
+{
+  "dataset": {
+    "exact_alignment_records": 130,
+    "repo_id": "bghira/minimax-music3-rvq-reverse-distillation",
+    "revision": "5029b1e7f1bbfbf028b76b38564fecccda94a111",
+    "split": "holdout"
+  },
+  "format": "simpletuner-minimaxmusic-replay-comparison-v1",
+  "metric": "per-record mean stitched condition-embedding cosine",
+  "models": [
+    {
+      "checkpoint": {
+        "checkpoint_bytes": 163955786,
+        "checkpoint_sha256": "fd7bfefb197cb7b65a6b6d9b1373c1772a15979f1a4b67d820764d610f0fb9a5",
+        "checkpoint_url": "https://www.serveurperso.com/ia/ssd/workspace/git/minimaxmusic.cpp/training/checkpoints/v1/best.pt",
+        "epoch": 17,
+        "reported_validation_loss": 5.292583839459852,
+        "source_commit": "d19efe9f94e41ac4c900aa30d56fe90c8dac7ef1",
+        "source_url": "https://github.com/ServeurpersoCom/minimaxmusic.cpp/commit/d19efe9f94e41ac4c900aa30d56fe90c8dac7ef1"
+      },
+      "model_id": "serveurperso-v1",
+      "parameter_count": 40978944,
+      "predicted_code_condition_cosine": {
+        "count": 130,
+        "max": 0.7290002703666687,
+        "mean": 0.6633286081827604,
+        "min": 0.6026018261909485,
+        "quantiles": {
+          "0.0": 0.6026018261909485,
+          "0.05": 0.6280524551868438,
+          "0.25": 0.6483299732208252,
+          "0.5": 0.6645751297473907,
+          "0.75": 0.6784153282642365,
+          "0.95": 0.6963281780481339,
+          "1.0": 0.7290002703666687
+        },
+        "std": 0.02217457298933777
+      },
+      "true_code_control": {
+        "count": 130,
+        "max": 0.999957263469696,
+        "mean": 0.9999072317893688,
+        "min": 0.9983662366867065,
+        "quantiles": {
+          "0.0": 0.9983662366867065,
+          "0.05": 0.999821075797081,
+          "0.25": 0.9999240636825562,
+          "0.5": 0.9999381899833679,
+          "0.75": 0.9999445825815201,
+          "0.95": 0.9999500513076782,
+          "1.0": 0.999957263469696
+        },
+        "std": 0.00015022638681815362
+      }
+    },
+    {
+      "checkpoint": {
+        "checkpoint": "full/output/final",
+        "checkpoint_bytes": 163931632,
+        "checkpoint_sha256": "1cea88dddf7a48583c13ac661a548b7a3535b933a76356bb158420d24da32f22",
+        "global_step": 17640
+      },
+      "model_id": "simpletuner-v1",
+      "parameter_count": 40978944,
+      "predicted_code_condition_cosine": {
+        "count": 130,
+        "max": 0.8286042809486389,
+        "mean": 0.7624421353523548,
+        "min": 0.7123062014579773,
+        "quantiles": {
+          "0.0": 0.7123062014579773,
+          "0.05": 0.7345191180706024,
+          "0.25": 0.747437372803688,
+          "0.5": 0.7631425857543945,
+          "0.75": 0.7765287756919861,
+          "0.95": 0.7904498308897019,
+          "1.0": 0.8286042809486389
+        },
+        "std": 0.019549602827450108
+      },
+      "true_code_control": {
+        "count": 130,
+        "max": 0.999957263469696,
+        "mean": 0.9999072317893688,
+        "min": 0.9983662366867065,
+        "quantiles": {
+          "0.0": 0.9983662366867065,
+          "0.05": 0.999821075797081,
+          "0.25": 0.9999240636825562,
+          "0.5": 0.9999381899833679,
+          "0.75": 0.9999445825815201,
+          "0.95": 0.9999500513076782,
+          "1.0": 0.999957263469696
+        },
+        "std": 0.00015022638681815362
+      }
+    },
+    {
+      "checkpoint": {
+        "checkpoint": "v2/output/final",
+        "checkpoint_bytes": 618960456,
+        "checkpoint_sha256": "47dfffb7a76d95589752df0c59ff6eaf3496f6fcf5408d324218887e471023c7",
+        "global_step": 17660
+      },
+      "model_id": "simpletuner-v2",
+      "parameter_count": 154736064,
+      "predicted_code_condition_cosine": {
+        "count": 130,
+        "max": 0.8361330032348633,
+        "mean": 0.7698414775041433,
+        "min": 0.7206224203109741,
+        "quantiles": {
+          "0.0": 0.7206224203109741,
+          "0.05": 0.7429912149906158,
+          "0.25": 0.7548101097345352,
+          "0.5": 0.770626574754715,
+          "0.75": 0.7835157215595245,
+          "0.95": 0.7986360996961593,
+          "1.0": 0.8361330032348633
+        },
+        "std": 0.019062989675575924
+      },
+      "true_code_control": {
+        "count": 130,
+        "max": 0.999957263469696,
+        "mean": 0.9999072317893688,
+        "min": 0.9983662366867065,
+        "quantiles": {
+          "0.0": 0.9983662366867065,
+          "0.05": 0.999821075797081,
+          "0.25": 0.9999240636825562,
+          "0.5": 0.9999381899833679,
+          "0.75": 0.9999445825815201,
+          "0.95": 0.9999500513076782,
+          "1.0": 0.999957263469696
+        },
+        "std": 0.00015022638681815362
+      }
+    },
+    {
+      "checkpoint": {
+        "checkpoint": "v3/output/final",
+        "checkpoint_bytes": 618960456,
+        "checkpoint_sha256": "356e97fea65c486a28814dc778b52645807ab8f784c6021d2403f688a2bfe788",
+        "global_step": 17660
+      },
+      "model_id": "simpletuner-v3",
+      "parameter_count": 154736064,
+      "predicted_code_condition_cosine": {
+        "count": 130,
+        "max": 0.8356426954269409,
+        "mean": 0.7702590204202212,
+        "min": 0.7236603498458862,
+        "quantiles": {
+          "0.0": 0.7236603498458862,
+          "0.05": 0.7415853887796402,
+          "0.25": 0.7547859847545624,
+          "0.5": 0.7704743146896362,
+          "0.75": 0.7837727516889572,
+          "0.95": 0.8004916191101075,
+          "1.0": 0.8356426954269409
+        },
+        "std": 0.019273954474224628
+      },
+      "true_code_control": {
+        "count": 130,
+        "max": 0.999957263469696,
+        "mean": 0.9999072317893688,
+        "min": 0.9983662366867065,
+        "quantiles": {
+          "0.0": 0.9983662366867065,
+          "0.05": 0.999821075797081,
+          "0.25": 0.9999240636825562,
+          "0.5": 0.9999381899833679,
+          "0.75": 0.9999445825815201,
+          "0.95": 0.9999500513076782,
+          "1.0": 0.999957263469696
+        },
+        "std": 0.00015022638681815362
+      }
+    },
+    {
+      "checkpoint": {
+        "checkpoint": "v4/output/final",
+        "checkpoint_bytes": 676055232,
+        "checkpoint_sha256": "e8fa93a7db2e5a090442d9492ea3f9f07664ce5bdb2ea7890559a4408f049744",
+        "global_step": 17660
+      },
+      "model_id": "simpletuner-v4",
+      "parameter_count": 169008576,
+      "predicted_code_condition_cosine": {
+        "count": 130,
+        "max": 0.897747814655304,
+        "mean": 0.8747961479883928,
+        "min": 0.8091457486152649,
+        "quantiles": {
+          "0.0": 0.8091457486152649,
+          "0.05": 0.8453496336936951,
+          "0.25": 0.8675179928541183,
+          "0.5": 0.879057914018631,
+          "0.75": 0.8857778161764145,
+          "0.95": 0.8938389658927918,
+          "1.0": 0.897747814655304
+        },
+        "std": 0.015662041602848998
+      },
+      "true_code_control": {
+        "count": 130,
+        "max": 0.999957263469696,
+        "mean": 0.9999072317893688,
+        "min": 0.9983662366867065,
+        "quantiles": {
+          "0.0": 0.9983662366867065,
+          "0.05": 0.999821075797081,
+          "0.25": 0.9999240636825562,
+          "0.5": 0.9999381899833679,
+          "0.75": 0.9999445825815201,
+          "0.95": 0.9999500513076782,
+          "1.0": 0.999957263469696
+        },
+        "std": 0.00015022638681815362
+      }
+    }
+  ]
+}

experiment-summary.json

@@ -0,0 +1,71 @@
+{
+  "encoders": [
+    {
+      "bytes": 163931632,
+      "condition_replay_checkpoint": "final",
+      "condition_replay_mean": 0.7624421353523548,
+      "config_filename": "encoders/minimax_music3_rvq_encoder_v1_41m_independent_heads.json",
+      "filename": "encoders/minimax_music3_rvq_encoder_v1_41m_independent_heads.safetensors",
+      "parameter_count": 40978944,
+      "sha256": "9683ba0a719fc8f462fb0013a413e22e59e3b29f50d90f34b25658fb827adb6e",
+      "source_checkpoint_path": "checkpoint-17500",
+      "source_repo_id": "SimpleTuner/open-rvq-encoder-minimax-music3-41m-v1",
+      "source_revision": "e9fa3bba9b5bb3b76e7fd6705f9ca5dbb7638ec4",
+      "version": "v1"
+    },
+    {
+      "bytes": 618960456,
+      "condition_replay_checkpoint": "final",
+      "condition_replay_mean": 0.7698414775041433,
+      "config_filename": "encoders/minimax_music3_rvq_encoder_v2_155m_wide_independent_heads.json",
+      "filename": "encoders/minimax_music3_rvq_encoder_v2_155m_wide_independent_heads.safetensors",
+      "parameter_count": 154736064,
+      "sha256": "47dfffb7a76d95589752df0c59ff6eaf3496f6fcf5408d324218887e471023c7",
+      "source_checkpoint_path": "final",
+      "source_repo_id": "SimpleTuner/open-rvq-encoder-minimax-music3-155m-v2",
+      "source_revision": "f317d72466b98cad097c1e6f4dd379f7e344c939",
+      "version": "v2"
+    },
+    {
+      "bytes": 618960456,
+      "condition_replay_checkpoint": "final",
+      "condition_replay_mean": 0.7702590204202212,
+      "config_filename": "encoders/minimax_music3_rvq_encoder_v3_155m_mert_aligned_independent_heads.json",
+      "filename": "encoders/minimax_music3_rvq_encoder_v3_155m_mert_aligned_independent_heads.safetensors",
+      "parameter_count": 154736064,
+      "sha256": "356e97fea65c486a28814dc778b52645807ab8f784c6021d2403f688a2bfe788",
+      "source_checkpoint_path": "final",
+      "source_repo_id": "SimpleTuner/open-rvq-encoder-minimax-music3-155m-v3",
+      "source_revision": "7c525705817d8878da02ed41429cad027552a568",
+      "version": "v3"
+    },
+    {
+      "bytes": 676055232,
+      "condition_replay_checkpoint": "final",
+      "condition_replay_mean": 0.8747961479883928,
+      "config_filename": "encoders/minimax_music3_rvq_encoder_v4_169m_autoregressive_depth_recommended.json",
+      "filename": "encoders/minimax_music3_rvq_encoder_v4_169m_autoregressive_depth_recommended.safetensors",
+      "parameter_count": 169008576,
+      "sha256": "e8fa93a7db2e5a090442d9492ea3f9f07664ce5bdb2ea7890559a4408f049744",
+      "source_checkpoint_path": "final",
+      "source_repo_id": "SimpleTuner/open-rvq-encoder-minimax-music3-169m-v4",
+      "source_revision": "b9a9165b99d12cb48cac812776938c8d9332dd59",
+      "version": "v4"
+    }
+  ],
+  "format": "simpletuner-minimax-music3-rvq-collection-v1",
+  "metric": {
+    "dataset_repo_id": "bghira/minimax-music3-rvq-reverse-distillation",
+    "dataset_revision": "5029b1e7f1bbfbf028b76b38564fecccda94a111",
+    "name": "per-record mean stitched condition-embedding cosine",
+    "records": 130,
+    "serveurperso_v1_mean": 0.6633286081827604,
+    "split": "holdout",
+    "true_code_control_mean": 0.9999072317893688
+  },
+  "notes": [
+    "The v1 packaged checkpoint is step 17500; its replay metric was measured on v1 final.",
+    "V2, v3, and v4 replay metrics were measured on the packaged final checkpoints."
+  ],
+  "recommended_version": "v4"
+}

minimax_music3_reference_adapter.py

@@ -0,0 +1,710 @@
+# Copyright 2026 SimpleTuner contributors
+#
+# 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.
+
+"""Reference-audio conditioning for MiniMax Music 3.
+
+This file is self-contained. It loads a released SimpleTuner RVQ encoder,
+encodes 44.1 kHz audio with the original DAV encoder, predicts eight RVQ codes
+per 25 Hz frame, and teacher-forces those codes through the official MiniMax
+Music 3 language-model path.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import math
+import re
+from contextlib import nullcontext
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn.utils import weight_norm
+
+
+COLLECTION_REPO_ID = "SimpleTuner/open-rvq-encoder-minimax-music3"
+OFFICIAL_MODEL_REPO_ID = "MiniMaxAI/MiniMax-Music3"
+DEFAULT_ENCODER_FILE = (
+    "encoders/minimax_music3_rvq_encoder_v4_169m_autoregressive_depth_recommended.safetensors"
+)
+DAV_FILE = "dav.pth"
+SAMPLE_RATE = 44_100
+DAV_HOP_SAMPLES = 512
+FRAME_RATE = 25
+LATENT_RATE_NUM = 441
+LATENT_RATE_DEN = 128
+AUDIO_CODE_OFFSET = 151_675
+AUDIO_END_TOKEN_ID = 151_670
+AUDIO_CFG_TOKEN_ID = 151_654
+SEMANTIC_VOCAB_SIZE = 16_384
+AR_CFG_SCALE = 1.5
+AR_TOP_K = 50
+LM_BLOCK_FRAMES = 256
+
+_SPECIAL_TAG_RE = re.compile(r"<\|([^|]*)\|>")
+_LEADING_TAGS_RE = re.compile(r"^[ \t]*((?:\[[^\]]+\][ \t]*)+)")
+
+
+@dataclasses.dataclass(frozen=True)
+class RVQEncoderConfig:
+    latent_channels: int = 128
+    codebook_vocab_sizes: tuple[int, ...] = (16_384, 1024, 1024, 1024, 1024, 1024, 1024, 1024)
+    d_model: int = 512
+    num_layers: int = 8
+    num_heads: int = 8
+    ff_mult: int = 4
+    dropout: float = 0.1
+    max_position_embeddings: int = 128
+    conv_dilations: tuple[int, ...] = (1, 3, 9)
+    mup: bool = False
+    mup_output_mult: float = 1.0
+    mup_readout_zero_init: bool = False
+    mup_attention_multiplier: float = 8.0
+    depth_decoder: bool = False
+    depth_decoder_dim: int = 512
+    depth_decoder_layers: int = 2
+    depth_decoder_heads: int = 8
+    depth_decoder_ff_mult: int = 4
+    depth_decoder_dropout: float = 0.1
+
+    @classmethod
+    def from_dict(cls, values: dict[str, Any]) -> "RVQEncoderConfig":
+        normalized = dict(values)
+        for key in ("codebook_vocab_sizes", "conv_dilations"):
+            if key in normalized:
+                normalized[key] = tuple(normalized[key])
+        known = {field.name for field in dataclasses.fields(cls)}
+        unknown = sorted(set(normalized) - known)
+        if unknown:
+            raise ValueError(f"Unknown RVQ encoder configuration fields: {unknown}")
+        return cls(**normalized)
+
+
+class RVQResBlock(nn.Module):
+    def __init__(self, dim: int, dilation: int):
+        super().__init__()
+        self.norm = nn.GroupNorm(1, dim)
+        self.conv1 = nn.Conv1d(dim, dim, kernel_size=3, padding=dilation, dilation=dilation)
+        self.conv2 = nn.Conv1d(dim, dim, kernel_size=1)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        residual = self.conv1(F.gelu(self.norm(hidden_states)))
+        return hidden_states + self.conv2(F.gelu(residual))
+
+
+class RVQMuTransformerEncoderLayer(nn.Module):
+    def __init__(self, config: RVQEncoderConfig):
+        super().__init__()
+        if config.d_model % config.num_heads:
+            raise ValueError("d_model must be divisible by num_heads")
+        self.num_heads = config.num_heads
+        self.head_dim = config.d_model // config.num_heads
+        self.attention_multiplier = config.mup_attention_multiplier
+        self.norm1 = nn.LayerNorm(config.d_model)
+        self.norm2 = nn.LayerNorm(config.d_model)
+        self.q_proj = nn.Linear(config.d_model, config.d_model)
+        self.k_proj = nn.Linear(config.d_model, config.d_model)
+        self.v_proj = nn.Linear(config.d_model, config.d_model)
+        self.out_proj = nn.Linear(config.d_model, config.d_model)
+        self.linear1 = nn.Linear(config.d_model, config.d_model * config.ff_mult)
+        self.linear2 = nn.Linear(config.d_model * config.ff_mult, config.d_model)
+        self.dropout = nn.Dropout(config.dropout)
+        self.attn_dropout = nn.Dropout(config.dropout)
+
+    def _split_heads(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch, frames, _ = hidden_states.shape
+        return hidden_states.view(batch, frames, self.num_heads, self.head_dim).transpose(1, 2)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        normalized = self.norm1(hidden_states)
+        query = self._split_heads(self.q_proj(normalized))
+        key = self._split_heads(self.k_proj(normalized))
+        value = self._split_heads(self.v_proj(normalized))
+        scores = torch.matmul(query, key.transpose(-2, -1)) * (self.attention_multiplier / self.head_dim)
+        probabilities = self.attn_dropout(F.softmax(scores.float(), dim=-1).to(query.dtype))
+        attended = torch.matmul(probabilities, value).transpose(1, 2).contiguous()
+        attended = attended.view(hidden_states.shape)
+        hidden_states = hidden_states + self.dropout(self.out_proj(attended))
+        feedforward = self.linear2(self.dropout(F.gelu(self.linear1(self.norm2(hidden_states)))))
+        return hidden_states + self.dropout(feedforward)
+
+
+class RVQDepthDecoderLayer(nn.Module):
+    def __init__(self, config: RVQEncoderConfig):
+        super().__init__()
+        dim = config.depth_decoder_dim
+        if dim % config.depth_decoder_heads:
+            raise ValueError("depth_decoder_dim must be divisible by depth_decoder_heads")
+        self.num_heads = config.depth_decoder_heads
+        self.head_dim = dim // config.depth_decoder_heads
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.q_proj = nn.Linear(dim, dim)
+        self.k_proj = nn.Linear(dim, dim)
+        self.v_proj = nn.Linear(dim, dim)
+        self.out_proj = nn.Linear(dim, dim)
+        self.linear1 = nn.Linear(dim, dim * config.depth_decoder_ff_mult)
+        self.linear2 = nn.Linear(dim * config.depth_decoder_ff_mult, dim)
+        self.dropout = nn.Dropout(config.depth_decoder_dropout)
+        self.attn_dropout = nn.Dropout(config.depth_decoder_dropout)
+
+    def _split_heads(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch, depth, _ = hidden_states.shape
+        return hidden_states.view(batch, depth, self.num_heads, self.head_dim).transpose(1, 2)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        normalized = self.norm1(hidden_states)
+        query = self._split_heads(self.q_proj(normalized))
+        key = self._split_heads(self.k_proj(normalized))
+        value = self._split_heads(self.v_proj(normalized))
+        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        depth = hidden_states.shape[1]
+        scores = scores.masked_fill(torch.ones((depth, depth), dtype=torch.bool, device=scores.device).triu(1), -torch.inf)
+        probabilities = self.attn_dropout(F.softmax(scores.float(), dim=-1).to(query.dtype))
+        attended = torch.matmul(probabilities, value).transpose(1, 2).contiguous()
+        attended = attended.view(hidden_states.shape)
+        hidden_states = hidden_states + self.dropout(self.out_proj(attended))
+        feedforward = self.linear2(self.dropout(F.gelu(self.linear1(self.norm2(hidden_states)))))
+        return hidden_states + self.dropout(feedforward)
+
+
+class MiniMaxMusicRVQDepthDecoder(nn.Module):
+    def __init__(self, config: RVQEncoderConfig):
+        super().__init__()
+        self.config = config
+        self.context_projection = nn.Linear(config.d_model, config.depth_decoder_dim, bias=False)
+        self.prior_embeddings = nn.ModuleList(
+            nn.Embedding(vocab_size, config.depth_decoder_dim) for vocab_size in config.codebook_vocab_sizes[:-1]
+        )
+        self.position = nn.Parameter(torch.zeros(1, len(config.codebook_vocab_sizes), config.depth_decoder_dim))
+        self.layers = nn.ModuleList(RVQDepthDecoderLayer(config) for _ in range(config.depth_decoder_layers))
+        self.norm = nn.LayerNorm(config.depth_decoder_dim)
+        self.heads = nn.ModuleList(
+            nn.Linear(config.depth_decoder_dim, vocab_size) for vocab_size in config.codebook_vocab_sizes[1:]
+        )
+
+    def _decode(self, sequence: torch.Tensor) -> torch.Tensor:
+        hidden_states = sequence + self.position[:, : sequence.shape[1]].to(sequence.dtype)
+        for layer in self.layers:
+            hidden_states = layer(hidden_states)
+        return self.norm(hidden_states)
+
+    def forward(self, frame_context: torch.Tensor, semantic_codes: torch.Tensor) -> list[torch.Tensor]:
+        batch, frames, _ = frame_context.shape
+        sequence = torch.cat(
+            (
+                self.context_projection(frame_context).flatten(0, 1).unsqueeze(1),
+                self.prior_embeddings[0](semantic_codes).flatten(0, 1).unsqueeze(1),
+            ),
+            dim=1,
+        )
+        logits = []
+        for acoustic_index, head in enumerate(self.heads):
+            head_logits = head(self._decode(sequence)[:, -1]).view(batch, frames, -1)
+            logits.append(head_logits)
+            if acoustic_index + 1 < len(self.heads):
+                selected = head_logits.argmax(dim=-1)
+                prior = self.prior_embeddings[acoustic_index + 1](selected).flatten(0, 1).unsqueeze(1)
+                sequence = torch.cat((sequence, prior), dim=1)
+        return logits
+
+
+class MiniMaxMusicRVQEncoder(nn.Module):
+    def __init__(self, config: RVQEncoderConfig):
+        super().__init__()
+        self.config = config
+        self.conv_in = nn.Conv1d(config.latent_channels, config.d_model, kernel_size=7, padding=3)
+        self.blocks = nn.ModuleList(RVQResBlock(config.d_model, dilation) for dilation in config.conv_dilations)
+        self.position = nn.Parameter(torch.zeros(1, config.max_position_embeddings, config.d_model))
+        if config.mup:
+            self.transformer = nn.ModuleList(RVQMuTransformerEncoderLayer(config) for _ in range(config.num_layers))
+        else:
+            layer = nn.TransformerEncoderLayer(
+                d_model=config.d_model,
+                nhead=config.num_heads,
+                dim_feedforward=config.d_model * config.ff_mult,
+                dropout=config.dropout,
+                activation="gelu",
+                batch_first=True,
+                norm_first=True,
+            )
+            self.transformer = nn.TransformerEncoder(layer, config.num_layers)
+        self.norm_out = nn.LayerNorm(config.d_model)
+        readout_sizes = config.codebook_vocab_sizes[:1] if config.depth_decoder else config.codebook_vocab_sizes
+        self.heads = nn.ModuleList(nn.Linear(config.d_model, vocab_size) for vocab_size in readout_sizes)
+        self.depth_decoder = MiniMaxMusicRVQDepthDecoder(config) if config.depth_decoder else None
+
+    def forward(self, latents: torch.Tensor, pool: torch.Tensor) -> list[torch.Tensor]:
+        hidden_states = self.conv_in(latents.transpose(1, 2))
+        for block in self.blocks:
+            hidden_states = block(hidden_states)
+        hidden_states = torch.bmm(pool.to(hidden_states.dtype), hidden_states.transpose(1, 2))
+        hidden_states = hidden_states + self.position[:, : pool.shape[1]].to(hidden_states.dtype)
+        layers = self.transformer if isinstance(self.transformer, nn.ModuleList) else self.transformer.layers
+        for layer in layers:
+            hidden_states = layer(hidden_states)
+        hidden_states = self.norm_out(hidden_states)
+        semantic = self.heads[0](hidden_states)
+        if self.depth_decoder is None:
+            return [head(hidden_states) for head in self.heads]
+        semantic_codes = semantic.argmax(dim=-1)
+        return [semantic, *self.depth_decoder(hidden_states, semantic_codes)]
+
+
+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:
+        return hidden_states + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * hidden_states).pow(2)
+
+
+class DAVResidualUnit(nn.Module):
+    def __init__(self, dim: int, dilation: int):
+        super().__init__()
+        self.block = nn.Sequential(
+            MiniMaxMusic3Snake1d(dim),
+            weight_norm(nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=3 * dilation)),
+            MiniMaxMusic3Snake1d(dim),
+            weight_norm(nn.Conv1d(dim, dim, kernel_size=1)),
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        residual = self.block(hidden_states)
+        if residual.shape[-1] != hidden_states.shape[-1]:
+            padding = (hidden_states.shape[-1] - residual.shape[-1]) // 2
+            hidden_states = hidden_states[..., padding : hidden_states.shape[-1] - padding]
+        return hidden_states + residual
+
+
+class DAVEncoderBlock(nn.Module):
+    def __init__(self, dim: int, stride: int):
+        super().__init__()
+        self.block = nn.Sequential(
+            DAVResidualUnit(dim // 2, 1),
+            DAVResidualUnit(dim // 2, 3),
+            DAVResidualUnit(dim // 2, 9),
+            MiniMaxMusic3Snake1d(dim // 2),
+            weight_norm(
+                nn.Conv1d(dim // 2, dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2))
+            ),
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.block(hidden_states)
+
+
+class DAVEncoder(nn.Module):
+    def __init__(self, encoder_dim: int = 64, rates: tuple[int, ...] = (2, 4, 8, 8), latent_dim: int = 1024):
+        super().__init__()
+        layers: list[nn.Module] = [weight_norm(nn.Conv1d(1, encoder_dim, kernel_size=7, padding=3))]
+        for stride in rates:
+            encoder_dim *= 2
+            layers.append(DAVEncoderBlock(encoder_dim, stride))
+        layers.extend((MiniMaxMusic3Snake1d(encoder_dim), weight_norm(nn.Conv1d(encoder_dim, latent_dim, 3, padding=1))))
+        self.block = nn.Sequential(*layers)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.block(hidden_states)
+
+
+class DAVEncoderOnly(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.encoder = DAVEncoder()
+        self.mean_proj = nn.Conv1d(1024, 64, kernel_size=1)
+
+    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
+        if waveform.ndim == 2:
+            waveform = waveform.unsqueeze(0)
+        if waveform.ndim != 3 or waveform.shape[1] not in (1, 2):
+            raise ValueError("waveform must have shape [channels, samples] or [batch, channels, samples]")
+        if waveform.shape[1] == 1:
+            waveform = waveform.repeat(1, 2, 1)
+        remainder = waveform.shape[-1] % DAV_HOP_SAMPLES
+        if remainder:
+            waveform = F.pad(waveform, (0, DAV_HOP_SAMPLES - remainder))
+        batch, _, samples = waveform.shape
+        hidden_states = self.encoder(waveform.reshape(batch * 2, 1, samples))
+        return self.mean_proj(hidden_states).reshape(batch, 128, -1)
+
+
+def frame_latent_starts(frame_count: int) -> list[int]:
+    """Map continuous 25 Hz frames to the 44.1 kHz / 512 DAV timeline."""
+    if frame_count <= 0:
+        raise ValueError("frame_count must be positive")
+    return [(index * LATENT_RATE_NUM) // LATENT_RATE_DEN for index in range(frame_count + 1)]
+
+
+def build_pool_matrix(bounds: list[int]) -> torch.Tensor:
+    if len(bounds) < 2:
+        raise ValueError("at least two frame boundaries are required")
+    origin = bounds[0]
+    local = [value - origin for value in bounds]
+    pool = torch.zeros((len(local) - 1, local[-1]), dtype=torch.float32)
+    for index, (start, end) in enumerate(zip(local[:-1], local[1:])):
+        if end <= start:
+            raise ValueError(f"invalid DAV span [{start}, {end}) for frame {index}")
+        pool[index, start:end] = 1.0 / (end - start)
+    return pool
+
+
+def _clean_caption(caption: str) -> str:
+    def rewrite(match: re.Match) -> str:
+        parts = match.group(1).strip().split(None, 1)
+        return f"{parts[0]} is {parts[1]}" if len(parts) == 2 else parts[0]
+
+    text = _SPECIAL_TAG_RE.sub(rewrite, caption)
+    lines = []
+    for line in text.splitlines():
+        line = re.sub(r"^\s{0,3}#{1,6}\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
+        lines.append(re.sub(r"(?<!\*)\*([^*\n]+)\*(?!\*)", r"\1", line).rstrip())
+    text = re.sub(r"^\s*[-*_]{3,}\s*$", "", "\n".join(lines), flags=re.MULTILINE)
+    return re.sub(r"\n{2,}", "\n", text.replace("• ", "").replace("    ", ""))
+
+
+def _normalize_lyrics(lyrics: str) -> str:
+    lines = []
+    for line in lyrics.split("\n"):
+        match = _LEADING_TAGS_RE.match(line)
+        lines.append(match.group(1).strip() if match else line)
+    text = "\n".join(lines).replace("] ", "]\n").replace(" [", "\n[").replace(" ^ ", "\n")
+    text = re.sub(r"\[([^\]]+)\]", lambda match: f"[{match.group(1).lower()}]", text)
+    return f"[start]\n{text}"
+
+
+def build_text_ids(tokenizer, prompt: str, lyrics: str, device: torch.device) -> torch.Tensor:
+    if not prompt.strip() or not lyrics.strip():
+        raise ValueError("prompt and lyrics must be non-empty")
+    text = (
+        f"<|im_start|><|caption_start|>{_clean_caption(prompt)}<|caption_end|>"
+        f"<|lyrics_start|>{_normalize_lyrics(lyrics)}<|lyrics_end|><|im_end|><|audio_start|>"
+    )
+    input_ids = tokenizer(text, return_tensors="pt")["input_ids"]
+    unconditional = input_ids.clone()
+    unconditional[:, 1:-2] = AUDIO_CFG_TOKEN_ID
+    return torch.cat((input_ids, unconditional), dim=0).to(device)
+
+
+def _sample_top_k(logits: torch.Tensor, generator: torch.Generator | None) -> torch.Tensor:
+    values = torch.nan_to_num(logits.float(), nan=-1e9, posinf=1e9, neginf=-1e9)
+    threshold = torch.topk(values, min(AR_TOP_K, values.shape[-1]), dim=-1).values[..., -1, None]
+    probabilities = torch.softmax(values.masked_fill(values < threshold, -torch.inf), dim=-1)
+    sample_device = generator.device if generator is not None else probabilities.device
+    return torch.multinomial(probabilities.to(sample_device), 1, generator=generator).squeeze(-1).to(values.device)
+
+
+def _official_depth_hidden(language_model, depth_decoder, hidden: torch.Tensor, codes: torch.Tensor) -> torch.Tensor:
+    sequence = [depth_decoder.projection(hidden).unsqueeze(1)]
+    semantic = language_model.model.embed_tokens(codes[:, 0] + AUDIO_CODE_OFFSET)
+    sequence.append(depth_decoder.projection(semantic).unsqueeze(1))
+    hidden_parts = []
+    for index in range(1, codes.shape[1]):
+        depth_hidden = depth_decoder(torch.cat(sequence, dim=1))[:, -1]
+        hidden_parts.append(depth_hidden)
+        if index < codes.shape[1] - 1:
+            embedding = depth_decoder.audio_embeddings(codes[:, index] + (index - 1) * 1024)
+            sequence.append(depth_decoder.projection(embedding).unsqueeze(1))
+    return torch.cat(hidden_parts, dim=-1)
+
+
+def _embed_official_codes(language_model, depth_decoder, codes: torch.Tensor) -> torch.Tensor:
+    semantic = language_model.model.embed_tokens(codes[..., 0] + AUDIO_CODE_OFFSET)
+    offsets = torch.arange(codes.shape[-1] - 1, device=codes.device) * 1024
+    acoustic = depth_decoder.audio_embeddings(codes[..., 1:] + offsets).sum(dim=-2)
+    return (semantic + acoustic.to(semantic.dtype)) * (codes.shape[-1] ** -0.5)
+
+
+@torch.inference_mode()
+def replay_codes_diffusers(
+    pipeline,
+    codes: torch.Tensor,
+    *,
+    prompt: str,
+    lyrics: str,
+    generator: torch.Generator | None = None,
+) -> torch.Tensor:
+    """Replay predicted codes through official Diffusers LM components."""
+    if codes.ndim != 2 or codes.shape[1] != 8 or codes.shape[0] == 0:
+        raise ValueError("codes must have shape [frames, 8]")
+    language_model = pipeline.language_model
+    depth_decoder = pipeline.rvq_depth_decoder
+    device = next(language_model.parameters()).device
+    codes = codes.to(device=device, dtype=torch.long)
+    text_ids = build_text_ids(pipeline.tokenizer, prompt, lyrics, device)
+    text_output = language_model.model(inputs_embeds=language_model.model.embed_tokens(text_ids), use_cache=True)
+    past = text_output.past_key_values
+    last_hidden = text_output.last_hidden_state[:, -1]
+
+    vocab_mask = torch.ones(language_model.config.vocab_size, dtype=torch.bool, device=device)
+    vocab_mask[AUDIO_CODE_OFFSET : AUDIO_CODE_OFFSET + SEMANTIC_VOCAB_SIZE] = False
+    vocab_mask[AUDIO_END_TOKEN_ID] = False
+    logits = language_model.lm_head(last_hidden).float().masked_fill(vocab_mask, -torch.inf)
+    conditioned, unconditioned = logits[:1], logits[1:2]
+    guided = unconditioned + (conditioned - unconditioned) * AR_CFG_SCALE
+    threshold = torch.topk(conditioned, AR_TOP_K, dim=-1).values[..., -1, None]
+    warmup_token = _sample_top_k(guided.masked_fill(conditioned < threshold, -torch.inf), generator)
+    if int(warmup_token.item()) == AUDIO_END_TOKEN_ID:
+        raise ValueError("the selected seed ended during the required AR warm-up frame")
+    warmup_semantic = (warmup_token - AUDIO_CODE_OFFSET).repeat(2)
+
+    sequence = [depth_decoder.projection(last_hidden).unsqueeze(1)]
+    semantic_embed = language_model.model.embed_tokens(warmup_semantic + AUDIO_CODE_OFFSET)
+    sequence.append(depth_decoder.projection(semantic_embed).unsqueeze(1))
+    warmup_codes = [warmup_semantic]
+    for index in range(1, 8):
+        depth_hidden = depth_decoder(torch.cat(sequence, dim=1))[:, -1]
+        depth_logits = depth_decoder.audio_heads[index - 1](depth_hidden).float()
+        depth_guided = depth_logits[1:2] + (depth_logits[:1] - depth_logits[1:2]) * AR_CFG_SCALE
+        code = _sample_top_k(depth_guided, generator).repeat(2)
+        warmup_codes.append(code)
+        if index < 7:
+            embedding = depth_decoder.audio_embeddings(code + (index - 1) * 1024)
+            sequence.append(depth_decoder.projection(embedding).unsqueeze(1))
+    warmup_codes = torch.stack(warmup_codes, dim=1)
+    output = language_model.model(
+        inputs_embeds=_embed_official_codes(language_model, depth_decoder, warmup_codes).unsqueeze(1),
+        past_key_values=past,
+        use_cache=True,
+    )
+    past = output.past_key_values
+    first_hidden = output.last_hidden_state[:, -1]
+
+    hidden_parts = [first_hidden.unsqueeze(1)]
+    if codes.shape[0] > 1:
+        feedback_codes = codes[:-1].unsqueeze(0).repeat(2, 1, 1)
+        feedback = _embed_official_codes(language_model, depth_decoder, feedback_codes)
+        for start in range(0, feedback.shape[1], LM_BLOCK_FRAMES):
+            output = language_model.model(
+                inputs_embeds=feedback[:, start : start + LM_BLOCK_FRAMES],
+                past_key_values=past,
+                use_cache=True,
+            )
+            past = output.past_key_values
+            hidden_parts.append(output.last_hidden_state)
+    global_hidden = torch.cat(hidden_parts, dim=1)
+    repeated_codes = codes.unsqueeze(0).repeat(2, 1, 1).reshape(-1, 8)
+    depth_hidden = _official_depth_hidden(
+        language_model,
+        depth_decoder,
+        global_hidden.reshape(-1, global_hidden.shape[-1]),
+        repeated_codes,
+    ).view(2, codes.shape[0], -1)
+    return torch.cat((global_hidden[:1], depth_hidden[:1]), dim=-1).cpu()
+
+
+class MiniMaxMusic3ReferenceAdapter:
+    def __init__(self, dav_encoder: DAVEncoderOnly, rvq_encoder: MiniMaxMusicRVQEncoder):
+        self.dav_encoder = dav_encoder.eval()
+        self.rvq_encoder = rvq_encoder.eval()
+
+    @classmethod
+    def from_files(cls, encoder_file: str | Path, config_file: str | Path, dav_file: str | Path):
+        import json
+
+        from safetensors.torch import load_file
+
+        config = RVQEncoderConfig.from_dict(json.loads(Path(config_file).read_text(encoding="utf-8")))
+        rvq_encoder = MiniMaxMusicRVQEncoder(config)
+        rvq_encoder.load_state_dict(load_file(str(encoder_file)), strict=True)
+        dav_encoder = DAVEncoderOnly()
+        checkpoint = torch.load(dav_file, map_location="cpu", weights_only=True)
+        dav_state = {key: value for key, value in checkpoint.items() if key.startswith(("encoder.", "mean_proj."))}
+        dav_encoder.load_state_dict(dav_state, strict=True)
+        return cls(dav_encoder, rvq_encoder)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        repo_id: str = COLLECTION_REPO_ID,
+        *,
+        encoder_file: str = DEFAULT_ENCODER_FILE,
+        revision: str | None = None,
+        cache_dir: str | None = None,
+        official_model_revision: str | None = None,
+    ):
+        from huggingface_hub import hf_hub_download
+
+        encoder_path = hf_hub_download(repo_id, encoder_file, revision=revision, cache_dir=cache_dir)
+        config_path = hf_hub_download(
+            repo_id,
+            str(Path(encoder_file).with_suffix(".json")),
+            revision=revision,
+            cache_dir=cache_dir,
+        )
+        dav_path = hf_hub_download(
+            OFFICIAL_MODEL_REPO_ID,
+            DAV_FILE,
+            revision=official_model_revision,
+            cache_dir=cache_dir,
+        )
+        return cls.from_files(encoder_path, config_path, dav_path)
+
+    @staticmethod
+    def _resample(waveform: torch.Tensor, sample_rate: int) -> torch.Tensor:
+        if waveform.ndim == 1:
+            waveform = waveform.unsqueeze(0)
+        if waveform.ndim == 3:
+            if waveform.shape[0] != 1:
+                raise ValueError("reference audio batch size must be one")
+            waveform = waveform[0]
+        if waveform.ndim != 2:
+            raise ValueError("waveform must have shape [samples], [channels, samples], or [1, channels, samples]")
+        if sample_rate != SAMPLE_RATE:
+            import torchaudio
+
+            waveform = torchaudio.functional.resample(waveform, sample_rate, SAMPLE_RATE)
+        return waveform.float()
+
+    @torch.inference_mode()
+    def predict_codes(
+        self,
+        waveform: torch.Tensor,
+        sample_rate: int,
+        *,
+        device: str | torch.device | None = None,
+        encoder_dtype: torch.dtype | None = None,
+    ) -> torch.Tensor:
+        waveform = self._resample(waveform, sample_rate)
+        original_samples = waveform.shape[-1]
+        if original_samples < SAMPLE_RATE // FRAME_RATE:
+            raise ValueError("reference audio must contain at least one 25 Hz frame")
+        device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+        encoder_dtype = encoder_dtype or (torch.bfloat16 if device.type == "cuda" else torch.float32)
+
+        self.dav_encoder.to(device=device, dtype=torch.float32)
+        latents = self.dav_encoder(waveform.unsqueeze(0).to(device=device, dtype=torch.float32))[0].cpu()
+        self.dav_encoder.to("cpu")
+
+        frame_count = int(original_samples * FRAME_RATE // SAMPLE_RATE)
+        bounds = frame_latent_starts(frame_count)
+        while frame_count and bounds[-1] > latents.shape[-1]:
+            frame_count -= 1
+            bounds = frame_latent_starts(frame_count)
+        if not frame_count:
+            raise ValueError("DAV encoding produced no complete reference frames")
+
+        window_size = self.rvq_encoder.config.max_position_embeddings
+        regular_starts = list(range(0, max(frame_count - window_size + 1, 0), window_size))
+        if frame_count >= window_size:
+            regular_starts.append(frame_count - window_size if not regular_starts else regular_starts[-1])
+            regular_starts = sorted(set(regular_starts))
+            tail = frame_count - window_size
+            if tail not in regular_starts:
+                regular_starts.append(tail)
+        else:
+            regular_starts = [0]
+
+        predictions = torch.empty((frame_count, 8), dtype=torch.long)
+        assigned = torch.zeros(frame_count, dtype=torch.bool)
+        self.rvq_encoder.to(device=device, dtype=encoder_dtype)
+        autocast = torch.autocast(device.type, dtype=encoder_dtype) if device.type == "cuda" else nullcontext()
+        with autocast:
+            for frame_start in regular_starts:
+                frame_end = min(frame_start + window_size, frame_count)
+                local_bounds = bounds[frame_start : frame_end + 1]
+                latent_start, latent_end = local_bounds[0], local_bounds[-1]
+                window_latents = latents[:, latent_start:latent_end].transpose(0, 1).to(device, dtype=encoder_dtype)
+                pool = build_pool_matrix(local_bounds).to(device)
+                logits = self.rvq_encoder(window_latents.unsqueeze(0), pool.unsqueeze(0))
+                predicted = torch.stack([head.argmax(dim=-1)[0] for head in logits], dim=-1).cpu()
+                take = ~assigned[frame_start:frame_end]
+                predictions[frame_start:frame_end][take] = predicted[take]
+                assigned[frame_start:frame_end][take] = True
+        self.rvq_encoder.to("cpu")
+        if not assigned.all():
+            raise RuntimeError("RVQ window inference did not cover every reference frame")
+        return predictions
+
+    def encode_reference(
+        self,
+        pipeline,
+        waveform: torch.Tensor,
+        sample_rate: int,
+        *,
+        prompt: str,
+        lyrics: str,
+        generator: torch.Generator | None = None,
+        device: str | torch.device | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        codes = self.predict_codes(waveform, sample_rate, device=device)
+        frame_hiddens = replay_codes_diffusers(
+            pipeline,
+            codes,
+            prompt=prompt,
+            lyrics=lyrics,
+            generator=generator,
+        )
+        return frame_hiddens, codes
+
+
+def install_diffusers_reference_adapter() -> None:
+    """Allow the official modular pipeline to accept precomputed frame_hiddens."""
+    from diffusers.modular_pipelines.minimax_music3.encoders import (
+        MiniMaxMusic3AutoregressiveStep,
+        MiniMaxMusic3TokenizeStep,
+    )
+    from diffusers.modular_pipelines.modular_pipeline_utils import InputParam
+
+    if getattr(MiniMaxMusic3TokenizeStep, "_simpletuner_reference_patch", False):
+        return
+
+    original_tokenize_call = MiniMaxMusic3TokenizeStep.__call__
+    original_generate_call = MiniMaxMusic3AutoregressiveStep.__call__
+    original_tokenize_inputs = MiniMaxMusic3TokenizeStep.inputs.fget
+    original_generate_inputs = MiniMaxMusic3AutoregressiveStep.inputs.fget
+
+    def tokenize_inputs(self):
+        inputs = original_tokenize_inputs(self)
+        for value in inputs:
+            if value.name in {"prompt", "lyrics"}:
+                value.required = False
+        inputs.append(InputParam("frame_hiddens", default=None, type_hint=torch.Tensor))
+        return inputs
+
+    def generate_inputs(self):
+        inputs = original_generate_inputs(self)
+        for value in inputs:
+            if value.name == "text_ids":
+                value.required = False
+        inputs.append(InputParam("frame_hiddens", default=None, type_hint=torch.Tensor))
+        return inputs
+
+    def tokenize_call(self, components, state):
+        block_state = self.get_block_state(state)
+        if getattr(block_state, "frame_hiddens", None) is not None:
+            return components, state
+        return original_tokenize_call(self, components, state)
+
+    def generate_call(self, components, state):
+        block_state = self.get_block_state(state)
+        if getattr(block_state, "frame_hiddens", None) is not None:
+            return components, state
+        return original_generate_call(self, components, state)
+
+    MiniMaxMusic3TokenizeStep.inputs = property(tokenize_inputs)
+    MiniMaxMusic3AutoregressiveStep.inputs = property(generate_inputs)
+    MiniMaxMusic3TokenizeStep.__call__ = tokenize_call
+    MiniMaxMusic3AutoregressiveStep.__call__ = generate_call
+    MiniMaxMusic3TokenizeStep._simpletuner_reference_patch = True
+    MiniMaxMusic3AutoregressiveStep._simpletuner_reference_patch = True