# Copyright 2026 The MiniMax Team and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import numpy as np import torch from ...configuration_utils import FrozenDict from ...guiders import ClassifierFreeGuidance from ...models import MiniMaxMusic3ConditionEncoder, MiniMaxMusic3Transformer1DModel from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import logging from ...utils.torch_utils import randn_tensor from ..modular_pipeline import ( BlockState, LoopSequentialPipelineBlocks, ModularPipelineBlocks, PipelineState, ) from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .before_denoise import _CHUNK_FRAMES from .modular_pipeline import MiniMaxMusic3ModularPipeline logger = logging.get_logger(__name__) # pylint: disable=invalid-name # Neighboring windows overlap by ~344 latent frames (100-frame hop at ~3.445 latents per frame); only the first 172 # of them are spliced and blended, from the previous window's carry spanning latent frames [L - 344, L - 172). _OVERLAP_LATENT_LENGTH = 172 class MiniMaxMusic3ChunkConditionStep(ModularPipelineBlocks): model_name = "minimax-music3" @property def description(self) -> str: return ( "Chunk conditioning step that projects the window's per-frame hidden states onto the Flow-VAE latent " "timeline and splices in the previous window's conditioning over the overlapping latent frames." ) @property def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("condition_encoder", MiniMaxMusic3ConditionEncoder), ComponentSpec("transformer", MiniMaxMusic3Transformer1DModel), ] @property def inputs(self) -> list[InputParam]: return [ InputParam( "frame_hiddens", required=True, type_hint=torch.Tensor, description="Per-frame hidden states generated by the autoregressive step.", ), InputParam( "chunk_starts", required=True, type_hint=list, description="Frame index at which each 200-frame denoising window starts.", ), ] @torch.no_grad() def __call__(self, components: MiniMaxMusic3ModularPipeline, block_state: BlockState, k: int): device = components._execution_device chunk_start = block_state.chunk_starts[k] chunk_end = min(chunk_start + _CHUNK_FRAMES, block_state.frame_hiddens.shape[1]) condition = components.condition_encoder(block_state.frame_hiddens[:, chunk_start:chunk_end].to(device)) condition = condition.to(components.transformer.dtype) overlap = 0 if block_state.previous_latent is not None: overlap = min(block_state.previous_latent.shape[-1], condition.shape[1]) condition[:, :overlap] = block_state.previous_condition[:, :overlap] block_state.condition = condition block_state.overlap = overlap return components, block_state class MiniMaxMusic3ChunkPrepareLatentsStep(ModularPipelineBlocks): model_name = "minimax-music3" @property def description(self) -> str: return ( "Chunk latent preparation step that draws the window's initial noise and snapshots the noise over the " "overlapping latent frames as the blending prompt." ) @property def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("transformer", MiniMaxMusic3Transformer1DModel)] @property def inputs(self) -> list[InputParam]: return [InputParam.template("generator")] @torch.no_grad() def __call__(self, components: MiniMaxMusic3ModularPipeline, block_state: BlockState, k: int): device = components._execution_device latents = randn_tensor( (1, components.num_channels_latents, block_state.condition.shape[1]), generator=block_state.generator, device=device, dtype=block_state.condition.dtype, ) block_state.noise_prompt = latents[..., : block_state.overlap].clone() if block_state.overlap > 0 else None block_state.latents = latents return components, block_state class MiniMaxMusic3ChunkSetTimestepsStep(ModularPipelineBlocks): model_name = "minimax-music3" @property def description(self) -> str: return "Chunk scheduler step that resets the flow-matching sigma schedule for the current window." @property def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] @property def inputs(self) -> list[InputParam]: return [ InputParam( "num_inference_steps", default=30, type_hint=int, description="Number of flow-matching Euler steps per chunk.", ), ] @torch.no_grad() def __call__(self, components: MiniMaxMusic3ModularPipeline, block_state: BlockState, k: int): device = components._execution_device sigmas = np.linspace(1.0, 1.0 / block_state.num_inference_steps, block_state.num_inference_steps) components.scheduler.set_timesteps(sigmas=sigmas, device=device) block_state.timesteps = components.scheduler.timesteps return components, block_state class MiniMaxMusic3ChunkDenoiseInner(ModularPipelineBlocks): model_name = "minimax-music3" @property def description(self) -> str: return ( "Inner denoising loop that flow-matches one window's latents over the scheduler timesteps. The guider " "manages the conditional/unconditional transformer passes (the unconditional conditioning is all zeros), " "and the overlapping latent frames are blended toward the previous window's trailing latents at every " "step so neighboring windows share their boundary." ) @property def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("transformer", MiniMaxMusic3Transformer1DModel), ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), ComponentSpec( "guider", ClassifierFreeGuidance, config=FrozenDict({"guidance_scale": 1.7}), default_creation_method="from_config", ), ] @property def inputs(self) -> list[InputParam]: return [ InputParam( "num_inference_steps", default=30, type_hint=int, description="Number of flow-matching Euler steps per chunk.", ), ] @torch.no_grad() def __call__(self, components: MiniMaxMusic3ModularPipeline, block_state: BlockState, k: int): latents = block_state.latents timesteps = block_state.timesteps overlap = block_state.overlap # The unconditional branch conditions on zeros, not on a re-encoded empty prompt. guider_inputs = { "encoder_hidden_states": (block_state.condition, torch.zeros_like(block_state.condition)), } for i, t in enumerate(timesteps): if overlap > 0: time_value = t.to(latents.dtype) latents[..., :overlap] = (1.0 - (1.0 - 1e-6) * time_value) * block_state.noise_prompt + ( time_value * block_state.previous_latent[..., :overlap] ) # The transformer consumes the scheduler timestep directly: flow-matching time in [0, 1], 0 = noise. timestep = t.expand(latents.shape[0]).to(latents.dtype) components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) guider_state = components.guider.prepare_inputs(guider_inputs) for guider_state_batch in guider_state: components.guider.prepare_models(components.transformer) cond_kwargs = {key: getattr(guider_state_batch, key) for key in guider_inputs} guider_state_batch.noise_pred = components.transformer( hidden_states=latents, timestep=timestep, return_dict=False, **cond_kwargs, )[0] components.guider.cleanup_models(components.transformer) velocity = components.guider(guider_state)[0] latents = components.scheduler.step(velocity, t, latents, return_dict=False)[0] block_state.progress_bar.update() block_state.latents = latents return components, block_state class MiniMaxMusic3ChunkUpdateStep(ModularPipelineBlocks): model_name = "minimax-music3" @property def description(self) -> str: return ( "Post-denoising update step that restores the previous window's latents over the overlap, appends the " "window's latents to the chunk list, and carries the trailing latents and conditioning to the next window." ) @torch.no_grad() def __call__(self, components: MiniMaxMusic3ModularPipeline, block_state: BlockState, k: int): latents = block_state.latents if block_state.overlap > 0: latents[..., : block_state.overlap] = block_state.previous_latent[..., : block_state.overlap] overlap_start = max(0, latents.shape[-1] - 2 * _OVERLAP_LATENT_LENGTH) overlap_end = max(overlap_start, latents.shape[-1] - _OVERLAP_LATENT_LENGTH) block_state.previous_latent = latents[..., overlap_start:overlap_end] block_state.previous_condition = block_state.condition[:, overlap_start:overlap_end] block_state.latent_chunks.append(latents) return components, block_state class MiniMaxMusic3ChunkLoopWrapper(LoopSequentialPipelineBlocks): model_name = "minimax-music3" @property def description(self) -> str: return ( "Pipeline block that iterates over the 200-frame denoising windows. At each window it runs sub-blocks " "for conditioning, latent preparation, scheduler reset, denoising, and the overlap carry update." ) @property def loop_inputs(self) -> list[InputParam]: return [ InputParam( "chunk_starts", required=True, type_hint=list, description="Frame index at which each 200-frame denoising window starts.", ), ] @property def loop_intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "latent_chunks", type_hint=list, description="List of per-window denoised latent tensors (uncropped).", ), ] @torch.no_grad() def __call__(self, components: MiniMaxMusic3ModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) block_state.latent_chunks = [] block_state.previous_latent = None block_state.previous_condition = None num_chunks = len(block_state.chunk_starts) with self.progress_bar(total=num_chunks * block_state.num_inference_steps) as progress_bar: block_state.progress_bar = progress_bar for k in range(num_chunks): components, block_state = self.loop_step(components, block_state, k=k) block_state.progress_bar = None self.set_block_state(state, block_state) return components, state class MiniMaxMusic3ChunkDenoiseStep(MiniMaxMusic3ChunkLoopWrapper): block_classes = [ MiniMaxMusic3ChunkConditionStep, MiniMaxMusic3ChunkPrepareLatentsStep, MiniMaxMusic3ChunkSetTimestepsStep, MiniMaxMusic3ChunkDenoiseInner, MiniMaxMusic3ChunkUpdateStep, ] block_names = ["condition", "prepare_latents", "set_timesteps", "denoise_inner", "update_chunk"] @property def description(self) -> str: return ( "Chunk denoise step that iterates over the 200-frame denoising windows.\n" "At each window: condition -> prepare_latents -> set_timesteps -> denoise_inner -> update_chunk." )