from __future__ import annotations import inspect import json import os import re from pathlib import Path from typing import Any, Sequence import numpy as np import torch from transformers import AutoTokenizer from transformers.feature_extraction_utils import BatchFeature from transformers.processing_utils import ProcessorMixin def _clean_text(text: str) -> str: return " ".join(str(text).strip().split()) def _as_list(value: Any, batch_size: int, name: str) -> list[Any]: if isinstance(value, (str, Path)) or value is None or np.isscalar(value): return [value] * batch_size if isinstance(value, torch.Tensor) and value.ndim <= 2: return [value] if batch_size == 1 else list(value) if isinstance(value, np.ndarray) and value.ndim <= 2: return [value] if batch_size == 1 else list(value) values = list(value) if len(values) != batch_size: raise ValueError(f"{name} must contain {batch_size} items, got {len(values)}") return values def _pad_1d(rows: list[torch.Tensor], pad_value: int) -> tuple[torch.Tensor, torch.Tensor]: max_len = max((row.numel() for row in rows), default=0) values = torch.full((len(rows), max_len), pad_value, dtype=torch.long) mask = torch.zeros((len(rows), max_len), dtype=torch.long) for idx, row in enumerate(rows): length = row.numel() values[idx, :length] = row mask[idx, :length] = 1 return values, mask class ArkttsProcessor(ProcessorMixin): attributes = ["tokenizer"] tokenizer_class = ("PreTrainedTokenizerFast", "PreTrainedTokenizer") valid_kwargs = ["num_codebooks", "semantic_begin_id", "audio_sampling_rate"] def __init__( self, tokenizer, num_codebooks: int = 10, semantic_begin_id: int = 151678, audio_sampling_rate: int = 44100, **kwargs, ): super().__init__(tokenizer) self.num_codebooks = int(num_codebooks) self.semantic_begin_id = int(semantic_begin_id) self.audio_sampling_rate = int(audio_sampling_rate) @classmethod def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs) -> "ArkttsProcessor": trust_remote_code = bool(kwargs.pop("trust_remote_code", False)) shared_names = { "cache_dir", "force_download", "local_files_only", "token", "revision", "subfolder" } shared = {key: kwargs[key] for key in list(kwargs) if key in shared_names} config = {} local_config = os.path.join(str(pretrained_model_name_or_path), "processor_config.json") if os.path.isfile(local_config): with open(local_config, "r", encoding="utf-8") as handle: config = json.load(handle) tokenizer = AutoTokenizer.from_pretrained( pretrained_model_name_or_path, use_fast=True, trust_remote_code=trust_remote_code, fix_mistral_regex=False, **shared, ) return cls( tokenizer=tokenizer, num_codebooks=config.get("num_codebooks", 10), semantic_begin_id=config.get("semantic_begin_id", 151678), audio_sampling_rate=config.get("audio_sampling_rate", 44100), ) def _encode(self, text: str) -> torch.Tensor: encode_kwargs = {"add_special_tokens": False} if "allowed_special" in inspect.signature(self.tokenizer.encode).parameters: encode_kwargs["allowed_special"] = "all" return torch.tensor(self.tokenizer.encode(text, **encode_kwargs), dtype=torch.long) @staticmethod def _format_reference_text(text: str) -> str: cleaned = _clean_text(text) if re.search(r"<\|speaker:\d+\|>", cleaned): return cleaned return f"<|speaker:0|>{cleaned}" def _prompt_segments(self, text: str, reference_text: str | None, has_reference: bool): target = _clean_text(text) if not target: raise ValueError("text must not be empty") def encode_parts(parts: list[str]) -> torch.Tensor: return torch.cat([self._encode(part) for part in parts]) if not has_reference: full = encode_parts([ "<|im_start|>system\n", "convert the provided text to speech", "<|im_end|>\n", "<|im_start|>user\n", target, "<|im_end|>\n", "<|im_start|>assistant\n<|voice|>", ]) return full, self._encode("") if not reference_text: raise ValueError("reference_text is required when a reference voice is provided") prefix = encode_parts([ "<|im_start|>system\n", "convert the provided text to speech reference to the following:\n\nText:\n", self._format_reference_text(reference_text), "\n\nSpeech:\n", ]) suffix = encode_parts([ "<|im_end|>\n", "<|im_start|>user\n", target, "<|im_end|>\n", "<|im_start|>assistant\n<|voice|>", ]) return prefix, suffix def _load_audio(self, value: Any, sampling_rate: int | None) -> torch.Tensor: source_rate = sampling_rate if isinstance(value, (str, Path)): try: import soundfile as sf except ImportError as exc: raise ImportError("soundfile is required for reference audio paths") from exc array, source_rate = sf.read(str(value), dtype="float32", always_2d=True) array = array.mean(axis=1) audio = torch.from_numpy(np.asarray(array, dtype=np.float32)) else: if isinstance(value, dict): source_rate = value.get("sampling_rate", source_rate) value = value.get("array") if isinstance(value, (tuple, list)) and len(value) == 2 and np.isscalar(value[1]): value, source_rate = value audio = torch.as_tensor(value, dtype=torch.float32) if audio.ndim == 2: audio = audio.mean(dim=0) if audio.ndim != 1: raise ValueError(f"reference audio must be mono or channels-first, got {tuple(audio.shape)}") if audio.numel() == 0: raise ValueError("reference audio must not be empty") if source_rate is None: raise ValueError("sampling_rate is required for reference audio arrays") if int(source_rate) != self.audio_sampling_rate: try: from torchaudio.functional import resample except ImportError as exc: raise ImportError("torchaudio is required to resample reference audio") from exc audio = resample(audio, int(source_rate), self.audio_sampling_rate) return audio.contiguous() def __call__( self, text: str | Sequence[str], reference_text: str | Sequence[str] | None = None, reference_audio: Any = None, reference_codes: Any = None, sampling_rate: int | Sequence[int] | None = None, return_tensors: str = "pt", **kwargs, ) -> BatchFeature: if kwargs: raise TypeError(f"Unexpected processor arguments: {sorted(kwargs)}") if return_tensors != "pt": raise ValueError("ArkttsProcessor currently supports return_tensors='pt' only") texts = [text] if isinstance(text, str) else list(text) if not texts: raise ValueError("text batch must not be empty") batch_size = len(texts) ref_texts = _as_list(reference_text, batch_size, "reference_text") if reference_audio is not None and reference_codes is not None: raise ValueError("Provide reference_audio or reference_codes, not both") has_reference = reference_audio is not None or reference_codes is not None prefix_rows, suffix_rows = zip(*[ self._prompt_segments(item, ref_texts[idx], has_reference) for idx, item in enumerate(texts) ]) prefix_ids, prefix_mask = _pad_1d(list(prefix_rows), self.tokenizer.pad_token_id) suffix_ids, suffix_mask = _pad_1d(list(suffix_rows), self.tokenizer.pad_token_id) data: dict[str, torch.Tensor] = { "prefix_input_ids": prefix_ids, "prefix_attention_mask": prefix_mask, "suffix_input_ids": suffix_ids, "suffix_attention_mask": suffix_mask, } if reference_codes is not None: code_items = _as_list(reference_codes, batch_size, "reference_codes") loaded = [] for item in code_items: if isinstance(item, (str, Path)): item = np.load(str(item)) codes = torch.as_tensor(item, dtype=torch.long) if codes.ndim != 2 or codes.shape[0] != self.num_codebooks or codes.shape[1] == 0: raise ValueError( f"reference codes must have shape [{self.num_codebooks}, T>0], got {tuple(codes.shape)}" ) if codes.min() < 0 or codes.max() >= 4096: raise ValueError("reference codes must be in [0, 4095]") loaded.append(codes) max_frames = max(item.shape[1] for item in loaded) padded = torch.full((batch_size, self.num_codebooks, max_frames), -1, dtype=torch.long) lengths = torch.empty(batch_size, dtype=torch.long) for idx, codes in enumerate(loaded): lengths[idx] = codes.shape[1] padded[idx, :, : codes.shape[1]] = codes data["reference_codes"] = padded data["reference_code_lengths"] = lengths if reference_audio is not None: audio_items = _as_list(reference_audio, batch_size, "reference_audio") rate_items = _as_list(sampling_rate, batch_size, "sampling_rate") loaded_audio = [self._load_audio(item, rate_items[idx]) for idx, item in enumerate(audio_items)] max_samples = max(item.numel() for item in loaded_audio) padded_audio = torch.zeros((batch_size, 1, max_samples), dtype=torch.float32) lengths = torch.empty(batch_size, dtype=torch.long) for idx, audio in enumerate(loaded_audio): lengths[idx] = audio.numel() padded_audio[idx, 0, : audio.numel()] = audio data["reference_audio_values"] = padded_audio data["reference_audio_lengths"] = lengths return BatchFeature(data=data) @property def model_input_names(self) -> list[str]: return [ "prefix_input_ids", "prefix_attention_mask", "suffix_input_ids", "suffix_attention_mask", "reference_codes", "reference_code_lengths", "reference_audio_values", "reference_audio_lengths", ] def batch_decode(self, *args, **kwargs): return self.tokenizer.batch_decode(*args, **kwargs) def decode(self, *args, **kwargs): return self.tokenizer.decode(*args, **kwargs) __all__ = ["ArkttsProcessor"]