(
model: Any,
processor: Any,
request: dict[str, Any],
args: argparse.Namespace,
device: str,
)
| 349 | |
| 350 | |
| 351 | def run_once( |
| 352 | model: Any, |
| 353 | processor: Any, |
| 354 | request: dict[str, Any], |
| 355 | args: argparse.Namespace, |
| 356 | device: str, |
| 357 | ) -> tuple[np.ndarray, int, float, dict[str, Any]]: |
| 358 | text = str(request.get("text", "")) |
| 359 | if not text: |
| 360 | raise RuntimeError("VibeVoice warmbench request missing text") |
| 361 | voice_samples = [str(resolve_repo_path(Path(path))) for path in request.get("voice_samples", [])] |
| 362 | if not voice_samples: |
| 363 | voice_samples = [str(resolve_repo_path(Path(path))) for path in (args.voice_samples or DEFAULT_VOICES)] |
| 364 | seed = int(request.get("seed", args.seed)) |
| 365 | ddpm_steps = int(request.get("ddpm_steps", args.ddpm_steps)) |
| 366 | max_new_tokens = int(request.get("max_new_tokens", args.max_new_tokens)) |
| 367 | max_length_times = float(request.get("max_length_times", args.max_length_times)) |
| 368 | cfg_scale = float(request.get("cfg_scale", args.cfg_scale)) |
| 369 | prompt_noise_file = str(request.get("prompt_noise_file", "")) or args.prompt_noise_file |
| 370 | noise_file = str(request.get("diffusion_noise_file", "")) or args.noise_file |
| 371 | seed_all(seed, args.backend) |
| 372 | model.set_ddpm_inference_steps(ddpm_steps) |
| 373 | inputs = processor( |
| 374 | text=text, |
| 375 | voice_samples=voice_samples or None, |
| 376 | padding=True, |
| 377 | return_tensors="pt", |
| 378 | return_attention_mask=True, |
| 379 | ) |
| 380 | inputs = move_batch_to_device(dict(inputs), device) |
| 381 | acoustic_vae_dim = int(model.model.config.acoustic_vae_dim) |
| 382 | with controlled_generation_noise(prompt_noise_file, noise_file, acoustic_vae_dim): |
| 383 | started = time.perf_counter() |
| 384 | outputs = model.generate( |
| 385 | **inputs, |
| 386 | max_new_tokens=max_new_tokens if max_new_tokens > 0 else None, |
| 387 | cfg_scale=cfg_scale, |
| 388 | tokenizer=processor.tokenizer, |
| 389 | generation_config={"do_sample": False}, |
| 390 | verbose=False, |
| 391 | show_progress_bar=False, |
| 392 | is_prefill=not args.disable_prefill, |
| 393 | max_length_times=max_length_times, |
| 394 | ) |
| 395 | if args.backend == "cuda": |
| 396 | torch.cuda.synchronize(args.device) |
| 397 | ended = time.perf_counter() |
| 398 | audio = extract_audio(outputs) |
| 399 | input_tokens = int(inputs["input_ids"].shape[-1]) |
| 400 | output_tokens = int(outputs.sequences.shape[-1]) |
| 401 | generated_token_ids = [int(token) for token in outputs.sequences[0, input_tokens:].detach().cpu().tolist()] |
| 402 | reach_max_step = getattr(outputs, "reach_max_step_sample", None) |
| 403 | reached_max_step = bool(reach_max_step.detach().cpu().any().item()) if torch.is_tensor(reach_max_step) else False |
| 404 | metrics = { |
| 405 | "input_tokens": input_tokens, |
| 406 | "output_tokens": output_tokens, |
| 407 | "generated_tokens": output_tokens - input_tokens, |
| 408 | "reached_max_step": reached_max_step, |
no test coverage detected