()
| 301 | |
| 302 | |
| 303 | def main() -> int: |
| 304 | args = parse_args() |
| 305 | if args.family != "ace_step": |
| 306 | raise RuntimeError(f"unsupported ACE-Step warmbench family: {args.family}") |
| 307 | |
| 308 | backend = normalize_device(args.backend, args.device) |
| 309 | checkpoint_dir = Path(args.checkpoint_dir) |
| 310 | if not checkpoint_dir.is_dir(): |
| 311 | raise RuntimeError(f"ACE-Step checkpoint dir not found: {checkpoint_dir}") |
| 312 | |
| 313 | os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") |
| 314 | os.environ.setdefault("ACESTEP_DISABLE_TQDM", "1") |
| 315 | torch.manual_seed(0) |
| 316 | np.random.seed(0) |
| 317 | torch.set_num_threads(max(1, args.threads)) |
| 318 | torch.backends.cuda.matmul.allow_tf32 = False |
| 319 | torch.backends.cudnn.allow_tf32 = False |
| 320 | torch.backends.cudnn.benchmark = False |
| 321 | install_controlled_vae_encode_noise(args.noise_file) |
| 322 | |
| 323 | requests = load_requests(args) |
| 324 | init_llm = any(request_needs_llm(request) for request in requests) |
| 325 | dit_handler, llm_handler = initialize_handlers( |
| 326 | checkpoint_dir, |
| 327 | args.config_path, |
| 328 | args.lm_model_path, |
| 329 | backend, |
| 330 | args.lm_backend, |
| 331 | args.lm_dtype, |
| 332 | init_llm, |
| 333 | ) |
| 334 | |
| 335 | output_root = Path(args.output_dir) if args.output_dir else DEFAULT_OUTPUT_ROOT |
| 336 | timing_lines: list[str] = [ |
| 337 | f"ace_step.backend {backend}", |
| 338 | f"ace_step.lm_backend {args.lm_backend}", |
| 339 | f"ace_step.lm_dtype {args.lm_dtype}", |
| 340 | f"ace_step.lm_initialized {1 if init_llm else 0}", |
| 341 | "ace_step.python_tf32_disabled 1", |
| 342 | "ace_step.python_tqdm_disabled 1", |
| 343 | ] |
| 344 | |
| 345 | if requests: |
| 346 | for warmup_index in range(max(0, args.warmup)): |
| 347 | _, warmup_timing = run_request( |
| 348 | dit_handler, |
| 349 | llm_handler, |
| 350 | requests[0], |
| 351 | output_root / "warmup" / f"{warmup_index:02d}", |
| 352 | backend, |
| 353 | args.noise_file, |
| 354 | ) |
| 355 | timing_lines.extend(warmup_timing) |
| 356 | |
| 357 | steps: list[dict[str, Any]] = [] |
| 358 | for request_index, request in enumerate(requests): |
| 359 | total_ms = 0.0 |
| 360 | last_step: dict[str, Any] | None = None |
no test coverage detected