Mirror processing_moss_tts.py

processing_moss_tts.py

@@ -0,0 +1,914 @@
+# coding=utf-8
+"""Processor for the MOSS-TTS-Local-Transformer-v1.5 HuggingFace release."""
+
+from __future__ import annotations
+
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
+
+import torch
+import torchaudio
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoTokenizer,
+    BatchFeature,
+    PreTrainedTokenizerBase,
+    ProcessorMixin,
+    logging,
+    processing_utils,
+)
+
+from .configuration_moss_tts import MossTTSLocalConfig
+
+
+if hasattr(processing_utils, "MODALITY_TO_BASE_CLASS_MAPPING"):
+    processing_utils.MODALITY_TO_BASE_CLASS_MAPPING["audio_tokenizer"] = "PreTrainedModel"
+else:
+    processing_utils.AUTO_TO_BASE_CLASS_MAPPING["AutoModel"] = "PreTrainedModel"
+logger = logging.get_logger(__name__)
+
+AUDIO_PLACEHOLDER = "<|audio|>"
+USER_ROLE_PREFIX = "user\n"
+USER_TEMPLATE_REFERENCE_PREFIX = (
+    "<user_inst>\n"
+    "- Reference(s):\n"
+)
+USER_TEMPLATE_AFTER_REFERENCE_SUFFIX = (
+    "\n"
+    "- Text:\n"
+)
+USER_TEMPLATE_SUFFIX = "\n</user_inst>"
+ASSISTANT_TURN_PREFIX = "\n"
+ASSISTANT_ROLE_PREFIX = "assistant\n"
+USER_MESSAGE_FIELDS = (
+    "text",
+    "reference",
+    "instruction",
+    "tokens",
+    "quality",
+    "sound_event",
+    "ambient_sound",
+    "language",
+)
+
+
+def _normalize_template_value(value: Any) -> str:
+    if value is None:
+        return "None"
+    resolved = str(value).strip()
+    return resolved or "None"
+
+
+def _render_user_prompt_after_reference(
+    language_code: object | None = None,
+    prompt_fields: Optional[Dict[str, Any]] = None,
+) -> str:
+    fields = dict(prompt_fields or {})
+    return (
+        "\n- Instruction:\n"
+        + _normalize_template_value(fields.get("instruction"))
+        + "\n- Tokens:\n"
+        + _normalize_template_value(fields.get("tokens"))
+        + "\n- Quality:\n"
+        + _normalize_template_value(fields.get("quality"))
+        + "\n- Sound Event:\n"
+        + _normalize_template_value(fields.get("sound_event"))
+        + "\n- Ambient Sound:\n"
+        + _normalize_template_value(fields.get("ambient_sound"))
+        + "\n- Language:\n"
+        + _normalize_template_value(fields.get("language", language_code))
+        + USER_TEMPLATE_AFTER_REFERENCE_SUFFIX
+    )
+
+
+@dataclass
+class Message:
+    def to_dict(self) -> Dict[str, Any]:
+        raise NotImplementedError
+
+
+@dataclass
+class UserMessage(Message):
+    text: Optional[str] = None
+    reference: Optional[List[Optional[Union[str, os.PathLike, torch.Tensor]]]] = None
+    instruction: Optional[str] = None
+    tokens: Optional[int] = None
+    quality: Optional[str] = None
+    sound_event: Optional[str] = None
+    ambient_sound: Optional[str] = None
+    language: Optional[str] = None
+
+    def __post_init__(self) -> None:
+        template = """<user_inst>
+- Reference(s):
+{reference}
+- Instruction:
+{instruction}
+- Tokens:
+{tokens}
+- Quality:
+{quality}
+- Sound Event:
+{sound_event}
+- Ambient Sound:
+{ambient_sound}
+- Language:
+{language}
+- Text:
+{text}
+</user_inst>"""
+
+        audio_codes_list: list[Union[str, os.PathLike, torch.Tensor]] = []
+        if self.reference is None:
+            reference = "None"
+        else:
+            reference_items: list[str] = []
+            for speaker_idx, speaker_reference in enumerate(self.reference):
+                if speaker_reference is None:
+                    continue
+                # Keep raw audio placeholders directly under "- Reference(s):".
+                # Speaker labels such as "[S1]:" change the token sequence and
+                # can affect voice-clone conditioning.
+                reference_items.append(AUDIO_PLACEHOLDER)
+                audio_codes_list.append(speaker_reference)
+            reference = "\n".join(reference_items) if reference_items else "None"
+
+        self._content = (
+            template.replace("{reference}", str(reference))
+            .replace("{instruction}", str(self.instruction))
+            .replace("{tokens}", str(self.tokens))
+            .replace("{quality}", str(self.quality))
+            .replace("{sound_event}", str(self.sound_event))
+            .replace("{ambient_sound}", str(self.ambient_sound))
+            .replace("{language}", str(self.language))
+            .replace("{text}", str(self.text))
+        )
+        self._audio_codes_list = audio_codes_list
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "role": "user",
+            "content": self._content,
+            "audio_codes_list": self._audio_codes_list,
+            "text": self.text,
+            "instruction": self.instruction,
+            "tokens": self.tokens,
+            "quality": self.quality,
+            "sound_event": self.sound_event,
+            "ambient_sound": self.ambient_sound,
+            "language": self.language,
+        }
+
+
+@dataclass
+class AssistantMessage(Message):
+    audio_codes_list: List[Union[str, os.PathLike, torch.Tensor]]
+    content: str = AUDIO_PLACEHOLDER
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "role": "assistant",
+            "content": self.content,
+            "audio_codes_list": self.audio_codes_list,
+        }
+
+
+class MossTTSLocalProcessor(ProcessorMixin):
+    attributes = ["tokenizer"]
+    tokenizer_class = "AutoTokenizer"
+    audio_tokenizer_class = "AutoModel"
+
+    tokenizer: PreTrainedTokenizerBase
+    audio_tokenizer: Any
+
+    def __init__(
+        self,
+        tokenizer: PreTrainedTokenizerBase,
+        audio_tokenizer: Any = None,
+        model_config: Optional[MossTTSLocalConfig] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(tokenizer=tokenizer, audio_tokenizer=audio_tokenizer, **kwargs)
+        self.tokenizer = tokenizer
+        self.audio_tokenizer = audio_tokenizer
+        self.model_config = model_config or MossTTSLocalConfig()
+
+        def _id_to_token(token_id: int) -> str:
+            token = tokenizer.convert_ids_to_tokens(int(token_id))
+            if isinstance(token, list):
+                return token[0] if token else ""
+            return cast(str, token)
+
+        self.audio_user_slot_token = _id_to_token(self.model_config.audio_user_slot_token_id)
+        self.audio_assistant_slot_token = _id_to_token(self.model_config.audio_assistant_slot_token_id)
+        self.audio_start_token = _id_to_token(self.model_config.audio_start_token_id)
+        self.audio_end_token = _id_to_token(self.model_config.audio_end_token_id)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        trust_remote_code = kwargs.pop("trust_remote_code", True)
+        kwargs.pop("_from_auto", None)
+        codec_path = kwargs.pop("codec_path", None)
+        codec_weight_dtype = kwargs.pop("codec_weight_dtype", "bf16")
+        codec_compute_dtype = kwargs.pop("codec_compute_dtype", None)
+        codec_attention_implementation = kwargs.pop("codec_attention_implementation", None)
+
+        codec_kwargs: Dict[str, Any] = {}
+        if codec_weight_dtype is not None:
+            # Default to bf16 codec weights; callers can pass codec_weight_dtype="fp32" to override.
+            codec_kwargs["codec_weight_dtype"] = codec_weight_dtype
+
+        model_ref = Path(str(pretrained_model_name_or_path))
+        model_ref_or_name = model_ref if model_ref.exists() else pretrained_model_name_or_path
+        model_config = cast(
+            MossTTSLocalConfig,
+            AutoConfig.from_pretrained(
+                model_ref_or_name,
+                *args,
+                trust_remote_code=trust_remote_code,
+                **kwargs,
+            ),
+        )
+
+        if codec_path is None:
+            try:
+                processor_dict, _ = cls.get_processor_dict(
+                    pretrained_model_name_or_path,
+                    **dict(kwargs),
+                )
+                codec_path = processor_dict.get("audio_tokenizer_name_or_path")
+                audio_tokenizer_dict = processor_dict.get("audio_tokenizer", {})
+                if isinstance(audio_tokenizer_dict, dict):
+                    codec_path = audio_tokenizer_dict.get("audio_tokenizer_name_or_path") or codec_path
+            except Exception:
+                codec_path = None
+        if codec_path is None:
+            codec_path = getattr(model_config, "audio_tokenizer_name_or_path", None)
+        if codec_path is None:
+            codec_path = "OpenMOSS-Team/MOSS-Audio-Tokenizer-v2"
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_ref_or_name,
+            *args,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
+        audio_tokenizer = AutoModel.from_pretrained(
+            codec_path,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+            **codec_kwargs,
+        )
+        if codec_compute_dtype is not None and hasattr(audio_tokenizer, "set_compute_dtype"):
+            audio_tokenizer.set_compute_dtype(codec_compute_dtype)
+        if codec_attention_implementation is not None and hasattr(
+            audio_tokenizer, "set_attention_implementation"
+        ):
+            audio_tokenizer.set_attention_implementation(codec_attention_implementation)
+        return cls(
+            tokenizer=tokenizer,
+            audio_tokenizer=audio_tokenizer,
+            model_config=model_config,
+            **kwargs,
+        )
+
+    @staticmethod
+    def build_user_message(
+        text: Optional[str] = None,
+        reference: Optional[List[Optional[Union[str, os.PathLike, torch.Tensor]]]] = None,
+        instruction: Optional[str] = None,
+        tokens: Optional[int] = None,
+        quality: Optional[str] = None,
+        sound_event: Optional[str] = None,
+        ambient_sound: Optional[str] = None,
+        language: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        if reference is not None and not isinstance(reference, list):
+            reference = [reference]
+        return UserMessage(
+            text=text,
+            reference=reference,
+            instruction=instruction,
+            tokens=tokens,
+            quality=quality,
+            sound_event=sound_event,
+            ambient_sound=ambient_sound,
+            language=language,
+        ).to_dict()
+
+    @staticmethod
+    def build_assistant_message(
+        audio_codes_list: List[Union[str, os.PathLike, torch.Tensor]],
+        content: str = AUDIO_PLACEHOLDER,
+    ) -> Dict[str, Any]:
+        return AssistantMessage(audio_codes_list=audio_codes_list, content=content).to_dict()
+
+    def _assert_fixed_nq(self, n_vq: Optional[int]) -> int:
+        config_nq = int(self.model_config.n_vq)
+        if n_vq is not None and int(n_vq) != config_nq:
+            raise ValueError(
+                "This MOSS-TTS-Local-Transformer-v1.5 release uses the RVQ depth stored in the model config. "
+                f"Expected n_vq={config_nq}, got {int(n_vq)}."
+            )
+        return config_nq
+
+    def _encode_text(self, text: str) -> list[int]:
+        try:
+            return list(self.tokenizer.encode(text, add_special_tokens=False))
+        except TypeError:
+            return list(self.tokenizer.encode(text))
+
+    def _build_text_rows(self, token_ids: Sequence[int], *, device: Optional[torch.device] = None) -> torch.Tensor:
+        rows = torch.full(
+            (len(token_ids), int(self.model_config.n_vq) + 1),
+            int(self.model_config.audio_pad_token_id),
+            dtype=torch.long,
+            device=device,
+        )
+        if token_ids:
+            rows[:, 0] = torch.tensor([int(token_id) for token_id in token_ids], dtype=torch.long, device=rows.device)
+        return rows
+
+    def _build_audio_rows(self, audio_tokens: torch.Tensor, slot_token_id: int) -> torch.Tensor:
+        rows = torch.full(
+            (int(audio_tokens.shape[0]), int(self.model_config.n_vq) + 1),
+            int(self.model_config.audio_pad_token_id),
+            dtype=torch.long,
+            device=audio_tokens.device,
+        )
+        if rows.shape[0] > 0:
+            rows[:, 0] = int(slot_token_id)
+            rows[:, 1:] = audio_tokens.to(dtype=torch.long)
+        return rows
+
+    def _user_prompt_prefix_ids(self) -> list[int]:
+        return (
+            [int(self.model_config.im_start_token_id)]
+            + self._encode_text(USER_ROLE_PREFIX)
+            + self._encode_text(USER_TEMPLATE_REFERENCE_PREFIX)
+        )
+
+    def _user_prompt_after_reference_ids(
+        self,
+        language_code: object | None,
+        prompt_fields: Optional[Dict[str, Any]],
+    ) -> list[int]:
+        return self._encode_text(
+            _render_user_prompt_after_reference(
+                language_code=language_code,
+                prompt_fields=prompt_fields,
+            )
+        )
+
+    def _assistant_prompt_prefix_ids(self) -> list[int]:
+        return (
+            self._encode_text(USER_TEMPLATE_SUFFIX)
+            + [int(self.model_config.im_end_token_id)]
+            + self._encode_text(ASSISTANT_TURN_PREFIX)
+            + [int(self.model_config.im_start_token_id)]
+            + self._encode_text(ASSISTANT_ROLE_PREFIX)
+        )
+
+    def _prompt_fields_from_user_message(self, message: Dict[str, Any]) -> dict[str, Any]:
+        fields = {}
+        for key in ("instruction", "tokens", "quality", "sound_event", "ambient_sound"):
+            if key in message and message.get(key) is not None:
+                fields[key] = message.get(key)
+        if "language" in message and message.get("language") is not None:
+            fields["language"] = message.get("language")
+        return fields
+
+    def _build_generation_or_voice_clone_codes(
+        self,
+        message: Dict[str, Any],
+        n_vq: int,
+    ) -> torch.Tensor:
+        if "text" not in message:
+            raise ValueError("Direct MOSS-TTS-Local-Transformer-v1.5 generation requires messages built by build_user_message(...).")
+        text = "" if message.get("text") is None else str(message.get("text"))
+        prompt_fields = self._prompt_fields_from_user_message(message)
+        language_code = message.get("language")
+        audio_codes_list = self._resolve_audio_items(message.get("audio_codes_list", []), n_vq)
+        text_token_ids = self._encode_text(text)
+
+        if audio_codes_list:
+            parts: list[torch.Tensor] = [self._build_text_rows(
+                self._user_prompt_prefix_ids(),
+                device=audio_codes_list[0].device,
+            )]
+            for reference_codes in audio_codes_list:
+                parts.append(self._build_text_rows([int(self.model_config.audio_start_token_id)], device=reference_codes.device))
+                parts.append(self._build_audio_rows(reference_codes, int(self.model_config.audio_user_slot_token_id)))
+                parts.append(self._build_text_rows([int(self.model_config.audio_end_token_id)], device=reference_codes.device))
+            parts.append(
+                self._build_text_rows(
+                    self._user_prompt_after_reference_ids(language_code, prompt_fields)
+                    + text_token_ids
+                    + self._assistant_prompt_prefix_ids()
+                    + [int(self.model_config.audio_start_token_id)],
+                    device=audio_codes_list[0].device,
+                )
+            )
+            return torch.cat(parts, dim=0)
+
+        prompt_token_ids = (
+            self._user_prompt_prefix_ids()
+            + self._encode_text("None")
+            + self._user_prompt_after_reference_ids(language_code, prompt_fields)
+            + text_token_ids
+            + self._assistant_prompt_prefix_ids()
+            + [int(self.model_config.audio_start_token_id)]
+        )
+        return self._build_text_rows(prompt_token_ids)
+
+    def _build_continuation_codes(
+        self,
+        conversation: list[Dict[str, Any]],
+        n_vq: int,
+    ) -> torch.Tensor:
+        if len(conversation) < 2:
+            raise ValueError("continuation mode requires a user message followed by an assistant audio message.")
+        user_message = conversation[-2]
+        assistant_message = conversation[-1]
+        if user_message.get("role") != "user" or assistant_message.get("role") != "assistant":
+            raise ValueError("continuation mode requires the last two messages to be user, assistant.")
+        if "text" not in user_message:
+            raise ValueError("Direct MOSS-TTS-Local-Transformer-v1.5 continuation requires user messages built by build_user_message(...).")
+
+        text = "" if user_message.get("text") is None else str(user_message.get("text"))
+        prompt_fields = self._prompt_fields_from_user_message(user_message)
+        language_code = user_message.get("language")
+        prompt_token_ids = (
+            self._user_prompt_prefix_ids()
+            + self._encode_text("None")
+            + self._user_prompt_after_reference_ids(language_code, prompt_fields)
+            + self._encode_text(text)
+            + self._assistant_prompt_prefix_ids()
+            + [int(self.model_config.audio_start_token_id)]
+        )
+        audio_codes_list = self._resolve_audio_items(assistant_message.get("audio_codes_list", []), n_vq)
+        if not audio_codes_list:
+            return self._build_text_rows(prompt_token_ids)
+        if len(audio_codes_list) != 1:
+            raise ValueError("The MOSS-TTS-Local-Transformer-v1.5 continuation path expects exactly one prompt audio item.")
+        prompt_audio_codes = audio_codes_list[0]
+        return torch.cat(
+            [
+                self._build_text_rows(prompt_token_ids, device=prompt_audio_codes.device),
+                self._build_audio_rows(prompt_audio_codes, int(self.model_config.audio_assistant_slot_token_id)),
+            ],
+            dim=0,
+        )
+
+    def _try_build_direct_codes(
+        self,
+        conversation: list[Dict[str, Any]],
+        mode: str,
+        n_vq: int,
+    ) -> Optional[torch.Tensor]:
+        if mode == "generation" and len(conversation) == 1 and conversation[-1].get("role") == "user":
+            if "text" in conversation[-1]:
+                return self._build_generation_or_voice_clone_codes(conversation[-1], n_vq)
+            return None
+        if mode == "continuation" and len(conversation) >= 2:
+            if "text" in conversation[-2]:
+                return self._build_continuation_codes(conversation, n_vq)
+            return None
+        return None
+
+    def __call__(self, *args, **kwargs) -> BatchFeature:
+        conversations = args[0] if args else kwargs.pop("conversations")
+        mode: str = kwargs.pop("mode", "generation")
+        apply_chat_template: bool = kwargs.pop("apply_chat_template", True)
+        n_vq = self._assert_fixed_nq(kwargs.pop("n_vq", None))
+
+        kwargs.pop("return_tensors", None)
+        kwargs.pop("padding", None)
+        kwargs.pop("truncation", None)
+
+        if mode not in {"generation", "continuation", "computing_loss"}:
+            raise ValueError(f"Unsupported mode: {mode}")
+        if isinstance(conversations, (Message, dict)):
+            conversations = [conversations]
+        elif isinstance(conversations, list) and conversations and all(
+            isinstance(item, (Message, dict)) for item in conversations
+        ):
+            conversations = [conversations]
+
+        input_ids_list: list[torch.Tensor] = []
+        for conversation in conversations:
+            if isinstance(conversation, (Message, dict)):
+                conversation = [conversation]
+            conversation = [self._normalize_message(message) for message in conversation]
+
+            if (mode == "generation") ^ (conversation[-1]["role"] == "user"):
+                raise ValueError("generation mode must end with a user message.")
+            if mode == "continuation" and conversation[-1]["role"] != "assistant":
+                raise ValueError("continuation mode must end with an assistant message.")
+
+            direct_codes = self._try_build_direct_codes(conversation, mode, n_vq)
+            if direct_codes is not None:
+                input_ids_list.append(direct_codes)
+                continue
+
+            unified_parts = []
+            for message_idx, message in enumerate(conversation):
+                content = str(message["content"])
+                if apply_chat_template:
+                    add_generation_prompt = mode == "generation" and message_idx == len(conversation) - 1
+                    try:
+                        content = self.tokenizer.apply_chat_template(
+                            [{"role": message["role"], "content": content}],
+                            add_generation_prompt=add_generation_prompt,
+                            tokenize=False,
+                        )
+                    except Exception:
+                        logger.warning("apply_chat_template failed; falling back to raw message content.")
+
+                raw_audio_items = message.get("audio_codes_list", [])
+                audio_codes_list = self._resolve_audio_items(raw_audio_items, n_vq)
+                unified_parts.append(
+                    self._get_unified_codes(
+                        role=message["role"],
+                        content=content,
+                        audio_codes_list=audio_codes_list,
+                        truncation=(mode == "continuation"),
+                    )
+                )
+
+            unified_codes = torch.cat(unified_parts, dim=0)
+            if mode == "generation":
+                audio_start_row = torch.full(
+                    (1, n_vq + 1),
+                    int(self.model_config.audio_pad_token_id),
+                    dtype=unified_codes.dtype,
+                    device=unified_codes.device,
+                )
+                audio_start_row[:, 0] = int(self.model_config.audio_start_token_id)
+                unified_codes = torch.cat([unified_codes, audio_start_row], dim=0)
+            input_ids_list.append(unified_codes)
+
+        return BatchFeature(data=self._pad(input_ids_list))
+
+    def _normalize_message(self, message: Union[Message, Dict[str, Any]]) -> Dict[str, Any]:
+        if isinstance(message, Message):
+            return message.to_dict()
+        if not isinstance(message, dict):
+            raise TypeError("Each message must be a Message or dict.")
+        if "content" in message and "audio_codes_list" in message:
+            return message
+        role = message.get("role")
+        if role == "user":
+            return self.build_user_message(**{key: message.get(key) for key in USER_MESSAGE_FIELDS})
+        if role == "assistant":
+            return self.build_assistant_message(
+                audio_codes_list=message.get("audio_codes_list", []),
+                content=message.get("content", AUDIO_PLACEHOLDER),
+            )
+        raise ValueError(f"Unsupported role: {role}")
+
+    def _resolve_audio_items(
+        self,
+        raw_audio_items: list[Any],
+        n_vq: int,
+    ) -> list[torch.Tensor]:
+        if not raw_audio_items:
+            return []
+        resolved: list[Optional[torch.Tensor]] = [None] * len(raw_audio_items)
+        paths: list[str] = []
+        path_positions: list[int] = []
+        for index, item in enumerate(raw_audio_items):
+            if isinstance(item, torch.Tensor):
+                if item.ndim != 2 or int(item.shape[1]) != n_vq:
+                    raise ValueError(f"audio code tensor must have shape [T, {n_vq}], got {tuple(item.shape)}.")
+                resolved[index] = item.to(dtype=torch.long).cpu()
+            elif isinstance(item, (str, os.PathLike)):
+                paths.append(str(item))
+                path_positions.append(index)
+            else:
+                raise TypeError("Audio items must be tensors or path-like values.")
+        if paths:
+            encoded = self.encode_audios_from_path(paths, n_vq=n_vq)
+            for position, codes in zip(path_positions, encoded):
+                resolved[position] = codes
+        return [cast(torch.Tensor, item) for item in resolved]
+
+    def _pad(self, input_ids_list: list[torch.Tensor]) -> Dict[str, torch.Tensor]:
+        device = input_ids_list[0].device
+        lengths = torch.tensor([item.shape[0] for item in input_ids_list], device=device)
+        padded = torch.nn.utils.rnn.pad_sequence(
+            input_ids_list,
+            batch_first=True,
+            padding_value=int(self.model_config.audio_pad_token_id),
+            padding_side="left",
+        )
+        left_pad_mask = (padded.shape[1] - lengths).unsqueeze(1) > torch.arange(
+            padded.shape[1],
+            device=device,
+        ).unsqueeze(0)
+        padded[..., 0][left_pad_mask] = int(self.model_config.pad_token_id)
+        attention_mask = torch.zeros(padded.shape[:2], dtype=torch.bool, device=device)
+        attention_mask[~left_pad_mask] = True
+        return {"input_ids": padded, "attention_mask": attention_mask}
+
+    @staticmethod
+    def _replace_audio_placeholders(
+        content: str,
+        lengths: list[int],
+        slot_token: str,
+        audio_start_token: str,
+        audio_end_token: str,
+    ) -> str:
+        placeholder_count = content.count(AUDIO_PLACEHOLDER)
+        if placeholder_count != len(lengths):
+            raise ValueError(
+                f"Number of {AUDIO_PLACEHOLDER} ({placeholder_count}) does not match "
+                f"audio item count ({len(lengths)})."
+            )
+        lengths_iter = iter(lengths)
+
+        def replacer(_: re.Match) -> str:
+            length = int(next(lengths_iter))
+            if length <= 0:
+                return f"{audio_start_token}{audio_end_token}"
+            return f"{audio_start_token}{slot_token * length}{audio_end_token}"
+
+        return re.sub(re.escape(AUDIO_PLACEHOLDER), replacer, content)
+
+    def _get_unified_codes(
+        self,
+        role: str,
+        content: str,
+        audio_codes_list: list[torch.Tensor],
+        truncation: bool,
+    ) -> torch.Tensor:
+        n_vq = int(self.model_config.n_vq)
+        slot_token = self.audio_user_slot_token if role == "user" else self.audio_assistant_slot_token
+        content = self._replace_audio_placeholders(
+            content=content,
+            lengths=[int(codes.shape[0]) for codes in audio_codes_list],
+            slot_token=slot_token,
+            audio_start_token=self.audio_start_token,
+            audio_end_token=self.audio_end_token,
+        )
+        text_codes = torch.tensor(
+            self.tokenizer.encode(content),
+            dtype=torch.long,
+            device=audio_codes_list[0].device if audio_codes_list else None,
+        )
+
+        audio_start_indices = torch.where(text_codes == int(self.model_config.audio_start_token_id))[0]
+        audio_end_indices = torch.where(text_codes == int(self.model_config.audio_end_token_id))[0]
+        if len(audio_start_indices) != len(audio_codes_list) or len(audio_end_indices) != len(audio_codes_list):
+            raise ValueError("Audio placeholders do not match the encoded audio spans.")
+
+        if not audio_codes_list:
+            audio_codes = torch.full(
+                (len(text_codes), n_vq),
+                int(self.model_config.audio_pad_token_id),
+                dtype=torch.long,
+                device=text_codes.device,
+            )
+        else:
+            pieces: list[torch.Tensor] = []
+            prefix_idx = 0
+            for start_t, end_t, codes in zip(audio_start_indices, audio_end_indices, audio_codes_list):
+                start_idx = int(start_t.item())
+                end_idx = int(end_t.item())
+                pad_before = torch.full(
+                    (start_idx - prefix_idx + 1, n_vq),
+                    int(self.model_config.audio_pad_token_id),
+                    dtype=torch.long,
+                    device=codes.device,
+                )
+                pieces.extend([pad_before, codes.to(dtype=torch.long)])
+                prefix_idx = end_idx
+            if truncation:
+                trailing = torch.zeros(
+                    (0, n_vq),
+                    dtype=torch.long,
+                    device=audio_codes_list[0].device,
+                )
+            else:
+                last_end = int(audio_end_indices[-1].item())
+                trailing = torch.full(
+                    (len(text_codes) - last_end, n_vq),
+                    int(self.model_config.audio_pad_token_id),
+                    dtype=torch.long,
+                    device=audio_codes_list[0].device,
+                )
+            pieces.append(trailing)
+            audio_codes = torch.cat(pieces, dim=0)
+
+        if text_codes.shape[0] != audio_codes.shape[0]:
+            min_len = min(text_codes.shape[0], audio_codes.shape[0])
+            text_codes = text_codes[:min_len]
+            audio_codes = audio_codes[:min_len]
+        return torch.cat([text_codes.unsqueeze(1), audio_codes], dim=1)
+
+    def _parse_text_codes(self, start_length: int, text_codes: torch.LongTensor) -> str:
+        text = cast(str, self.tokenizer.decode(text_codes))
+        prefix = cast(str, self.tokenizer.decode(text_codes[:start_length]))
+        text = text[len(prefix):]
+        audio_pattern = re.compile(
+            rf"(?:{re.escape(self.audio_start_token)})?"
+            rf"(?:{re.escape(self.audio_assistant_slot_token)})*"
+            rf"{re.escape(self.audio_end_token)}"
+        )
+        return audio_pattern.sub(
+            lambda match: AUDIO_PLACEHOLDER if self.audio_assistant_slot_token in match.group(0) else "",
+            text,
+        )
+
+    def _parse_audio_codes(
+        self,
+        start_length: int,
+        audio_codes: torch.LongTensor,
+        *,
+        return_stereo: bool = True,
+    ) -> list[torch.Tensor]:
+        is_pad = audio_codes.eq(int(self.model_config.audio_pad_token_id)).all(dim=1)
+        non_pad = ~is_pad
+        if not bool(non_pad.any().item()):
+            return []
+        idx = torch.nonzero(non_pad).squeeze(1)
+        breaks = torch.where(idx[1:] != idx[:-1] + 1)[0] + 1
+        segment_indices = [idx] if breaks.numel() == 0 else list(torch.tensor_split(idx, breaks.cpu().tolist()))
+        code_segments = [audio_codes[segment] for segment in segment_indices]
+        decoded = self.decode_audio_codes(code_segments, return_stereo=return_stereo)
+
+        if start_length > 0 and code_segments and decoded:
+            first_code_length = int(code_segments[0].shape[0])
+            if first_code_length > 0:
+                trim_ratio = max(0.0, min(float(start_length) / float(first_code_length), 1.0))
+                if trim_ratio >= 1.0:
+                    decoded = decoded[1:]
+                elif trim_ratio > 0.0:
+                    trim_samples = int(decoded[0].shape[-1] * trim_ratio)
+                    decoded[0] = decoded[0][..., trim_samples:]
+        return decoded
+
+    def decode(self, output: Any, *, return_stereo: bool = True) -> list[Optional[AssistantMessage]]:
+        generated_messages: list[Optional[AssistantMessage]] = []
+        for start_length, generation_ids in output:
+            content = self._parse_text_codes(int(start_length), generation_ids[:, 0])
+            audio_codes_list = self._parse_audio_codes(
+                int(start_length),
+                generation_ids[:, 1:],
+                return_stereo=return_stereo,
+            )
+            if content == "":
+                generated_messages.append(None)
+            else:
+                generated_messages.append(
+                    AssistantMessage(
+                        content=content,
+                        audio_codes_list=cast(list[Union[str, torch.Tensor]], audio_codes_list),
+                    )
+                )
+        return generated_messages
+
+    @staticmethod
+    def loudness_normalize(
+        wav: torch.Tensor,
+        target_dbfs: float = -20.0,
+        gain_range: tuple[float, float] = (-3.0, 3.0),
+    ) -> torch.Tensor:
+        wav = wav.to(torch.float32)
+        if wav.numel() == 0:
+            return wav
+        current_dbfs = 10.0 * torch.log10(torch.mean(wav**2) + 1e-9)
+        gain = max(gain_range[0], min(float(target_dbfs - current_dbfs), gain_range[1]))
+        return wav * (10.0 ** (gain / 20.0))
+
+    def _get_audio_tokenizer_device(self) -> torch.device:
+        audio_tokenizer = getattr(self, "audio_tokenizer", None)
+        if audio_tokenizer is None:
+            raise RuntimeError("audio_tokenizer is not set.")
+        try:
+            return next(audio_tokenizer.parameters()).device
+        except StopIteration:
+            return torch.device("cpu")
+
+    def encode_audios_from_wav(
+        self,
+        wav_list: Union[torch.Tensor, list[torch.Tensor]],
+        sampling_rate: int,
+        n_vq: Optional[int] = None,
+    ) -> list[torch.Tensor]:
+        n_vq = self._assert_fixed_nq(n_vq)
+        if self.audio_tokenizer is None:
+            raise RuntimeError("audio_tokenizer is not set.")
+        if isinstance(wav_list, torch.Tensor):
+            wav_list = [wav_list]
+        target_sr = int(self.model_config.sampling_rate)
+        device = self._get_audio_tokenizer_device()
+        prepared = []
+        for wav in wav_list:
+            if wav.ndim == 1:
+                wav = wav.unsqueeze(0)
+            if wav.shape[0] == 1:
+                wav = wav.repeat(2, 1)
+            elif wav.shape[0] > 2:
+                wav = wav[:2]
+            if int(sampling_rate) != target_sr:
+                wav = torchaudio.functional.resample(wav, int(sampling_rate), target_sr)
+            prepared.append(self.loudness_normalize(wav).to(device))
+
+        if hasattr(self.audio_tokenizer, "batch_encode"):
+            encoded = self.audio_tokenizer.batch_encode(prepared, num_quantizers=n_vq)
+            audio_codes = encoded.audio_codes
+            audio_lengths = encoded.audio_codes_lengths
+        else:
+            max_len = max(int(wav.shape[-1]) for wav in prepared)
+            input_values = torch.zeros(len(prepared), 1, max_len, dtype=torch.float32, device=device)
+            padding_mask = torch.zeros(len(prepared), max_len, dtype=torch.bool, device=device)
+            for index, wav in enumerate(prepared):
+                input_values[index, 0, : wav.shape[-1]] = wav
+                padding_mask[index, : wav.shape[-1]] = True
+            encoded = self.audio_tokenizer.encode(
+                input_values,
+                padding_mask=padding_mask,
+                num_quantizers=n_vq,
+                return_dict=True,
+            )
+            audio_codes = encoded.audio_codes
+            audio_lengths = encoded.audio_codes_lengths
+
+        if audio_codes is None or audio_lengths is None:
+            raise RuntimeError("audio_tokenizer did not return audio_codes/audio_codes_lengths.")
+        result = []
+        for index in range(int(audio_codes.shape[1])):
+            length = int(audio_lengths[index].item())
+            result.append(audio_codes[:, index, :length].transpose(0, 1).contiguous().cpu().long())
+        return result
+
+    def encode_audios_from_path(
+        self,
+        wav_path_list: Union[str, os.PathLike, list[Union[str, os.PathLike]]],
+        n_vq: Optional[int] = None,
+    ) -> list[torch.Tensor]:
+        if isinstance(wav_path_list, (str, os.PathLike)):
+            wav_path_list = [wav_path_list]
+        wavs = []
+        target_sr = int(self.model_config.sampling_rate)
+        for wav_path in wav_path_list:
+            wav, sr = torchaudio.load(str(wav_path))
+            if int(sr) != target_sr:
+                wav = torchaudio.functional.resample(wav, int(sr), target_sr)
+            wavs.append(wav)
+        return self.encode_audios_from_wav(wavs, target_sr, n_vq=n_vq)
+
+    def decode_audio_codes(
+        self,
+        audio_tokens_list: Union[torch.Tensor, list[torch.Tensor]],
+        *,
+        return_stereo: bool = True,
+    ) -> list[torch.Tensor]:
+        if self.audio_tokenizer is None:
+            raise RuntimeError("audio_tokenizer is not set.")
+        if isinstance(audio_tokens_list, torch.Tensor):
+            audio_tokens_list = [audio_tokens_list]
+        if not audio_tokens_list:
+            return []
+
+        n_vq = int(self.model_config.n_vq)
+        device = self._get_audio_tokenizer_device()
+        codes_list = [
+            codes[:, :n_vq].transpose(0, 1).contiguous().to(device=device, dtype=torch.long)
+            for codes in audio_tokens_list
+        ]
+        max_len = max(int(codes.shape[1]) for codes in codes_list)
+        audio_codes = torch.zeros(n_vq, len(codes_list), max_len, device=device, dtype=torch.long)
+        padding_mask = torch.zeros(len(codes_list), max_len, device=device, dtype=torch.bool)
+        for index, codes in enumerate(codes_list):
+            length = int(codes.shape[1])
+            audio_codes[:, index, :length] = codes
+            padding_mask[index, :length] = True
+
+        decoded = self.audio_tokenizer.decode(
+            audio_codes,
+            padding_mask=padding_mask,
+            num_quantizers=n_vq,
+            return_dict=True,
+            chunk_duration=8,
+        )
+        audio = decoded.audio
+        audio_lengths = decoded.audio_lengths
+        if audio is None or audio_lengths is None:
+            raise RuntimeError("audio_tokenizer.decode did not return audio/audio_lengths.")
+        wavs = []
+        for index in range(int(audio.shape[0])):
+            length = int(audio_lengths[index].item())
+            wav = audio[index, :, :length].contiguous().cpu().to(torch.float32)
+            if not return_stereo:
+                if wav.shape[0] == 1:
+                    wav = wav.squeeze(0)
+                else:
+                    wav = wav.mean(dim=0)
+            wavs.append(wav)
+        return wavs