[Diffusion] Add mixed-resolution benchmark support (for #20762) (#20863)

Signed-off-by: Fengyuan Yu <15fengyuan@gmail.com>
Co-authored-by: Fengyuan Yu <15fengyuan@gmail.com>
Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
Fengyuan Yu
2026-04-22 09:22:19 +03:00
committed by GitHub
co-authored by Fengyuan Yu ronnie_zheng
parent e39f0f4ff3
commit 5c245d978f
3 changed files with 120 additions and 18 deletions
@@ -29,7 +29,7 @@ import dataclasses
import json
import time
from dataclasses import dataclass
from typing import Any, Dict, List, Tuple
from typing import Any, Dict, List, Optional, Tuple
import torch
from tqdm import tqdm
@@ -81,6 +81,8 @@ class BenchArgs:
task_name: str = "unknown"
num_prompts: int = 10
batch_size: int = 1
random_request_config: str = None
random_request_seed: int = 42
# Benchmark Execution
skip_warmup: bool = False
@@ -151,6 +153,23 @@ class BenchArgs:
help="Batch size per generation call (currently only bs=1 is supported)",
)
parser.add_argument(
"--random-request-config",
type=str,
default=None,
help=(
"JSON string defining random request profiles. "
"Each profile may contain: width, height, num_inference_steps, etc. "
"The 'weight' field controls sampling probability (relative weight)."
),
)
parser.add_argument(
"--random-request-seed",
type=int,
default=42,
help="Random seed for sampling request profiles (default: 42).",
)
# Benchmark Execution
parser.add_argument(
"--skip-warmup", action="store_true", help="Skip warmup batch"
@@ -186,17 +205,22 @@ def generate_batch(
engine: DiffGenerator,
bench_args: BenchArgs,
prompts: List[str],
user_sampling_params: Dict[str, Any],
user_sampling_params: List[Dict[str, Any]],
) -> BatchOutput:
"""Generate batch of images/videos synchronously."""
assert len(user_sampling_params) == len(prompts), (
f"user_sampling_params length ({len(user_sampling_params)}) must match "
f"prompts length ({len(prompts)})"
)
output = BatchOutput()
start_time = time.perf_counter()
torch.cuda.reset_peak_memory_stats()
for prompt in prompts:
for prompt, params in zip(prompts, user_sampling_params):
try:
sampling_params_kwargs = dict(user_sampling_params)
sampling_params_kwargs = dict(params)
sampling_params_kwargs["prompt"] = prompt
result = engine.generate(sampling_params_kwargs=sampling_params_kwargs)
@@ -227,6 +251,7 @@ def calculate_metrics(
total_duration: float,
resolution: Tuple[int, int, int],
num_requests: int,
all_sampling_params: Optional[List[Dict[str, Any]]] = None,
) -> Dict[str, Any]:
"""Calculate generation-specific throughput metrics."""
successful = [o for o in outputs if o.success]
@@ -235,8 +260,15 @@ def calculate_metrics(
peak_memory = max((o.peak_memory_mb for o in outputs), default=0)
width, height, frames = resolution
pixels_per_sample = width * height * frames
total_pixels = num_success * pixels_per_sample
if all_sampling_params:
total_pixels = sum(
p.get("width", width)
* p.get("height", height)
* p.get("num_frames", frames)
for p in all_sampling_params[:num_success]
)
else:
total_pixels = num_success * width * height * frames
metrics = {
"num_requests": num_requests,
@@ -272,6 +304,11 @@ def throughput_test(
engine = initialize_engine(server_args)
if bench_args.random_request_config and bench_args.dataset != "random":
raise ValueError(
"--random-request-config can only be used with --dataset random"
)
logger.info(f"Loading {bench_args.dataset} dataset...")
if bench_args.dataset == "vbench":
bench_args.task_name = engine.server_args.pipeline_config.task_type
@@ -281,7 +318,7 @@ def throughput_test(
else:
raise ValueError(f"Unknown dataset: {bench_args.dataset}")
sampling_params = {
_sampling_params = {
"guidance_scale": bench_args.guidance_scale,
"num_inference_steps": bench_args.num_inference_steps,
"height": bench_args.height,
@@ -290,18 +327,29 @@ def throughput_test(
"seed": bench_args.seed,
}
if bench_args.disable_safety_checker:
sampling_params["safety_checker"] = None
_sampling_params["safety_checker"] = None
total_count = min(bench_args.num_prompts, len(dataset))
all_prompts = [dataset[i].prompt for i in range(total_count)]
if bench_args.random_request_config:
all_sampling_params = []
for i in range(total_count):
params = dict(_sampling_params)
params.update(dataset.get_sampling_params(i))
all_sampling_params.append(params)
else:
all_sampling_params = [_sampling_params] * total_count
if not bench_args.skip_warmup:
logger.info("Running warmup batch...")
warmup_count = min(bench_args.batch_size, len(dataset))
warmup_prompts = [dataset[i].prompt for i in range(warmup_count)]
generate_batch(engine, bench_args, warmup_prompts, sampling_params)
warmup_count = min(bench_args.batch_size, total_count)
warmup_prompts = all_prompts[:warmup_count]
warmup_sampling_params = all_sampling_params[:warmup_count]
generate_batch(engine, bench_args, warmup_prompts, warmup_sampling_params)
logger.info(f"Running benchmark with {bench_args.num_prompts} prompts...")
outputs: List[BatchOutput] = []
total_count = min(bench_args.num_prompts, len(dataset))
all_prompts = [dataset[i].prompt for i in range(total_count)]
start_time = time.perf_counter()
@@ -315,9 +363,10 @@ def throughput_test(
for batch_start in range(0, total_count, bench_args.batch_size):
batch_end = min(batch_start + bench_args.batch_size, total_count)
batch_prompts = all_prompts[batch_start:batch_end]
batch_sampling_params = all_sampling_params[batch_start:batch_end]
batch_output = generate_batch(
engine, bench_args, batch_prompts, sampling_params
engine, bench_args, batch_prompts, batch_sampling_params
)
outputs.append(batch_output)
@@ -332,6 +381,7 @@ def throughput_test(
total_duration,
resolution=resolution,
num_requests=total_count,
all_sampling_params=all_sampling_params,
)
display_results(
@@ -191,6 +191,8 @@ async def async_request_image_sglang(
if input.width and input.height:
payload["size"] = f"{input.width}x{input.height}"
if input.num_inference_steps:
payload["num_inference_steps"] = input.num_inference_steps
# Merge extra parameters
payload.update(input.extra_body)
@@ -299,6 +301,8 @@ async def async_request_video_sglang(
payload["size"] = f"{input.width}x{input.height}"
if input.num_frames:
payload["num_frames"] = input.num_frames
if input.num_inference_steps:
payload["num_inference_steps"] = input.num_inference_steps
if input.fps:
payload["fps"] = input.fps
@@ -519,6 +523,11 @@ async def benchmark(args):
setattr(args, "task_name", task_name)
if args.random_request_config and args.dataset != "random":
raise ValueError(
"--random-request-config can only be used with --dataset random"
)
if args.dataset == "vbench":
dataset = VBenchDataset(args, api_url, args.model)
elif args.dataset == "random":
@@ -720,6 +729,25 @@ if __name__ == "__main__":
)
parser.add_argument("--width", type=int, default=None, help="Image/Video width.")
parser.add_argument("--height", type=int, default=None, help="Image/Video height.")
parser.add_argument(
"--random-request-config",
type=str,
default=None,
help=(
"JSON string defining random request profiles. "
"Each profile may contain: width, height, num_inference_steps, etc. "
"The 'weight' field controls sampling probability (relative weight). "
"Example: "
'[{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},'
'{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
),
)
parser.add_argument(
"--random-request-seed",
type=int,
default=42,
help="Random seed for sampling request profiles (default: 42).",
)
parser.add_argument(
"--num-frames", type=int, default=None, help="Number of frames (for video)."
)
@@ -1,6 +1,7 @@
import glob
import json
import os
import random
import re
import subprocess
import uuid
@@ -286,16 +287,39 @@ class RandomDataset(BaseDataset):
super().__init__(args, api_url, model)
self.num_prompts = args.num_prompts or 100
self.random_request_config = args.random_request_config
if self.random_request_config:
self.random_request_config = json.loads(self.random_request_config)
weights = [p.pop("weight") for p in self.random_request_config]
seed = args.random_request_seed
rng = random.Random(seed)
self._sampled_requests = rng.choices(
self.random_request_config, weights=weights, k=self.num_prompts
)
else:
self._sampled_requests = None
def get_sampling_params(self, idx: int) -> dict:
"""Return the per-request sampling profile dict, or empty dict if not mix-diffusion."""
if self._sampled_requests:
return self._sampled_requests[idx]
return {}
def __len__(self) -> int:
return self.num_prompts
def __getitem__(self, idx: int) -> RequestFuncInput:
profile = self._sampled_requests[idx] if self._sampled_requests else {}
return RequestFuncInput(
prompt=f"Random prompt {idx} for benchmarking diffusion models",
api_url=self.api_url,
model=self.model,
width=self.args.width,
height=self.args.height,
num_frames=self.args.num_frames,
fps=self.args.fps,
width=profile.get("width", self.args.width),
height=profile.get("height", self.args.height),
num_frames=profile.get("num_frames", self.args.num_frames),
num_inference_steps=profile.get(
"num_inference_steps", self.args.num_inference_steps
),
fps=profile.get("fps", self.args.fps),
)