drama3_engine.py

"""Fish Audio Drama 3 on one GPU, streamed, inside this process.

This replaces the SGLang Omni server the endpoint ran. It accepts the same
/v1/audio/speech request body that server.py builds (`engine_request`) and yields the
same events, so the web server, its silence guard and its retries are unchanged.

Numerics are those of the release and of SGLang Omni:

- Model: the official fish-speech PyTorch implementation (vendored under fish_speech/),
  with the checkpoint's block-FP8 MoE experts run as released (W8A8 Triton kernel) and
  every other weight in bfloat16. Nothing is quantized further.
- Prompts: SGLang Omni's reference-bracket-v1 layout, plus the described-character
  and directing-history layouts of drama3_space_ext.py (Figures 11-13 of the report).
- Sampling: the request's temperature / top-p / top-k over end-of-speech and the 4,096
  semantic codes, repetition-aware resampling of the main token, the same codebook
  heads. Only the 4,097 lm_head rows that can be sampled are computed.
- Audio: references encode under bfloat16 autocast as in SGLang Omni's preprocessing;
  the codec decodes SGLang Omni's stream windows (10, 10, 12, 32, then 64 frames, each
  with 32 frames of left context) in float32.

Speed comes from CUDA-graph decode (one graph replay per frame), and streaming each
window as soon as its frames exist.
"""

from __future__ import annotations

import base64
import collections
import hashlib
import io
import json
import math
import os
import time
from collections.abc import Iterator, Mapping
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F

from fish_speech.codec import load_codec
from fish_speech.models.text2semantic.errors import InvalidInputError
from fish_speech.models.text2semantic.fish_omni.backbone import (
    CHECKPOINT_BACKBONE_PREFIX,
    CHECKPOINT_LM_HEAD_KEY,
    FishOmniBackbone,
    load_backbone_weights,
)
from fish_speech.models.text2semantic.fish_omni.config import FishOmniConfig
from fish_speech.models.text2semantic.fish_omni.generate import (
    SPEAKER_TAG_PATTERN,
    UNICODE_WHITE_SPACE,
    Drama3DecodeGraphs,
    ensure_drama3_caches,
    validate_drama3_text,
    with_speaker_tag,
)
from fish_speech.models.text2semantic.fish_omni.model import (
    AUDIO_DECODER_PREFIXES,
    FishOmniModel,
    iter_checkpoint_tensors,
)
from fish_speech.models.text2semantic.fish_omni.tower import (
    RAS_TEMPERATURE,
    RAS_TOP_K,
    RAS_TOP_P,
    RAS_WINDOW_FRAMES,
    FishOmniAudioDecoder,
    load_audio_decoder_weights,
    sample_top_k_top_p,
)

MODEL_ID = os.environ.get("MODEL_ID", "fishaudio/drama-3")
# The revision the H200 endpoint pinned (endpoint/deployment.json of the original Space).
MODEL_REVISION = os.environ.get("MODEL_REVISION", "e0a9218aa48888270f06ef1ed9ea67622c0e17d2")
DTYPE = torch.bfloat16
SAMPLE_RATE = 44100
FRAME_SAMPLES = 2048

# SGLang Omni's prompt literals (sglang_omni/models/fishaudio_drama3/prompt.py) and the
# scene instruction of drama3_space_ext.py.
TTS_INSTRUCTION = "convert the provided text to speech"
REFERENCE_INSTRUCTION = "convert the provided text to speech reference to the following:"
MESSAGE_END = "<|im_end|>\n"
GENERATION_START = "<|im_start|>assistant\n<|voice|>"
SCENE_INSTRUCTION = (
    "Voice every line of the script in order, with <|audio_separator|> "
    "between consecutive lines."
)
ENVELOPE_MARKER = "@@drama3-space@@"

# SGLang Omni defaults (fishaudio_drama3/config.py).
MAX_PROMPT_ROWS = 8192
MAX_NEW_TOKENS_CAP = 1200
DEFAULT_TEMPERATURE = 1.0
DEFAULT_TOP_P = 0.9
DEFAULT_TOP_K = 30
STREAM_STRIDE = 10
STREAM_FOLLOWUP_STRIDE_RAMP = (10, 12, 32)
STREAM_FOLLOWUP_STRIDE = 64
# SGLang Omni decodes each window with 32 frames of left context. The codec's
# receptive field is longer (its window-limited transformers stack to hundreds of
# frames): 32 frames leaves the streamed audio ~3.5 dB SNR from a whole-take decode,
# with discontinuities at window seams; 320 frames reaches ~50 dB, the float16 floor.
STREAM_OVERLAP_FRAMES = int(os.environ.get("DRAMA3_STREAM_OVERLAP", 320))
MAX_REFERENCE_SECONDS = 60.0

# Prompt rows per prefill forward. Larger chunks launch fewer kernels; the rows and
# their causal attention are the same.
PREFILL_CHUNK_ROWS = int(os.environ.get("DRAMA3_PREFILL_CHUNK", 512))
REFERENCE_CACHE_SIZE = 64
# KV buckets prepared at startup: first takes (≤1024 rows) and retakes with history.
WARM_UP_BUCKETS = (1024, 2048)


class RequestError(ValueError):
    """A request the model refuses (HTTP 400 on the SGLang Omni server)."""


def log(message: str) -> None:
    print(f"[engine] {message}", flush=True)


# ─── Loading ──────────────────────────────────────────────────────────────────


def checkpoint_path() -> Path:
    from huggingface_hub import snapshot_download

    # The files README.md's preload_from_hub bakes into the Space image: use them
    # without a network round trip, and download only if they are missing.
    kwargs = dict(revision=MODEL_REVISION, allow_patterns=["*.json", "*.safetensors", "codec.pth"])
    try:
        return Path(snapshot_download(MODEL_ID, local_files_only=True, **kwargs))
    except Exception as error:
        log(f"{MODEL_ID}@{MODEL_REVISION[:8]} is not preloaded ({type(error).__name__}); downloading 41 GB")
        return Path(snapshot_download(MODEL_ID, **kwargs))


@dataclass
class Weights:
    """The model and codec, on CPU or (on ZeroGPU) packed for the GPU; GPU-side state
    such as caches and CUDA graphs lives in Engine, one per process."""

    config: FishOmniConfig
    model: FishOmniModel
    codec: torch.nn.Module

    @property
    def candidate_ids(self) -> torch.Tensor:
        return self.model.backbone.candidate_ids

    @property
    def candidate_head(self) -> torch.Tensor:
        return self.model.backbone.candidate_head


def load_weights(checkpoint_dir: Path | None = None, device: str = "cuda") -> Weights:
    """Build on the meta device, fill on CPU (FP8 experts stay FP8), then move."""
    import drama3_fast

    started = time.perf_counter()
    checkpoint_dir = Path(checkpoint_dir or checkpoint_path())
    log(f"loading weights from {checkpoint_dir}")
    with open(checkpoint_dir / "config.json", encoding="utf-8") as f:
        config = FishOmniConfig.from_dict(json.load(f))
    with torch.device("meta"):
        backbone = FishOmniBackbone(config.text, fp8_experts=config.fp8_experts, dtype=DTYPE)
        audio_decoder = FishOmniAudioDecoder(config, dtype=DTYPE)
    backbone.materialize(["cpu"])
    load_backbone_weights(
        backbone, iter_checkpoint_tensors(checkpoint_dir, skip_prefixes=AUDIO_DECODER_PREFIXES)
    )
    audio_decoder.to_empty(device="cpu")
    load_audio_decoder_weights(
        audio_decoder,
        iter_checkpoint_tensors(
            checkpoint_dir, skip_prefixes=(CHECKPOINT_BACKBONE_PREFIX, CHECKPOINT_LM_HEAD_KEY)
        ),
    )
    from tokenizers import Tokenizer

    tokenizer = Tokenizer.from_file(str(checkpoint_dir / "tokenizer.json"))
    model = FishOmniModel(config, backbone, audio_decoder, tokenizer).eval().requires_grad_(False)
    # The main token is drawn over end-of-speech and the semantic band only: keep the
    # lm_head rows of those ids, in the order the tower samples them.
    band = torch.arange(config.semantic_start_token_id, config.semantic_end_token_id + 1)
    candidate_ids = torch.cat((torch.tensor([config.eos_token_id]), band)).sort().values
    backbone.register_buffer("candidate_ids", candidate_ids, persistent=False)
    backbone.register_buffer(
        "candidate_head", backbone.lm_head.weight.detach()[candidate_ids].clone(), persistent=False
    )
    del backbone.lm_head  # 1 GB the decode step no longer reads
    drama3_fast.fuse_projections(model)
    drama3_fast.install_decode_attention()
    codec = load_codec(str(checkpoint_dir / "codec.pth"), torch.float32).requires_grad_(False)
    # Encoding stays fp32 (under bfloat16 autocast, as SGLang Omni's preprocessing);
    # the convolutional decoder runs in float16, as SGLang Omni's vocoder.
    codec.decoder.half()
    log(f"loaded {MODEL_ID}@{MODEL_REVISION[:8]} on CPU in {time.perf_counter() - started:.0f}s")
    if device != "cpu":
        model.to(device)
        codec.to(device)
    return Weights(config, model, codec)


# ─── Prompts ──────────────────────────────────────────────────────────────────


@dataclass(frozen=True)
class Character:
    name: str
    description: str
    scene: str | None


@dataclass(frozen=True)
class Reference:
    transcript: str | None
    codes: torch.Tensor  # [num_codebooks, frames], CPU int64


def reference_text_block(transcript: str) -> str:
    """The speaker tag follows "Text: " directly; leading white space moves behind it."""
    tagged = with_speaker_tag(transcript)
    body = tagged.lstrip(UNICODE_WHITE_SPACE)
    tag = SPEAKER_TAG_PATTERN.match(body)
    assert tag is not None
    leading = tagged[: len(tagged) - len(body)]
    return tag.group(0) + leading + body[tag.end():]


def turn_text(text: str, character: Character | None) -> str:
    if character is not None:
        return f"{character.name}: {text}"
    return with_speaker_tag(text)


def build_segments(
    text: str,
    voices: list[Reference],
    history: list[tuple[str, torch.Tensor]],
    character: Character | None,
) -> list[str | torch.Tensor]:
    """SGLang Omni's stock layout, or drama3_space_ext's described / directing one."""
    segments: list[str | torch.Tensor] = ["<|im_start|>system\n"]
    if character is not None:
        lines = [f"Scene: {character.scene}"] if character.scene else []
        lines += ["Characters:", f"- {character.name}: {character.description}", SCENE_INSTRUCTION]
        segments.append("\n".join(lines))
    else:
        segments.append(REFERENCE_INSTRUCTION if voices else TTS_INSTRUCTION)
        for index, voice in enumerate(voices, 1):
            segments += [
                f"\n\n[Reference {index}]\nText: {reference_text_block(voice.transcript)}\nSpeech: ",
                voice.codes,
            ]
    segments.append(MESSAGE_END)
    for earlier_text, codes in history:
        segments += [
            "<|im_start|>user\n", turn_text(earlier_text, character), MESSAGE_END,
            GENERATION_START, codes, MESSAGE_END,
        ]
    segments += ["<|im_start|>user\n", turn_text(text, character), MESSAGE_END, GENERATION_START]
    return segments


def segment_rows(segments, tokenizer, config: FishOmniConfig) -> torch.Tensor:
    """Text rows are [token, 0, ...]; audio rows [c0 + semantic_begin, c0..cN-1].

    Every literal is tokenized on its own: tokenizing the joined string changes ids
    at the seams.
    """
    num_codebooks = config.audio_decoder.num_codebooks
    codebook_size = config.audio_decoder.vocab_size
    rows = []
    for segment in segments:
        if isinstance(segment, str):
            ids = torch.tensor(tokenizer.encode(segment).ids, dtype=torch.long)
            text_rows = torch.zeros((len(ids), 1 + num_codebooks), dtype=torch.long)
            text_rows[:, 0] = ids
            rows.append(text_rows)
        else:
            codes = segment.to(device="cpu", dtype=torch.long)
            if codes.ndim != 2 or codes.shape[0] != num_codebooks or codes.shape[1] == 0:
                raise RequestError(f"audio codes must be [{num_codebooks}, frames]; got {tuple(codes.shape)}")
            if codes.min() < 0 or codes.max() >= codebook_size:
                raise RequestError(f"audio codes must be in [0, {codebook_size})")
            rows.append(torch.cat((codes[:1] + config.semantic_start_token_id, codes)).T)
    return torch.cat(rows)


def checked_text(value: object, location: str) -> str:
    if not isinstance(value, str):
        raise RequestError(f"{location} must be a string")
    try:
        validate_drama3_text(value, location)
    except InvalidInputError as error:
        raise RequestError(str(error)) from error
    other = {int(m.group(1)) for m in SPEAKER_TAG_PATTERN.finditer(value)} - {0}
    if other:
        raise RequestError(f"only speaker 0 is supported; got speaker tags {sorted(other)}")
    return value


def single_line(value: object, location: str, *, optional: bool = False) -> str | None:
    if value is None and optional:
        return None
    if not isinstance(value, str) or not value.strip(UNICODE_WHITE_SPACE):
        raise RequestError(f"{location} must be a non-empty string")
    if "\n" in value or "\r" in value:
        raise RequestError(f"{location} must be a single line")
    return checked_text(value, location).strip(UNICODE_WHITE_SPACE)


@dataclass(frozen=True)
class SpeechRequest:
    text: str
    voice_count: int
    references: list[Mapping]  # voice clips first, then one clip per history turn
    history: list[str]
    character: Character | None
    temperature: float
    top_p: float
    top_k: int
    max_new_tokens: int


def number(params: Mapping, name: str, default: float) -> float:
    value = params.get(name)
    if value is None:
        return default
    if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
        raise RequestError(f"{name} must be a number")
    return float(value)


def parse_request(body: Mapping) -> SpeechRequest:
    """The /v1/audio/speech fields server.py sends, checked like SGLang Omni does."""
    text = body.get("input")
    references = list(body.get("references") or [])
    if not isinstance(text, str):
        raise RequestError("input must be a string")
    temperature = number(body, "temperature", DEFAULT_TEMPERATURE)
    top_p = number(body, "top_p", DEFAULT_TOP_P)
    top_k = body.get("top_k", DEFAULT_TOP_K)
    if isinstance(top_k, bool) or not isinstance(top_k, int) or top_k < 1:
        raise RequestError("top_k must be a positive integer")
    if not 0 < top_p <= 1:
        raise RequestError("top_p must be in (0, 1]")
    if temperature < 0:
        raise RequestError("temperature must be non-negative")
    max_new_tokens = body.get("max_new_tokens") or MAX_NEW_TOKENS_CAP
    if isinstance(max_new_tokens, bool) or not isinstance(max_new_tokens, int) or max_new_tokens < 1:
        raise RequestError("max_new_tokens must be a positive integer")
    max_new_tokens = min(max_new_tokens, MAX_NEW_TOKENS_CAP)
    sampling = dict(temperature=temperature, top_p=top_p, top_k=top_k, max_new_tokens=max_new_tokens)

    if not text.startswith(ENVELOPE_MARKER):
        checked_text(text, "input")
        for index, reference in enumerate(references):
            checked_text(reference.get("text") or "", f"references.{index}")
        return SpeechRequest(text=text, voice_count=len(references), references=references,
                             history=[], character=None, **sampling)

    try:
        spec = json.loads(text[len(ENVELOPE_MARKER):])
    except json.JSONDecodeError as error:
        raise RequestError(f"malformed envelope: {error}") from error
    if not isinstance(spec, Mapping):
        raise RequestError("the envelope must be a JSON object")
    line = checked_text(spec.get("text"), "text")
    history = spec.get("history") or []
    if not isinstance(history, list) or not all(isinstance(t, str) for t in history):
        raise RequestError("envelope history must be a list of strings")
    for index, earlier in enumerate(history):
        checked_text(earlier, f"history.{index}")
    voice_count = spec.get("voice_references", 0)
    if (isinstance(voice_count, bool) or not isinstance(voice_count, int) or voice_count < 0
            or voice_count + len(history) != len(references)):
        raise RequestError(
            f"voice_references {voice_count!r} and {len(history)} history turns do not "
            f"account for {len(references)} references"
        )
    for index, reference in enumerate(references[:voice_count]):
        checked_text(reference.get("text"), f"references.{index}")
    character = None
    raw = spec.get("character")
    if raw is not None:
        if not isinstance(raw, Mapping):
            raise RequestError("envelope character must be an object")
        if voice_count:
            raise RequestError("a described character cannot also have voice references")
        name = single_line(raw.get("name"), "character.name")
        if ":" in name:
            raise RequestError("character.name must not contain ':'")
        character = Character(
            name=name,
            description=single_line(raw.get("description"), "character.description"),
            scene=single_line(raw.get("scene"), "character.scene", optional=True),
        )
    return SpeechRequest(text=line, voice_count=voice_count, references=references,
                         history=list(history), character=character, **sampling)


def reference_audio(reference: Mapping) -> np.ndarray:
    """Mono float32 samples at 44.1 kHz from a base64 clip (FLAC or WAV)."""
    import soundfile as sf

    data = reference.get("data") or reference.get("base64")
    if not isinstance(data, str):
        raise RequestError("a reference needs base64 audio in data")
    try:
        samples, rate = sf.read(io.BytesIO(base64.b64decode(data)), dtype="float32", always_2d=True)
    except Exception as error:
        raise RequestError(f"could not decode reference audio: {error}") from error
    samples = samples.mean(axis=1)
    if rate != SAMPLE_RATE:
        import librosa

        samples = librosa.resample(samples, orig_sr=rate, target_sr=SAMPLE_RATE)
    if samples.size > MAX_REFERENCE_SECONDS * SAMPLE_RATE:
        raise RequestError(f"reference audio is longer than {MAX_REFERENCE_SECONDS:g} s")
    if samples.size < FRAME_SAMPLES:
        raise RequestError("reference audio is too short")
    return np.ascontiguousarray(samples, dtype=np.float32)


# ─── GPU state ────────────────────────────────────────────────────────────────


def decode_candidates(tower, candidates, candidate_logits, hidden, positions, frames, history,
                      *, temperature, top_k, top_p) -> torch.Tensor:
    """FishOmniAudioDecoder.decode_step, given the logits of the candidate ids."""
    first = candidates[sample_top_k_top_p(candidate_logits, temperature, top_k, top_p, check_finite=False)]
    retry = candidates[
        sample_top_k_top_p(candidate_logits, RAS_TEMPERATURE, RAS_TOP_K, RAS_TOP_P, check_finite=False)
    ]
    repeated = (history[-RAS_WINDOW_FRAMES:] == first[:, None]).any(-1)
    repeated &= tower.is_semantic(first)
    main = torch.where(repeated, retry, first)
    c0 = (main - tower.semantic_begin).clamp(0, tower.config.vocab_size - 1)
    semantic = tower.is_semantic(main)
    hidden = tower.run_shared(hidden, c0, semantic, positions, frames)
    codes = [c0]
    for layer, head in zip(tower.residual_layers, tower.codebook_heads):
        hidden = layer(hidden, codes[-1], semantic, positions)
        codes.append(sample_top_k_top_p(head(hidden), temperature, top_k, top_p, check_finite=False))
    return torch.stack((main, *codes), dim=-1)


class DecodeGraphs(Drama3DecodeGraphs):
    """CUDA graphs of one whole decode step, bookkeeping included.

    Like upstream Drama3DecodeGraphs (lm_head -> tower sampling -> backbone), with:
    - only the 4,097 lm_head rows that can be sampled (end-of-speech and the semantic
      band): the full head is 252,288 x 2,048, 1 GB read per frame for logits that
      are never used;
    - the next step's inputs, the frame log and the stop status written inside the
      graph, so the host can queue step n + 1 before it reads step n's status.
    """

    def __init__(self, model: FishOmniModel, candidate_ids: torch.Tensor, candidate_head: torch.Tensor) -> None:
        super().__init__(model)
        device = self.hidden.device
        self.candidate_ids = candidate_ids
        self.candidate_head = candidate_head
        self.out_frames = torch.zeros((MAX_NEW_TOKENS_CAP + 1, self.frame.shape[1]), dtype=torch.long, device=device)
        self.out_index = torch.zeros(1, dtype=torch.long, device=device)
        self.status = torch.zeros(2, dtype=torch.long, device=device)

    def run(self, kv_len: int) -> None:
        backbone, tower = self.model.backbone, self.model.audio_decoder
        logits = F.linear(self.hidden, self.candidate_head)
        finite = torch.isfinite(logits).all()
        sampled = decode_candidates(
            tower, self.candidate_ids, logits, self.hidden, self.position, self.frame, self.history,
            temperature=self.temperature, top_k=self.top_k, top_p=self.top_p,
        )
        embeddings = tower.embed_backbone(sampled, backbone.embed(sampled[:, 0]))
        hidden = backbone(inputs_embeds=embeddings[None], input_pos=self.position + 1, static_kv_len=kv_len)
        self.out_frames.index_copy_(0, self.out_index.clamp(max=MAX_NEW_TOKENS_CAP), sampled)
        self.status.copy_(torch.stack((sampled[0, 0], finite.long())))
        self.hidden.copy_(hidden[0])
        self.frame.copy_(sampled)
        self.history.copy_(torch.cat((self.history[1:], sampled[0, :1])))
        self.position.add_(1)
        self.out_index.add_(1)

    def start(self, hidden, position, prompt_frames, temperature, top_k, top_p) -> None:
        super().start(hidden, position, prompt_frames, temperature, top_k, top_p)
        self.out_index.zero_()

    def replay(self, position: int) -> None:
        self.graphs[next(b for b in self.buckets if b >= position + 2)].replay()


@dataclass
class PendingAudio:
    done: torch.cuda.Event
    pcm: torch.Tensor  # pinned int16 on the host


class Engine:
    """Generation on the GPU that holds `weights` (one ZeroGPU worker, or a local GPU)."""

    def __init__(self, weights: Weights) -> None:
        import drama3_fast

        if drama3_fast.enabled():
            # In the GPU process: on ZeroGPU, torch._dynamo must first be imported
            # after ZeroGPU has restored the real torch.cuda functions.
            drama3_fast.compile_layers(weights.model)
        self.weights = weights
        self.model = weights.model
        self.config = weights.config
        self.codec = weights.codec
        self.tokenizer = self.model.tokenizer
        self.reference_codes: collections.OrderedDict[str, torch.Tensor] = collections.OrderedDict()
        self.graphs: DecodeGraphs | None = None
        self.vocoder_stream: torch.cuda.Stream | None = None

    @property
    def device(self) -> torch.device:
        return self.model.audio_decoder.device

    # Audio ------------------------------------------------------------------

    @torch.inference_mode()
    def encode(self, reference: Mapping) -> torch.Tensor:
        """Codes [num_codebooks, frames] of a reference clip, cached by content.

        As in SGLang Omni's preprocessing: fp32 weights under bfloat16 autocast."""
        key = hashlib.sha1(str(reference.get("data") or reference.get("base64")).encode()).hexdigest()
        if (codes := self.reference_codes.get(key)) is not None:
            self.reference_codes.move_to_end(key)
            return codes
        samples = torch.from_numpy(reference_audio(reference)).to(self.device)
        audio = samples.clamp(-1.0, 1.0)[None, None]
        lengths = torch.tensor([audio.shape[-1]], device=self.device)
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=self.device.type == "cuda"):
            indices, _ = self.codec.encode(audio, lengths)
        codes = indices[0].long().cpu()
        self.reference_codes[key] = codes
        while len(self.reference_codes) > REFERENCE_CACHE_SIZE:
            self.reference_codes.popitem(last=False)
        return codes

    @torch.inference_mode()
    def vocode(self, codes: torch.Tensor) -> torch.Tensor:
        """Waveform of codes [num_codebooks, frames]: the quantizer in fp32, then the
        convolutional decoder in float16 as SGLang Omni's vocoder runs it."""
        latent = self.codec.quantizer.decode(codes[None].contiguous())
        decoder_dtype = next(self.codec.decoder.parameters()).dtype
        return self.codec.decoder(latent.to(decoder_dtype))[0, 0].float()

    # Prompt -----------------------------------------------------------------

    def prompt(self, request: SpeechRequest) -> torch.Tensor:
        codes = [self.encode(reference) for reference in request.references]
        voices = [
            Reference(transcript=reference.get("text"), codes=code)
            for reference, code in zip(request.references[: request.voice_count], codes)
        ]
        history = list(zip(request.history, codes[request.voice_count:]))
        rows = segment_rows(build_segments(request.text, voices, history, request.character),
                            self.tokenizer, self.config)
        if len(rows) > MAX_PROMPT_ROWS:
            raise RequestError(
                f"a prompt of {len(rows)} rows exceeds the maximum allowed length of {MAX_PROMPT_ROWS} rows"
            )
        return rows

    # Generation --------------------------------------------------------------

    @torch.inference_mode()
    def warm_up(self, reference: Mapping | None) -> None:
        """Compile and capture the common KV buckets and run a short take, so the
        first visitors do not wait for it (compiled kernels are cached on disk for
        the other GPU workers too)."""
        started = time.perf_counter()
        for bucket in WARM_UP_BUCKETS:
            self.graph_for(bucket)
        body = {"input": "Places, everyone.", "references": [reference] if reference else None,
                "max_new_tokens": 2 * STREAM_STRIDE, "temperature": 1.0, "top_p": 0.9, "top_k": 4097}
        for _ in self.speech(body):
            pass
        log(f"warmed up in {time.perf_counter() - started:.1f}s")
        properties = torch.cuda.get_device_properties(self.device)
        replay = self.graphs.graphs[WARM_UP_BUCKETS[0]]
        torch.cuda.synchronize(self.device)
        timer = time.perf_counter()
        for _ in range(30):
            replay.replay()
        torch.cuda.synchronize(self.device)
        log(f"{properties.name}, {properties.multi_processor_count} SMs, "
            f"{properties.total_memory / 2**30:.0f} GiB, torch {torch.__version__}: one decode step "
            f"{(time.perf_counter() - timer) / 30 * 1000:.2f} ms on the GPU; "
            f"{torch.cuda.max_memory_allocated(self.device) / 2**30:.1f} GiB peak")

    def decode_graphs(self) -> DecodeGraphs:
        if self.graphs is None:
            model = self.model
            ensure_drama3_caches(model, DecodeGraphs.graph_capacity(model))
            self.graphs = DecodeGraphs(model, self.weights.candidate_ids, self.weights.candidate_head)
            model.decode_graphs = self.graphs
        return self.graphs

    def graph_for(self, kv_len: int) -> None:
        """Capture one KV bucket's graph on first use: a fresh worker captures only
        the buckets its requests reach. Capturing first compiles the layers."""
        graphs = self.decode_graphs()
        if kv_len not in graphs.graphs:
            started = time.perf_counter()
            with torch.cuda.device(self.device):
                graphs.capture(kv_len)
            log(f"decode graph for {kv_len} cache rows ready in {time.perf_counter() - started:.1f}s")

    @torch.inference_mode()
    def prefill(self, prompt: torch.Tensor, request: SpeechRequest) -> DecodeGraphs:
        model = self.model
        backbone, tower = model.backbone, model.audio_decoder
        graphs = self.decode_graphs()
        prompt_rows = len(prompt)
        capacity = prompt_rows + request.max_new_tokens
        if capacity > graphs.graph_capacity(model):
            raise RequestError("the prompt and audio limit exceed the context")
        needed = next(b for b in graphs.buckets if b >= capacity)
        for bucket in graphs.buckets:
            if prompt_rows + 1 <= bucket <= needed:
                self.graph_for(bucket)
        device = tower.device
        frames = prompt.to(device)
        positions = torch.arange(prompt_rows, device=device)
        for start in range(0, prompt_rows, PREFILL_CHUNK_ROWS):
            end = min(start + PREFILL_CHUNK_ROWS, prompt_rows)
            chunk = frames[start:end]
            embeddings = tower.embed_backbone(chunk, backbone.embed(chunk[:, 0]))
            hidden = backbone(inputs_embeds=embeddings[None], input_pos=positions[start:end])[0]
            paired_end = min(end, prompt_rows - 1)
            if paired_end > start:
                tower.prefill(hidden[: paired_end - start], positions[start:paired_end],
                              frames[start:paired_end], frames[start + 1: paired_end + 1])
        graphs.start(hidden[-1:], prompt_rows - 1, frames, request.temperature, request.top_k, request.top_p)
        return graphs

    @torch.inference_mode()
    def frames(self, prompt: torch.Tensor, request: SpeechRequest) -> Iterator[int]:
        """Yield the number of frames in graphs.out_frames after each step; return
        the finish reason. Step n + 1 is queued before step n's status is read."""
        graphs = self.prefill(prompt, request)
        first_position = len(prompt) - 1
        eos = self.config.eos_token_id
        status = torch.empty(2, dtype=torch.long, pin_memory=True)
        ready = torch.cuda.Event()
        graphs.replay(first_position)
        for step in range(request.max_new_tokens):
            status.copy_(graphs.status, non_blocking=True)
            ready.record()
            if step + 1 < request.max_new_tokens:
                graphs.replay(first_position + step + 1)
            ready.synchronize()
            main, finite = status.tolist()
            if not finite:
                raise FloatingPointError("non-finite logits; refusing to sample invalid audio")
            if main == eos:
                return "stop"
            self.frame_ready = ready
            yield step + 1
        return "length"

    @torch.inference_mode()
    def queue_window(self, start: int, end: int, emitted: int) -> PendingAudio:
        """Vocode frames start..end on the side stream once they exist; the samples of
        frames before `emitted` are context only."""
        graphs = self.graphs
        if self.vocoder_stream is None:
            self.vocoder_stream = torch.cuda.Stream(device=self.device)
        stream = self.vocoder_stream
        stream.wait_event(self.frame_ready)
        with torch.cuda.stream(stream):
            codes = graphs.out_frames[start:end, 1:].T.contiguous()
            audio = self.vocode(codes)[(emitted - start) * FRAME_SAMPLES:]
            pcm = (audio.clamp(-1.0, 1.0) * 32767.0).to(torch.int16)
            host = torch.empty(pcm.shape, dtype=torch.int16, pin_memory=True)
            host.copy_(pcm, non_blocking=True)
            done = torch.cuda.Event()
            done.record(stream)
        return PendingAudio(done, host)

    def speech(self, body: Mapping, *, deadline: float | None = None) -> Iterator[dict]:
        """Events of one take: speech.audio.delta (PCM16 bytes) ... speech.audio.done."""
        started = time.perf_counter()
        timings: dict = {}
        try:
            request = parse_request(body)
            prompt = self.prompt(request)
        except RequestError as error:
            yield {"type": "error", "message": f"Drama 3 request: {error}"}
            return
        timings["prompt_ms"] = round((time.perf_counter() - started) * 1000)
        timings["prompt_rows"] = len(prompt)

        produced = emitted = 0
        window_end, windows = STREAM_STRIDE, 0
        pending: collections.deque[PendingAudio] = collections.deque()
        finish = "length"
        generator = self.frames(prompt, request)
        decode_started = None

        def window(end: int) -> PendingAudio:
            return self.queue_window(max(0, emitted - STREAM_OVERLAP_FRAMES), end, emitted)

        try:
            while True:
                try:
                    produced = next(generator)
                except StopIteration as stop:
                    finish = stop.value
                    break
                if decode_started is None:
                    decode_started = time.perf_counter()
                    timings["prefill_ms"] = round((decode_started - started) * 1000) - timings["prompt_ms"]
                if produced == window_end:
                    pending.append(window(produced))
                    emitted = produced
                    ramp = STREAM_FOLLOWUP_STRIDE_RAMP
                    window_end += ramp[windows] if windows < len(ramp) else STREAM_FOLLOWUP_STRIDE
                    windows += 1
                while pending and pending[0].done.query():
                    timings.setdefault("first_audio_ms", round((time.perf_counter() - started) * 1000))
                    yield {"type": "speech.audio.delta", "pcm": pending.popleft().pcm.numpy().tobytes()}
                if deadline is not None and time.time() > deadline:
                    finish = "length"
                    break
        finally:
            generator.close()
        decode_s = time.perf_counter() - (decode_started or started)
        if produced > emitted:
            pending.append(window(produced))
            emitted = produced
        while pending:
            item = pending.popleft()
            item.done.synchronize()
            timings.setdefault("first_audio_ms", round((time.perf_counter() - started) * 1000))
            yield {"type": "speech.audio.delta", "pcm": item.pcm.numpy().tobytes()}
        timings.update(
            frames=produced,
            decode_frames_per_s=round(produced / decode_s, 1) if produced and decode_s > 0 else None,
            total_ms=round((time.perf_counter() - started) * 1000),
        )
        yield {"type": "speech.audio.done", "finish_reason": finish, "timings": timings, "frames": produced}