MiniMax Music 3

"""MiniMax Music 3, generated from a declarative template."""

from __future__ import annotations

import base64
import tempfile

from pydantic import BaseModel, ConfigDict, Field

from runware_serverless import serve

MODEL = "MiniMaxAI/MiniMax-Music3"
REVISION = "1b525ccb685ae809f33027c51c351eeb5833cf72"
CACHE = "/weights"
DEFAULT_DURATION = 60


class Request(BaseModel):
    # An input this model does not read is a line the caller believed was doing
    # something, so it is refused rather than dropped. Without this, pydantic
    # ignores it: a width sent to a model of fixed size, or a misspelled seed,
    # both come back as a result that silently answered a different request.
    model_config = ConfigDict(extra="forbid")

    prompt: str = Field(description="The style, instrumentation and mood to write in.")
    lyrics: str = Field(default="", description="The words to sing. Left empty, the track is instrumental.")
    duration: float = Field(
        default=DEFAULT_DURATION,
        ge=10,
        le=240,
        description="How long the track runs, in seconds.",
    )
    seed: int | None = Field(default=None, description="Repeat a result by repeating its seed.")


class Generated(BaseModel):
    audio: str = Field(description="The result as a base64 WAV, stereo.")
    duration: float
    sample_rate: int
    seed: int | None


@serve
class GeneratedMusic:
    def load(self) -> None:
        if hasattr(self, "pipeline"):
            return

        # Imported here rather than at the top of the file. The build imports
        # this file with the platform's own dependencies, so nothing above
        # exists until a worker runs.
        import torch
        from diffusers import ModularPipeline

        self.torch = torch

        self.pipeline = ModularPipeline.from_pretrained(
            MODEL, revision=REVISION, cache_dir=CACHE
        )
        # Every component stays on the card. A worker's host RAM is half its
        # card's VRAM, so parking anything on the host runs out of memory long
        # before the card does, and the kill arrives as the container's exit 137
        # rather than as a CUDA error.
        self.pipeline.load_components(dtype=torch.bfloat16, device_map="cuda")

        # load_components reports a component it could not build as a warning
        # and leaves it None, so a wrong argument above does not fail the load:
        # it surfaces one request later as an attribute error on None, far from
        # its cause. Fail here instead, naming what is missing. Every declared
        # component is meant to load, because this call names no workflow.
        missing = self.pipeline.null_component_names
        if missing:
            raise RuntimeError(
                "these components did not load, and the warnings above this line"
                " say why: " + ", ".join(missing)
            )

    def predict(self, request: Request) -> Generated:
        import soundfile

        generator = None
        if request.seed is not None:
            generator = self.torch.Generator().manual_seed(request.seed)

        results = self.pipeline(
            prompt=request.prompt,
            lyrics=request.lyrics,
            audio_duration=request.duration,
            generator=generator,
            output_type="np",
        )

        track = _audios(results)[0]
        rate = self.pipeline.sampling_rate

        # soundfile writes frames by channel, and the pipeline hands back
        # channels by frames, so the track arrives transposed or it is written as
        # one channel of twice the length.
        with tempfile.NamedTemporaryFile(suffix=".wav") as handle:
            soundfile.write(handle.name, track.T, rate, format="WAV")
            handle.seek(0)
            encoded = base64.b64encode(handle.read()).decode("ascii")

        # What came out, not what was asked for: the model rounds the length to
        # its own frame rate.
        return Generated(
            audio=encoded,
            duration=track.shape[-1] / rate,
            sample_rate=rate,
            seed=request.seed,
        )


def _audios(results: object) -> object:
    """The waveform, however the pipeline hands it back."""
    if isinstance(results, dict):
        return results["audios"]
    if hasattr(results, "audios"):
        return results.audios

    return results

configVersion: 1
id: minimax-music3
name: MiniMax Music 3
source:
  type: code
  entry: model.py
compute:
  gpuType: l40s
build:
  requirements:
    - torch==2.14.1
    - torchvision==0.29.1
    - diffusers==0.42.1
    - transformers==5.21.0
    - accelerate==1.15.0
    - soundfile==0.14.0
volumes:
  - /weights