413 lines
16 KiB
Python
Executable File
413 lines
16 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Generate a deterministic VBench-v1 suite through four MiniMax-H3 services.
|
|
|
|
The program deliberately separates generation from VBench scoring. It creates a
|
|
flat ``{prompt}-{sample_index}.mp4`` directory accepted by VBench standard mode,
|
|
records a resumable JSONL audit trail, and balances multi-label VBench dimensions
|
|
across replicas with a deterministic greedy assignment.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import shutil
|
|
import statistics
|
|
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import requests
|
|
|
|
|
|
ALL_DIMENSIONS = [
|
|
"subject_consistency",
|
|
"background_consistency",
|
|
"temporal_flickering",
|
|
"motion_smoothness",
|
|
"dynamic_degree",
|
|
"aesthetic_quality",
|
|
"imaging_quality",
|
|
"object_class",
|
|
"multiple_objects",
|
|
"human_action",
|
|
"color",
|
|
"spatial_relationship",
|
|
"scene",
|
|
"temporal_style",
|
|
"appearance_style",
|
|
"overall_consistency",
|
|
]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Job:
|
|
request_id: str
|
|
prompt: str
|
|
dimensions: tuple[str, ...]
|
|
sample_index: int
|
|
seed: int
|
|
replica_index: int
|
|
port: int
|
|
|
|
@property
|
|
def filename(self) -> str:
|
|
return f"{self.prompt}-{self.sample_index}.mp4"
|
|
|
|
|
|
def atomic_copy(source: Path, destination: Path) -> None:
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = destination.with_name(destination.name + ".partial")
|
|
shutil.copy2(source, temporary)
|
|
os.replace(temporary, destination)
|
|
|
|
|
|
def load_prompts(metadata_path: Path, dimensions: set[str]) -> list[dict[str, Any]]:
|
|
raw = json.loads(metadata_path.read_text(encoding="utf-8"))
|
|
merged: dict[str, set[str]] = {}
|
|
order: list[str] = []
|
|
for item in raw:
|
|
prompt = str(item["prompt_en"])
|
|
item_dimensions = set(item.get("dimension", [])) & dimensions
|
|
if not item_dimensions:
|
|
continue
|
|
if prompt not in merged:
|
|
merged[prompt] = set()
|
|
order.append(prompt)
|
|
merged[prompt].update(item_dimensions)
|
|
return [{"prompt": prompt, "dimensions": sorted(merged[prompt])} for prompt in order]
|
|
|
|
|
|
def coverage_subset(prompts: list[dict[str, Any]], max_prompts: int) -> list[dict[str, Any]]:
|
|
if max_prompts <= 0 or max_prompts >= len(prompts):
|
|
return prompts
|
|
remaining = list(prompts)
|
|
selected: list[dict[str, Any]] = []
|
|
counts = {dimension: 0 for dimension in ALL_DIMENSIONS}
|
|
while remaining and len(selected) < max_prompts:
|
|
# Prefer prompts covering the least represented dimensions; stable order
|
|
# is the final tie-breaker, so repeated runs produce the same suite.
|
|
best_index = min(
|
|
range(len(remaining)),
|
|
key=lambda index: (
|
|
sum(counts.get(dim, 0) for dim in remaining[index]["dimensions"]),
|
|
-len(remaining[index]["dimensions"]),
|
|
index,
|
|
),
|
|
)
|
|
item = remaining.pop(best_index)
|
|
selected.append(item)
|
|
for dimension in item["dimensions"]:
|
|
counts[dimension] += 1
|
|
return selected
|
|
|
|
|
|
def build_jobs(args: argparse.Namespace) -> list[Job]:
|
|
dimensions = set(args.dimensions.split(","))
|
|
unknown = dimensions - set(ALL_DIMENSIONS)
|
|
if unknown:
|
|
raise SystemExit(f"unknown VBench dimensions: {sorted(unknown)}")
|
|
prompts = coverage_subset(load_prompts(args.metadata, dimensions), args.max_prompts)
|
|
ports = [int(value) for value in args.ports.split(",") if value]
|
|
if not ports:
|
|
raise SystemExit("at least one port is required")
|
|
|
|
total_counts = [0] * len(ports)
|
|
dimension_counts = [{dimension: 0 for dimension in dimensions} for _ in ports]
|
|
jobs: list[Job] = []
|
|
for prompt_index, item in enumerate(prompts):
|
|
for sample_index in range(args.samples_per_prompt):
|
|
item_dimensions = tuple(item["dimensions"])
|
|
rotation = (prompt_index + sample_index) % len(ports)
|
|
candidates = list(range(len(ports)))
|
|
candidates.sort(
|
|
key=lambda replica: (
|
|
sum(dimension_counts[replica][dim] for dim in item_dimensions),
|
|
total_counts[replica],
|
|
(replica - rotation) % len(ports),
|
|
)
|
|
)
|
|
replica = candidates[0]
|
|
for dimension in item_dimensions:
|
|
dimension_counts[replica][dimension] += 1
|
|
total_counts[replica] += 1
|
|
jobs.append(
|
|
Job(
|
|
request_id=f"p{prompt_index:04d}-s{sample_index}",
|
|
prompt=item["prompt"],
|
|
dimensions=item_dimensions,
|
|
sample_index=sample_index,
|
|
seed=args.seed + prompt_index * args.samples_per_prompt + sample_index,
|
|
replica_index=replica,
|
|
port=ports[replica],
|
|
)
|
|
)
|
|
return jobs
|
|
|
|
|
|
def make_payload(args: argparse.Namespace, job: Job, steps: int) -> dict[str, Any]:
|
|
return {
|
|
"model": args.model,
|
|
"prompt": job.prompt,
|
|
"num_outputs_per_prompt": 1,
|
|
"num_inference_steps": steps,
|
|
"flow_shift": args.flow_shift,
|
|
"audio_flow_shift": args.audio_flow_shift,
|
|
"seed": job.seed,
|
|
"task": "t2va",
|
|
"conditions": [],
|
|
"target": {
|
|
"short_edge": args.short_edge,
|
|
"aspect_ratio": args.aspect_ratio,
|
|
"duration_seconds": args.duration_seconds,
|
|
},
|
|
}
|
|
|
|
|
|
def submit_and_wait(
|
|
session: requests.Session, args: argparse.Namespace, job: Job, steps: int
|
|
) -> tuple[dict[str, Any], dict[str, Any]]:
|
|
response = session.post(
|
|
f"http://{args.host}:{job.port}/v1/videos",
|
|
json=make_payload(args, job, steps),
|
|
timeout=args.submit_timeout,
|
|
)
|
|
if response.status_code != 200:
|
|
raise RuntimeError(f"submit HTTP {response.status_code}: {response.text[:1000]}")
|
|
status = response.json()
|
|
video_id = status.get("id")
|
|
if not video_id:
|
|
raise RuntimeError(f"submit response has no id: {status}")
|
|
deadline = time.monotonic() + args.request_timeout
|
|
while status.get("status") not in {"completed", "failed"}:
|
|
if time.monotonic() >= deadline:
|
|
raise TimeoutError(f"video job {video_id} exceeded {args.request_timeout}s")
|
|
time.sleep(args.poll_interval)
|
|
response = session.get(
|
|
f"http://{args.host}:{job.port}/v1/videos/{video_id}",
|
|
timeout=args.poll_timeout,
|
|
)
|
|
if response.status_code != 200:
|
|
raise RuntimeError(f"poll HTTP {response.status_code}: {response.text[:1000]}")
|
|
status = response.json()
|
|
if status.get("status") != "completed":
|
|
raise RuntimeError(f"job failed: {status.get('error') or status}")
|
|
return status, {"video_id": video_id}
|
|
|
|
|
|
def run_job(session: requests.Session, args: argparse.Namespace, job: Job) -> dict[str, Any]:
|
|
started_epoch = time.time()
|
|
started = time.monotonic()
|
|
result: dict[str, Any] = {
|
|
**job.__dict__,
|
|
"dimensions": list(job.dimensions),
|
|
"filename": job.filename,
|
|
"started_at_epoch": started_epoch,
|
|
"success": False,
|
|
"error": None,
|
|
}
|
|
try:
|
|
status, extra = submit_and_wait(session, args, job, args.num_inference_steps)
|
|
source_value = status.get("file_path")
|
|
if not source_value:
|
|
raise RuntimeError(f"completed response has no file_path: {status}")
|
|
source = Path(source_value)
|
|
if not source.is_file():
|
|
raise FileNotFoundError(f"server output does not exist: {source}")
|
|
destination = args.videos_dir / job.filename
|
|
atomic_copy(source, destination)
|
|
result.update(extra)
|
|
result.update(
|
|
{
|
|
"success": True,
|
|
"server_file_path": str(source),
|
|
"saved_file_path": str(destination),
|
|
"inference_time_s": status.get("inference_time_s"),
|
|
"peak_memory_mb": status.get("peak_memory_mb"),
|
|
"file_size_bytes": destination.stat().st_size,
|
|
}
|
|
)
|
|
except Exception as exc: # Preserve the rest of a long run and make it resumable.
|
|
result["error"] = f"{type(exc).__name__}: {exc}"
|
|
result["latency_s"] = time.monotonic() - started
|
|
result["finished_at_epoch"] = time.time()
|
|
return result
|
|
|
|
|
|
def completed_request_ids(results_path: Path, videos_dir: Path) -> set[str]:
|
|
completed: set[str] = set()
|
|
if not results_path.is_file():
|
|
return completed
|
|
for line in results_path.read_text(encoding="utf-8").splitlines():
|
|
try:
|
|
row = json.loads(line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
video = videos_dir / str(row.get("filename", ""))
|
|
if row.get("success") and video.is_file() and video.stat().st_size > 0:
|
|
completed.add(str(row["request_id"]))
|
|
return completed
|
|
|
|
|
|
def write_manifest(jobs: list[Job], path: Path) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with path.open("w", encoding="utf-8") as handle:
|
|
for job in jobs:
|
|
row = {**job.__dict__, "dimensions": list(job.dimensions), "filename": job.filename}
|
|
handle.write(json.dumps(row, ensure_ascii=False) + "\n")
|
|
|
|
|
|
def run_worker(
|
|
replica: int,
|
|
jobs: list[Job],
|
|
args: argparse.Namespace,
|
|
output_handle: Any,
|
|
output_lock: threading.Lock,
|
|
) -> int:
|
|
failures = 0
|
|
if not jobs:
|
|
return failures
|
|
with requests.Session() as session:
|
|
if args.warmup_requests:
|
|
warmup = jobs[0]
|
|
for index in range(args.warmup_requests):
|
|
started = time.monotonic()
|
|
try:
|
|
submit_and_wait(session, args, warmup, args.warmup_inference_steps)
|
|
print(
|
|
f"replica={replica} warmup={index + 1}/{args.warmup_requests} "
|
|
f"latency={time.monotonic() - started:.2f}s success=true",
|
|
flush=True,
|
|
)
|
|
except Exception as exc:
|
|
raise RuntimeError(f"replica {replica} warmup failed: {exc}") from exc
|
|
for index, job in enumerate(jobs, start=1):
|
|
result = run_job(session, args, job)
|
|
with output_lock:
|
|
output_handle.write(json.dumps(result, ensure_ascii=False) + "\n")
|
|
output_handle.flush()
|
|
failures += int(not result["success"])
|
|
print(
|
|
f"replica={replica} request={index}/{len(jobs)} id={job.request_id} "
|
|
f"success={str(result['success']).lower()} latency={result['latency_s']:.2f}s "
|
|
f"error={result['error']}",
|
|
flush=True,
|
|
)
|
|
return failures
|
|
|
|
|
|
def summarize(jobs: list[Job], results_path: Path, output_path: Path) -> dict[str, Any]:
|
|
latest: dict[str, dict[str, Any]] = {}
|
|
if results_path.is_file():
|
|
for line in results_path.read_text(encoding="utf-8").splitlines():
|
|
try:
|
|
row = json.loads(line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
latest[str(row.get("request_id"))] = row
|
|
rows = [latest[job.request_id] for job in jobs if job.request_id in latest]
|
|
success = [row for row in rows if row.get("success")]
|
|
latencies = [float(row["latency_s"]) for row in success]
|
|
by_replica: dict[str, Any] = {}
|
|
for replica in sorted({job.replica_index for job in jobs}):
|
|
replica_jobs = [job for job in jobs if job.replica_index == replica]
|
|
replica_rows = [row for row in rows if int(row["replica_index"]) == replica]
|
|
by_dimension = {
|
|
dimension: sum(dimension in job.dimensions for job in replica_jobs)
|
|
for dimension in ALL_DIMENSIONS
|
|
if any(dimension in job.dimensions for job in replica_jobs)
|
|
}
|
|
by_replica[str(replica)] = {
|
|
"expected": len(replica_jobs),
|
|
"recorded": len(replica_rows),
|
|
"completed": sum(bool(row.get("success")) for row in replica_rows),
|
|
"dimension_jobs": by_dimension,
|
|
}
|
|
summary = {
|
|
"expected": len(jobs),
|
|
"recorded_latest": len(rows),
|
|
"completed": len(success),
|
|
"failed_or_missing": len(jobs) - len(success),
|
|
"latency_mean_s": statistics.fmean(latencies) if latencies else None,
|
|
"latency_median_s": statistics.median(latencies) if latencies else None,
|
|
"by_replica": by_replica,
|
|
}
|
|
output_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
return summary
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--metadata", type=Path, required=True)
|
|
parser.add_argument("--videos-dir", type=Path, required=True)
|
|
parser.add_argument("--results", type=Path, required=True)
|
|
parser.add_argument("--manifest", type=Path, required=True)
|
|
parser.add_argument("--summary", type=Path, required=True)
|
|
parser.add_argument("--host", default="127.0.0.1")
|
|
parser.add_argument("--ports", default="30010,30020,30030,30040")
|
|
parser.add_argument("--model", default="/data/hf_models/MiniMax-H3")
|
|
parser.add_argument("--dimensions", default=",".join(ALL_DIMENSIONS))
|
|
parser.add_argument("--samples-per-prompt", type=int, default=1)
|
|
parser.add_argument("--max-prompts", type=int, default=0)
|
|
parser.add_argument("--num-inference-steps", type=int, default=20)
|
|
parser.add_argument("--short-edge", type=int, default=768)
|
|
parser.add_argument("--duration-seconds", type=float, default=5.0)
|
|
parser.add_argument("--aspect-ratio", default="16:9")
|
|
parser.add_argument("--flow-shift", type=float, default=12.0)
|
|
parser.add_argument("--audio-flow-shift", type=float, default=3.0)
|
|
parser.add_argument("--seed", type=int, default=1101)
|
|
parser.add_argument("--warmup-requests", type=int, default=1)
|
|
parser.add_argument("--warmup-inference-steps", type=int, default=5)
|
|
parser.add_argument("--submit-timeout", type=float, default=120.0)
|
|
parser.add_argument("--poll-timeout", type=float, default=30.0)
|
|
parser.add_argument("--poll-interval", type=float, default=1.0)
|
|
parser.add_argument("--request-timeout", type=float, default=3600.0)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
if not args.metadata.is_file():
|
|
raise SystemExit(f"VBench metadata missing: {args.metadata}")
|
|
if args.samples_per_prompt < 1 or args.samples_per_prompt > 5:
|
|
raise SystemExit("samples-per-prompt must be between 1 and 5")
|
|
args.videos_dir.mkdir(parents=True, exist_ok=True)
|
|
args.results.parent.mkdir(parents=True, exist_ok=True)
|
|
jobs = build_jobs(args)
|
|
write_manifest(jobs, args.manifest)
|
|
completed = completed_request_ids(args.results, args.videos_dir)
|
|
pending = [job for job in jobs if job.request_id not in completed]
|
|
print(
|
|
f"suite prompts={len({job.prompt for job in jobs})} jobs={len(jobs)} "
|
|
f"completed={len(completed)} pending={len(pending)} replicas={len(set(job.port for job in jobs))}",
|
|
flush=True,
|
|
)
|
|
grouped: dict[int, list[Job]] = {}
|
|
for job in pending:
|
|
grouped.setdefault(job.replica_index, []).append(job)
|
|
failures = 0
|
|
lock = threading.Lock()
|
|
with args.results.open("a", encoding="utf-8", buffering=1) as output:
|
|
with ThreadPoolExecutor(max_workers=max(1, len(grouped))) as executor:
|
|
futures = {
|
|
executor.submit(run_worker, replica, shard, args, output, lock): replica
|
|
for replica, shard in grouped.items()
|
|
}
|
|
for future in as_completed(futures):
|
|
try:
|
|
failures += future.result()
|
|
except Exception as exc:
|
|
print(f"replica={futures[future]} fatal={type(exc).__name__}: {exc}", flush=True)
|
|
failures += len(grouped[futures[future]])
|
|
summary = summarize(jobs, args.results, args.summary)
|
|
print(json.dumps(summary, ensure_ascii=False, indent=2), flush=True)
|
|
return int(failures > 0 or summary["failed_or_missing"] > 0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|