| |
| """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() |
|
|
|
|
