()
| 499 | |
| 500 | |
| 501 | def main() -> int: |
| 502 | args = parse_args() |
| 503 | if args.family != "vibevoice": |
| 504 | raise RuntimeError(f"unsupported VibeVoice warmbench family: {args.family}") |
| 505 | |
| 506 | model_root = resolve_repo_path(args.model).resolve() |
| 507 | if not (model_root / "config.json").is_file(): |
| 508 | raise RuntimeError(f"missing VibeVoice model config: {model_root / 'config.json'}") |
| 509 | |
| 510 | VibeVoiceForConditionalGenerationInference, VibeVoiceProcessor, module_path = load_reference_symbols(args.reference_root) |
| 511 | |
| 512 | device = select_device(args) |
| 513 | dtype = torch.float32 if args.backend == "cpu" else torch_dtype(args.dtype) |
| 514 | tokenizer_kwargs: dict[str, Any] = {"local_files_only": args.local_files_only} |
| 515 | if args.tokenizer_model: |
| 516 | tokenizer_kwargs["language_model_pretrained_name"] = str(resolve_repo_path(Path(args.tokenizer_model))) |
| 517 | processor = VibeVoiceProcessor.from_pretrained(str(model_root), **tokenizer_kwargs) |
| 518 | |
| 519 | model = VibeVoiceForConditionalGenerationInference.from_pretrained( |
| 520 | str(model_root), |
| 521 | torch_dtype=dtype, |
| 522 | attn_implementation=args.attn_implementation, |
| 523 | device_map=("cuda" if args.backend == "cuda" else "cpu"), |
| 524 | ) |
| 525 | model.eval() |
| 526 | patch_transformers_cache_signature(model) |
| 527 | requests = load_requests(args) |
| 528 | if not requests: |
| 529 | raise RuntimeError("VibeVoice warmbench request sequence is empty") |
| 530 | warmup_request = dict(requests[0]) |
| 531 | if args.warmup_text: |
| 532 | warmup_request["text"] = args.warmup_text |
| 533 | for request in requests: |
| 534 | voice_samples = request.get("voice_samples", args.voice_samples or [str(path) for path in DEFAULT_VOICES]) |
| 535 | if not voice_samples: |
| 536 | raise RuntimeError("VibeVoice warmbench request requires voice_samples") |
| 537 | for voice_sample in [str(resolve_repo_path(Path(path))) for path in voice_samples]: |
| 538 | if not Path(voice_sample).is_file(): |
| 539 | raise RuntimeError(f"missing VibeVoice voice sample: {voice_sample}") |
| 540 | for voice_sample in warmup_request.get("voice_samples", args.voice_samples or [str(path) for path in DEFAULT_VOICES]): |
| 541 | voice_sample = str(resolve_repo_path(Path(voice_sample))) |
| 542 | if not Path(voice_sample).is_file(): |
| 543 | raise RuntimeError(f"missing VibeVoice voice sample: {voice_sample}") |
| 544 | audio_out_dir = args.output_dir if args.output_dir is not None else args.audio_out_dir |
| 545 | |
| 546 | warmup_outputs: list[tuple[np.ndarray, int, float, dict[str, Any]]] = [] |
| 547 | for _ in range(max(0, args.warmup)): |
| 548 | warmup_outputs.append(run_once(model, processor, warmup_request, args, device)) |
| 549 | |
| 550 | last_outputs: list[tuple[np.ndarray, int, float, dict[str, Any]]] = [] |
| 551 | if args.batch: |
| 552 | current_batch: list[tuple[np.ndarray, int, float, dict[str, Any]]] | None = None |
| 553 | for _ in range(max(1, args.iterations)): |
| 554 | current_batch = run_batch(model, processor, requests, args, device) |
| 555 | assert current_batch is not None |
| 556 | last_outputs = current_batch |
| 557 | else: |
| 558 | for request in requests: |
no test coverage detected