| """Command-line orchestration: tokenize, release the encoder, then render.""" |
|
|
| import argparse |
| import gc |
| import os |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| from . import __version__ |
| from .audio import load_audio, save_wav_atomic |
| from .quantizer import WannabeTokenizer |
| from .renderer import YuE2Renderer |
|
|
|
|
| def positive_int(value: str) -> int: |
| try: |
| parsed = int(value) |
| except ValueError as exc: |
| raise argparse.ArgumentTypeError("must be an integer") from exc |
| if str(parsed) != value or parsed < 1: |
| raise argparse.ArgumentTypeError("must be a positive integer without extra characters") |
| return parsed |
|
|
|
|
| def save_tokens_atomic(path: Path, tokens: np.ndarray) -> None: |
| temporary = path.with_name(path.name + ".tmp") |
| try: |
| with temporary.open("wb") as f: |
| np.save(f, tokens, allow_pickle=False) |
| os.replace(temporary, path) |
| finally: |
| if temporary.exists(): |
| temporary.unlink() |
|
|
|
|
| def build_parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser( |
| description="Tokenize one audio file with WannabeTokenizer and reconstruct it with YuE2." |
| ) |
| parser.add_argument("input", type=Path, help="Input audio file") |
| parser.add_argument("--tokenizer", required=True, type=Path, help="Unified WannabeTokenizer .safetensors file") |
| parser.add_argument("--yue2-checkpoint", required=True, type=Path, help="YuE2 bf16 all-in-one .safetensors file") |
| parser.add_argument("--output-dir", type=Path, default=Path("output"), help="Destination folder (default: output)") |
| parser.add_argument("--device", default="auto", help="Torch device, for example auto, cuda, cuda:0, or cpu") |
| parser.add_argument("--seed", type=int, default=0, help="NAR noise seed (default: 0)") |
| parser.add_argument("--steps", type=positive_int, default=32, help="Flow-matching steps (default: 32)") |
| parser.add_argument("--batch-windows", type=positive_int, default=8, help="Tokenizer windows per batch (default: 8)") |
| parser.add_argument("--overwrite", action="store_true", help="Replace existing tokens.npy and reconstructed.wav") |
| parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}") |
| return parser |
|
|
|
|
| def main() -> None: |
| args = build_parser().parse_args() |
| audio_path = args.input.expanduser() |
| tokenizer_path = args.tokenizer.expanduser() |
| yue2_path = args.yue2_checkpoint.expanduser() |
| output_dir = args.output_dir.expanduser() |
| for label, path in ( |
| ("Input audio", audio_path), |
| ("Tokenizer checkpoint", tokenizer_path), |
| ("YuE2 checkpoint", yue2_path), |
| ): |
| if not path.is_file(): |
| raise FileNotFoundError(f"{label} not found: {path}") |
|
|
| device = "cuda" if args.device == "auto" and torch.cuda.is_available() else args.device |
| if device == "auto": |
| device = "cpu" |
| if device.startswith("cuda") and not torch.cuda.is_available(): |
| raise RuntimeError("A CUDA device was requested, but CUDA is not available") |
| if device == "cpu": |
| print("Warning: YuE2 reconstruction on CPU is very slow and requires substantial RAM.") |
|
|
| tokens_path = output_dir / "tokens.npy" |
| audio_output_path = output_dir / "reconstructed.wav" |
| outputs = (tokens_path, audio_output_path) |
| if not args.overwrite: |
| existing = [str(path) for path in outputs if path.exists()] |
| if existing: |
| raise FileExistsError("Output already exists (use --overwrite to replace): " + ", ".join(existing)) |
| output_dir.mkdir(parents=True, exist_ok=True) |
|
|
| waveform = load_audio(audio_path) |
| print(f"Tokenizing on {device}...") |
| tokenizer = WannabeTokenizer(tokenizer_path, device=device) |
| tokens = tokenizer.tokenize(waveform, batch_windows=args.batch_windows) |
| save_tokens_atomic(tokens_path, tokens) |
| print(f"Wrote {tokens_path} ({len(tokens)} frames, {len(tokens) / 25:.2f} seconds)") |
|
|
| |
| del tokenizer |
| gc.collect() |
| if device.startswith("cuda"): |
| torch.cuda.empty_cache() |
|
|
| print(f"Reconstructing on {device} with seed {args.seed} and {args.steps} steps...") |
| renderer = YuE2Renderer(yue2_path, device=device) |
| reconstructed = renderer.reconstruct(tokens, seed=args.seed, steps=args.steps) |
| save_wav_atomic(audio_output_path, reconstructed) |
| print(f"Wrote {audio_output_path} (48 kHz stereo WAV)") |
|
|
