diff --git a/docs/source/openvino/models.mdx b/docs/source/openvino/models.mdx index 913828fb0c..e6beba2384 100644 --- a/docs/source/openvino/models.mdx +++ b/docs/source/openvino/models.mdx @@ -34,6 +34,7 @@ Here is the list of the supported architectures : - CLIP - CamemBERT - ChatGLM (ChatGLM2, ChatGLM3, GLM4) +- Chatterbox (text-to-speech, English and multilingual) - CodeGen - CodeGen2 - Cohere diff --git a/optimum/commands/export/openvino.py b/optimum/commands/export/openvino.py index bc63bb760f..ab81186aca 100644 --- a/optimum/commands/export/openvino.py +++ b/optimum/commands/export/openvino.py @@ -85,7 +85,7 @@ def parse_args_openvino(parser: "ArgumentParser"): optional_group.add_argument( "--library", type=str, - choices=["transformers", "diffusers", "timm", "sentence_transformers", "open_clip", "kokoro"], + choices=["transformers", "diffusers", "timm", "sentence_transformers", "open_clip", "kokoro", "chatterbox"], default=None, help="The library used to load the model before export. If not provided, will attempt to infer the local checkpoint's library", ) diff --git a/optimum/exporters/openvino/__main__.py b/optimum/exporters/openvino/__main__.py index d2bfda8db2..8889510cb0 100644 --- a/optimum/exporters/openvino/__main__.py +++ b/optimum/exporters/openvino/__main__.py @@ -38,6 +38,7 @@ is_transformers_version, ) from optimum.intel.utils.modeling_utils import ( + _ChatterboxForTextToSpeech, _infer_library_from_model_name_or_path, _KokoroForTextToSpeech, _OpenClipForZeroShotImageClassification, @@ -87,7 +88,7 @@ def infer_task( if task == "auto": if library_name == "open_clip": task = "zero-shot-image-classification" - elif library_name == "kokoro": + elif library_name in ("kokoro", "chatterbox"): task = "text-to-audio" else: try: @@ -518,6 +519,10 @@ def bitnet_load_hook(self, state_dict, prefix, *args, **kwargs): model = _OpenClipForZeroShotImageClassification.from_pretrained(model_name_or_path, cache_dir=cache_dir) elif library_name == "kokoro": model = _KokoroForTextToSpeech.from_pretrained(model_name_or_path, cache_dir=cache_dir, token=token) + elif library_name == "chatterbox": + model = _ChatterboxForTextToSpeech.from_pretrained( + model_name_or_path, cache_dir=cache_dir, token=token, **loading_kwargs + ) else: # remote code models like phi3_v internvl2, minicpmv, internvl2, nanollava, maira2 should be loaded using AutoModelForCausalLM and not AutoModelForImageTextToText # TODO: use config.auto_map to load remote code models instead (for other models we can directly use config.architectures) diff --git a/optimum/exporters/openvino/chatterbox.py b/optimum/exporters/openvino/chatterbox.py new file mode 100644 index 0000000000..e8353609fe --- /dev/null +++ b/optimum/exporters/openvino/chatterbox.py @@ -0,0 +1,381 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""OpenVINO export helpers for the ResembleAI Chatterbox TTS model. + +Chatterbox is a multi-stage text-to-speech pipeline. For OpenVINO inference it is +decomposed into three exported submodels plus a set of precomputed assets: + +* ``openvino_t3`` -- the token-to-token model: a Llama backbone with a speech-token + head, exported as a stateful decoder that consumes ``inputs_embeds`` (the speech/text + embedding tables and the conditioning prefix are precomputed assets applied in Python). +* ``openvino_flow`` -- the whole S3Gen flow (conformer encoder + Euler ODE solver + + flow-matching estimator), exported as a single graph with the diffusion noise as an + explicit input so that generation is deterministic and loop-free. +* ``openvino_hifigan`` -- the HiFiGAN vocoder turning mel-spectrograms into a waveform. + ``torch.istft`` is replaced by a traceable overlap-add inverse STFT. + +The assets (``chatterbox_assets.safetensors`` + ``chatterbox_config.json``) carry the +embedding tables, the built-in voice conditioning prefix and the S3Gen reference dict. +""" + +import json +import logging +from pathlib import Path +from typing import Dict + +import torch +import torch.nn.functional as F + +from optimum.exporters.openvino.base import OpenVINOConfig +from optimum.exporters.openvino.patching_utils import ModelPatcher +from optimum.utils.input_generators import DummyInputGenerator +from optimum.utils.normalized_config import NormalizedConfig + + +logger = logging.getLogger(__name__) + + +# Speech tokens >= SPEECH_VOCAB_SIZE are special tokens that must not reach the vocoder. +SPEECH_VOCAB_SIZE = 6561 + + +def manual_istft( + real: torch.Tensor, img: torch.Tensor, n_fft: int, hop_len: int, window: torch.Tensor +) -> torch.Tensor: + """Traceable inverse STFT via ``irfft`` + overlap-add. + + Mirrors ``torch.istft`` (center=True, Hann window) but only uses operations that the + OpenVINO PyTorch frontend can convert. ``real``/``img`` have shape ``[B, n_fft//2+1, T]``. + """ + complex_spec = torch.complex(real, img) + frames = torch.fft.irfft(complex_spec, n=n_fft, dim=1) # [B, n_fft, T] + frames = frames * window.view(1, -1, 1) + eye = torch.eye(n_fft, device=frames.device, dtype=frames.dtype).unsqueeze(1) # [n_fft, 1, n_fft] + out = F.conv_transpose1d(frames, eye, stride=hop_len) # [B, 1, L] + win_sq = (window**2).view(1, -1, 1).expand(1, -1, frames.shape[-1]) + norm = F.conv_transpose1d(win_sq, eye, stride=hop_len) + out = out / (norm + 1e-8) + pad = n_fft // 2 + return out[:, 0, pad:-pad] if pad > 0 else out[:, 0, :] + + +class ChatterboxFlowWrapper(torch.nn.Module): + """Wraps the S3Gen flow (``CausalMaskedDiffWithXvec``) into a single traceable graph. + + The diffusion noise is provided as an explicit input and the number of ODE timesteps is + fixed, so the Euler solver is unrolled at export time. + """ + + def __init__(self, flow: torch.nn.Module, n_timesteps: int = 10): + super().__init__() + self.flow = flow + self.n_timesteps = n_timesteps + cfm = flow.decoder + + # Replace the noise-sampling forward with one that consumes injected noise. + def cfm_forward( + mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, noised_mels=None, meanflow=False + ): + z = getattr(cfm, "_injected_noise", None) + if z is None: + z = torch.randn_like(mu) + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) + if cfm.t_scheduler == "cosine": + t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + return cfm.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond, meanflow=False), None + + cfm.forward = cfm_forward + + def forward(self, token, token_len, prompt_token, prompt_token_len, prompt_feat, embedding, noise): + self.flow.decoder._injected_noise = noise + feat, _ = self.flow.inference( + token=token, + token_len=token_len, + prompt_token=prompt_token, + prompt_token_len=prompt_token_len, + prompt_feat=prompt_feat, + prompt_feat_len=None, + embedding=embedding, + finalize=True, + n_timesteps=self.n_timesteps, + ) + return feat + + +class ChatterboxHifiganWrapper(torch.nn.Module): + """Wraps the HiFiGAN vocoder, replacing ``torch.istft`` with a traceable implementation.""" + + def __init__(self, hift: torch.nn.Module): + super().__init__() + self.hift = hift + n_fft = hift.istft_params["n_fft"] + hop = hift.istft_params["hop_len"] + window = hift.stft_window + + def patched_istft(magnitude, phase): + magnitude = torch.clip(magnitude, max=1e2) + real = magnitude * torch.cos(phase) + img = magnitude * torch.sin(phase) + return manual_istft(real, img, n_fft, hop, window.to(magnitude.device)) + + hift._istft = patched_istft + + def forward(self, speech_feat): + wav, _ = self.hift.inference(speech_feat=speech_feat) + return wav + + +class DummyChatterboxFlowInputGenerator(DummyInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "token", + "token_len", + "prompt_token", + "prompt_token_len", + "prompt_feat", + "embedding", + "noise", + ) + + def __init__(self, task: str, normalized_config: NormalizedConfig, **kwargs): + self.task = task + self.n_mels = getattr(normalized_config, "n_mels", 80) + self.token_mel_ratio = getattr(normalized_config, "token_mel_ratio", 2) + self.prompt_token_len = getattr(normalized_config, "prompt_token_len", 157) + self.prompt_feat_len = getattr(normalized_config, "prompt_feat_len", 314) + self.speaker_embedding_dim = getattr(normalized_config, "speaker_embedding_dim", 192) + self.seq_len = 64 + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + total_tokens = self.prompt_token_len + self.seq_len + mel_t = self.token_mel_ratio * total_tokens + if input_name == "token": + return self.random_float_tensor([1, self.seq_len], min_value=0, max_value=SPEECH_VOCAB_SIZE - 1) + if input_name == "token_len": + return torch.tensor([self.seq_len], dtype=torch.float32) + if input_name == "prompt_token": + return self.random_float_tensor([1, self.prompt_token_len], min_value=0, max_value=SPEECH_VOCAB_SIZE - 1) + if input_name == "prompt_token_len": + return torch.tensor([self.prompt_token_len], dtype=torch.float32) + if input_name == "prompt_feat": + return self.random_float_tensor([1, self.prompt_feat_len, self.n_mels]) + if input_name == "embedding": + return self.random_float_tensor([1, self.speaker_embedding_dim]) + if input_name == "noise": + return self.random_float_tensor([1, self.n_mels, mel_t]) + raise ValueError(f"Unsupported input {input_name} for DummyChatterboxFlowInputGenerator") + + +def _make_pad_mask_dynamic(lengths, max_len: int = 0): + """Drop-in for the S3Gen ``make_pad_mask`` that keeps the mask length dynamic. + + The original implementation calls ``lengths.max().item()`` which bakes the sequence + length into the traced graph as a constant. Deriving the range length from a tensor + instead keeps the exported model valid for arbitrary token lengths. + """ + lengths = lengths.long() + batch_size = lengths.size(0) + range_len = max_len if max_len > 0 else lengths.max() + seq_range = torch.arange(0, range_len, dtype=torch.int64, device=lengths.device) + seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, seq_range.shape[0]) + return seq_range_expand >= lengths.unsqueeze(-1) + + +class ChatterboxFlowModelPatcher(ModelPatcher): + """Patches non-traceable ops used by the S3Gen flow during export. + + Two substitutions are applied for the duration of the export: + + * ``torch.atleast_2d`` -> a ``reshape`` for 1D tensors (no OpenVINO conversion rule). + * ``make_pad_mask`` -> a variant whose mask length is derived from a tensor rather than + a Python ``int``, so the exported flow generalizes to arbitrary token lengths. + """ + + def __enter__(self): + super().__enter__() + self._orig_atleast_2d = torch.atleast_2d + torch.atleast_2d = lambda t: t.reshape(1, -1) if (torch.is_tensor(t) and t.ndim == 1) else t + + import chatterbox.models.s3gen.flow as flow_module + import chatterbox.models.s3gen.utils.mask as mask_module + + self._mask_modules = [mask_module, flow_module] + self._orig_make_pad_mask = mask_module.make_pad_mask + for mod in self._mask_modules: + if hasattr(mod, "make_pad_mask"): + mod.make_pad_mask = _make_pad_mask_dynamic + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + torch.atleast_2d = self._orig_atleast_2d + for mod in self._mask_modules: + if hasattr(mod, "make_pad_mask"): + mod.make_pad_mask = self._orig_make_pad_mask + + +class ChatterboxFlowOpenVINOConfig(OpenVINOConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyChatterboxFlowInputGenerator,) + NORMALIZED_CONFIG_CLASS = NormalizedConfig + _MODEL_PATCHER = ChatterboxFlowModelPatcher + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + return { + "token": {1: "token_length"}, + "token_len": {}, + "prompt_token": {1: "prompt_token_length"}, + "prompt_token_len": {}, + "prompt_feat": {1: "prompt_feat_length"}, + "embedding": {}, + "noise": {2: "mel_length"}, + } + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + return {"mel": {0: "batch_size", 2: "mel_out_length"}} + + +class DummyChatterboxHifiganInputGenerator(DummyInputGenerator): + SUPPORTED_INPUT_NAMES = ("speech_feat",) + + def __init__(self, task: str, normalized_config: NormalizedConfig, **kwargs): + self.task = task + self.n_mels = getattr(normalized_config, "n_mels", 80) + self.seq_len = 114 + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "speech_feat": + return self.random_float_tensor([1, self.n_mels, self.seq_len]) + raise ValueError(f"Unsupported input {input_name} for DummyChatterboxHifiganInputGenerator") + + +class ChatterboxHifiganOpenVINOConfig(OpenVINOConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyChatterboxHifiganInputGenerator,) + NORMALIZED_CONFIG_CLASS = NormalizedConfig + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + return {"speech_feat": {0: "batch_size", 2: "mel_length"}} + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + return {"waveform": {0: "batch_size", 1: "audio_length"}} + + +def build_chatterbox_t3_for_export(tts): + """Reconstruct a ``LlamaForCausalLM`` from the Chatterbox T3 backbone and speech head. + + Inference is driven by ``inputs_embeds``, so the (unused) ``embed_tokens`` layer keeps + its original vocabulary while the LM head is replaced by the speech-token head. + """ + from transformers import LlamaConfig, LlamaForCausalLM + + t3 = tts.t3 + speech_vocab = t3.hp.speech_tokens_dict_size + llama_cfg = LlamaConfig(**t3.cfg.to_dict()) + lm = LlamaForCausalLM(llama_cfg) + lm.model.load_state_dict(t3.tfmr.state_dict(), strict=True) + lm.lm_head = torch.nn.Linear(llama_cfg.hidden_size, speech_vocab, bias=t3.speech_head.bias is not None) + with torch.no_grad(): + lm.lm_head.weight.copy_(t3.speech_head.weight) + if t3.speech_head.bias is not None: + lm.lm_head.bias.copy_(t3.speech_head.bias) + lm.config.vocab_size = speech_vocab + lm.eval() + return lm + + +def get_t3_export_config(lm): + """Build the inputs_embeds stateful decoder export config for the T3 Llama backbone.""" + from optimum.exporters.tasks import TasksManager + + from .model_configs import LMInputEmbedsConfigHelper + + task = "text-generation-with-past" + export_config_class = TasksManager._SUPPORTED_MODEL_TYPE["llama"]["openvino"][task] + base_config = export_config_class(lm.config, task=task, use_past=True, use_past_in_inputs=True) + return LMInputEmbedsConfigHelper(base_config) + + +def get_chatterbox_models_for_export(model, n_cfm_timesteps: int = 10): + """Return ``{name: (submodel, export_config)}`` and the per-submodel statefulness list. + + ``model`` is the ``ChatterboxTTS`` object loaded by ``_ChatterboxForTextToSpeech``. + """ + config = model.config + + lm = build_chatterbox_t3_for_export(model) + t3_config = get_t3_export_config(lm) + + flow_wrapper = ChatterboxFlowWrapper(model.s3gen.flow, n_timesteps=n_cfm_timesteps).eval() + flow_config = ChatterboxFlowOpenVINOConfig(config, task="text-to-audio") + + hifigan_wrapper = ChatterboxHifiganWrapper(model.s3gen.mel2wav).eval() + hifigan_config = ChatterboxHifiganOpenVINOConfig(config, task="text-to-audio") + + models_and_export_configs = { + "t3": (lm, t3_config), + "flow": (flow_wrapper, flow_config), + "hifigan": (hifigan_wrapper, hifigan_config), + } + # Only the autoregressive T3 backbone benefits from a stateful (kv-cache) export. + stateful_submodels = [True, False, False] + return models_and_export_configs, stateful_submodels + + +def save_chatterbox_assets(model, output: Path, n_cfm_timesteps: int = 10): + """Save embedding tables, the built-in voice conditioning and config needed at inference.""" + from safetensors.torch import save_file + + t3 = model.t3 + config = model.config + + with torch.inference_mode(): + cond_prefix_emb = t3.prepare_conditioning(model.conds.t3).detach().clone() + + gen = model.conds.gen + assets = { + "speech_emb_weight": t3.speech_emb.weight.detach().clone(), + "speech_pos_emb_weight": t3.speech_pos_emb.emb.weight.detach().clone(), + "text_emb_weight": t3.text_emb.weight.detach().clone(), + "text_pos_emb_weight": t3.text_pos_emb.emb.weight.detach().clone(), + "cond_prefix_emb": cond_prefix_emb, + "gen_prompt_token": gen["prompt_token"].detach().to(torch.float32).clone(), + "gen_prompt_token_len": gen["prompt_token_len"].detach().to(torch.float32).clone(), + "gen_prompt_feat": gen["prompt_feat"].detach().to(torch.float32).clone(), + "gen_embedding": gen["embedding"].detach().to(torch.float32).clone(), + } + assets = {k: v.contiguous() for k, v in assets.items()} + save_file(assets, str(output / "chatterbox_assets.safetensors")) + + # Persist the inference metadata as a JSON config. + config_dict = {k: v for k, v in vars(config).items() if not k.startswith("_")} + config_dict["n_cfm_timesteps"] = n_cfm_timesteps + with open(output / "chatterbox_config.json", "w", encoding="utf-8") as f: + json.dump(config_dict, f, indent=2) + + # Save the tokenizer (and Chinese Cangjie table for the multilingual model) so the + # text front-end is available at inference time. + ckpt_dir = getattr(model, "_chatterbox_ckpt_dir", None) + if ckpt_dir is not None: + import shutil + + tokenizer_file = getattr(config, "tokenizer_file", "tokenizer.json") + tokenizer_files = [tokenizer_file] + if getattr(config, "multilingual", False): + tokenizer_files.append("Cangjie5_TC.json") + for fname in tokenizer_files: + src = Path(ckpt_dir) / fname + if src.is_file(): + shutil.copy(src, output / fname) diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py index dbe8e403cd..ade90d07f0 100644 --- a/optimum/exporters/openvino/convert.py +++ b/optimum/exporters/openvino/convert.py @@ -587,7 +587,7 @@ def export_from_model( ) library_name = _infer_library_from_model_or_model_class(model) - if library_name not in ("open_clip", "kokoro"): + if library_name not in ("open_clip", "kokoro", "chatterbox"): TasksManager.standardize_model_attributes(model, library_name=library_name) if hasattr(model.config, "export_model_type") and model.config.export_model_type is not None: @@ -602,7 +602,9 @@ def export_from_model( custom_architecture = library_name == "transformers" and model_type not in TasksManager._SUPPORTED_MODEL_TYPE - if task is not None and task != "auto": + if library_name == "chatterbox": + task = "text-to-audio" + elif task is not None and task != "auto": task = TasksManager.map_from_synonym(task) else: if library_name == "kokoro": @@ -691,7 +693,15 @@ def export_from_model( model, library_name, task, preprocessors, custom_export_configs, fn_get_submodels ) - if library_name == "diffusers": + if library_name == "chatterbox": + from .chatterbox import get_chatterbox_models_for_export + + n_cfm_timesteps = getattr(model.config, "n_cfm_timesteps", 10) + export_config = None + models_and_export_configs, stateful_submodels = get_chatterbox_models_for_export( + model, n_cfm_timesteps=n_cfm_timesteps + ) + elif library_name == "diffusers": export_config, models_and_export_configs = get_diffusion_models_for_export_ext(model, exporter="openvino") stateful_submodels = False elif stateful and is_encoder_decoder and not custom_architecture: @@ -717,7 +727,12 @@ def export_from_model( ) logging.disable(logging.NOTSET) - if library_name == "open_clip": + if library_name == "chatterbox": + from .chatterbox import save_chatterbox_assets + + save_chatterbox_assets(model, output, n_cfm_timesteps=getattr(model.config, "n_cfm_timesteps", 10)) + files_subpaths = ["openvino_" + model_name + ".xml" for model_name in models_and_export_configs.keys()] + elif library_name == "open_clip": if hasattr(model.config, "save_pretrained"): model.config.save_pretrained(output) diff --git a/optimum/exporters/openvino/utils.py b/optimum/exporters/openvino/utils.py index ca7b675b15..ddab2ecc43 100644 --- a/optimum/exporters/openvino/utils.py +++ b/optimum/exporters/openvino/utils.py @@ -547,7 +547,9 @@ def set_simplified_chat_template(ov_tokenizer_model, processor_chat_template=Non def allow_skip_tracing_check(library_name, model_type): - if library_name in ["diffusers", "sentence_transformers"]: + # Chatterbox submodels (S3Gen flow / HiFiGAN) draw random noise internally, so the + # TorchScript trace-check comparison is expected to differ and is not meaningful. + if library_name in ["diffusers", "sentence_transformers", "chatterbox"]: return True return model_type in SKIP_CHECK_TRACE_MODELS diff --git a/optimum/intel/openvino/modeling_text2speech.py b/optimum/intel/openvino/modeling_text2speech.py index 494f277fda..0ed247954d 100644 --- a/optimum/intel/openvino/modeling_text2speech.py +++ b/optimum/intel/openvino/modeling_text2speech.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import logging import os from pathlib import Path @@ -20,6 +21,7 @@ import numpy as np import openvino import torch +import torch.nn.functional as F from huggingface_hub import hf_hub_download from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE from torch import nn @@ -162,6 +164,21 @@ class OVModelForTextToSpeechSeq2Seq(OVModelForSeq2SeqLM): @classmethod def from_pretrained(cls, model_id, **kwargs): + # Chatterbox models do not ship a config.json recognizable by AutoConfig; their + # metadata lives in `chatterbox_config.json`. When such a model is detected we + # bypass the generic library/config inference (which requires a config.json) and + # route directly to the Chatterbox implementation. + if kwargs.get("config") is None and not kwargs.get("export", False): + config = _try_load_chatterbox_config( + model_id, + cache_dir=kwargs.get("cache_dir", HUGGINGFACE_HUB_CACHE), + token=kwargs.get("token"), + revision=kwargs.get("revision"), + local_files_only=kwargs.get("local_files_only", False), + ) + if config is not None: + return _OVModelForChatterboxTextToSpeech._from_pretrained(model_id, config, **kwargs) + # For Kokoro models, load config via PretrainedConfig since AutoConfig # does not recognize the "kokoro" model_type. if kwargs.get("config") is None: @@ -197,10 +214,14 @@ def _from_pretrained( ): if getattr(config, "model_type", None) == "kokoro": return _OVModelForKokoroTextToSpeech._from_pretrained(model_id, config, **kwargs) + elif getattr(config, "model_type", None) == "chatterbox": + return _OVModelForChatterboxTextToSpeech._from_pretrained(model_id, config, **kwargs) elif getattr(config, "architectures", None) and "SpeechT5ForTextToSpeech" in config.architectures: return _OVModelForSpeechT5ForTextToSpeech._from_pretrained(model_id, config, **kwargs) else: - raise ValueError(f"{getattr(config, 'model_type')} are not supported text-to-audio model using OpenVINO") + raise ValueError( + f"{getattr(config, 'architectures', None)} are not supported text-to-audio model using OpenVINO" + ) def reshape(self, *args, **kwargs): logger.warning("Static shapes are not supported for this model.") @@ -890,3 +911,500 @@ def preprocess_input( "segments": preprocessed_segments, "speed": speed, } + + +def _try_load_chatterbox_config( + model_id, cache_dir=HUGGINGFACE_HUB_CACHE, token=None, revision=None, local_files_only=False +): + """Load the Chatterbox inference config from an exported model directory or hub repo. + + Returns a ``PretrainedConfig`` with ``model_type == "chatterbox"`` if a + ``chatterbox_config.json`` is found, otherwise ``None``. + """ + config_path = None + model_path = Path(model_id) + if model_path.is_dir(): + candidate = model_path / "chatterbox_config.json" + if candidate.is_file(): + config_path = str(candidate) + else: + try: + config_path = hf_hub_download( + repo_id=str(model_id), + filename="chatterbox_config.json", + cache_dir=cache_dir, + token=token, + revision=revision, + local_files_only=local_files_only, + ) + except Exception: + config_path = None + + if config_path is None: + return None + + with open(config_path, "r", encoding="utf-8") as f: + config_dict = json.load(f) + config = PretrainedConfig() + for key, value in config_dict.items(): + setattr(config, key, value) + config.model_type = "chatterbox" + config.export_model_type = "chatterbox" + return config + + +def _chatterbox_punc_norm(text: str) -> str: + """Light text normalization matching the original Chatterbox front-end.""" + if len(text) == 0: + return "You need to add some text for me to talk." + if text[0].islower(): + text = text[0].upper() + text[1:] + text = " ".join(text.split()) + punc_to_replace = [ + ("...", ", "), + ("…", ", "), + (":", ","), + (" - ", ", "), + (";", ", "), + ("—", "-"), + ("–", "-"), + (" ,", ","), + ("“", '"'), + ("”", '"'), + ("‘", "'"), + ("’", "'"), + ] + for old, new in punc_to_replace: + text = text.replace(old, new) + text = text.rstrip(" ") + sentence_enders = {".", "!", "?", "-", ","} + if not any(text.endswith(p) for p in sentence_enders): + text += "." + return text + + +class _OVModelForChatterboxTextToSpeech(OVBaseModel): + """OpenVINO inference for the ResembleAI Chatterbox TTS model. + + The pipeline runs three exported submodels with thin PyTorch glue: + + * ``t3`` -- stateful autoregressive Llama decoder that consumes ``inputs_embeds`` + and produces speech-token logits. Embedding tables and the built-in voice + conditioning prefix are stored as assets and applied in Python. + * ``flow`` -- whole S3Gen flow (token -> mel) with the diffusion noise as an input. + * ``hifigan`` -- vocoder (mel -> waveform). + """ + + export_feature = "text-to-audio" + auto_model_class = AutoModelForTextToSpectrogram + + SPEECH_VOCAB_SIZE = 6561 + + def __init__(self, t3, flow, hifigan, config: PretrainedConfig = None, **kwargs): + self.config = config + self.model_save_dir = kwargs.get("model_save_dir", None) + self._device = kwargs.get("device", "CPU").upper() + self.ov_config = kwargs.get("ov_config") or {} + self.is_dynamic = True + self._compile_only = kwargs.get("compile_only", False) + self.generation_config = kwargs.get("generation_config", None) + self._openvino_config = None + + self._t3_model = t3 + self._flow_model = flow + self._hifigan_model = hifigan + self.request_t3 = None + self.request_flow = None + self.request_hifigan = None + + # Assets and tokenizer + self._assets = kwargs.get("assets", {}) + self._tokenizer = kwargs.get("tokenizer", None) + self.preprocessors = kwargs.get("preprocessors", []) + self.multilingual = bool(getattr(config, "multilingual", False)) + + if kwargs.get("compile", True) and not self._compile_only: + self.compile() + + @staticmethod + def _core_for_models(): + return openvino.Core() + + @staticmethod + def _load_tokenizer(tokenizer_path, multilingual: bool = False): + """Load the Chatterbox text front-end. + + Prefers the original ``MTLTokenizer``/``EnTokenizer`` classes (they implement + language-specific normalization and the ``[lang]`` token prefix), and falls back + to a raw ``tokenizers.Tokenizer`` if the ``chatterbox`` package is unavailable. + """ + if tokenizer_path is None or not os.path.isfile(tokenizer_path): + logger.warning("Chatterbox tokenizer file not found; text preprocessing will be unavailable.") + return None + try: + if multilingual: + from chatterbox.models.tokenizers import MTLTokenizer + + return MTLTokenizer(str(tokenizer_path)) + from chatterbox.models.tokenizers import EnTokenizer + + return EnTokenizer(str(tokenizer_path)) + except Exception as e: + logger.warning( + f"Could not load the Chatterbox tokenizer class ({e}). Falling back to a raw tokenizer; " + "language-specific preprocessing will not be applied." + ) + try: + from tokenizers import Tokenizer + + return Tokenizer.from_file(str(tokenizer_path)) + except Exception as e2: + logger.warning(f"Could not load Chatterbox tokenizer: {e2}") + return None + + def _save_config(self, save_directory): + # The Chatterbox metadata is persisted as `chatterbox_config.json` rather than the + # standard config.json (handled in `_save_pretrained`). + pass + + def _save_pretrained(self, save_directory: Union[str, Path]): + """Persist the OpenVINO submodels and the Chatterbox assets to ``save_directory``.""" + import shutil + + save_directory = Path(save_directory) + models = { + "openvino_t3.xml": self._t3_model, + "openvino_flow.xml": self._flow_model, + "openvino_hifigan.xml": self._hifigan_model, + } + for file_name, ov_model in models.items(): + openvino.save_model(ov_model, save_directory / file_name) + + # Copy the non-IR artifacts (assets, config and tokenizer) from the source dir. + if self.model_save_dir is not None: + src = Path(self.model_save_dir) + for extra in ("chatterbox_assets.safetensors", "chatterbox_config.json", "tokenizer.json"): + src_path = src / extra + if src_path.is_file(): + shutil.copy(src_path, save_directory / extra) + + @classmethod + def _from_pretrained( + cls, + model_id: Union[str, Path], + config: "PretrainedConfig", + token: Optional[Union[bool, str]] = None, + revision: Optional[str] = None, + force_download: bool = False, + cache_dir: str = HUGGINGFACE_HUB_CACHE, + local_files_only: bool = False, + load_in_8bit: bool = False, + quantization_config: Union[OVWeightQuantizationConfig, Dict] = None, + trust_remote_code: bool = False, + **kwargs, + ): + from safetensors.torch import load_file + + device = kwargs.pop("device", "CPU") + ov_config = kwargs.pop("ov_config", None) + compile_only = kwargs.pop("compile_only", False) + enable_compilation = kwargs.pop("compile", True) + + file_names = { + "t3": "openvino_t3.xml", + "flow": "openvino_flow.xml", + "hifigan": "openvino_hifigan.xml", + } + tokenizer_file = getattr(config, "tokenizer_file", "tokenizer.json") + extra_files = ["chatterbox_assets.safetensors", tokenizer_file] + if getattr(config, "multilingual", False): + extra_files.append("Cangjie5_TC.json") + + if os.path.isdir(model_id): + model_save_dir = Path(model_id) + resolved = {k: os.path.join(model_id, v) for k, v in file_names.items()} + for extra in extra_files: + resolved[extra] = os.path.join(model_id, extra) + else: + resolved = {} + for name, file_name in {**file_names, **{e: e for e in extra_files}}.items(): + for suffix in [".bin"] if name in file_names else []: + bin_name = file_name.replace(".xml", suffix) + hf_hub_download( + repo_id=str(model_id), + filename=bin_name, + token=token, + revision=revision, + cache_dir=cache_dir, + force_download=force_download, + local_files_only=local_files_only, + ) + resolved[name] = hf_hub_download( + repo_id=str(model_id), + filename=file_name, + token=token, + revision=revision, + cache_dir=cache_dir, + force_download=force_download, + local_files_only=local_files_only, + ) + model_save_dir = Path(resolved["t3"]).parent + + t3_model = cls._core_for_models().read_model(resolved["t3"]) + flow_model = cls._core_for_models().read_model(resolved["flow"]) + hifigan_model = cls._core_for_models().read_model(resolved["hifigan"]) + + assets = load_file(resolved["chatterbox_assets.safetensors"]) + + tokenizer = cls._load_tokenizer( + resolved.get(tokenizer_file), multilingual=getattr(config, "multilingual", False) + ) + + return cls( + t3=t3_model, + flow=flow_model, + hifigan=hifigan_model, + config=config, + device=device, + ov_config=ov_config, + model_save_dir=model_save_dir, + compile_only=compile_only, + compile=enable_compilation, + assets=assets, + tokenizer=tokenizer, + ) + + def compile(self): + core = self._core_for_models() + if self.request_t3 is None: + self.request_t3 = core.compile_model(self._t3_model, self._device, self.ov_config).create_infer_request() + if self.request_flow is None: + self.request_flow = core.compile_model(self._flow_model, self._device, self.ov_config) + if self.request_hifigan is None: + self.request_hifigan = core.compile_model(self._hifigan_model, self._device, self.ov_config) + + def clear_requests(self): + self.request_t3 = None + self.request_flow = None + self.request_hifigan = None + + def reshape(self, *args, **kwargs): + logger.warning("Static shapes are not supported for Chatterbox model.") + return self + + def can_generate(self) -> bool: + return True + + # ------------------------------------------------------------------ tokenization + def _text_to_tokens(self, text: str, language_id: Optional[str] = None) -> torch.Tensor: + if self._tokenizer is None: + raise ValueError( + "The Chatterbox tokenizer is not available. Make sure the tokenizer file is present " + "in the model directory." + ) + if self.multilingual and language_id is None: + language_id = "en" + + # The original Chatterbox tokenizer classes expose `text_to_tokens`, applying + # language-specific normalization and the `[lang]` prefix for the multilingual model. + if hasattr(self._tokenizer, "text_to_tokens"): + if self.multilingual: + tokens = self._tokenizer.text_to_tokens(text, language_id=language_id.lower()) + else: + tokens = self._tokenizer.text_to_tokens(text) + return tokens.to(dtype=torch.long) + + # Raw `tokenizers.Tokenizer` fallback (English only, no language handling). + text = text.replace(" ", "[SPACE]") + ids = self._tokenizer.encode(text).ids + return torch.tensor(ids, dtype=torch.long).unsqueeze(0) + + def preprocess_input(self, text: str, language_id: Optional[str] = None, **kwargs) -> dict: + """Normalize and tokenize text into ``input_ids`` ready for ``generate``. + + Args: + text: The input text to synthesize. + language_id: Two-letter language code (e.g. ``"ru"``, ``"fr"``, ``"zh"``) for the + multilingual model. Ignored by the English-only model. + """ + text = _chatterbox_punc_norm(text) + return {"input_ids": self._text_to_tokens(text, language_id=language_id)} + + # ------------------------------------------------------------------ T3 stage + def _emb(self, weight_key: str, idx: torch.Tensor) -> torch.Tensor: + return F.embedding(idx, self._assets[weight_key]) + + def _run_t3( + self, + input_ids: torch.Tensor, + max_new_tokens: int = 1000, + temperature: float = 0.8, + cfg_weight: float = 0.5, + repetition_penalty: float = 1.2, + min_p: float = 0.05, + top_p: float = 1.0, + ) -> torch.Tensor: + from transformers.generation.logits_process import ( + MinPLogitsWarper, + RepetitionPenaltyLogitsProcessor, + TopPLogitsWarper, + ) + + cfg = self.config + sot, eot = cfg.start_text_token, cfg.stop_text_token + sst, est = cfg.start_speech_token, cfg.stop_speech_token + + if input_ids.dim() == 1: + input_ids = input_ids.unsqueeze(0) + + # CFG: duplicate the sequence (conditional + unconditional). + text_tokens = torch.cat([input_ids, input_ids], dim=0) + text_tokens = F.pad(text_tokens, (1, 0), value=sot) + text_tokens = F.pad(text_tokens, (0, 1), value=eot) + + text_e = self._emb("text_emb_weight", text_tokens) + # The unconditional (CFG) branch zeros the token embedding BEFORE the positional + # embedding is added, so the position information is preserved (matches the reference). + text_e[1].zero_() + text_e = text_e + self._assets["text_pos_emb_weight"][: text_tokens.shape[1]] + + cond = self._assets["cond_prefix_emb"].expand(2, -1, -1) + # The reference pipeline prefixes the speech part with an initial start-of-speech + # token (inside `prepare_input_embeds`) and then concatenates an explicit BOS token, + # so the prefill contains two identical speech-token embeddings at position 0. + bos = torch.full((2, 1), sst, dtype=torch.long) + bos_e = self._emb("speech_emb_weight", bos) + self._assets["speech_pos_emb_weight"][0:1] + embeds = torch.cat([cond, text_e, bos_e, bos_e], dim=1) + + rep = RepetitionPenaltyLogitsProcessor(penalty=float(repetition_penalty)) + topp = TopPLogitsWarper(top_p=top_p) + minp = MinPLogitsWarper(min_p=min_p) + + self.request_t3.reset_state() + + # The stateful decoder derives RoPE positions from the attention-mask length, which + # must cover the full sequence seen so far (past cache + current tokens). + past_len = 0 + + def run(ie): + nonlocal past_len + am = np.ones((ie.shape[0], ie.shape[1] + past_len), dtype=np.int64) + res = self.request_t3.infer( + { + "inputs_embeds": ie.numpy().astype(np.float32), + "attention_mask": am, + "beam_idx": np.arange(ie.shape[0], dtype=np.int32), + } + ) + past_len += ie.shape[1] + return torch.from_numpy(next(iter(res.values()))) + + logits = run(embeds) + generated = bos[:1].clone() # track only the conditional batch + predicted = [] + for i in range(max_new_tokens): + step = logits[:, -1, :] + cond_logits, uncond_logits = step[0:1], step[1:2] + scaled = cond_logits + cfg_weight * (cond_logits - uncond_logits) + ids = generated + scaled = rep(ids, scaled) + if temperature != 1.0: + scaled = scaled / temperature + scaled = minp(ids, scaled) + scaled = topp(ids, scaled) + probs = torch.softmax(scaled, dim=-1) + next_token = torch.multinomial(probs, num_samples=1) + predicted.append(next_token) + generated = torch.cat([generated, next_token], dim=1) + if next_token.view(-1).item() == est: + break + next_e = self._emb("speech_emb_weight", next_token) + self._assets["speech_pos_emb_weight"][i + 1 : i + 2] + next_e = torch.cat([next_e, next_e]) # CFG + logits = run(next_e) + + self.request_t3.reset_state() + if not predicted: + return torch.zeros((1, 0), dtype=torch.long) + return torch.cat(predicted, dim=1) + + # ------------------------------------------------------------------ flow + vocoder + def _run_flow_and_vocoder(self, speech_tokens: torch.Tensor) -> torch.Tensor: + st = speech_tokens.view(-1) + st = st[st < self.SPEECH_VOCAB_SIZE] + token = st.unsqueeze(0).to(torch.float32) + token_len = torch.tensor([token.shape[1]], dtype=torch.float32) + + prompt_token = self._assets["gen_prompt_token"] + prompt_token_len = self._assets["gen_prompt_token_len"] + prompt_feat = self._assets["gen_prompt_feat"] + embedding = self._assets["gen_embedding"] + + token_mel_ratio = getattr(self.config, "token_mel_ratio", 2) + n_mels = getattr(self.config, "n_mels", 80) + total_tokens = prompt_token.shape[1] + token.shape[1] + mel_t = token_mel_ratio * total_tokens + noise = torch.randn(1, n_mels, mel_t) + + flow_out = self.request_flow( + [ + token.numpy(), + token_len.numpy(), + prompt_token.numpy(), + prompt_token_len.numpy(), + prompt_feat.numpy(), + embedding.numpy(), + noise.numpy(), + ] + ) + mel = flow_out[0] + wav = self.request_hifigan([mel])[0] + return torch.from_numpy(wav) + + # ------------------------------------------------------------------ public API + @torch.no_grad() + def generate( + self, + input_ids: Optional[Union[torch.Tensor, np.ndarray]] = None, + text: Optional[str] = None, + language_id: Optional[str] = None, + max_new_tokens: int = 1000, + temperature: float = 0.8, + cfg_weight: float = 0.5, + repetition_penalty: float = 1.2, + min_p: float = 0.05, + top_p: float = 1.0, + **kwargs, + ) -> torch.FloatTensor: + """Generate an audio waveform from token ids or raw text. + + Args: + input_ids: Tokenized input from ``preprocess_input``. If not provided, ``text`` + is tokenized internally. + text: Raw text to synthesize (used when ``input_ids`` is ``None``). + language_id: Two-letter language code (e.g. ``"ru"``, ``"fr"``, ``"zh"``) for the + multilingual model; ignored by the English-only model. Used only when the + text is tokenized internally (i.e. ``input_ids`` is ``None``). + """ + if input_ids is None: + if text is None: + raise ValueError("Either `input_ids` or `text` must be provided.") + input_ids = self.preprocess_input(text, language_id=language_id)["input_ids"] + if isinstance(input_ids, np.ndarray): + input_ids = torch.from_numpy(input_ids) + + speech_tokens = self._run_t3( + input_ids, + max_new_tokens=max_new_tokens, + temperature=temperature, + cfg_weight=cfg_weight, + repetition_penalty=repetition_penalty, + min_p=min_p, + top_p=top_p, + ) + if speech_tokens.shape[1] == 0: + raise RuntimeError("Chatterbox T3 produced no speech tokens for the given input.") + return self._run_flow_and_vocoder(speech_tokens) + + @property + def sampling_rate(self) -> int: + return int(getattr(self.config, "sampling_rate", 24000)) diff --git a/optimum/intel/utils/import_utils.py b/optimum/intel/utils/import_utils.py index b7b833a473..d9475acb2a 100644 --- a/optimum/intel/utils/import_utils.py +++ b/optimum/intel/utils/import_utils.py @@ -116,6 +116,13 @@ _kokoro_version = importlib_metadata.version("kokoro") except importlib_metadata.PackageNotFoundError: _kokoro_available = False +_chatterbox_available = importlib.util.find_spec("chatterbox") is not None +_chatterbox_version = "N/A" +if _chatterbox_available: + try: + _chatterbox_version = importlib_metadata.version("chatterbox-tts") + except importlib_metadata.PackageNotFoundError: + _chatterbox_available = False _safetensors_version = "N/A" _safetensors_available = importlib.util.find_spec("safetensors") is not None @@ -302,6 +309,10 @@ def is_kokoro_available(): return _kokoro_available +def is_chatterbox_available(): + return _chatterbox_available + + def is_safetensors_available(): return _safetensors_available diff --git a/optimum/intel/utils/modeling_utils.py b/optimum/intel/utils/modeling_utils.py index c1ac1edbc6..bcf3bd259a 100644 --- a/optimum/intel/utils/modeling_utils.py +++ b/optimum/intel/utils/modeling_utils.py @@ -165,6 +165,20 @@ def _is_kokoro_model( return False +def _is_chatterbox_model(all_files: list) -> bool: + """Detect ResembleAI Chatterbox TTS models. + + The Chatterbox repository does not ship a ``config.json`` recognizable by + ``AutoConfig``. Instead it is identified by its characteristic checkpoint + files (the T3 token-to-token model, the S3Gen vocoder and the voice encoder). + """ + file_names = {Path(f).name for f in all_files} + has_t3 = "t3_cfg.safetensors" in file_names or "t3_cfg.pt" in file_names + has_s3gen = "s3gen.safetensors" in file_names or "s3gen.pt" in file_names + has_ve = "ve.safetensors" in file_names or "ve.pt" in file_names + return has_t3 and has_s3gen and has_ve + + def _infer_library_from_model_name_or_path( model_name_or_path: Union[str, Path], subfolder: str = "", @@ -179,6 +193,8 @@ def _infer_library_from_model_name_or_path( library_name = "open_clip" elif _is_kokoro_model(model_name_or_path, all_files, cache_dir=cache_dir, token=token): library_name = "kokoro" + elif _is_chatterbox_model(all_files): + library_name = "chatterbox" else: library_name = TasksManager._infer_library_from_model_name_or_path( model_name_or_path=model_name_or_path, cache_dir=cache_dir @@ -197,6 +213,8 @@ def _infer_library_from_model_or_model_class( library_name = "open_clip" elif model.__module__.startswith("kokoro") or getattr(model, "_kokoro_model", False): library_name = "kokoro" + elif model.__module__.startswith("chatterbox") or getattr(model, "_chatterbox_model", False): + library_name = "chatterbox" elif model.__module__.startswith("optimum"): # for wrapped models like timm in optimum.intel.openvino.modeling_timm library_name = TasksManager._infer_library_from_model_or_model_class(model=model.model) @@ -499,3 +517,131 @@ def from_pretrained( model.config = config return model + + +class _ChatterboxForTextToSpeech: + """Wrapper for loading a ResembleAI Chatterbox TTS model for OpenVINO export. + + Chatterbox is a multi-stage TTS pipeline (a T3 token-to-token model with a Llama + backbone, the S3Gen flow-matching vocoder and a voice encoder) that is not natively + recognized by ``AutoModel``. This wrapper loads the original ``ChatterboxTTS`` object + and attaches a ``PretrainedConfig`` carrying the metadata required by the export and + inference code, so it conforms to optimum-intel expectations. + """ + + # File sets identifying each Chatterbox variant. + _ENGLISH_FILES = ["ve.safetensors", "t3_cfg.safetensors", "s3gen.safetensors", "tokenizer.json", "conds.pt"] + _MULTILINGUAL_FILES = [ + "ve.pt", + "t3_mtl23ls_v2.safetensors", + "s3gen.pt", + "grapheme_mtl_merged_expanded_v1.json", + "conds.pt", + "Cangjie5_TC.json", + ] + + @classmethod + def from_pretrained( + cls, + model_name_or_path: Union[str, Path], + cache_dir: str = HUGGINGFACE_HUB_CACHE, + token: Optional[Union[bool, str]] = None, + multilingual: Optional[bool] = None, + **kwargs, + ): + try: + from chatterbox.mtl_tts import ChatterboxMultilingualTTS + from chatterbox.tts import ChatterboxTTS + except ImportError: + raise ImportError( + "To load a Chatterbox TTS model, the `chatterbox-tts` package is required. " + "Please install it with `pip install chatterbox-tts`." + ) + + model_path = Path(model_name_or_path) + + # Decide which variant to load (English vs multilingual 23-language model). + # For local directories the variant is inferred from the available checkpoint files; + # for remote repositories the multilingual model is used by default. + if multilingual is None: + if model_path.is_dir(): + has_english = (model_path / "t3_cfg.safetensors").is_file() + has_multilingual = (model_path / "t3_mtl23ls_v2.safetensors").is_file() + # Prefer multilingual unless only the English checkpoint is present. + multilingual = has_multilingual or not has_english + else: + multilingual = True + + required_files = cls._MULTILINGUAL_FILES if multilingual else cls._ENGLISH_FILES + + if model_path.is_dir(): + ckpt_dir = model_path + else: + # Download the required checkpoint files and resolve their common parent dir. + from huggingface_hub import hf_hub_download + + local_path = None + for fpath in required_files: + local_path = hf_hub_download( + repo_id=str(model_name_or_path), filename=fpath, cache_dir=cache_dir, token=token + ) + ckpt_dir = Path(local_path).parent + + loader = ChatterboxMultilingualTTS if multilingual else ChatterboxTTS + tts = loader.from_local(ckpt_dir, device="cpu") + + config = cls._build_config(tts, multilingual=multilingual) + tts.config = config + tts._chatterbox_model = True + tts._chatterbox_ckpt_dir = str(ckpt_dir) + tts._chatterbox_multilingual = multilingual + return tts + + @staticmethod + def _build_config(tts, multilingual: bool = False) -> PretrainedConfig: + hp = tts.t3.hp + llama_cfg = tts.t3.cfg + flow = tts.s3gen.flow + cfm = flow.decoder + hift = tts.s3gen.mel2wav + + config = PretrainedConfig() + config.model_type = "chatterbox" + config.export_model_type = "chatterbox" + + # T3 (token-to-token) parameters + config.hidden_size = llama_cfg.hidden_size + config.speech_vocab_size = hp.speech_tokens_dict_size + config.start_text_token = hp.start_text_token + config.stop_text_token = hp.stop_text_token + config.start_speech_token = hp.start_speech_token + config.stop_speech_token = hp.stop_speech_token + config.speech_cond_prompt_len = hp.speech_cond_prompt_len + config.is_multilingual = hp.is_multilingual + config.llama_config = llama_cfg.to_dict() + + # Tokenizer / front-end metadata. The multilingual model uses MTLTokenizer with a + # language token prepended to the text, the English model uses EnTokenizer. + config.multilingual = bool(multilingual) + config.tokenizer_file = "grapheme_mtl_merged_expanded_v1.json" if multilingual else "tokenizer.json" + + # S3Gen flow parameters + config.n_mels = flow.output_size + config.token_mel_ratio = flow.token_mel_ratio + config.pre_lookahead_len = flow.pre_lookahead_len + config.inference_cfg_rate = cfm.inference_cfg_rate + config.t_scheduler = cfm.t_scheduler + config.n_cfm_timesteps = 10 + config.speaker_embedding_dim = flow.spk_embed_affine_layer.in_features + + # Built-in voice conditionals geometry + gen = tts.conds.gen + config.prompt_token_len = int(gen["prompt_token"].shape[1]) + config.prompt_feat_len = int(gen["prompt_feat"].shape[1]) + + # HiFiGAN vocoder parameters + config.istft_n_fft = hift.istft_params["n_fft"] + config.istft_hop_len = hift.istft_params["hop_len"] + config.sampling_rate = tts.sr + + return config diff --git a/tests/openvino/chatterbox_test_utils.py b/tests/openvino/chatterbox_test_utils.py new file mode 100644 index 0000000000..1ac14dc840 --- /dev/null +++ b/tests/openvino/chatterbox_test_utils.py @@ -0,0 +1,149 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Helpers to build a tiny, randomly-initialized Chatterbox TTS model for tests. + +The real ``ResembleAI/chatterbox`` checkpoint is ~2 GB, so the tests construct a small +variant at test time instead of downloading it: the T3 Llama backbone is reduced to a +couple of layers, the S3Gen flow estimator is shrunk, and the conditioning is synthetic. +The S3Gen conformer encoder and HiFiGAN vocoder keep their real (offline-built) shapes +because their internal dimensions are tightly coupled. Only the tokenizer file (small) is +fetched from the hub so the text front-end matches the real model. +""" + +import contextlib +import copy + +import torch + + +def _tiny_llama_config_patch(): + """Context manager that temporarily shrinks the T3 Llama backbone configuration.""" + from chatterbox.models.t3 import llama_configs + + @contextlib.contextmanager + def _patch(): + original = llama_configs.LLAMA_CONFIGS["Llama_520M"] + tiny = copy.deepcopy(original) + # Keep hidden_size/head_dim intact (the Perceiver resampler and embedding tables + # assume hidden_size=1024); only reduce depth and the MLP width. + tiny.update(intermediate_size=256, num_hidden_layers=2) + llama_configs.LLAMA_CONFIGS["Llama_520M"] = tiny + try: + yield + finally: + llama_configs.LLAMA_CONFIGS["Llama_520M"] = original + + return _patch() + + +@contextlib.contextmanager +def _tiny_s3gen_estimator_patch(): + """Temporarily shrink the S3Gen flow estimator (the heaviest exported S3Gen part).""" + import chatterbox.models.s3gen.s3gen as s3gen_module + + original = s3gen_module.ConditionalDecoder + + def tiny_decoder(*args, **kwargs): + kwargs.update(channels=[64], n_blocks=1, num_mid_blocks=2, num_heads=2, attention_head_dim=32) + return original(*args, **kwargs) + + s3gen_module.ConditionalDecoder = tiny_decoder + try: + yield + finally: + s3gen_module.ConditionalDecoder = original + + +def build_tiny_chatterbox(tmp_dir, multilingual: bool = False, seed: int = 0): + """Build a tiny random-weight Chatterbox model and return it ready for export. + + Args: + tmp_dir: A directory used to stage the (small) tokenizer file. + multilingual: Whether to build the multilingual variant (multilingual tokenizer + + larger text vocabulary). + seed: RNG seed for reproducible random weights. + + Returns: + A ``ChatterboxTTS``-like object with a ``config`` and synthetic ``conds`` attached, + suitable for ``optimum.exporters.openvino.export_from_model``. + """ + from pathlib import Path + + from chatterbox.models.s3gen import S3Gen + from chatterbox.models.t3 import T3 + from chatterbox.models.t3.modules.cond_enc import T3Cond + from chatterbox.models.t3.modules.t3_config import T3Config + from chatterbox.models.voice_encoder import VoiceEncoder + from chatterbox.tts import ChatterboxTTS, Conditionals + from huggingface_hub import hf_hub_download + + from optimum.intel.utils.modeling_utils import _ChatterboxForTextToSpeech + + torch.manual_seed(seed) + tmp_dir = Path(tmp_dir) + + with _tiny_llama_config_patch(), _tiny_s3gen_estimator_patch(): + hp = T3Config.multilingual() if multilingual else T3Config.english_only() + t3 = T3(hp).eval() + s3gen = S3Gen().eval() + ve = VoiceEncoder().eval() + + # The text front-end must match the real model; fetch only the small tokenizer file. + repo_id = "ResembleAI/chatterbox" + tokenizer_file = "grapheme_mtl_merged_expanded_v1.json" if multilingual else "tokenizer.json" + tokenizer_path = hf_hub_download(repo_id, tokenizer_file) + if multilingual: + from chatterbox.models.tokenizers import MTLTokenizer + + # The multilingual tokenizer loads the Cangjie table from the file's directory. + hf_hub_download(repo_id, "Cangjie5_TC.json") + tokenizer = MTLTokenizer(tokenizer_path) + else: + from chatterbox.models.tokenizers import EnTokenizer + + tokenizer = EnTokenizer(tokenizer_path) + + # Synthetic built-in voice conditioning (T3Cond + S3Gen reference dict). + speaker_emb = torch.randn(1, hp.speaker_embed_size) + cond_prompt = torch.randint(0, 100, (1, hp.speech_cond_prompt_len)) + t3_cond = T3Cond( + speaker_emb=speaker_emb, + cond_prompt_speech_tokens=cond_prompt, + emotion_adv=0.5 * torch.ones(1, 1, 1), + ) + n_prompt = 8 + gen = { + "prompt_token": torch.randint(0, 6561, (1, n_prompt)), + "prompt_token_len": torch.tensor([n_prompt]), + "prompt_feat": torch.randn(1, 2 * n_prompt, 80), + "prompt_feat_len": None, + "embedding": torch.randn(1, s3gen.flow.spk_embed_affine_layer.in_features), + } + + tts = ChatterboxTTS(t3, s3gen, ve, tokenizer, device="cpu", conds=Conditionals(t3_cond, gen)) + tts.config = _ChatterboxForTextToSpeech._build_config(tts, multilingual=multilingual) + tts._chatterbox_model = True + + # Stage the tokenizer files where the asset saver expects them. + ckpt_dir = tmp_dir / "tiny_chatterbox_ckpt" + ckpt_dir.mkdir(parents=True, exist_ok=True) + import shutil + + shutil.copy(tokenizer_path, ckpt_dir / tokenizer_file) + if multilingual: + shutil.copy(hf_hub_download(repo_id, "Cangjie5_TC.json"), ckpt_dir / "Cangjie5_TC.json") + tts._chatterbox_ckpt_dir = str(ckpt_dir) + + return tts diff --git a/tests/openvino/test_export.py b/tests/openvino/test_export.py index 6201757474..1d1462405e 100644 --- a/tests/openvino/test_export.py +++ b/tests/openvino/test_export.py @@ -13,6 +13,7 @@ # limitations under the License. +import importlib.util import unittest from pathlib import Path @@ -425,3 +426,39 @@ def test_export_custom_model(self): ov_outputs = ov_model(**tokens) self.assertTrue(torch.allclose(ov_outputs.token_embeddings, model_outputs.token_embeddings, atol=1e-4)) self.assertTrue(torch.allclose(ov_outputs.sentence_embedding, model_outputs.sentence_embedding, atol=1e-4)) + + +@unittest.skipUnless( + importlib.util.find_spec("chatterbox") is not None, + "chatterbox-tts package is not installed", +) +class ChatterboxExportModelTest(unittest.TestCase): + def _check_chatterbox_export(self, output_path: Path): + # Three OpenVINO submodels are produced: the T3 token-to-token model, the S3Gen + # flow (token -> mel) and the HiFiGAN vocoder (mel -> waveform). + for name in ("t3", "flow", "hifigan"): + self.assertTrue((output_path / f"openvino_{name}.xml").exists(), f"missing openvino_{name}.xml") + self.assertTrue((output_path / f"openvino_{name}.bin").exists(), f"missing openvino_{name}.bin") + # Assets (embedding tables, built-in voice conditioning) and metadata. + self.assertTrue((output_path / "chatterbox_assets.safetensors").exists()) + self.assertTrue((output_path / "chatterbox_config.json").exists()) + + def test_chatterbox_export(self): + from chatterbox_test_utils import build_tiny_chatterbox + + with TemporaryDirectory() as tmpdirname: + model = build_tiny_chatterbox(tmpdirname, multilingual=False) + output_path = Path(tmpdirname) / "ov" + export_from_model(model=model, output=output_path, task="text-to-audio") + self._check_chatterbox_export(output_path) + self.assertTrue((output_path / "tokenizer.json").exists()) + + def test_chatterbox_multilingual_export(self): + from chatterbox_test_utils import build_tiny_chatterbox + + with TemporaryDirectory() as tmpdirname: + model = build_tiny_chatterbox(tmpdirname, multilingual=True) + output_path = Path(tmpdirname) / "ov" + export_from_model(model=model, output=output_path, task="text-to-audio") + self._check_chatterbox_export(output_path) + self.assertTrue((output_path / "grapheme_mtl_merged_expanded_v1.json").exists()) diff --git a/tests/openvino/test_seq2seq.py b/tests/openvino/test_seq2seq.py index db500d5438..0af44e8c84 100644 --- a/tests/openvino/test_seq2seq.py +++ b/tests/openvino/test_seq2seq.py @@ -14,6 +14,7 @@ import copy import gc +import importlib.util import os import unittest from tempfile import TemporaryDirectory @@ -1302,6 +1303,74 @@ def test_compare_to_kokoro(self): gc.collect() +@unittest.skipUnless( + importlib.util.find_spec("chatterbox") is not None, + "chatterbox-tts package is not installed", +) +class OVModelForChatterboxTextToSpeechIntegrationTest(unittest.TestCase): + OVMODEL_CLASS = OVModelForTextToSpeechSeq2Seq + + def _export_tiny_model(self, tmpdirname, multilingual=False): + # Build a tiny random-weight Chatterbox at test time instead of downloading the + # ~2 GB checkpoint, and export it to OpenVINO. + from chatterbox_test_utils import build_tiny_chatterbox + + from optimum.exporters.openvino import export_from_model + + model = build_tiny_chatterbox(tmpdirname, multilingual=multilingual) + output_path = os.path.join(tmpdirname, "ov") + export_from_model(model=model, output=output_path, task="text-to-audio") + return output_path + + def test_export_and_inference(self): + set_seed(SEED) + with TemporaryDirectory() as tmpdirname: + output_path = self._export_tiny_model(tmpdirname) + ov_model = self.OVMODEL_CLASS.from_pretrained(output_path, device=OPENVINO_DEVICE) + + from optimum.intel.openvino.modeling_text2speech import _OVModelForChatterboxTextToSpeech + + self.assertIsInstance(ov_model, _OVModelForChatterboxTextToSpeech) + self.assertIsInstance(ov_model.config, PretrainedConfig) + self.assertEqual(ov_model.config.model_type, "chatterbox") + self.assertFalse(ov_model.multilingual) + + inputs = ov_model.preprocess_input("Hello, this is a test.") + self.assertIn("input_ids", inputs) + waveform = ov_model.generate(**inputs, max_new_tokens=120) + + self.assertIsInstance(waveform, torch.Tensor) + self.assertEqual(waveform.dim(), 2) + self.assertGreater(waveform.shape[-1], 0) + # The synthesized audio must be finite and non-silent. + self.assertTrue(torch.isfinite(waveform).all()) + self.assertGreater(float(waveform.abs().max()), 1e-3) + + del ov_model + gc.collect() + + def test_multilingual_export_and_inference(self): + set_seed(SEED) + with TemporaryDirectory() as tmpdirname: + output_path = self._export_tiny_model(tmpdirname, multilingual=True) + ov_model = self.OVMODEL_CLASS.from_pretrained(output_path, device=OPENVINO_DEVICE) + + self.assertTrue(ov_model.multilingual) + # Russian input via the `language_id` argument, mirroring the original model API. + waveform = ov_model.generate( + text="Всем привет. Вы используете оптимум-интел.", + language_id="ru", + max_new_tokens=120, + ) + self.assertEqual(waveform.dim(), 2) + self.assertGreater(waveform.shape[-1], 0) + self.assertTrue(torch.isfinite(waveform).all()) + self.assertGreater(float(waveform.abs().max()), 1e-3) + + del ov_model + gc.collect() + + class OVModelForPix2StructIntegrationTest(OVSeq2SeqTestMixin): SUPPORTED_ARCHITECTURES = ["pix2struct"] TASK = "image-to-text" # is it fine as well with visual-question-answering?