modeling_arktts.py

from __future__ import annotations

import math
import shutil
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.nn.attention import SDPBackend, sdpa_kernel
from transformers import GenerationConfig, PreTrainedModel
from transformers.generation import (
    LogitsProcessor,
    LogitsProcessorList,
    StoppingCriteriaList,
)
from transformers.modeling_outputs import ModelOutput
from transformers.utils.hub import cached_file

from .configuration_arktts import ArkttsConfig


@dataclass
class ArkttsModelOutput(ModelOutput):
    logits: Optional[Tensor] = None
    hidden_states: Optional[Tensor] = None
    codebook_logits: Optional[Tensor] = None


@dataclass
class ArkttsGenerateOutput(ModelOutput):
    codes: Optional[Tensor] = None
    code_lengths: Optional[Tensor] = None
    finished: Optional[Tensor] = None


class ArkttsSemanticLogitsProcessor(LogitsProcessor):
    def __init__(self, semantic_begin_id: int, semantic_end_id: int, eos_token_id: int):
        self.semantic_begin_id = int(semantic_begin_id)
        self.semantic_end_id = int(semantic_end_id)
        self.eos_token_id = int(eos_token_id)

    def __call__(self, input_ids: Tensor, scores: Tensor) -> Tensor:
        filtered = torch.full_like(scores, float("-inf"))
        filtered[:, self.semantic_begin_id : self.semantic_end_id + 1] = scores[
            :, self.semantic_begin_id : self.semantic_end_id + 1
        ]
        filtered[:, self.eos_token_id] = scores[:, self.eos_token_id]
        return filtered


class ArkttsLegacyTopKTopPLogitsProcessor(LogitsProcessor):
    """Matches the candidate filtering order used by the original inference code."""

    def __init__(self, top_k: int, top_p: float):
        if top_k <= 0:
            raise ValueError("top_k must be positive")
        if not 0.0 < top_p <= 1.0:
            raise ValueError("top_p must be in (0, 1]")
        self.top_k = int(top_k)
        self.top_p = float(top_p)

    def __call__(self, input_ids: Tensor, scores: Tensor) -> Tensor:
        sorted_scores, sorted_indices = torch.sort(scores, descending=True, dim=-1)
        cumulative = torch.cumsum(torch.softmax(sorted_scores, dim=-1), dim=-1)
        positions = torch.arange(sorted_scores.shape[-1], device=scores.device)
        threshold = torch.tensor(self.top_p, dtype=cumulative.dtype, device=cumulative.device)
        remove_sorted = (cumulative > threshold) | (positions >= self.top_k)
        remove_sorted[..., 0] = False
        remove = torch.zeros_like(remove_sorted).scatter(1, sorted_indices, remove_sorted)
        return scores.masked_fill(remove, float("-inf"))


class ArkttsKVCache(nn.Module):
    def __init__(
        self,
        batch_size: int,
        max_length: int,
        heads: int,
        head_dim: int,
        dtype,
        return_full: bool = False,
    ):
        super().__init__()
        shape = (batch_size, heads, max_length, head_dim)
        self.register_buffer("keys", torch.zeros(shape, dtype=dtype), persistent=False)
        self.register_buffer("values", torch.zeros(shape, dtype=dtype), persistent=False)
        self.return_full = bool(return_full)

    def update(self, cache_position: Tensor, keys: Tensor, values: Tensor):
        self.keys[:, :, cache_position] = keys
        self.values[:, :, cache_position] = values
        end = self.keys.shape[2] if self.return_full else int(cache_position[-1].item()) + 1
        return self.keys[:, :, :end], self.values[:, :, :end]


class ArkttsRMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float):
        super().__init__()
        self.eps = float(eps)
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: Tensor) -> Tensor:
        normalized = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
        return normalized.to(x.dtype) * self.weight


def _precompute_rope(length: int, head_dim: int, base: float) -> Tensor:
    frequencies = 1.0 / (
        base ** (torch.arange(0, head_dim, 2).float()[: head_dim // 2] / head_dim)
    )
    phases = torch.outer(torch.arange(length), frequencies)
    values = torch.polar(torch.ones_like(phases), phases)
    return torch.stack((values.real, values.imag), dim=-1).to(torch.bfloat16)


def _apply_rope(x: Tensor, rope: Tensor) -> Tensor:
    shaped = x.float().reshape(*x.shape[:-1], -1, 2)
    if rope.ndim == 3:
        rope = rope[None, :, None]
    elif rope.ndim == 4:
        rope = rope[:, :, None]
    else:
        raise ValueError(f"Unexpected RoPE shape: {tuple(rope.shape)}")
    output = torch.stack(
        (
            shaped[..., 0] * rope[..., 0] - shaped[..., 1] * rope[..., 1],
            shaped[..., 1] * rope[..., 0] + shaped[..., 0] * rope[..., 1],
        ),
        dim=-1,
    )
    return output.flatten(3).to(x.dtype)


class ArkttsAttention(nn.Module):
    def __init__(
        self,
        dim: int,
        n_head: int,
        n_local_heads: int,
        head_dim: int,
        qkv_bias: bool,
        output_bias: bool,
        qk_norm: bool,
        norm_eps: float,
        dropout: float,
        use_sdpa: bool,
    ):
        super().__init__()
        total = (n_head + 2 * n_local_heads) * head_dim
        self.wqkv = nn.Linear(dim, total, bias=qkv_bias)
        self.wo = nn.Linear(n_head * head_dim, dim, bias=output_bias)
        self.n_head = int(n_head)
        self.n_local_heads = int(n_local_heads)
        self.head_dim = int(head_dim)
        self.dropout = float(dropout)
        self.use_sdpa = bool(use_sdpa)
        self.qk_norm = bool(qk_norm)
        if self.qk_norm:
            self.q_norm = ArkttsRMSNorm(head_dim, norm_eps)
            self.k_norm = ArkttsRMSNorm(head_dim, norm_eps)
        self.kv_cache: Optional[ArkttsKVCache] = None

    def forward(
        self,
        x: Tensor,
        rope: Tensor,
        attention_mask: Optional[Tensor],
        cache_position: Optional[Tensor] = None,
    ) -> Tensor:
        batch, length, _ = x.shape
        query_size = self.n_head * self.head_dim
        kv_size = self.n_local_heads * self.head_dim
        query, key, value = self.wqkv(x).split((query_size, kv_size, kv_size), dim=-1)
        query = query.view(batch, length, self.n_head, self.head_dim)
        key = key.view(batch, length, self.n_local_heads, self.head_dim)
        value = value.view(batch, length, self.n_local_heads, self.head_dim)
        if self.qk_norm:
            query = self.q_norm(query)
            key = self.k_norm(key)
        query = _apply_rope(query, rope).transpose(1, 2)
        key = _apply_rope(key, rope).transpose(1, 2)
        value = value.transpose(1, 2)
        if self.kv_cache is not None:
            if cache_position is None:
                raise ValueError("cache_position is required when KV cache is enabled")
            key, value = self.kv_cache.update(cache_position, key, value)
        repeats = self.n_head // self.n_local_heads
        key = key.repeat_interleave(repeats, dim=1)
        value = value.repeat_interleave(repeats, dim=1)
        if self.use_sdpa:
            output = F.scaled_dot_product_attention(
                query,
                key,
                value,
                attn_mask=attention_mask,
                dropout_p=self.dropout if self.training else 0.0,
            )
        else:
            scores = query @ key.transpose(-2, -1) / math.sqrt(self.head_dim)
            if attention_mask is not None:
                scores = scores.masked_fill(~attention_mask, float("-inf"))
            probabilities = torch.softmax(scores, dim=-1)
            if self.training and self.dropout:
                probabilities = F.dropout(probabilities, p=self.dropout)
            output = probabilities @ value
        output = output.transpose(1, 2).contiguous().view(batch, length, query_size)
        return self.wo(output)


class ArkttsFeedForward(nn.Module):
    def __init__(self, dim: int, intermediate_size: int):
        super().__init__()
        self.w1 = nn.Linear(dim, intermediate_size, bias=False)
        self.w2 = nn.Linear(intermediate_size, dim, bias=False)
        self.w3 = nn.Linear(dim, intermediate_size, bias=False)

    def forward(self, x: Tensor) -> Tensor:
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class ArkttsTransformerBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        intermediate_size: int,
        n_head: int,
        n_local_heads: int,
        head_dim: int,
        qkv_bias: bool,
        output_bias: bool,
        qk_norm: bool,
        norm_eps: float,
        dropout: float,
        use_sdpa: bool,
    ):
        super().__init__()
        self.attention = ArkttsAttention(
            dim, n_head, n_local_heads, head_dim, qkv_bias, output_bias,
            qk_norm, norm_eps, dropout, use_sdpa,
        )
        self.feed_forward = ArkttsFeedForward(dim, intermediate_size)
        self.ffn_norm = ArkttsRMSNorm(dim, norm_eps)
        self.attention_norm = ArkttsRMSNorm(dim, norm_eps)

    def forward(self, x, rope, attention_mask, cache_position=None):
        hidden = x + self.attention(self.attention_norm(x), rope, attention_mask, cache_position)
        return hidden + self.feed_forward(self.ffn_norm(hidden))


class ArkttsModel(PreTrainedModel):
    config_class = ArkttsConfig
    base_model_prefix = ""
    main_input_name = "input_ids"
    _no_split_modules = ["ArkttsTransformerBlock"]
    _supports_sdpa = True

    def __init__(self, config: ArkttsConfig):
        super().__init__(config)
        self.embeddings = nn.Embedding(config.vocab_size, config.dim)
        self.codebook_embeddings = nn.Embedding(config.codebook_size * config.num_codebooks, config.dim)
        self.layers = nn.ModuleList([
            ArkttsTransformerBlock(
                config.dim, config.intermediate_size, config.n_head, config.n_local_heads,
                config.head_dim, config.attention_qkv_bias, config.attention_o_bias,
                config.attention_qk_norm, config.norm_eps, config.dropout, True,
            )
            for _ in range(config.n_layer)
        ])
        self.norm = ArkttsRMSNorm(config.dim, config.norm_eps)
        self.fast_project_in = (
            nn.Linear(config.dim, config.fast_dim)
            if config.fast_dim != config.dim else nn.Identity()
        )
        self.fast_embeddings = nn.Embedding(config.codebook_size, config.fast_dim)
        self.fast_layers = nn.ModuleList([
            ArkttsTransformerBlock(
                config.fast_dim, config.fast_intermediate_size, config.fast_n_head,
                config.fast_n_local_heads, config.fast_head_dim,
                config.fast_attention_qkv_bias, config.fast_attention_o_bias,
                config.fast_attention_qk_norm, config.norm_eps, config.dropout, False,
            )
            for _ in range(config.n_fast_layer)
        ])
        self.fast_norm = ArkttsRMSNorm(config.fast_dim, config.norm_eps)
        self.fast_output = nn.Linear(config.fast_dim, config.codebook_size, bias=False)
        self.register_buffer(
            "freqs_cis", _precompute_rope(config.max_seq_len, config.head_dim, config.rope_base),
            persistent=False,
        )
        self.register_buffer(
            "fast_freqs_cis",
            _precompute_rope(config.num_codebooks, config.fast_head_dim, config.rope_base),
            persistent=False,
        )
        self.__dict__["_arktts_codec"] = None
        self.post_init()

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if isinstance(module, nn.Linear) and module.bias is not None:
                module.bias.data.zero_()

    def get_input_embeddings(self):
        return self.embeddings

    def set_input_embeddings(self, value):
        self.embeddings = value

    def get_output_embeddings(self):
        return None

    def _embed(self, input_ids: Tensor) -> Tensor:
        codebook_embeds = []
        for index in range(self.config.num_codebooks):
            codebook_embeds.append(
                self.codebook_embeddings(input_ids[:, index + 1] + index * self.config.codebook_size)
            )
        codebook_sum = torch.stack(codebook_embeds, dim=1).sum(dim=1)
        semantic = (input_ids[:, 0] >= self.config.semantic_begin_id) & (
            input_ids[:, 0] <= self.config.semantic_end_id
        )
        codebook_sum = torch.where(semantic.unsqueeze(-1), codebook_sum, 0.0)
        return self.embeddings(input_ids[:, 0]) + codebook_sum

    @staticmethod
    def _causal_mask(attention_mask: Tensor, query_positions: Tensor, key_length: int) -> Tensor:
        if attention_mask.shape[1] < key_length:
            attention_mask = F.pad(
                attention_mask,
                (0, key_length - attention_mask.shape[1]),
                value=0,
            )
        key_positions = torch.arange(key_length, device=attention_mask.device)
        causal = key_positions[None, :] <= query_positions[:, None]
        return causal[None, None] & attention_mask[:, None, None, :key_length].bool()

    def forward(
        self,
        input_ids: Tensor,
        attention_mask: Optional[Tensor] = None,
        labels: Optional[Tensor] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ):
        del labels, output_hidden_states, kwargs
        if input_ids.ndim != 3 or input_ids.shape[1] != self.config.num_codebooks + 1:
            raise ValueError(
                f"input_ids must have shape [B, {self.config.num_codebooks + 1}, T]"
            )
        batch, _, length = input_ids.shape
        if attention_mask is None:
            attention_mask = torch.ones((batch, length), dtype=torch.long, device=input_ids.device)
        position_ids = attention_mask.long().cumsum(-1).sub(1).clamp_min(0)
        rope = self.freqs_cis[position_ids]
        mask = self._causal_mask(
            attention_mask,
            torch.arange(length, device=input_ids.device),
            length,
        )
        hidden = self._embed(input_ids)
        for layer in self.layers:
            hidden = layer(hidden, rope, mask)
        normalized = self.norm(hidden)
        logits = F.linear(normalized, self.embeddings.weight)
        output = ArkttsModelOutput(logits=logits, hidden_states=hidden)
        if return_dict is False:
            return (logits, hidden)
        return output

    def _setup_generation_caches(self, batch_size: int, max_length: int, dtype):
        del max_length
        for layer in self.layers:
            layer.attention.kv_cache = ArkttsKVCache(
                batch_size,
                self.config.max_seq_len,
                self.config.n_local_heads,
                self.config.head_dim,
                dtype,
                return_full=True,
            ).to(self.device)
        for layer in self.fast_layers:
            layer.attention.kv_cache = ArkttsKVCache(
                batch_size, self.config.num_codebooks, self.config.fast_n_local_heads,
                self.config.fast_head_dim, dtype, return_full=True,
            ).to(self.device)

    def _slow_step(
        self,
        input_ids: Tensor,
        cache_position: Tensor,
        position_ids: Tensor,
        attention_mask: Tensor,
    ) -> tuple[Tensor, Tensor]:
        hidden = self._embed(input_ids)
        rope = self.freqs_cis[position_ids]
        mask = self._causal_mask(
            attention_mask,
            cache_position,
            self.config.max_seq_len,
        )
        for layer in self.layers:
            hidden = layer(hidden, rope, mask, cache_position)
        hidden = hidden[:, -1:]
        normalized = self.norm(hidden)
        logits = F.linear(normalized, self.embeddings.weight)[:, -1]
        fast_hidden = normalized if self.config.norm_fastlayer_input else hidden
        return logits, fast_hidden

    def _fast_step(self, hidden: Tensor, position: int) -> Tensor:
        cache_position = torch.tensor([position], device=hidden.device, dtype=torch.long)
        rope = self.fast_freqs_cis[cache_position]
        key_mask = torch.ones(
            (hidden.shape[0], self.config.num_codebooks),
            device=hidden.device,
            dtype=torch.bool,
        )
        mask = self._causal_mask(key_mask, cache_position, self.config.num_codebooks)
        for layer in self.fast_layers:
            hidden = layer(hidden, rope, mask, cache_position)
        return self.fast_output(self.fast_norm(hidden))[:, -1]

    @staticmethod
    def _as_processor_list(value) -> LogitsProcessorList:
        if value is None:
            return LogitsProcessorList()
        if isinstance(value, LogitsProcessorList):
            return value
        return LogitsProcessorList(value)

    @staticmethod
    def _sample(scores: Tensor, generator=None) -> Tensor:
        probabilities = torch.softmax(scores, dim=-1)
        random = torch.rand(
            probabilities.shape,
            dtype=probabilities.dtype,
            device=probabilities.device,
            generator=generator,
        )
        noise = -torch.log(random)
        return torch.argmax(probabilities / noise, dim=-1)

    def _processed_scores(
        self,
        input_ids: Tensor,
        scores: Tensor,
        processors: LogitsProcessorList,
        top_k: int,
        top_p: float,
        temperature: float,
    ) -> Tensor:
        scores = processors(input_ids, scores)
        scores = ArkttsLegacyTopKTopPLogitsProcessor(top_k, top_p)(input_ids, scores)
        temperature_value = torch.tensor(
            temperature, dtype=scores.dtype, device=scores.device
        ).clamp_min(1e-5)
        return scores / temperature_value

    def _sample_semantic(
        self,
        history: Tensor,
        logits: Tensor,
        custom_processors: LogitsProcessorList,
        top_k: int,
        top_p: float,
        temperature: float,
        previous: Optional[Tensor],
        do_sample: bool,
        generator=None,
    ) -> Tensor:
        processors = LogitsProcessorList([
            ArkttsSemanticLogitsProcessor(
                self.config.semantic_begin_id,
                self.config.semantic_end_id,
                self.config.eos_token_id,
            ),
            *custom_processors,
        ])
        regular_scores = self._processed_scores(
            history, logits, processors, top_k, top_p, temperature
        )
        if not do_sample:
            return regular_scores.argmax(dim=-1)
        normal = self._sample(regular_scores, generator=generator)
        high_scores = self._processed_scores(
            history,
            logits,
            processors,
            top_k,
            self.config.ras_top_p,
            self.config.ras_temperature,
        )
        high = self._sample(high_scores, generator=generator)
        if previous is None:
            return normal
        repeated = (previous == normal[:, None]).any(dim=1)
        semantic = (normal >= self.config.semantic_begin_id) & (
            normal <= self.config.semantic_end_id
        )
        return torch.where(repeated & semantic, high, normal)

    def _generate_codebooks(
        self,
        slow_hidden: Tensor,
        semantic: Tensor,
        processors: LogitsProcessorList,
        top_k: int,
        top_p: float,
        temperature: float,
        do_sample: bool,
        generator=None,
    ) -> Tensor:
        hidden = self.fast_project_in(slow_hidden)
        self._fast_step(hidden, 0)
        current = (semantic - self.config.semantic_begin_id).clamp(0, self.config.codebook_size - 1)
        codebooks = [current]
        fast_history = current[:, None]
        hidden = self.fast_embeddings(current)[:, None]
        for position in range(1, self.config.num_codebooks):
            scores = self._fast_step(hidden, position)
            scores = self._processed_scores(
                fast_history, scores, processors, top_k, top_p, temperature
            )
            current = self._sample(scores, generator=generator) if do_sample else scores.argmax(dim=-1)
            codebooks.append(current)
            fast_history = torch.cat((fast_history, current[:, None]), dim=1)
            hidden = self.fast_embeddings(current)[:, None]
        return torch.stack(codebooks, dim=1)

    def _prepare_prompt(
        self,
        input_ids=None,
        attention_mask=None,
        prefix_input_ids=None,
        prefix_attention_mask=None,
        suffix_input_ids=None,
        suffix_attention_mask=None,
        reference_codes=None,
        reference_code_lengths=None,
        reference_audio_values=None,
        reference_audio_lengths=None,
    ):
        if input_ids is not None:
            if input_ids.ndim != 3 or input_ids.shape[1] != self.config.num_codebooks + 1:
                raise ValueError("Direct input_ids must have shape [B, num_codebooks + 1, T]")
            if attention_mask is None:
                attention_mask = torch.ones(
                    input_ids.shape[0], input_ids.shape[-1], dtype=torch.long, device=input_ids.device
                )
            return input_ids.to(self.device), attention_mask.to(self.device)
        if prefix_input_ids is None or suffix_input_ids is None:
            raise ValueError("Processor output or direct input_ids is required")
        prefix_input_ids = prefix_input_ids.to(self.device)
        suffix_input_ids = suffix_input_ids.to(self.device)
        prefix_attention_mask = prefix_attention_mask.to(self.device)
        suffix_attention_mask = suffix_attention_mask.to(self.device)
        if reference_audio_values is not None:
            if reference_codes is not None:
                raise ValueError("Provide reference audio or reference codes, not both")
            reference_codes, reference_code_lengths = self.encode_audio(
                reference_audio_values.to(self.device), reference_audio_lengths.to(self.device)
            )
        if reference_codes is not None:
            reference_codes = reference_codes.to(self.device)
            reference_code_lengths = reference_code_lengths.to(self.device)
        batch_size = prefix_input_ids.shape[0]
        rows = []
        for batch_index in range(batch_size):
            prefix = prefix_input_ids[batch_index, prefix_attention_mask[batch_index].bool()]
            suffix = suffix_input_ids[batch_index, suffix_attention_mask[batch_index].bool()]
            if reference_codes is None:
                semantic_row = torch.cat((prefix, suffix))
                values = torch.zeros(
                    (self.config.num_codebooks + 1, semantic_row.numel()),
                    dtype=torch.long,
                    device=self.device,
                )
                values[0] = semantic_row
            else:
                length = int(reference_code_lengths[batch_index])
                codes = reference_codes[batch_index, :, :length].long()
                semantic_codes = codes[0] + self.config.semantic_begin_id
                semantic_row = torch.cat((prefix, semantic_codes, suffix))
                values = torch.zeros(
                    (self.config.num_codebooks + 1, semantic_row.numel()),
                    dtype=torch.long,
                    device=self.device,
                )
                values[0] = semantic_row
                values[1:, prefix.numel() : prefix.numel() + length] = codes
            rows.append(values)
        max_length = max(row.shape[1] for row in rows)
        prompt = torch.zeros(
            (batch_size, self.config.num_codebooks + 1, max_length),
            dtype=torch.long,
            device=self.device,
        )
        prompt[:, 0] = self.config.pad_token_id
        prompt_mask = torch.zeros((batch_size, max_length), dtype=torch.long, device=self.device)
        for batch_index, row in enumerate(rows):
            start = max_length - row.shape[1]
            prompt[batch_index, :, start:] = row
            prompt_mask[batch_index, start:] = 1
        return prompt, prompt_mask

    @torch.inference_mode()
    def generate(
        self,
        input_ids: Optional[Tensor] = None,
        attention_mask: Optional[Tensor] = None,
        prefix_input_ids: Optional[Tensor] = None,
        prefix_attention_mask: Optional[Tensor] = None,
        suffix_input_ids: Optional[Tensor] = None,
        suffix_attention_mask: Optional[Tensor] = None,
        reference_codes: Optional[Tensor] = None,
        reference_code_lengths: Optional[Tensor] = None,
        reference_audio_values: Optional[Tensor] = None,
        reference_audio_lengths: Optional[Tensor] = None,
        generation_config: Optional[GenerationConfig] = None,
        logits_processor=None,
        codebook_logits_processor=None,
        stopping_criteria=None,
        max_new_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        top_k: Optional[int] = None,
        do_sample: Optional[bool] = None,
        return_dict_in_generate: bool = False,
        generator=None,
        **kwargs,
    ):
        if kwargs:
            raise TypeError(f"Unexpected generation arguments: {sorted(kwargs)}")
        config = generation_config or getattr(self, "generation_config", GenerationConfig())
        config_max_new = getattr(config, "max_new_tokens", None)
        config_temperature = getattr(config, "temperature", None)
        config_top_p = getattr(config, "top_p", None)
        config_top_k = getattr(config, "top_k", None)
        max_new_tokens = int(max_new_tokens if max_new_tokens is not None else (config_max_new or 512))
        temperature = float(temperature if temperature is not None else (config_temperature or 0.7))
        top_p = float(top_p if top_p is not None else (config_top_p or 0.9))
        top_k = int(top_k if top_k is not None else (config_top_k or 50))
        do_sample = bool(do_sample if do_sample is not None else getattr(config, "do_sample", True))
        if max_new_tokens <= 0:
            raise ValueError("max_new_tokens must be positive")
        if temperature <= 0:
            raise ValueError("temperature must be positive")
        prompt, prompt_mask = self._prepare_prompt(
            input_ids, attention_mask, prefix_input_ids, prefix_attention_mask,
            suffix_input_ids, suffix_attention_mask, reference_codes,
            reference_code_lengths, reference_audio_values, reference_audio_lengths,
        )
        batch_size, _, prompt_width = prompt.shape
        if prompt_width >= self.config.max_seq_len:
            raise ValueError(
                f"Prompt length {prompt_width} must be smaller than {self.config.max_seq_len}"
            )
        max_new_tokens = min(max_new_tokens, self.config.max_seq_len - prompt_width)
        self._setup_generation_caches(
            batch_size, prompt_width + max_new_tokens, next(self.parameters()).dtype
        )
        semantic_processors = self._as_processor_list(logits_processor)
        codebook_processors = self._as_processor_list(codebook_logits_processor)
        criteria = stopping_criteria or StoppingCriteriaList()
        if not isinstance(criteria, StoppingCriteriaList):
            criteria = StoppingCriteriaList(criteria)
        cache_position = torch.arange(prompt_width, device=self.device, dtype=torch.long)
        position_ids = prompt_mask.cumsum(-1).sub(1).clamp_min(0)
        logits, slow_hidden = self._slow_step(
            prompt, cache_position, position_ids, prompt_mask
        )
        semantic_history = prompt[:, 0]
        prompt_lengths = prompt_mask.sum(-1)
        previous = None
        finished = torch.zeros(batch_size, dtype=torch.bool, device=self.device)
        code_lengths = torch.zeros(batch_size, dtype=torch.long, device=self.device)
        generated_frames = []

        for step in range(max_new_tokens):
            active_before = ~finished
            semantic = self._sample_semantic(
                semantic_history, logits, semantic_processors, top_k, top_p,
                temperature, previous, do_sample, generator,
            )
            codebooks = self._generate_codebooks(
                slow_hidden, semantic, codebook_processors, top_k, top_p,
                temperature, do_sample, generator,
            )
            emitted = active_before & (semantic != self.config.eos_token_id)
            frame = torch.where(emitted[:, None], codebooks, -1)
            generated_frames.append(frame)
            code_lengths += emitted.long()
            semantic_history = torch.cat((semantic_history, semantic[:, None]), dim=1)

            if previous is None:
                previous = torch.zeros(
                    (batch_size, self.config.ras_window_size),
                    dtype=torch.long,
                    device=self.device,
                )
            else:
                previous = previous.roll(-1, dims=1)
                previous[:, -1] = semantic
            finished |= semantic.eq(self.config.eos_token_id)
            if criteria:
                stopped = criteria(semantic_history, logits)
                if not isinstance(stopped, Tensor):
                    stopped = torch.full_like(finished, bool(stopped))
                finished |= stopped.to(device=self.device, dtype=torch.bool)
            if finished.all():
                break

            next_column = torch.cat((semantic[:, None], codebooks), dim=1).unsqueeze(-1)
            new_valid = active_before.long()[:, None]
            prompt_mask = torch.cat((prompt_mask, new_valid), dim=1)
            physical_position = torch.tensor([prompt_width + step], device=self.device)
            token_position = (prompt_lengths + step)[:, None]
            with sdpa_kernel(SDPBackend.MATH):
                logits, slow_hidden = self._slow_step(
                    next_column, physical_position, token_position, prompt_mask
                )

        if generated_frames:
            codes = torch.stack(generated_frames, dim=2)
            max_valid = int(code_lengths.max().item()) if code_lengths.numel() else 0
            codes = codes[:, :, :max_valid]
        else:
            codes = torch.empty(
                (batch_size, self.config.num_codebooks, 0), dtype=torch.long, device=self.device
            )
        result = ArkttsGenerateOutput(codes=codes, code_lengths=code_lengths, finished=finished)
        return result if return_dict_in_generate else codes

    def _codec_path(self) -> str:
        source = str(getattr(self.config, "_name_or_path", ""))
        local = Path(source)
        if local.is_dir() and (local / self.config.codec_filename).is_file():
            return str(local / self.config.codec_filename)
        resolved = cached_file(source, self.config.codec_filename)
        if resolved is None:
            raise FileNotFoundError(f"Could not resolve {self.config.codec_filename} from {source}")
        return resolved

    def save_pretrained(self, save_directory, *args, **kwargs):
        result = super().save_pretrained(save_directory, *args, **kwargs)
        source = Path(self._codec_path()).resolve()
        destination = Path(save_directory) / self.config.codec_filename
        destination.parent.mkdir(parents=True, exist_ok=True)
        if source != destination.resolve():
            shutil.copy2(source, destination)
        return result

    def load_codec(self, device=None, dtype=None):
        codec = self.__dict__.get("_arktts_codec")
        target_device = torch.device(device) if device is not None else self.device
        target_dtype = dtype or self.dtype
        if target_device.type == "cpu":
            target_dtype = torch.float32
        if codec is None:
            from .modeling_arktts_codec import ArkttsCodec

            codec = ArkttsCodec(self.config)
            state = torch.load(self._codec_path(), map_location="cpu", weights_only=True)
            if "state_dict" in state:
                state = state["state_dict"]
            if any("generator." in key for key in state):
                state = {
                    key.replace("generator.", ""): value
                    for key, value in state.items() if "generator." in key
                }
            state = {
                key: value
                for key, value in state.items()
                if not key.endswith(("freqs_cis", "causal_mask"))
            }
            codec.load_state_dict(state, strict=True)
            codec.eval()
            self.__dict__["_arktts_codec"] = codec
        codec.to(device=target_device, dtype=target_dtype)
        return codec

    @torch.inference_mode()
    def encode_audio(self, audio_values: Tensor, audio_lengths: Optional[Tensor] = None):
        codec = self.load_codec(device=audio_values.device)
        audio_values = audio_values.to(dtype=next(codec.parameters()).dtype)
        return codec.encode(audio_values, audio_lengths)

    @torch.inference_mode()
    def decode_audio(self, codes: Tensor):
        if codes.ndim == 2:
            codes = codes.unsqueeze(0)
        if codes.ndim != 3 or codes.shape[1] != self.config.num_codebooks:
            raise ValueError(f"codes must have shape [B, {self.config.num_codebooks}, T]")
        codec = self.load_codec(device=codes.device)
        waveforms = []
        lengths = []
        for item in codes:
            valid = (item >= 0).all(dim=0)
            length = int(valid.sum().item())
            if length == 0:
                waveform = torch.empty(0, device=codes.device, dtype=torch.float32)
            else:
                waveform = codec.decode(item[:, :length].unsqueeze(0))[0, 0].float()
            waveforms.append(waveform)
            lengths.append(waveform.numel())
        max_length = max(lengths, default=0)
        padded = torch.zeros((len(waveforms), max_length), dtype=torch.float32, device=codes.device)
        for index, waveform in enumerate(waveforms):
            padded[index, : waveform.numel()] = waveform
        return padded, torch.tensor(lengths, dtype=torch.long, device=codes.device)

    @torch.inference_mode()
    def generate_audio(self, **kwargs):
        codes = self.generate(**kwargs)
        waveforms, lengths = self.decode_audio(codes)
        return waveforms, lengths, codes


__all__ = [
    "ArkttsConfig",
    "ArkttsGenerateOutput",
    "ArkttsLegacyTopKTopPLogitsProcessor",
    "ArkttsModel",
    "ArkttsModelOutput",
    "ArkttsSemanticLogitsProcessor",
]