Keep only periodic c0 reference constraints

README.md

@@ -27,21 +27,16 @@ The project started with a 41M-parameter, single-GPU community proof at **0.6633
 
 Use **v4** unless reproducing an experiment.
 
+The released ComfyUI and Diffusers adapters support one reference-generation method. Every fifth generated semantic `c0` token is restricted to the encoder's top-5 candidates by default. MiniMax chooses the token and generates all acoustic codebooks. The interval is configurable from 1 through 10.
+
 ## Files
 
-| Version | Parameters | Experiment | Replay cosine |
+| File | Parameters | Experiment | Replay cosine |
 |---|---:|---|---:|
-| v1 | 40,978,944 | Baseline; eight independent heads | 0.7624* |
-| v2 | 154,736,064 | Wider shared encoder | 0.7698 |
-| v3 | 154,736,064 | v2 plus training-only MERT alignment | 0.7703 |
-| **v4** | 169,008,576 | Causal acoustic decoder across codebook depth | **0.8748** |
-
-Weight files in [`encoders/`](encoders/):
-
-- **v1**: `minimax_music3_rvq_encoder_v1_41m_independent_heads.safetensors`
-- **v2**: `minimax_music3_rvq_encoder_v2_155m_wide_independent_heads.safetensors`
-- **v3**: `minimax_music3_rvq_encoder_v3_155m_mert_aligned_independent_heads.safetensors`
-- **v4**: `minimax_music3_rvq_encoder_v4_169m_autoregressive_depth_recommended.safetensors`
+| `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/).
 
@@ -136,7 +131,14 @@ 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**.
+Restart ComfyUI. Upload a reference audio file and select v4 in **MiniMax Music3 RVQ Reference Encoder Loader**.
+
+- [`comfyui_workflow_example.json`](comfyui_workflow_example.json) constrains every fifth generated semantic `c0` token to the encoder's top-5 candidates.
+- `reference_interval=1` constrains every frame.
+- `reference_interval=5` is the tested default.
+- `reference_interval=10` constrains every tenth frame and gives the language model more freedom.
+
+The seven acoustic codebooks are generated by MiniMax. They are not copied from the reference. Describe the target arrangement in `caption`. Provide the desired sectioned lyrics. Prompt adherence and audio quality vary; this is not a general audio-to-audio conversion system.
 
 The node package reads the encoder files directly from this clone. They can instead be placed in:
 
@@ -144,7 +146,7 @@ The node package reads the encoder files directly from this clone. They can inst
 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.
+The interval-5 graph was verified with a 30-second reference, five Euler steps, the pruned int8 text encoder, the fp16 diffusion model, and the v4 encoder. It produced a 29.99-second stereo 44.1 kHz FLAC. Use 30 diffusion steps for normal output.
 
 ## Diffusers
 
@@ -190,6 +192,7 @@ frame_hiddens, predicted_codes = adapter.encode_reference(
     lyrics="[instrumental]",
     generator=generator,
     device="cuda",
+    reference_interval=5,
 )
 
 result = pipe(
@@ -200,7 +203,7 @@ result = pipe(
 )
 ```
 
-The patch only adds a precomputed-`frame_hiddens` bypass to the modular pipeline. It does not replace MiniMax model code.
+`reference_interval` accepts integers from 1 through 10. The patch only adds a precomputed-`frame_hiddens` bypass to the modular pipeline. It does not replace MiniMax model code.
 
 ## Limits
 
@@ -209,8 +212,8 @@ The patch only adds a precomputed-`frame_hiddens` bypass to the modular pipeline
 - 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.
+- Reference generation still needs the official MiniMax language model and RVQ depth decoder.
+- The released integration uses the RVQ encoder's top-5 semantic candidates. Other candidate counts are not exposed.
 
 ## Credits
 

comfyui_open_rvq/nodes.py

@@ -15,9 +15,9 @@ 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.ldm.minimax_music.ar import AUDIO_FRAMES_PER_SECOND, CFG_SCALE, CFG_TOP_K, MiniMaxMusic3AR, sample_topk
+from comfy.ldm.minimax_music.prompt import AUDIO_CODE_OFFSET
 from comfy.text_encoders.minimax_music import MiniMaxMusic3TEModel
 
 
@@ -29,30 +29,18 @@ from minimax_music3_reference_adapter import MiniMaxMusic3ReferenceAdapter  # no
 
 
 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)
+REFERENCE_CANDIDATE_COUNT = 5
+DEFAULT_REFERENCE_INTERVAL = 5
+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(
@@ -71,25 +59,59 @@ def _run_depth(model: MiniMaxMusic3AR, device, execution_dtype, core):
     comfy.model_prefetch.prefetch_queue_pop(queue, device, None)
 
 
+def _sample_semantic_candidates(model, hidden, semantic_candidates, cfg_scale, top_k, generator):
+    if model.model.pruned_lm_head:
+        logits = model.model.lm_head_pruned(hidden).float()
+        token_candidates = semantic_candidates + 1
+        offset = 1
+    else:
+        logits = model.model.lm_head(hidden).float()
+        token_candidates = semantic_candidates + AUDIO_CODE_OFFSET
+        offset = AUDIO_CODE_OFFSET
+    token_candidates = token_candidates.to(device=logits.device, dtype=torch.long).unsqueeze(0)
+    conditioned = logits[:1].gather(-1, token_candidates)
+    unconditioned = logits[1:2].gather(-1, token_candidates)
+    guided = unconditioned + (conditioned - unconditioned) * cfg_scale
+    selected_index = sample_topk(guided, min(top_k, guided.shape[-1]), generator).view(1, 1)
+    return token_candidates.gather(-1, selected_index).squeeze(-1) - offset
+
+
 @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}]")
+def replay_reference_codes(
+    model,
+    input_ids,
+    semantic_candidates,
+    seed,
+    device,
+    cfg_scale,
+    top_k,
+    reference_interval,
+):
+    if semantic_candidates.ndim != 2 or semantic_candidates.shape[0] == 0:
+        raise ValueError("semantic candidates must have shape [frames, candidates]")
+    if semantic_candidates.shape[1] != REFERENCE_CANDIDATE_COUNT:
+        raise ValueError(f"semantic candidates must contain exactly {REFERENCE_CANDIDATE_COUNT} codes per frame")
+    if not 1 <= reference_interval <= 10:
+        raise ValueError("reference_interval must be between 1 and 10")
+
+    frame_count = semantic_candidates.shape[0]
     prompt_tokens = int(input_ids.shape[1])
     input_ids = input_ids.to(device)
-    codes = codes.to(device=device, dtype=torch.long)
+    semantic_candidates = semantic_candidates.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
+    from comfy.ldm.minimax_music.prompt import 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)
+
+    def embed_text(token_ids):
+        if model.model.pruned_embedding:
+            return model.model.embed_tokens_prefill(token_ids, out_dtype=execution_dtype)
+        return model.model.embed_tokens(token_ids, out_dtype=execution_dtype)
+
+    past = model.model.init_kv_cache(2, prompt_tokens + frame_count + 2, device, execution_dtype)
+    output = model.model(None, embeds=embed_text(text_ids), past_key_values=past, dtype=execution_dtype)
     last_hidden, past = output[0][:, -1], output[2]
     from comfy.ldm.minimax_music.ar import derive_seed
 
@@ -123,19 +145,41 @@ def replay_reference_codes(model: MiniMaxMusic3AR, input_ids, codes, seed, devic
     output = model.model(None, embeds=feedback, past_key_values=past, dtype=execution_dtype)
     last_hidden, past = output[0][:, -1], output[2]
 
+    all_semantic_candidates = torch.arange(16384, device=device, dtype=torch.long)
     frames = []
-    for frame_index in range(codes.shape[0]):
+    for frame_index in range(frame_count):
         comfy.model_management.throw_exception_if_processing_interrupted()
-        frame_codes = codes[frame_index].unsqueeze(0).repeat(2, 1)
+        frame_candidates = (
+            semantic_candidates[frame_index]
+            if frame_index % reference_interval == 0
+            else all_semantic_candidates
+        )
+        c0 = _sample_semantic_candidates(
+            model,
+            last_hidden,
+            frame_candidates,
+            cfg_scale,
+            top_k,
+            generator,
+        ).repeat(2)
+        c0_embed = model._embed_c0(c0, execution_dtype)
         result: dict[str, torch.Tensor] = {}
 
-        def teacher_core():
-            result["hidden"] = _teacher_depth_hidden(model, last_hidden, frame_codes, execution_dtype)
+        def depth_core():
+            result["codes"], result["hidden"] = model._depth_codes(
+                last_hidden,
+                c0,
+                c0_embed,
+                generator,
+                execution_dtype,
+                cfg_scale,
+                top_k,
+            )
 
-        _run_depth(model, device, execution_dtype, teacher_core)
+        _run_depth(model, device, execution_dtype, depth_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)
+        if frame_index + 1 < frame_count:
+            feedback = model._embed_audio_frame(result["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)
@@ -145,19 +189,20 @@ _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:
+    semantic_candidates = token_weight_pairs.get("minimax_reference_semantic_candidates")
+    if semantic_candidates 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,
+        semantic_candidates,
         int(token_weight_pairs["seed"]),
         self.execution_device,
         float(token_weight_pairs["cfg_scale"]),
         int(token_weight_pairs["top_k"]),
+        int(token_weight_pairs["minimax_reference_interval"]),
     )
     return hidden.unsqueeze(0), None, {}
 
@@ -170,7 +215,11 @@ if not getattr(MiniMaxMusic3TEModel, "_simpletuner_reference_patch", False):
 class MiniMaxMusic3RVQReferenceEncoderLoader:
     @classmethod
     def INPUT_TYPES(cls):
-        encoders = [name for name in folder_paths.get_filename_list(ENCODER_FOLDER) if name.lower().endswith(".safetensors")]
+        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,)}}
 
@@ -199,6 +248,7 @@ class MiniMaxMusic3ReferenceAudioEncode:
                 "caption": ("STRING", {"multiline": True, "dynamicPrompts": True}),
                 "lyrics": ("STRING", {"multiline": True, "dynamicPrompts": True}),
                 "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
+                "reference_interval": ("INT", {"default": DEFAULT_REFERENCE_INTERVAL, "min": 1, "max": 10, "step": 1}),
             },
             "optional": {
                 "cfg_scale": ("FLOAT", {"default": CFG_SCALE, "min": 0.0, "max": 100.0, "step": 0.1}),
@@ -211,10 +261,22 @@ class MiniMaxMusic3ReferenceAudioEncode:
     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(
+    def encode(
+        self,
+        clip,
+        reference_encoder,
+        audio,
+        caption,
+        lyrics,
+        seed,
+        reference_interval=DEFAULT_REFERENCE_INTERVAL,
+        cfg_scale=CFG_SCALE,
+        top_k=CFG_TOP_K,
+    ):
+        codes, semantic_candidates = reference_encoder.predict_codes_with_semantic_candidates(
             audio["waveform"],
             int(audio["sample_rate"]),
+            semantic_top_k=REFERENCE_CANDIDATE_COUNT,
             device=comfy.model_management.get_torch_device(),
         )
         tokens = clip.tokenize(
@@ -225,7 +287,8 @@ class MiniMaxMusic3ReferenceAudioEncode:
             cfg_scale=cfg_scale,
             top_k=top_k,
         )
-        tokens["minimax_reference_codes"] = codes
+        tokens["minimax_reference_semantic_candidates"] = semantic_candidates
+        tokens["minimax_reference_interval"] = reference_interval
         conditioning = clip.encode_from_tokens_scheduled(tokens)
         for cond in conditioning:
             hidden = cond[0]

comfyui_workflow_example.json

@@ -369,7 +369,7 @@
       ],
       "size": [
         520,
-        640
+        560
       ],
       "flags": {},
       "order": 8,
@@ -418,6 +418,14 @@
             "name": "seed"
           }
         },
+        {
+          "localized_name": "reference_interval",
+          "name": "reference_interval",
+          "type": "INT",
+          "widget": {
+            "name": "reference_interval"
+          }
+        },
         {
           "localized_name": "cfg_scale",
           "name": "cfg_scale",
@@ -459,9 +467,10 @@
         "Node name for S&R": "MiniMaxMusic3ReferenceAudioEncode"
       },
       "widgets_values": [
-        "rock",
+        "1980s heavy metal cover, distorted electric guitars, driving bass, double-kick drums, dramatic arena production",
         "[instrumental]",
         299871780,
+        5,
         1.5,
         50
       ]

minimax_music3_reference_adapter.py

@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Reference-audio conditioning for MiniMax Music 3.
+"""Reference-audio constrained generation 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.
+per 25 Hz frame, and periodically constrains the official MiniMax Music 3
+language model to the encoder's top-5 semantic candidates.
 """
 
 from __future__ import annotations
@@ -52,7 +52,8 @@ AUDIO_CFG_TOKEN_ID = 151_654
 SEMANTIC_VOCAB_SIZE = 16_384
 AR_CFG_SCALE = 1.5
 AR_TOP_K = 50
-LM_BLOCK_FRAMES = 256
+REFERENCE_CANDIDATE_COUNT = 5
+DEFAULT_REFERENCE_INTERVAL = 5
 
 _SPECIAL_TAG_RE = re.compile(r"<\|([^|]*)\|>")
 _LEADING_TAGS_RE = re.compile(r"^[ \t]*((?:\[[^\]]+\][ \t]*)+)")
@@ -406,26 +407,43 @@ def build_text_ids(tokenizer, prompt: str, lyrics: str, device: torch.device) ->
     return torch.cat((input_ids, unconditional), dim=0).to(device)
 
 
-def _sample_top_k(logits: torch.Tensor, generator: torch.Generator | None) -> torch.Tensor:
+def _sample_top_k(
+    logits: torch.Tensor,
+    generator: torch.Generator | None,
+    top_k: int = AR_TOP_K,
+) -> 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]
+    threshold = torch.topk(values, min(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:
+def _sample_official_depth_codes(
+    language_model,
+    depth_decoder,
+    hidden: torch.Tensor,
+    semantic_codes: torch.Tensor,
+    generator: torch.Generator | None,
+    cfg_scale: float = AR_CFG_SCALE,
+    top_k: int = AR_TOP_K,
+) -> tuple[torch.Tensor, torch.Tensor]:
     sequence = [depth_decoder.projection(hidden).unsqueeze(1)]
-    semantic = language_model.model.embed_tokens(codes[:, 0] + AUDIO_CODE_OFFSET)
+    semantic = language_model.model.embed_tokens(semantic_codes + AUDIO_CODE_OFFSET)
     sequence.append(depth_decoder.projection(semantic).unsqueeze(1))
+    codes = [semantic_codes]
     hidden_parts = []
-    for index in range(1, codes.shape[1]):
+    for index in range(1, 8):
         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)
+        hidden_parts.append(depth_hidden[:1])
+        logits = depth_decoder.audio_heads[index - 1](depth_hidden).float()
+        guided = logits[1:2] + (logits[:1] - logits[1:2]) * cfg_scale
+        code = _sample_top_k(guided, generator, top_k).repeat(2)
+        codes.append(code)
+        if index < 7:
+            embedding = depth_decoder.audio_embeddings(code + (index - 1) * 1024)
             sequence.append(depth_decoder.projection(embedding).unsqueeze(1))
-    return torch.cat(hidden_parts, dim=-1)
+    return torch.stack(codes, dim=1), torch.cat(hidden_parts, dim=-1)
 
 
 def _embed_official_codes(language_model, depth_decoder, codes: torch.Tensor) -> torch.Tensor:
@@ -435,22 +453,53 @@ def _embed_official_codes(language_model, depth_decoder, codes: torch.Tensor) ->
     return (semantic + acoustic.to(semantic.dtype)) * (codes.shape[-1] ** -0.5)
 
 
+def _sample_semantic_candidates_diffusers(
+    language_model,
+    hidden: torch.Tensor,
+    semantic_candidates: torch.Tensor,
+    cfg_scale: float,
+    top_k: int,
+    generator: torch.Generator | None,
+) -> torch.Tensor:
+    token_candidates = semantic_candidates.to(device=hidden.device, dtype=torch.long) + AUDIO_CODE_OFFSET
+    token_candidates = token_candidates.unsqueeze(0)
+    logits = language_model.lm_head(hidden).float()
+    conditioned = logits[:1].gather(-1, token_candidates)
+    unconditioned = logits[1:2].gather(-1, token_candidates)
+    guided = unconditioned + (conditioned - unconditioned) * cfg_scale
+    selected_index = _sample_top_k(guided, generator, min(top_k, guided.shape[-1])).view(1, 1)
+    return token_candidates.gather(-1, selected_index).squeeze(-1) - AUDIO_CODE_OFFSET
+
+
 @torch.inference_mode()
 def replay_codes_diffusers(
     pipeline,
     codes: torch.Tensor,
+    semantic_candidates: torch.Tensor,
     *,
     prompt: str,
     lyrics: str,
     generator: torch.Generator | None = None,
+    reference_interval: int = DEFAULT_REFERENCE_INTERVAL,
+    cfg_scale: float = AR_CFG_SCALE,
+    top_k: int = AR_TOP_K,
 ) -> torch.Tensor:
-    """Replay predicted codes through official Diffusers LM components."""
+    """Generate a fixed-length rollout with periodic top-5 semantic constraints."""
     if codes.ndim != 2 or codes.shape[1] != 8 or codes.shape[0] == 0:
         raise ValueError("codes must have shape [frames, 8]")
+    if semantic_candidates.ndim != 2 or semantic_candidates.shape != (codes.shape[0], REFERENCE_CANDIDATE_COUNT):
+        raise ValueError(
+            f"semantic_candidates must have shape [frames, {REFERENCE_CANDIDATE_COUNT}]"
+        )
+    if not 1 <= reference_interval <= 10:
+        raise ValueError("reference_interval must be between 1 and 10")
+    if not 1 <= top_k <= SEMANTIC_VOCAB_SIZE:
+        raise ValueError(f"top_k must be between 1 and {SEMANTIC_VOCAB_SIZE}")
     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)
+    semantic_candidates = semantic_candidates.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
@@ -461,56 +510,64 @@ def replay_codes_diffusers(
     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)
+    guided = unconditioned + (conditioned - unconditioned) * cfg_scale
+    threshold = torch.topk(conditioned, top_k, dim=-1).values[..., -1, None]
+    warmup_token = _sample_top_k(guided.masked_fill(conditioned < threshold, -torch.inf), generator, top_k)
     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)
+    warmup_codes, _ = _sample_official_depth_codes(
+        language_model,
+        depth_decoder,
+        last_hidden,
+        warmup_semantic,
+        generator,
+        cfg_scale=cfg_scale,
+        top_k=top_k,
+    )
     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):
+    last_hidden = output.last_hidden_state[:, -1]
+    all_semantic_candidates = torch.arange(SEMANTIC_VOCAB_SIZE, device=device, dtype=torch.long)
+    hidden_frames = []
+    for frame_index in range(codes.shape[0]):
+        frame_candidates = (
+            semantic_candidates[frame_index]
+            if frame_index % reference_interval == 0
+            else all_semantic_candidates
+        )
+        semantic_code = _sample_semantic_candidates_diffusers(
+            language_model,
+            last_hidden,
+            frame_candidates,
+            cfg_scale,
+            top_k,
+            generator,
+        ).repeat(2)
+        sampled_codes, depth_hidden = _sample_official_depth_codes(
+            language_model,
+            depth_decoder,
+            last_hidden,
+            semantic_code,
+            generator,
+            cfg_scale=cfg_scale,
+            top_k=top_k,
+        )
+        hidden_frames.append(torch.cat((last_hidden[:1], depth_hidden), dim=-1).cpu())
+        if frame_index + 1 < codes.shape[0]:
             output = language_model.model(
-                inputs_embeds=feedback[:, start : start + LM_BLOCK_FRAMES],
+                inputs_embeds=_embed_official_codes(language_model, depth_decoder, sampled_codes).unsqueeze(1),
                 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()
+            last_hidden = output.last_hidden_state[:, -1]
+    return torch.cat(hidden_frames, dim=0).unsqueeze(0)
 
 
 class MiniMaxMusic3ReferenceAdapter:
@@ -577,14 +634,19 @@ class MiniMaxMusic3ReferenceAdapter:
         return waveform.float()
 
     @torch.inference_mode()
-    def predict_codes(
+    def _predict_codes(
         self,
         waveform: torch.Tensor,
         sample_rate: int,
         *,
         device: str | torch.device | None = None,
         encoder_dtype: torch.dtype | None = None,
-    ) -> torch.Tensor:
+        semantic_top_k: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if semantic_top_k is not None and not 1 <= semantic_top_k <= self.rvq_encoder.config.codebook_vocab_sizes[0]:
+            raise ValueError(
+                f"semantic_top_k must be between 1 and {self.rvq_encoder.config.codebook_vocab_sizes[0]}"
+            )
         waveform = self._resample(waveform, sample_rate)
         original_samples = waveform.shape[-1]
         if original_samples < SAMPLE_RATE // FRAME_RATE:
@@ -616,6 +678,9 @@ class MiniMaxMusic3ReferenceAdapter:
             regular_starts = [0]
 
         predictions = torch.empty((frame_count, 8), dtype=torch.long)
+        semantic_candidates = (
+            torch.empty((frame_count, semantic_top_k), dtype=torch.long) if semantic_top_k is not None else None
+        )
         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()
@@ -628,14 +693,56 @@ class MiniMaxMusic3ReferenceAdapter:
                 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()
+                window_semantic_candidates = (
+                    logits[0].topk(semantic_top_k, dim=-1).indices[0].cpu()
+                    if semantic_top_k is not None
+                    else None
+                )
                 take = ~assigned[frame_start:frame_end]
                 predictions[frame_start:frame_end][take] = predicted[take]
+                if semantic_candidates is not None:
+                    semantic_candidates[frame_start:frame_end][take] = window_semantic_candidates[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, semantic_candidates
+
+    def predict_codes(
+        self,
+        waveform: torch.Tensor,
+        sample_rate: int,
+        *,
+        device: str | torch.device | None = None,
+        encoder_dtype: torch.dtype | None = None,
+    ) -> torch.Tensor:
+        predictions, _ = self._predict_codes(
+            waveform,
+            sample_rate,
+            device=device,
+            encoder_dtype=encoder_dtype,
+        )
         return predictions
 
+    def predict_codes_with_semantic_candidates(
+        self,
+        waveform: torch.Tensor,
+        sample_rate: int,
+        *,
+        semantic_top_k: int = 5,
+        device: str | torch.device | None = None,
+        encoder_dtype: torch.dtype | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        predictions, semantic_candidates = self._predict_codes(
+            waveform,
+            sample_rate,
+            device=device,
+            encoder_dtype=encoder_dtype,
+            semantic_top_k=semantic_top_k,
+        )
+        assert semantic_candidates is not None
+        return predictions, semantic_candidates
+
     def encode_reference(
         self,
         pipeline,
@@ -646,14 +753,26 @@ class MiniMaxMusic3ReferenceAdapter:
         lyrics: str,
         generator: torch.Generator | None = None,
         device: str | torch.device | None = None,
+        reference_interval: int = DEFAULT_REFERENCE_INTERVAL,
+        cfg_scale: float = AR_CFG_SCALE,
+        top_k: int = AR_TOP_K,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        codes = self.predict_codes(waveform, sample_rate, device=device)
+        codes, semantic_candidates = self.predict_codes_with_semantic_candidates(
+            waveform,
+            sample_rate,
+            semantic_top_k=REFERENCE_CANDIDATE_COUNT,
+            device=device,
+        )
         frame_hiddens = replay_codes_diffusers(
             pipeline,
             codes,
+            semantic_candidates,
             prompt=prompt,
             lyrics=lyrics,
             generator=generator,
+            reference_interval=reference_interval,
+            cfg_scale=cfg_scale,
+            top_k=top_k,
         )
         return frame_hiddens, codes