YuE2 3B

"""YuE2 3B, 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 = "m-a-p/YuE2-3B"
REVISION = "8f24312f187763b9854adeec87e437e2b03bbff1"
VAE = "m-a-p/YuE2-Vae"
VAE_REVISION = "152733a19ad43aa67e367f9b5503ef8075bb5126"
CACHE = "/weights"
DEFAULT_COT = "full"


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")

    style: str = Field(description="The genre, instrumentation, voice and mood to write in.")
    lyrics: str = Field(description="The words to sing, in sections such as [verse] and [chorus].")
    cot: str = Field(
        default=DEFAULT_COT,
        pattern="^(full|melody|off)$",
        description="How much is planned before the audio: full plans melody and chords, melody only the melody, off skips the plan.",
    )
    abc: str | None = Field(
        default=None,
        description="A melody to follow, as ABC notation. Pass back an edited score from an earlier result to render a revision of it, or a transcribed melody to cover an existing song.",
    )
    cfg_scale: float | None = Field(
        default=None,
        ge=1.0,
        le=3.0,
        description="How closely the result follows the style text. Left out, the model's own default applies.",
    )
    seed: int = Field(default=831001, description="Repeat a result by repeating its seed.")


class Generated(BaseModel):
    audio: str = Field(description="The song as a base64 FLAC, stereo at 48 kHz.")
    score: str | None = Field(description="The score the model planned and sang, as ABC notation. Edit it and send it back as abc to render a revision.")
    seconds: float
    sample_rate: int
    seed: int
    truncated: dict = Field(description="Whether the plan or the audio hit its token budget rather than finishing. True in either means the song is cut short.")


@serve
class GeneratedSong:
    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.
        from yue2.pipeline import YuE2Pipeline

        self.pipeline = YuE2Pipeline.from_pretrained(
            MODEL,
            revision=REVISION,
            vae=VAE,
            vae_revision=VAE_REVISION,
            cache_dir=CACHE,
            device="cuda",
            # Its bar of dots goes to stdout on every step, which the platform
            # collects as log lines.
            progress=False,
        )

    def predict(self, request: Request) -> Generated:
        song = self.pipeline(
            style=request.style,
            lyrics=request.lyrics,
            cot=request.cot,
            abc=request.abc,
            cfg_scale=request.cfg_scale,
            seed=request.seed,
        )

        # The pipeline's own writer, rather than soundfile directly: it knows
        # the array's layout and which subtype each container takes, and it is
        # the one place that stays right if either changes.
        with tempfile.NamedTemporaryFile(suffix=".flac") as handle:
            song.save(handle.name)
            handle.seek(0)
            encoded = base64.b64encode(handle.read()).decode("ascii")

        return Generated(
            audio=encoded,
            score=song.abc,
            seconds=len(song.audio) / song.sample_rate,
            sample_rate=song.sample_rate,
            seed=request.seed,
            truncated=song.truncated,
        )

configVersion: 1
id: yue2-3b
name: YuE2 3B
source:
  type: code
  entry: model.py
compute:
  gpuType: l40s
build:
  requirements:
    - yue2_infer@https://huggingface.co/m-a-p/YuE2-3B/resolve/main/yue2_infer-0.1.5-py3-none-any.whl
volumes:
  - /weights