Spaces:
Running
Running
| from __future__ import annotations | |
| import argparse | |
| import logging | |
| from pathlib import Path | |
| from typing import Optional, Sequence | |
| import torch | |
| from transformers import AutoModelForCausalLM | |
| from moss_tts_nano.defaults import ( | |
| DEFAULT_AUDIO_TOKENIZER_PATH, | |
| DEFAULT_CHECKPOINT_PATH, | |
| DEFAULT_OUTPUT_DIR, | |
| ) | |
| from text_normalization_pipeline import WeTextProcessingManager, prepare_tts_request_texts | |
| MOSS_AUDIO_TOKENIZER_TYPE = "moss-audio-tokenizer-nano" | |
| DEFAULT_AUDIO_TOKENIZER_PRETRAINED_NAME_OR_PATH = DEFAULT_AUDIO_TOKENIZER_PATH | |
| DEFAULT_OUTPUT_AUDIO_PATH = DEFAULT_OUTPUT_DIR / "infer_output.wav" | |
| def set_logging() -> None: | |
| logging.basicConfig( | |
| format="%(asctime)s %(levelname)s %(name)s: %(message)s", | |
| level=logging.INFO, | |
| ) | |
| def waiting_for_debug(ip: str, port: int) -> None: | |
| import debugpy | |
| logging.info("waiting for debugger attach at %s:%s", ip, port) | |
| debugpy.listen((ip, port)) | |
| debugpy.wait_for_client() | |
| def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: | |
| parser = argparse.ArgumentParser(description="Run MOSS-TTS-Nano inference from a HF-style checkpoint.") | |
| parser.add_argument( | |
| "--checkpoint", | |
| default=str(DEFAULT_CHECKPOINT_PATH), | |
| help="Checkpoint directory loadable by from_pretrained().", | |
| ) | |
| parser.add_argument( | |
| "--output-audio-path", | |
| default=str(DEFAULT_OUTPUT_AUDIO_PATH), | |
| help="Where to save the generated waveform.", | |
| ) | |
| text_group = parser.add_mutually_exclusive_group(required=True) | |
| text_group.add_argument("--text", help="Text to synthesize.") | |
| text_group.add_argument("--text-file", help="Path to a UTF-8 text file to synthesize.") | |
| prompt_text_group = parser.add_mutually_exclusive_group(required=False) | |
| prompt_text_group.add_argument("--prompt-text", help="Reference transcript used by continuation mode.") | |
| prompt_text_group.add_argument("--prompt-text-file", help="UTF-8 reference transcript file used by continuation mode.") | |
| parser.add_argument("--text-tokenizer-path", default=None, help="Override the checkpoint-bundled text tokenizer.") | |
| parser.add_argument( | |
| "--audio-tokenizer-pretrained-name-or-path", | |
| default=DEFAULT_AUDIO_TOKENIZER_PRETRAINED_NAME_OR_PATH, | |
| help="HF path or repo id for the audio tokenizer. Defaults to OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano.", | |
| ) | |
| parser.add_argument( | |
| "--mode", | |
| default="voice_clone", | |
| choices=("continuation", "voice_clone"), | |
| help=( | |
| "Inference mode. continuation: plain TTS or prompt_text + prompt_speech continuation; " | |
| "voice_clone: prompt_speech + target_text." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--prompt-audio-path", | |
| default=None, | |
| help="Reference speech used by continuation-with-prompt or voice_clone mode.", | |
| ) | |
| parser.add_argument( | |
| "--reference-audio-path", | |
| default=None, | |
| help="Compatibility alias for --prompt-audio-path.", | |
| ) | |
| parser.add_argument("--device", default="auto", help="Device to run on, for example auto/cpu/cuda/cuda:0.") | |
| parser.add_argument( | |
| "--dtype", | |
| default="auto", | |
| choices=("auto", "float32", "float16", "bfloat16"), | |
| help="Weights dtype after loading.", | |
| ) | |
| parser.add_argument( | |
| "--nq", | |
| type=int, | |
| default=None, | |
| help="Only use the first nq RVQ layers for prompt encode, model decoding, and audio decode.", | |
| ) | |
| parser.add_argument("--max-new-frames", type=int, default=375, help="Maximum number of audio frames to generate.") | |
| parser.add_argument( | |
| "--voice-clone-max-text-tokens", | |
| type=int, | |
| default=75, | |
| help=( | |
| "Only for voice_clone mode: pocket-tts style sentence chunking target token budget. " | |
| "Set <= 0 to disable chunking and keep single-pass generation." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--voice-clone-max-memory-per-sample-gb", | |
| type=float, | |
| default=1.0, | |
| help=( | |
| "Only for chunked voice_clone mode on CUDA: maximum free-memory budget used to size each sample " | |
| "when auto-selecting the chunk batch size." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--print-voice-clone-text-chunks", | |
| action="store_true", | |
| help="Print the effective voice_clone text chunks before generation.", | |
| ) | |
| parser.add_argument( | |
| "--do-sample", | |
| type=int, | |
| nargs="?", | |
| const=1, | |
| default=1, | |
| choices=[0, 1], | |
| help="Sample instead of greedy decoding. Accepts bare --do-sample or --do-sample 0/1.", | |
| ) | |
| parser.add_argument("--text-temperature", type=float, default=None, help="Text-layer sampling temperature. Default: 1.5.") | |
| parser.add_argument("--text-top-p", type=float, default=None, help="Text-layer top-p sampling. Default: 1.0.") | |
| parser.add_argument("--text-top-k", type=int, default=None, help="Text-layer top-k sampling. Default: 50.") | |
| parser.add_argument("--audio-temperature", type=float, default=None, help="Audio-layer sampling temperature. Default: 1.7.") | |
| parser.add_argument("--audio-top-p", type=float, default=None, help="Audio-layer top-p sampling. Default: 0.8.") | |
| parser.add_argument("--audio-top-k", type=int, default=None, help="Audio-layer top-k sampling. Default: 25.") | |
| parser.add_argument( | |
| "--audio-repetition-penalty", | |
| type=float, | |
| default=None, | |
| help="Audio-layer repetition penalty. Default: 1.0.", | |
| ) | |
| parser.add_argument( | |
| "--enable-wetext-processing", | |
| type=int, | |
| nargs="?", | |
| const=1, | |
| default=1, | |
| choices=[0, 1], | |
| help="Enable WeTextProcessing normalization before inference.", | |
| ) | |
| parser.add_argument( | |
| "--disable-wetext-processing", | |
| action="store_true", | |
| help="Disable WeTextProcessing normalization even if --enable-wetext-processing 1 is set.", | |
| ) | |
| parser.add_argument( | |
| "--enable-normalize-tts-text", | |
| "--enable-robust-text-normalization", | |
| dest="enable_normalize_tts_text", | |
| action="store_true", | |
| default=True, | |
| help="Enable the repository's normalize_tts_text robust cleanup before inference.", | |
| ) | |
| parser.add_argument( | |
| "--disable-normalize-tts-text", | |
| "--disable-robust-text-normalization", | |
| dest="disable_normalize_tts_text", | |
| action="store_true", | |
| help="Disable the repository's normalize_tts_text robust cleanup before inference.", | |
| ) | |
| parser.add_argument("--temperature", type=float, default=None, help=argparse.SUPPRESS) | |
| parser.add_argument("--top-k", type=int, default=None, help=argparse.SUPPRESS) | |
| parser.add_argument("--top-p", type=float, default=None, help=argparse.SUPPRESS) | |
| parser.add_argument("--repetition-penalty", type=float, default=None, help=argparse.SUPPRESS) | |
| parser.add_argument("--seed", type=int, default=None, help="Optional random seed for sampling.") | |
| parser.add_argument("--debug_ip", type=str, default="localhost") | |
| parser.add_argument("--debug_port", type=int, default=32431) | |
| parser.add_argument("--debug", type=int, default=0, help="Run with debug-friendly settings.") | |
| return parser.parse_args(argv) | |
| def resolve_text(args: argparse.Namespace) -> str: | |
| if args.text is not None: | |
| return args.text | |
| return Path(args.text_file).read_text(encoding="utf-8") | |
| def resolve_prompt_text(args: argparse.Namespace) -> Optional[str]: | |
| if args.prompt_text is not None: | |
| return args.prompt_text | |
| if args.prompt_text_file is not None: | |
| return Path(args.prompt_text_file).read_text(encoding="utf-8") | |
| return None | |
| def resolve_device(device_arg: str) -> torch.device: | |
| if device_arg == "auto": | |
| return torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| return torch.device(device_arg) | |
| def resolve_dtype(dtype_arg: str, device: torch.device) -> torch.dtype: | |
| if dtype_arg == "float32": | |
| return torch.float32 | |
| if dtype_arg == "float16": | |
| return torch.float16 | |
| if dtype_arg == "bfloat16": | |
| return torch.bfloat16 | |
| if device.type == "cuda": | |
| if hasattr(torch.cuda, "is_bf16_supported") and torch.cuda.is_bf16_supported(): | |
| return torch.bfloat16 | |
| return torch.float16 | |
| return torch.float32 | |
| def load_model(checkpoint: str, device: torch.device, dtype: torch.dtype): | |
| model = AutoModelForCausalLM.from_pretrained( | |
| checkpoint, | |
| trust_remote_code=True, | |
| ) | |
| model.to(device=device, dtype=dtype) | |
| model._set_attention_implementation("sdpa") | |
| model.eval() | |
| return model | |
| def resolve_sampling_kwargs(args: argparse.Namespace) -> dict[str, object]: | |
| text_temperature = 1.0 if args.text_temperature is None else float(args.text_temperature) | |
| text_top_p = 1.0 if args.text_top_p is None else float(args.text_top_p) | |
| text_top_k = 50 if args.text_top_k is None else int(args.text_top_k) | |
| audio_temperature = 0.8 if args.audio_temperature is None else float(args.audio_temperature) | |
| audio_top_p = 0.95 if args.audio_top_p is None else float(args.audio_top_p) | |
| audio_top_k = 25 if args.audio_top_k is None else int(args.audio_top_k) | |
| audio_repetition_penalty = ( | |
| 1.2 if args.audio_repetition_penalty is None else float(args.audio_repetition_penalty) | |
| ) | |
| if args.temperature is not None: | |
| if args.text_temperature is None: | |
| text_temperature = float(args.temperature) | |
| if args.audio_temperature is None: | |
| audio_temperature = float(args.temperature) | |
| if args.top_p is not None: | |
| if args.text_top_p is None: | |
| text_top_p = float(args.top_p) | |
| if args.audio_top_p is None: | |
| audio_top_p = float(args.top_p) | |
| if args.top_k is not None: | |
| if args.text_top_k is None: | |
| text_top_k = int(args.top_k) | |
| if args.audio_top_k is None: | |
| audio_top_k = int(args.top_k) | |
| if args.repetition_penalty is not None and args.audio_repetition_penalty is None: | |
| audio_repetition_penalty = float(args.repetition_penalty) | |
| return { | |
| "text_temperature": text_temperature, | |
| "text_top_p": text_top_p, | |
| "text_top_k": text_top_k, | |
| "audio_temperature": audio_temperature, | |
| "audio_top_p": audio_top_p, | |
| "audio_top_k": audio_top_k, | |
| "audio_repetition_penalty": audio_repetition_penalty, | |
| } | |
| def maybe_print_voice_clone_text_chunks( | |
| *, | |
| model, | |
| args: argparse.Namespace, | |
| text: str, | |
| ) -> None: | |
| if args.mode != "voice_clone" or not args.print_voice_clone_text_chunks: | |
| return | |
| text_tokenizer = model._load_text_tokenizer( | |
| text_tokenizer=None, | |
| text_tokenizer_path=args.text_tokenizer_path, | |
| ) | |
| split_chunks = model._split_text_into_best_sentences( | |
| text_tokenizer=text_tokenizer, | |
| text=text, | |
| max_tokens=args.voice_clone_max_text_tokens, | |
| ) | |
| effective_chunks = split_chunks if len(split_chunks) > 1 else [text] | |
| print("Voice clone text chunks") | |
| print("----------------------") | |
| print( | |
| f"max_tokens={args.voice_clone_max_text_tokens} " | |
| f"split_chunks={len(split_chunks)} effective_chunks={len(effective_chunks)}" | |
| ) | |
| for chunk_index, chunk_text in enumerate(effective_chunks, start=1): | |
| print(f"[chunk {chunk_index}]") | |
| print(chunk_text) | |
| print() | |
| def main(argv: Optional[Sequence[str]] = None) -> dict[str, object]: | |
| set_logging() | |
| args = parse_args(argv) | |
| if args.debug == 1: | |
| waiting_for_debug(args.debug_ip, args.debug_port) | |
| device = resolve_device(args.device) | |
| dtype = resolve_dtype(args.dtype, device) | |
| if args.seed is not None: | |
| torch.manual_seed(args.seed) | |
| if torch.cuda.is_available(): | |
| torch.cuda.manual_seed_all(args.seed) | |
| model = load_model(args.checkpoint, device=device, dtype=dtype) | |
| sampling_kwargs = resolve_sampling_kwargs(args) | |
| raw_text = resolve_text(args) | |
| raw_prompt_text = resolve_prompt_text(args) or "" | |
| enable_wetext_processing = bool(args.enable_wetext_processing) and not bool(args.disable_wetext_processing) | |
| enable_normalize_tts_text = bool(args.enable_normalize_tts_text) and not bool(args.disable_normalize_tts_text) | |
| text_normalizer_manager = None | |
| if enable_wetext_processing: | |
| text_normalizer_manager = WeTextProcessingManager() | |
| snapshot = text_normalizer_manager.ensure_ready() | |
| if not snapshot.ready: | |
| raise RuntimeError(snapshot.error or snapshot.message) | |
| logging.info("WeTextProcessing ready for infer.py status=%s", snapshot.message) | |
| prepared_texts = prepare_tts_request_texts( | |
| text=raw_text, | |
| prompt_text=raw_prompt_text, | |
| voice="", | |
| enable_wetext=enable_wetext_processing, | |
| enable_normalize_tts_text=enable_normalize_tts_text, | |
| text_normalizer_manager=text_normalizer_manager, | |
| ) | |
| text = str(prepared_texts["text"]) | |
| prompt_text = str(prepared_texts["prompt_text"]).strip() or None | |
| logging.info( | |
| "text normalization method=%s language=%s text_chars=%d prompt_chars=%d", | |
| prepared_texts["normalization_method"], | |
| prepared_texts["text_normalization_language"] or "n/a", | |
| len(text), | |
| len(prompt_text or ""), | |
| ) | |
| maybe_print_voice_clone_text_chunks(model=model, args=args, text=text) | |
| logging.info("running inference mode=%s", args.mode) | |
| result = model.inference( | |
| text=text, | |
| output_audio_path=args.output_audio_path, | |
| mode=args.mode, | |
| prompt_text=prompt_text, | |
| prompt_audio_path=args.prompt_audio_path, | |
| reference_audio_path=args.reference_audio_path, | |
| text_tokenizer_path=args.text_tokenizer_path, | |
| audio_tokenizer_type=MOSS_AUDIO_TOKENIZER_TYPE, | |
| audio_tokenizer_pretrained_name_or_path=args.audio_tokenizer_pretrained_name_or_path, | |
| device=device, | |
| nq=args.nq, | |
| max_new_frames=args.max_new_frames, | |
| voice_clone_max_text_tokens=args.voice_clone_max_text_tokens, | |
| voice_clone_max_memory_per_sample_gb=args.voice_clone_max_memory_per_sample_gb, | |
| do_sample=bool(args.do_sample), | |
| use_kv_cache=True, | |
| **sampling_kwargs, | |
| ) | |
| logging.info( | |
| "saved generated audio to %s sample_rate=%s frames=%s", | |
| result["audio_path"], | |
| result["sample_rate"], | |
| int(result["audio_token_ids"].shape[0]), | |
| ) | |
| return result | |
| if __name__ == "__main__": | |
| main() | |