File size: 2,535 Bytes
0584718 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | #!/usr/bin/env python3
"""Generate one zero-shot voice-cloned utterance with the local NVFP4 bundle."""
from __future__ import annotations
import argparse
import re
from pathlib import Path
import soundfile as sf
import torch
from huggingface_hub import snapshot_download
from chatterbox_flash.nvfp4 import (
BASE_ASSET_FILENAMES,
BASE_CHECKPOINT_REVISION,
BASE_REPO_ID,
load_prepacked_tts_local,
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--prompt-wav", type=Path, required=True)
parser.add_argument("--text", required=True)
parser.add_argument("--output", type=Path, default=Path("output.wav"))
parser.add_argument(
"--block-size",
type=int,
choices=(16, 24),
default=16,
help="16 is the quality profile; 24 is the optional faster profile",
)
parser.add_argument("--seed", type=int, default=4242)
return parser.parse_args()
def main() -> None:
args = parse_args()
if not args.prompt_wav.is_file():
raise SystemExit(f"Missing reference audio: {args.prompt_wav}")
bundle = Path(__file__).resolve().parent
base_dir = Path(snapshot_download(
repo_id=BASE_REPO_ID,
revision=BASE_CHECKPOINT_REVISION,
allow_patterns=list(BASE_ASSET_FILENAMES),
))
capability = torch.cuda.get_device_capability()
gpu_slug = re.sub(r"[^a-zA-Z0-9_.-]+", "-", torch.cuda.get_device_name())
autotune_cache = (
Path.home()
/ ".cache"
/ "chatterbox-flash-nvfp4"
/ f"flashinfer-{gpu_slug}-sm{capability[0]}{capability[1]}-d{args.block_size}.json"
)
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
tts = load_prepacked_tts_local(
bundle,
base_dir,
device="cuda",
drf_block_size=args.block_size,
autotune_cache=autotune_cache,
)
waveform = tts.generate(
args.text,
audio_prompt_path=args.prompt_wav,
exaggeration=0.5,
num_steps=10,
temperature=0.2,
time_shift_tau=0.5,
omnivoice_schedule_t_shift=0.5,
cfg_scale=1.0,
position_temperature=5.0,
n_cfm_timesteps=2,
backend="flashinfer",
use_cuda_graph=True,
)
args.output.parent.mkdir(parents=True, exist_ok=True)
sf.write(args.output, waveform.float().numpy(), tts.sr)
print(f"Wrote {args.output} at {tts.sr} Hz")
if __name__ == "__main__":
main()
|