Refactor processing_moss_tts.py: Improve type hints, enhance message classes, and streamline audio token handling

processing_moss_tts.py

@@ -14,14 +14,22 @@
 # limitations under the License.
 
 import os
-from typing import Dict, List, Optional, Tuple, Type, Union, Literal, Final
+from typing import Any, Dict, List, Optional, Tuple, Type, Union, Literal, Final, cast
 from dataclasses import dataclass
 from pathlib import Path
 import re
 import torchaudio
 
 import torch
-from transformers import PreTrainedTokenizerBase, BatchFeature, ProcessorMixin, logging, AutoConfig, AutoModel, AutoTokenizer
+from transformers import (
+    PreTrainedTokenizerBase,
+    BatchFeature,
+    ProcessorMixin,
+    logging,
+    AutoConfig,
+    AutoModel,
+    AutoTokenizer,
+)
 
 from .configuration_moss_tts import MossTTSDelayConfig
 
@@ -34,8 +42,8 @@ AUDIO_PLACEHOLDER = "<|audio|>"
 
 @dataclass
 class Message:
-    pass
-    
+    def to_dict(self) -> Dict[str, Any]:
+        raise NotImplementedError
 
 
 @dataclass
@@ -78,13 +86,16 @@ class UserMessage(Message):
                 if speaker_reference is not None:
                     reference.append(f"[S{speaker_idx}]:\n{AUDIO_PLACEHOLDER}")
             reference = "\n".join(reference)
-            audio_codes_list = [speaker_reference for speaker_reference in self.reference if speaker_reference is not None]
+            audio_codes_list = [
+                speaker_reference
+                for speaker_reference in self.reference
+                if speaker_reference is not None
+            ]
         else:
             raise TypeError("`reference` should be exactly a list when it is not None.")
 
         content = (
-            template
-            .replace("{reference}", str(reference))
+            template.replace("{reference}", str(reference))
             .replace("{instruction}", str(self.instruction))
             .replace("{tokens}", str(self.tokens))
             .replace("{quality}", str(self.quality))
@@ -101,22 +112,23 @@ class UserMessage(Message):
         return {
             "role": "user",
             "content": self._content,
-            "audio_codes_list": self._audio_codes_list
+            "audio_codes_list": self._audio_codes_list,
         }
 
 
 @dataclass
 class AssistantMessage(Message):
     audio_codes_list: List[Union[str, torch.Tensor]]
-    content: str = AUDIO_PLACEHOLDER 
+    content: str = AUDIO_PLACEHOLDER
 
     def to_dict(self):
         return {
             "role": "assistant",
             "content": self.content,
-            "audio_codes_list": self.audio_codes_list
+            "audio_codes_list": self.audio_codes_list,
         }
 
+
 USER_MESSAGE_FIELDS = (
     "text",
     "reference",
@@ -129,27 +141,25 @@ USER_MESSAGE_FIELDS = (
 )
 
 
-
-
-
-
-
 class MossTTSDelayProcessor(ProcessorMixin):
     tokenizer_class = "AutoTokenizer"
     audio_tokenizer_class = "AutoModel"
 
+    tokenizer: PreTrainedTokenizerBase
+    audio_tokenizer: Any
+
     def __init__(
         self,
         tokenizer: PreTrainedTokenizerBase,
-        audio_tokenizer: AutoModel = None,
+        audio_tokenizer: Any = None,
         model_config: Optional[MossTTSDelayConfig] = None,
-        **kwargs
+        **kwargs,
     ):
-        super().__init__(
-            tokenizer=tokenizer,
-            audio_tokenizer=audio_tokenizer,
-            **kwargs
-        )
+        super().__init__(tokenizer=tokenizer, audio_tokenizer=audio_tokenizer, **kwargs)
+
+        # Explicit assignments for type-checkers; ProcessorMixin sets these too.
+        self.tokenizer = tokenizer
+        self.audio_tokenizer = audio_tokenizer
         if model_config is None:
             model_config = MossTTSDelayConfig()
         self.model_config = model_config
@@ -158,68 +168,107 @@ class MossTTSDelayProcessor(ProcessorMixin):
         self.imend_token_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
         self.newline_token_id = 198
 
-        self.audio_user_slot_token = tokenizer.convert_ids_to_tokens(self.model_config.audio_user_slot_token_id)
-        self.audio_assistant_gen_slot_token = tokenizer.convert_ids_to_tokens(self.model_config.audio_assistant_gen_slot_token_id)
-        self.audio_assistant_delay_slot_token = tokenizer.convert_ids_to_tokens(self.model_config.audio_assistant_delay_slot_token_id)
-        self.audio_start_token = tokenizer.convert_ids_to_tokens(self.model_config.audio_start_token_id)
-        self.audio_end_token = tokenizer.convert_ids_to_tokens(self.model_config.audio_end_token_id)
+        def _id_to_token(token_id: int) -> str:
+            tok = tokenizer.convert_ids_to_tokens(int(token_id))
+            if isinstance(tok, list):
+                return tok[0] if len(tok) > 0 else ""
+            return cast(str, tok)
+
+        self.audio_user_slot_token = _id_to_token(
+            self.model_config.audio_user_slot_token_id
+        )
+        self.audio_assistant_gen_slot_token = _id_to_token(
+            self.model_config.audio_assistant_gen_slot_token_id
+        )
+        self.audio_assistant_delay_slot_token = _id_to_token(
+            self.model_config.audio_assistant_delay_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, trust_remote_code=True, **kwargs):
-        kwargs.pop("_from_auto")
+    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)
+
+        audio_tokenizer_name_or_path = kwargs.pop(
+            "codec_path", "OpenMOSS-Team/MOSS-Audio-Tokenizer"
+        )
+
         pretrained_model_name_or_path = Path(pretrained_model_name_or_path)
-        model_config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs)
-        tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs)
+        model_config = cast(
+            MossTTSDelayConfig,
+            AutoConfig.from_pretrained(
+                pretrained_model_name_or_path,
+                *args,
+                trust_remote_code=trust_remote_code,
+                **kwargs,
+            ),
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            *args,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
+        audio_tokenizer = AutoModel.from_pretrained(
+            audio_tokenizer_name_or_path,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
 
-        audio_tokenizer_name_or_path = kwargs.pop("codec_path", "OpenMOSS-Team/MOSS-Audio-Tokenizer")
-        assert isinstance(audio_tokenizer_name_or_path, str), f"Unsupported audio_tokenizer_path input format: {type(audio_tokenizer_name_or_path)}"
-        audio_tokenizer = AutoModel.from_pretrained(audio_tokenizer_name_or_path, trust_remote_code=trust_remote_code, **kwargs)
-        
         return cls(
             tokenizer=tokenizer,
             audio_tokenizer=audio_tokenizer,
             model_config=model_config,
-            **kwargs
+            **kwargs,
         )
-    
-    def __call__(
-        self,
-        conversations: Union[Message, Dict, List[Message], List[Dict], List[List[Message]], List[List[Dict]]],
-        mode: str = "generation",
-        apply_chat_template: bool = True,
-        n_vq: Optional[int] = None
-    ) -> BatchFeature:
-        
+
+    def __call__(self, *args, **kwargs) -> BatchFeature:
+        conversations = args[0] if len(args) > 0 else kwargs.pop("conversations")
+        mode: str = kwargs.pop("mode", "generation")
+        apply_chat_template: bool = kwargs.pop("apply_chat_template", True)
+        n_vq: Optional[int] = kwargs.pop("n_vq", None)
+
+        # Common ProcessorMixin kwargs that we ignore because we always return torch tensors.
+        kwargs.pop("return_tensors", None)
+        kwargs.pop("padding", None)
+        kwargs.pop("truncation", None)
+
         """
-        mode 只会在将 Message 转换为 to_dict 时起作用;
+        mode only works when a Message is converted to a dict.
         """
-        
+
         if mode not in {"generation", "continuation"}:
             raise RuntimeError
-        
+
         if isinstance(conversations, (Message, Dict)):
             conversations = [conversations]
-        
+
         truncation = False
         if mode == "continuation":
             truncation = True
-        
+
         input_ids_list = []
         for conversation in conversations:
             if isinstance(conversation, (Message, Dict)):
                 conversation = [conversation]
 
+            # Normalize early so downstream logic always deals with dict messages.
+            conversation = [self._normalize_message(m) for m in conversation]
+
             if (mode == "generation") ^ (len(conversation) % 2 != 0):
                 raise ValueError
 
-            if (mode == "generation") ^ (conversation[-1]['role'] == "user"):
+            if (mode == "generation") ^ (conversation[-1]["role"] == "user"):
                 raise ValueError
 
             unified_codes = []
             for message_idx, message in enumerate(conversation):
-                message = self._normalize_message(message)
                 if apply_chat_template:
-                    add_generation_prompt = mode == "generation" and message_idx == len(conversation) - 1
+                    add_generation_prompt = (
+                        mode == "generation" and message_idx == len(conversation) - 1
+                    )
                     try:
                         content = self.tokenizer.apply_chat_template(
                             [{"role": message["role"], "content": message["content"]}],
@@ -229,29 +278,76 @@ class MossTTSDelayProcessor(ProcessorMixin):
                     except TypeError:
                         try:
                             content = self.tokenizer.apply_chat_template(
-                                [{"role": message["role"], "content": message["content"]}],
+                                [
+                                    {
+                                        "role": message["role"],
+                                        "content": message["content"],
+                                    }
+                                ],
                                 add_generation_prompt=add_generation_prompt,
                             )
                         except Exception:
-                            logger.warning("apply_chat_template failed; fallback to raw content.")
+                            logger.warning(
+                                "apply_chat_template failed; fallback to raw content."
+                            )
                             content = message["content"]
                 else:
-                    content = message['content']
-                
-                audio_codes_list = []
-                for audio_codes in message["audio_codes_list"]:
-                    if isinstance(audio_codes, torch.Tensor):
-                        if n_vq is not None and audio_codes.shape[1] != n_vq:
-                            raise RuntimeError("audio_codes's n_vq is not equal to the parameter `n_vq`. Your can set the parameter `n_vq` as None if you have already tokenzied the wavs.")
-                    else:
-                        audio_codes = self.encode_audios_from_path(audio_codes, n_vq)[0]
-                    audio_codes_list.append(audio_codes)
-                unified_codes.append(self._get_unified_codes(message['role'], content, audio_codes_list, truncation))
+                    content = message["content"]
+
+                if not isinstance(content, str):
+                    content = str(content)
+
+                # Batch-encode all path-based references in one call when possible.
+                # This ensures we actually exercise audio_tokenizer.batch_encode for multi-reference prompts,
+                # instead of repeatedly calling it with batch=1.
+                raw_audio_items = message.get("audio_codes_list", [])
+
+                audio_codes_list: List[torch.Tensor] = []
+                if len(raw_audio_items) > 0:
+                    encoded_items: List[Optional[torch.Tensor]] = [None] * len(
+                        raw_audio_items
+                    )
+                    paths: List[str] = []
+                    path_positions: List[int] = []
+
+                    for idx, item in enumerate(raw_audio_items):
+                        if isinstance(item, torch.Tensor):
+                            if n_vq is not None and item.shape[1] != n_vq:
+                                raise RuntimeError(
+                                    "audio_codes's n_vq is not equal to the parameter `n_vq`. Your can set the parameter `n_vq` as None if you have already tokenzied the wavs."
+                                )
+                            encoded_items[idx] = item
+                            continue
+
+                        if isinstance(item, (str, os.PathLike)):
+                            paths.append(str(item))
+                            path_positions.append(idx)
+                            continue
+
+                        raise TypeError(
+                            "Each audio item must be a torch.Tensor of codes or a path-like string."
+                        )
+
+                    if len(paths) > 0:
+                        encoded_from_paths = self.encode_audios_from_path(paths, n_vq)
+                        if len(encoded_from_paths) != len(paths):
+                            raise RuntimeError(
+                                "encode_audios_from_path returned an unexpected number of items."
+                            )
+                        for pos, codes in zip(path_positions, encoded_from_paths):
+                            encoded_items[pos] = codes
+
+                    audio_codes_list = [cast(torch.Tensor, t) for t in encoded_items]
+                unified_codes.append(
+                    self._get_unified_codes(
+                        message["role"], content, audio_codes_list, truncation
+                    )
+                )
 
             unified_codes = torch.cat(unified_codes)
             input_ids_list.append(unified_codes)
 
-        return self._pad(input_ids_list)
+        return BatchFeature(data=self._pad(input_ids_list))
 
     @staticmethod
     def build_user_message(
@@ -310,14 +406,23 @@ class MossTTSDelayProcessor(ProcessorMixin):
     def _pad(self, input_ids_list: List[torch.Tensor]):
         device = input_ids_list[0].device
         lengths = torch.tensor([w.shape[0] for w in input_ids_list], device=device)
-        pad_input_ids = torch.nn.utils.rnn.pad_sequence(input_ids_list, batch_first=True, padding_value=self.model_config.audio_pad_code, padding_side="left")
-        other_channel_mask = (pad_input_ids.shape[1] - lengths).unsqueeze(1) > torch.arange(pad_input_ids.shape[1], device=device).unsqueeze(0)
+        pad_input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids_list,
+            batch_first=True,
+            padding_value=self.model_config.audio_pad_code,
+            padding_side="left",
+        )
+        other_channel_mask = (pad_input_ids.shape[1] - lengths).unsqueeze(
+            1
+        ) > torch.arange(pad_input_ids.shape[1], device=device).unsqueeze(0)
         pad_input_ids[..., 0][other_channel_mask] = self.model_config.pad_token_id
-        attention_mask = torch.zeros(pad_input_ids.shape[0], pad_input_ids.shape[1], device=device)
+        attention_mask = torch.zeros(
+            pad_input_ids.shape[0], pad_input_ids.shape[1], device=device
+        )
         attention_mask[~other_channel_mask] = 1
         attention_mask = attention_mask.bool()
         return {
-            "input_ids": pad_input_ids, # [batch_size, seqlen, n_vq]
+            "input_ids": pad_input_ids,  # [batch_size, seqlen, n_vq]
             "attention_mask": attention_mask,
         }
 
@@ -329,7 +434,7 @@ class MossTTSDelayProcessor(ProcessorMixin):
         gen_slot_token: str,
         delay_slot_token: str,
         audio_start_token: str,
-        audio_end_token: str
+        audio_end_token: str,
     ) -> str:
         if n_vq < 1:
             raise ValueError(f"n_vq must be >= 1, got {n_vq}")
@@ -371,7 +476,9 @@ class MossTTSDelayProcessor(ProcessorMixin):
             return content, audio_codes_list
 
         if len(matches) != len(audio_codes_list):
-            raise ValueError("Audio placeholders do not match the provided audio codes list.")
+            raise ValueError(
+                "Audio placeholders do not match the provided audio codes list."
+            )
 
         new_audio_codes_list = []
         new_parts = []
@@ -381,18 +488,20 @@ class MossTTSDelayProcessor(ProcessorMixin):
             j = i
             while (
                 j + 1 < len(matches)
-                and content[matches[j].end():matches[j + 1].start()].strip() == ""
+                and content[matches[j].end() : matches[j + 1].start()].strip() == ""
             ):
                 j += 1
 
-            new_parts.append(content[last_pos:matches[i].start()])
+            new_parts.append(content[last_pos : matches[i].start()])
             new_parts.append(AUDIO_PLACEHOLDER)
             last_pos = matches[j].end()
 
             if j == i:
                 new_audio_codes_list.append(audio_codes_list[i])
             else:
-                new_audio_codes_list.append(torch.cat(audio_codes_list[i:j + 1], dim=0))
+                new_audio_codes_list.append(
+                    torch.cat(audio_codes_list[i : j + 1], dim=0)
+                )
 
             i = j + 1
 
@@ -408,9 +517,9 @@ class MossTTSDelayProcessor(ProcessorMixin):
             dtype=codes.dtype,
         )
         for i in range(codes.shape[1]):
-            delayed_tokens[i: i + codes.shape[0], i] = codes[:, i]
+            delayed_tokens[i : i + codes.shape[0], i] = codes[:, i]
         return delayed_tokens
-    
+
     @staticmethod
     def apply_de_delay_pattern(delay_codes: torch.Tensor) -> torch.Tensor:
         tokens = torch.full(
@@ -420,11 +529,16 @@ class MossTTSDelayProcessor(ProcessorMixin):
             dtype=delay_codes.dtype,
         )
         for i in range(delay_codes.shape[1]):
-            tokens[:, i] = delay_codes[i: i + tokens.shape[0], i]
+            tokens[:, i] = delay_codes[i : i + tokens.shape[0], i]
         return tokens
 
-
-    def _get_unified_codes(self, role: str, content: str, audio_codes_list: List[Union[str, torch.Tensor]], truncation: bool) -> torch.Tensor:
+    def _get_unified_codes(
+        self,
+        role: str,
+        content: str,
+        audio_codes_list: List[torch.Tensor],
+        truncation: bool,
+    ) -> torch.Tensor:
         """
         此时的 content 已经是带上了对话格式
         """
@@ -452,12 +566,23 @@ class MossTTSDelayProcessor(ProcessorMixin):
             audio_start_token=self.audio_start_token,
             audio_end_token=self.audio_end_token,
         )
-        text_codes = torch.tensor(self.tokenizer.encode(content), device=audio_codes_list[0].device if audio_codes_list else None)
+        text_codes = torch.tensor(
+            self.tokenizer.encode(content),
+            device=audio_codes_list[0].device if audio_codes_list else None,
+        )
 
-        audio_start_indices = torch.where(text_codes == self.model_config.audio_start_token_id)[0]
-        audio_end_indices = torch.where(text_codes == 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 provided audio codes list.")
+        audio_start_indices = torch.where(
+            text_codes == self.model_config.audio_start_token_id
+        )[0]
+        audio_end_indices = torch.where(
+            text_codes == 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 provided audio codes list."
+            )
 
         delay_audio_codes_list = []
         if len(audio_codes_list) == 0:
@@ -469,8 +594,14 @@ class MossTTSDelayProcessor(ProcessorMixin):
             )
         else:
             prefix_idx = 0
-            for audio_start_idx, audio_end_idx, audio_codes in zip(audio_start_indices, audio_end_indices, audio_codes_list):
-                delay_audio_codes = self.apply_delay_pattern(audio_codes, self.model_config.audio_pad_code)
+            for audio_start_idx_t, audio_end_idx_t, audio_codes in zip(
+                audio_start_indices, audio_end_indices, audio_codes_list
+            ):
+                audio_start_idx = int(audio_start_idx_t.item())
+                audio_end_idx = int(audio_end_idx_t.item())
+                delay_audio_codes = self.apply_delay_pattern(
+                    audio_codes, self.model_config.audio_pad_code
+                )
                 pad_codes = torch.full(
                     (audio_start_idx - prefix_idx + 1, n_vq),
                     self.model_config.audio_pad_code,
@@ -481,10 +612,13 @@ class MossTTSDelayProcessor(ProcessorMixin):
                 prefix_idx = audio_end_idx
 
             if truncation:
-                delay_audio_codes_list[-1] = delay_audio_codes_list[-1][:-(n_vq - 1), :]
+                delay_audio_codes_list[-1] = delay_audio_codes_list[-1][
+                    : -(n_vq - 1), :
+                ]
             else:
+                last_audio_end_idx = int(audio_end_indices[-1].item())
                 pad_codes = torch.full(
-                    (len(text_codes) - audio_end_indices[-1], n_vq),
+                    (len(text_codes) - last_audio_end_idx, n_vq),
                     self.model_config.audio_pad_code,
                     device=audio_codes_list[0].device,
                     dtype=audio_codes_list[0].dtype,
@@ -492,34 +626,36 @@ class MossTTSDelayProcessor(ProcessorMixin):
                 delay_audio_codes_list.append(pad_codes)
 
             delay_audio_codes_list = torch.cat(delay_audio_codes_list)
-        
+
         if text_codes.shape[0] != delay_audio_codes_list.shape[0]:
-            text_codes = text_codes[:delay_audio_codes_list.shape[0]]
+            text_codes = text_codes[: delay_audio_codes_list.shape[0]]
 
-        unified_codes = torch.cat([text_codes.unsqueeze(1), delay_audio_codes_list], dim=1)
+        unified_codes = torch.cat(
+            [text_codes.unsqueeze(1), delay_audio_codes_list], dim=1
+        )
         return unified_codes
 
     def _parse_text_codes(self, start_length, text_codes):
-        text = self.tokenizer.decode(text_codes)
-        prefix = self.tokenizer.decode(text_codes[:start_length])
-        text = text[len(prefix):]
+        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'(?:{self.audio_start_token})?'
-            rf'(?:{self.audio_assistant_gen_slot_token})*'
-            rf'(?:{self.audio_assistant_delay_slot_token})*'
-            rf'{self.audio_end_token}'
+            rf"(?:{self.audio_start_token})?"
+            rf"(?:{self.audio_assistant_gen_slot_token})*"
+            rf"(?:{self.audio_assistant_delay_slot_token})*"
+            rf"{self.audio_end_token}"
         )
 
         def normalize_audio_segments(text: str) -> str:
             def repl(match: re.Match) -> str:
                 seg = match.group(0)
-                # 如果片段内包含至少一个 gen_slot，则替换为 <|audio|>
+                # Replace with <|audio|> if gen_slot is present in the segment;
                 if self.audio_assistant_gen_slot_token in seg:
                     return AUDIO_PLACEHOLDER
-                # 否则直接删除
+                # Otherwise, remove it.
                 return ""
-            
+
             return AUDIO_PATTERN.sub(repl, text)
 
         return normalize_audio_segments(text)
@@ -543,18 +679,21 @@ class MossTTSDelayProcessor(ProcessorMixin):
 
         audio_codes_list = [audio_codes[s] for s in segments_idx]
 
-        decoded_audio_list = []
-        for segment_codes in audio_codes_list:
-            decoded_segment = self.decode_audio_codes([segment_codes])
-            if len(decoded_segment) > 0:
-                decoded_audio_list.append(decoded_segment[0])
+        # Batch-decode all audio segments together.
+        decoded_audio_list = self.decode_audio_codes(audio_codes_list)
 
         # Keep codec causal context by decoding the whole first segment first,
         # then trim at waveform level according to start_length ratio.
-        if start_length > 0 and len(audio_codes_list) > 0 and len(decoded_audio_list) > 0:
+        if (
+            start_length > 0
+            and len(audio_codes_list) > 0
+            and len(decoded_audio_list) > 0
+        ):
             first_codes_length = audio_codes_list[0].shape[0]
             if first_codes_length > 0:
-                trim_ratio = max(0.0, min(float(start_length) / float(first_codes_length), 1.0))
+                trim_ratio = max(
+                    0.0, min(float(start_length) / float(first_codes_length), 1.0)
+                )
                 first_audio = decoded_audio_list[0]
                 if trim_ratio >= 1.0:
                     decoded_audio_list = decoded_audio_list[1:]
@@ -564,7 +703,6 @@ class MossTTSDelayProcessor(ProcessorMixin):
 
         return decoded_audio_list
 
-
     def decode(self, output: List[Tuple[int, torch.Tensor]]):
         """
         1. 这里不管怎样，都需要一个完整的 assistant generation ids;
@@ -574,56 +712,219 @@ class MossTTSDelayProcessor(ProcessorMixin):
         genearted_messages = []
         for start_length, generation_ids in output:
             content = self._parse_text_codes(start_length, generation_ids[:, 0])
-            audio_codes_list = self._parse_audio_codes(start_length, generation_ids[:, 1:])
+            audio_codes_list = self._parse_audio_codes(
+                start_length, generation_ids[:, 1:]
+            )
             if content == "":
                 message = None
             else:
                 message = AssistantMessage(
                     content=content,
-                    audio_codes_list=audio_codes_list
+                    audio_codes_list=cast(
+                        List[Union[str, torch.Tensor]], audio_codes_list
+                    ),
                 )
             genearted_messages.append(message)
         return genearted_messages
 
     @staticmethod
-    def loudness_normalize(wav: torch.Tensor, target_dbfs: float = -20, gain_range: tuple[float, float] = (-3.0, 3.0)) -> torch.Tensor:
+    def loudness_normalize(
+        wav: torch.Tensor,
+        target_dbfs: float = -20,
+        gain_range: tuple[float, float] = (-3.0, 3.0),
+    ) -> torch.Tensor:
         wav = wav.to(torch.float32)
-        if wav.numel() == 0: return wav
-        rms = torch.sqrt(torch.mean(wav ** 2))
-        current_dbfs = 20.0 * torch.log10(rms + 1e-9)
+        if wav.numel() == 0:
+            return wav
+        current_dbfs = 10.0 * torch.log10(torch.mean(wav**2) + 1e-9)
         gain = float(target_dbfs - current_dbfs)
         gain = max(gain_range[0], min(gain, gain_range[1]))
         factor = 10.0 ** (gain / 20.0)
         return wav * factor
 
-    def encode_audios_from_wav(self, wav_list: List[torch.Tensor], sampling_rate: int, n_vq: int = None):
+    def _get_audio_tokenizer_device(self) -> torch.device:
+        """Best-effort device inference for `self.audio_tokenizer`.
+
+        Notes:
+        - Old TAC wrapper exposed `.device`, but standard `torch.nn.Module` does not.
+        - New MossAudioTokenizerModel is a `PreTrainedModel`; parameters define its device.
+        """
+
+        audio_tokenizer = getattr(self, "audio_tokenizer", None)
+        if audio_tokenizer is None:
+            logger.warning(
+                "audio_tokenizer is not set on processor. Using CPU as default."
+            )
+            return torch.device("cpu")
+
+        device_attr = getattr(audio_tokenizer, "device", None)
+        if isinstance(device_attr, torch.device):
+            return device_attr
+
+        try:
+            return next(audio_tokenizer.parameters()).device
+        except StopIteration:
+            # No parameters (shouldn't happen for real models); default to CPU.
+            logger.warning(
+                "No parameters found on audio_tokenizer. Using CPU as default."
+            )
+            return torch.device("cpu")
+
+    def encode_audios_from_wav(
+        self,
+        wav_list: List[torch.Tensor],
+        sampling_rate: int,
+        n_vq: Optional[int] = None,
+    ):
+        if self.audio_tokenizer is None:
+            raise RuntimeError("audio_tokenizer is not set on processor.")
+        audio_tokenizer = self.audio_tokenizer
+
         if isinstance(wav_list, torch.Tensor):
             wav_list = [wav_list]
         wav_list_ = []
         resample = False
         if sampling_rate != self.model_config.sampling_rate:
             resample = True
+        device = self._get_audio_tokenizer_device()
         for wav in wav_list:
             if wav.shape[0] > 1:
                 wav = torch.mean(wav, dim=0, keepdim=True)
             if resample:
-                wav = torchaudio.functional.resample(waveform=wav, orig_freq=sampling_rate, new_freq=self.model_config.sampling_rate)
+                wav = torchaudio.functional.resample(
+                    waveform=wav,
+                    orig_freq=sampling_rate,
+                    new_freq=self.model_config.sampling_rate,
+                )
+            wav = wav.to(device)
             wav_list_.append(self.loudness_normalize(wav.squeeze(0)))
-        return self.audio_tokenizer.encode(wav_list_, n_vq)
 
-    def encode_audios_from_path(self, wav_path_list: List[str], n_vq: int = None):
+        # New MossAudioTokenizerModel API: prefer batch_encode(list[wav])
+        if hasattr(audio_tokenizer, "batch_encode"):
+            enc = audio_tokenizer.batch_encode(wav_list_, num_quantizers=n_vq)
+            audio_codes = enc.audio_codes  # (NQ, B, T)
+            audio_codes_lengths = enc.audio_codes_lengths  # (B,)
+        else:
+            # Fallback: use encode() with explicit padding.
+            max_len = max(int(wav.shape[-1]) for wav in wav_list_)
+            input_values = torch.zeros(
+                len(wav_list_), 1, max_len, device=device, dtype=torch.float32
+            )
+            padding_mask = torch.zeros(
+                len(wav_list_), max_len, device=device, dtype=torch.bool
+            )
+            for i, wav in enumerate(wav_list_):
+                this_len = int(wav.shape[-1])
+                input_values[i, 0, :this_len] = wav
+                padding_mask[i, :this_len] = True
+            enc = audio_tokenizer.encode(
+                input_values,
+                padding_mask=padding_mask,
+                num_quantizers=n_vq,
+                return_dict=True,
+            )
+            audio_codes = enc.audio_codes
+            audio_codes_lengths = enc.audio_codes_lengths
+
+        if audio_codes is None or audio_codes_lengths is None:
+            raise RuntimeError(
+                "audio_tokenizer.encode() returned empty outputs (audio_codes/audio_codes_lengths)."
+            )
+
+        # Keep processor's historical contract: list[Tensor] with shape (T, NQ)
+        # and on CPU (so downstream text/audio packing remains device-agnostic).
+        codes_list: List[torch.Tensor] = []
+        for i in range(int(audio_codes.shape[1])):
+            length_i = int(audio_codes_lengths[i].item())
+            codes_i = (
+                audio_codes[:, i, :length_i]
+                .transpose(0, 1)
+                .contiguous()
+                .to(torch.long)
+                .cpu()
+            )
+            codes_list.append(codes_i)
+        return codes_list
+
+    def encode_audios_from_path(
+        self, wav_path_list: Union[str, List[str]], n_vq: Optional[int] = None
+    ):
         if isinstance(wav_path_list, str):
             wav_path_list = [wav_path_list]
-        wav_list = []
-        sampling_rate = None
+
+        if len(wav_path_list) == 0:
+            raise ValueError("Empty wav_path_list")
+
+        # Load + (if needed) resample each wav independently, so callers can
+        # pass a heterogeneous batch of files while still benefiting from
+        # audio_tokenizer.batch_encode.
+        target_sr = int(self.model_config.sampling_rate)
+        wav_list: List[torch.Tensor] = []
         for wav_path in wav_path_list:
             wav, sr = torchaudio.load(wav_path)
-            if sampling_rate is None:
-                sampling_rate = sr
-            elif sampling_rate != sr:
-                raise ValueError("sampling_rate of audios in the same batch should be the same.")
+            if int(sr) != target_sr:
+                wav = torchaudio.functional.resample(
+                    waveform=wav,
+                    orig_freq=int(sr),
+                    new_freq=target_sr,
+                )
+            wav_list.append(wav)
+
+        return self.encode_audios_from_wav(wav_list, target_sr, n_vq)
+
+    def decode_audio_codes(
+        self, audio_tokens_list: Union[torch.Tensor, List[torch.Tensor]]
+    ):
+        if self.audio_tokenizer is None:
+            raise RuntimeError("audio_tokenizer is not set on processor.")
+        audio_tokenizer = self.audio_tokenizer
+
+        if isinstance(audio_tokens_list, torch.Tensor):
+            audio_tokens_list = [audio_tokens_list]
+        if len(audio_tokens_list) == 0:
+            return []
+
+        device = self._get_audio_tokenizer_device()
+
+        # Processor uses (T, NQ); MossAudioTokenizer expects (NQ, T) (or (NQ, B, T)).
+        codes_list = [
+            codes.transpose(0, 1).contiguous().to(device=device, dtype=torch.long)
+            for codes in audio_tokens_list
+        ]
+
+        if hasattr(audio_tokenizer, "batch_decode"):
+            dec = audio_tokenizer.batch_decode(codes_list)
+            audio = dec.audio  # (B, C, T)
+            audio_lengths = dec.audio_lengths  # (B,)
+        else:
+            # Fallback: pad to (NQ, B, T) + mask, then decode.
+            nq = int(codes_list[0].shape[0])
+            max_t = max(int(c.shape[1]) for c in codes_list)
+            audio_codes = torch.zeros(
+                nq, len(codes_list), max_t, device=device, dtype=torch.long
+            )
+            padding_mask = torch.zeros(
+                len(codes_list), max_t, device=device, dtype=torch.bool
+            )
+            for i, c in enumerate(codes_list):
+                t = int(c.shape[1])
+                audio_codes[:, i, :t] = c
+                padding_mask[i, :t] = True
+            dec = audio_tokenizer.decode(
+                audio_codes, padding_mask=padding_mask, return_dict=True
+            )
+            audio = dec.audio
+            audio_lengths = dec.audio_lengths
+
+        if audio is None or audio_lengths is None:
+            raise RuntimeError(
+                "audio_tokenizer.decode() returned empty outputs (audio/audio_lengths)."
+            )
+
+        # Return historical contract: list of 1D waveforms (T,)
+        wav_list: List[torch.Tensor] = []
+        for i in range(int(audio.shape[0])):
+            length_i = int(audio_lengths[i].item())
+            wav = audio[i, 0, :length_i].contiguous().to(torch.float32).cpu()
             wav_list.append(wav)
-        return self.encode_audios_from_wav(wav_list, sampling_rate, n_vq)
-        
-    def decode_audio_codes(self, audio_tokens_list: List[torch.Tensor]):
-        return self.audio_tokenizer.decode(audio_tokens_list)
+        return wav_list