374 lines
14 KiB
Python
374 lines
14 KiB
Python
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()
|