diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index 15188fde5..9e8c21b20 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -457,6 +457,51 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate + - label: "Diffusion · Helios Distilled Acc/Perf Compare" + commands: + - | + timeout 90m bash -c " + set -euxo pipefail + export VLLM_IMAGE_FETCH_TIMEOUT=60 + pytest -s -v tests/e2e/accuracy/test_helios_distilled.py::test_helios_distilled_pr_matches_main_acc_perf \ + -m 'full_model and benchmark' \ + --run-level 'full_model' + " + - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/helios_distilled_acc_perf/result.json" + - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/helios_distilled_acc_perf/artifacts/**/*.mp4" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate - label: "Diffusion · Qwen Image Test" commands: diff --git a/benchmarks/diffusion/backends.py b/benchmarks/diffusion/backends.py index b7e533b2d..cc87b1c9a 100644 --- a/benchmarks/diffusion/backends.py +++ b/benchmarks/diffusion/backends.py @@ -278,7 +278,7 @@ async def async_request_v1_videos( return output # invoke a poll request (GET /v1/videos/{video_id}) - poll_interval = 2.0 # Unit(s) + poll_interval = 0.2 # Unit(s) timeout_seconds = 600.0 deadline = time.perf_counter() + timeout_seconds job_url = f"{input.api_url}/{job_id}" diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py index c59844f65..ae7b62852 100644 --- a/benchmarks/diffusion/diffusion_benchmark_serving.py +++ b/benchmarks/diffusion/diffusion_benchmark_serving.py @@ -68,6 +68,7 @@ import glob import json import logging import os +import pathlib import random import tempfile import time @@ -598,6 +599,13 @@ class RandomDataset(BaseDataset): if self.enable_negative_prompt: extra_body["negative_prompt"] = f"Negative prompt {idx} for benchmarking diffusion models" + request_fields = { + "width", + "height", + "num_frames", + "num_inference_steps", + "fps", + } params = { "width": self.args.width, "height": self.args.height, @@ -607,9 +615,11 @@ class RandomDataset(BaseDataset): } if self._sampled_requests: profile = self._sampled_requests[idx] - params.update(profile) + params.update({key: value for key, value in profile.items() if key in request_fields}) + extra_body.update({key: value for key, value in profile.items() if key not in request_fields}) + prompt = self.args.random_prompt or f"Random prompt {idx} for benchmarking diffusion models" return RequestFuncInput( - prompt=f"Random prompt {idx} for benchmarking diffusion models", + prompt=prompt, api_url=self.api_url, model=self.model, seed=self.args.seed, @@ -869,6 +879,35 @@ def calculate_metrics( return metrics +def _response_suffix(endpoint: str) -> str: + if endpoint == "/v1/videos": + return ".mp4" + return ".json" + + +def _save_success_responses( + outputs: list[RequestFuncOutput], + requests_list: list[RequestFuncInput], + output_dir: str, + endpoint: str, +) -> list[str]: + path = pathlib.Path(output_dir).expanduser().resolve() + path.mkdir(parents=True, exist_ok=True) + saved_paths: list[str] = [] + suffix = _response_suffix(endpoint) + + for index, (request, output) in enumerate(zip(requests_list, outputs, strict=False)): + if not output.success: + continue + item_path = path / f"{index:05d}_{request.request_id}{suffix}" + if isinstance(output.response_body, bytes): + item_path.write_bytes(output.response_body) + else: + item_path.write_text(json.dumps(output.response_body, indent=2, sort_keys=True), encoding="utf-8") + saved_paths.append(str(item_path)) + return saved_paths + + def wait_for_service(base_url: str, timeout: int = 120) -> None: print(f"Waiting for service at {base_url}...") start_time = time.time() @@ -998,6 +1037,13 @@ async def benchmark(args): metrics["model"] = args.model metrics["dataset"] = args.dataset metrics["task"] = args.task + if args.save_response_dir: + metrics["saved_responses"] = _save_success_responses( + outputs, + requests_list, + args.save_response_dir, + args.endpoint, + ) print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="=")) @@ -1154,6 +1200,12 @@ if __name__ == "__main__": ) parser.add_argument("--fps", type=int, default=None, help="FPS (for video).") parser.add_argument("--output-file", type=str, default=None, help="Output JSON file for metrics.") + parser.add_argument( + "--save-response-dir", + type=str, + default=None, + help="Directory to save successful raw responses. Video responses are saved as MP4 files.", + ) parser.add_argument( "--slo", action="store_true", @@ -1189,6 +1241,12 @@ if __name__ == "__main__": '{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]' ), ) + parser.add_argument( + "--random-prompt", + type=str, + default=None, + help="Prompt to use for every synthetic random request. Defaults to generated placeholder prompts.", + ) parser.add_argument( "--num-input-images", type=int, diff --git a/benchmarks/diffusion/performance_dashboard/helios_distilled_serving_performance.md b/benchmarks/diffusion/performance_dashboard/helios_distilled_serving_performance.md new file mode 100644 index 000000000..bb6a1db49 --- /dev/null +++ b/benchmarks/diffusion/performance_dashboard/helios_distilled_serving_performance.md @@ -0,0 +1,222 @@ +# Helios-Distilled Serving Performance Dashboard + +This document describes how to deploy and benchmark **BestWishYsh/Helios-Distilled** using vLLM-Omni. It follows the same serving benchmark structure as the Qwen-Image and Wan2.2 dashboards. + +--- + +# 1. Overview + +Helios-Distilled is a text-to-video diffusion model. Helios uses an autoregressive chunking pattern that generates 33 frames per chunk. For stable comparisons, choose `num_frames` as a multiple of 33, such as 99. + +This document covers: + +* Service launch configuration +* Benchmark scripts and usage +* Dataset and workload settings +* Performance and accuracy comparison methodology +* Reproducibility guidelines + +--- + +# 2. Test Environment + +| Component | Specification | +|------------|----------------| +| GPU | NVIDIA H100 for CI accuracy/perf gate; local bring-up also used L20X | +| Diffusion Attention Backend | FlashAttention | +| Model | `BestWishYsh/Helios-Distilled` | + +--- + +# 3. Service Launch Configuration + +## 3.1 Basic Serving Command + +```bash +vllm serve BestWishYsh/Helios-Distilled --omni \ + --port 8091 +``` + +## 3.2 Key Parameters + +| Parameter | Description | +| --- | --- | +| `--port` | Serving port | +| `--max-num-seqs` | Maximum number of active sequences | +| `--enable-diffusion-pipeline-profiler` | Optional stage-level profiling | +| `--profiler-config` | Optional torch profiler configuration | + +Record these parameters when reporting performance results. + +--- + +# 4. Benchmark Scripts + +## 4.1 Online Serving Benchmark + +Use the shared diffusion serving benchmark, as with Qwen-Image and Wan2.2: + +```bash +python benchmarks/diffusion/diffusion_benchmark_serving.py \ + --base-url http://localhost:8091 \ + --endpoint /v1/videos \ + --model BestWishYsh/Helios-Distilled \ + --dataset random \ + --task t2v \ + --num-prompts 1 \ + --max-concurrency 1 \ + --disable-tqdm \ + --random-request-config '[ + { + "width": 640, + "height": 384, + "num_inference_steps": 50, + "num_frames": 99, + "fps": 16, + "guidance_scale": 1.0, + "is_enable_stage2": true, + "pyramid_num_stages": 3, + "pyramid_num_inference_steps_list": [1, 1, 1], + "is_amplify_first_chunk": false, + "weight": 1 + } + ]' +``` + +`--random-request-config` forwards standard fields such as `width`, `height`, `num_frames`, `fps`, and `num_inference_steps` as first-class request fields. Model-specific fields are forwarded in the request body, which is required for Helios stage-2 settings. + +## 4.2 Warmed Serving Benchmark + +For local bring-up, use the same shared serving benchmark with one warmup request and one measured request: + +```bash +python benchmarks/diffusion/diffusion_benchmark_serving.py \ + --base-url http://localhost:8091 \ + --endpoint /v1/videos \ + --model BestWishYsh/Helios-Distilled \ + --dataset random \ + --task t2v \ + --num-prompts 1 \ + --max-concurrency 1 \ + --warmup-requests 1 \ + --warmup-num-inference-steps 1 \ + --random-prompt "A cat wearing sunglasses dances on a beach at sunset, cinematic lighting." \ + --seed 42 \ + --disable-tqdm \ + --output-file metrics.json \ + --save-response-dir responses \ + --random-request-config '[ + { + "width": 640, + "height": 384, + "num_inference_steps": 50, + "num_frames": 99, + "fps": 16, + "guidance_scale": 1.0, + "is_enable_stage2": true, + "pyramid_num_stages": 3, + "pyramid_num_inference_steps_list": [1, 1, 1], + "is_amplify_first_chunk": false, + "weight": 1 + } + ]' +``` + +## 4.3 PR-vs-Main Accuracy and Performance Comparison + +The PR-vs-main accuracy/performance gate lives in the Helios accuracy pytest. It checks out `origin/main`, starts a server for each checkout, and drives both with `diffusion_benchmark_serving.py`: + +```bash +pytest -s -v tests/e2e/accuracy/test_helios_distilled.py::test_helios_distilled_pr_matches_main_acc_perf \ + -m 'full_model and benchmark' \ + --run-level 'full_model' +``` + +--- + +# 5. Dataset & Workload Settings + +## 5.1 Recommended Evaluation Configurations + +### Dataset A (99 frames, 384x640) + +* Dataset: `random` +* Task: t2v +* Concurrency: 1 +* Frame count: 99 (`33 x 3`) +* FPS: 16 +* Stage-2 steps: `[1, 1, 1]` + +```json +[ + { + "width": 640, + "height": 384, + "num_inference_steps": 50, + "num_frames": 99, + "fps": 16, + "guidance_scale": 1.0, + "is_enable_stage2": true, + "pyramid_num_stages": 3, + "pyramid_num_inference_steps_list": [1, 1, 1], + "is_amplify_first_chunk": false, + "weight": 1 + } +] +``` + +--- + +# 6. Performance Metrics + +| Metric | Description | Unit | +| --- | --- | --- | +| Mean Latency | Mean request latency | seconds | +| P99 Latency | P99 request latency | seconds | +| Generated FPS | `num_frames / measured_time_s` | frames/second | +| Candidate/reference latency ratio | PR latency divided by reference latency | ratio | + +--- + +# 7. Accuracy Metrics + +The PR-vs-main gate compares encoded video frames extracted from MP4 outputs: + +| Metric | Threshold | Direction | +| --- | --- | --- | +| PSNR | `>= 35 dB` | Higher is better | +| MAE | `<= 3.0 / 255` | Lower is better | +| Candidate/reference latency ratio | `<= 1.03` | Lower is better | + +Local L20X bring-up for this PR produced: + +| Variant | Measured Latency (ms) | Generated FPS | +| --- | --- | --- | +| `origin/main` | 5317.95 | 18.62 | +| Candidate PR | 5283.27 | 18.74 | + +Encoded-video similarity: + +| Metric | Value | +| --- | --- | +| PSNR | 38.66 dB | +| MAE | 1.89 / 255 | +| Cosine similarity | 0.99978 | + +--- + +# 8. Reproducibility Checklist + +To ensure consistent and comparable benchmark results: + +* Record GPU type +* Record vLLM and vLLM-Omni versions +* Record benchmark parameters: resolution, frame count, FPS, seed, concurrency, and stage-2 steps +* Use frame counts that are multiples of 33 +* Keep prompt, seed, resolution, and sampling parameters identical when comparing PR vs main +* Ensure no background workload on GPUs during testing +* Use the same MP4 export and frame extraction path for both variants + +--- + +This document serves as the Helios-Distilled serving performance reference under vLLM-Omni. diff --git a/tests/e2e/accuracy/test_helios_distilled.py b/tests/e2e/accuracy/test_helios_distilled.py new file mode 100644 index 000000000..23d835f23 --- /dev/null +++ b/tests/e2e/accuracy/test_helios_distilled.py @@ -0,0 +1,368 @@ +from __future__ import annotations + +import json +import math +import os +import shutil +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from tempfile import TemporaryDirectory +from typing import Any + +import numpy as np +import pytest +import torch +from PIL import Image + +from tests.e2e.accuracy.helpers import reset_artifact_dir +from tests.helpers.mark import hardware_test + +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + +REPO_ROOT = Path(__file__).resolve().parents[3] +MODEL_NAME = "BestWishYsh/Helios-Distilled" +PROMPT = "A cat wearing sunglasses dances on a beach at sunset, cinematic lighting." +HEIGHT = 384 +WIDTH = 640 +NUM_FRAMES = 99 +FPS = 16 +SEED = 42 +GUIDANCE_SCALE = 1.0 +PYRAMID_NUM_INFERENCE_STEPS = (1, 1, 1) +MIN_PSNR_DB = 35.0 +MAX_MAE = 3.0 +MAX_TIME_RATIO = 1.03 +WARMUP_REQUESTS = 1 +WARMUP_NUM_INFERENCE_STEPS = 1 + + +def _run(command: list[str], *, cwd: Path, env: dict[str, str] | None = None) -> subprocess.CompletedProcess[str]: + print("+ " + " ".join(command), flush=True) + return subprocess.run( + command, + cwd=str(cwd), + env=env, + check=True, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + + +def _read_json(path: Path) -> dict[str, Any]: + return json.loads(path.read_text(encoding="utf-8")) + + +def _open_port(host: str = "127.0.0.1") -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind((host, 0)) + return int(sock.getsockname()[1]) + + +def _wait_for_service(base_url: str, process: subprocess.Popen, *, timeout_s: float = 1200.0) -> None: + import requests + + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if process.poll() is not None: + raise RuntimeError(f"Server exited with code {process.returncode} before becoming ready.") + try: + response = requests.get(f"{base_url}/health", timeout=1) + if response.status_code == 200: + return + except requests.RequestException: + pass + time.sleep(2) + raise TimeoutError(f"Server at {base_url} did not become ready within {timeout_s:.0f}s.") + + +def _stop_process(process: subprocess.Popen) -> None: + if process.poll() is not None: + return + try: + os.killpg(process.pid, signal.SIGTERM) + process.wait(timeout=30) + except Exception: + os.killpg(process.pid, signal.SIGKILL) + process.wait(timeout=30) + + +def _extract_frames(video_path: Path, output_dir: Path) -> list[Image.Image]: + output_dir.mkdir(parents=True, exist_ok=True) + frame_pattern = output_dir / "frame_%05d.png" + _run( + [ + "ffmpeg", + "-hide_banner", + "-loglevel", + "error", + "-i", + str(video_path), + "-vsync", + "0", + str(frame_pattern), + ], + cwd=output_dir, + ) + frame_paths = sorted(output_dir.glob("frame_*.png")) + if not frame_paths: + raise RuntimeError(f"No frames extracted from {video_path}") + return [Image.open(path).convert("RGB").copy() for path in frame_paths] + + +def _image_metrics(reference: np.ndarray, candidate: np.ndarray) -> dict[str, float]: + diff = reference.astype(np.float64) - candidate.astype(np.float64) + mse = float(np.mean(diff * diff)) + mae = float(np.mean(np.abs(diff))) + max_abs = float(np.max(np.abs(diff))) + psnr_db = float("inf") if mse == 0.0 else 20.0 * math.log10(255.0) - 10.0 * math.log10(mse) + ref_flat = reference.astype(np.float64).reshape(-1) + cand_flat = candidate.astype(np.float64).reshape(-1) + ref_norm = float(np.linalg.norm(ref_flat)) + cand_norm = float(np.linalg.norm(cand_flat)) + if ref_norm == 0.0 and cand_norm == 0.0: + cosine = 1.0 + else: + cosine = float(np.dot(ref_flat, cand_flat) / (ref_norm * cand_norm)) + return { + "mae": mae, + "mse": mse, + "max_abs": max_abs, + "psnr_db": psnr_db, + "cosine_similarity": cosine, + } + + +def _video_metrics(reference_frames: list[Image.Image], candidate_frames: list[Image.Image]) -> dict[str, Any]: + if len(reference_frames) != len(candidate_frames): + raise ValueError(f"Frame count mismatch: reference={len(reference_frames)} candidate={len(candidate_frames)}") + + ref_stack = np.stack([np.asarray(frame, dtype=np.uint8) for frame in reference_frames], axis=0) + cand_stack = np.stack([np.asarray(frame, dtype=np.uint8) for frame in candidate_frames], axis=0) + mid = len(reference_frames) // 2 + return { + "num_frames": len(reference_frames), + "frame0_metrics": _image_metrics(ref_stack[0], cand_stack[0]), + "mid_frame_index": mid, + "mid_frame_metrics": _image_metrics(ref_stack[mid], cand_stack[mid]), + "all_frames_metrics": _image_metrics(ref_stack, cand_stack), + } + + +def _reference_repo(output_root: Path) -> Path: + configured = os.environ.get("VLLM_HELIOS_REFERENCE_REPO") + if configured: + reference = Path(configured).expanduser().resolve() + if not reference.exists(): + raise FileNotFoundError(f"VLLM_HELIOS_REFERENCE_REPO does not exist: {reference}") + return reference + + reference = output_root / "origin_main" + if reference.exists(): + subprocess.run(["git", "worktree", "remove", "--force", str(reference)], cwd=str(REPO_ROOT), check=False) + if reference.exists(): + shutil.rmtree(reference) + _run(["git", "fetch", "origin", "main"], cwd=REPO_ROOT) + _run(["git", "worktree", "add", "--detach", str(reference), "origin/main"], cwd=REPO_ROOT) + return reference + + +def _server_env(repo: Path) -> dict[str, str]: + env = dict(os.environ) + env["PYTHONPATH"] = str(repo) + (os.pathsep + env["PYTHONPATH"] if env.get("PYTHONPATH") else "") + env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + return env + + +def _benchmark_config() -> str: + return json.dumps( + [ + { + "width": WIDTH, + "height": HEIGHT, + "num_inference_steps": 50, + "num_frames": NUM_FRAMES, + "fps": FPS, + "guidance_scale": GUIDANCE_SCALE, + "is_enable_stage2": True, + "pyramid_num_stages": 3, + "pyramid_num_inference_steps_list": list(PYRAMID_NUM_INFERENCE_STEPS), + "is_amplify_first_chunk": False, + "weight": 1, + } + ] + ) + + +def _start_server(*, repo: Path, variant_dir: Path) -> tuple[subprocess.Popen, str]: + port = _open_port() + base_url = f"http://127.0.0.1:{port}" + log_file = (variant_dir / "server.log").open("w", encoding="utf-8") + command = [ + sys.executable, + "-m", + "vllm_omni.entrypoints.cli.main", + "serve", + MODEL_NAME, + "--host", + "127.0.0.1", + "--port", + str(port), + "--omni", + ] + print("+ " + " ".join(command), flush=True) + process = subprocess.Popen( + command, + cwd=str(repo), + env=_server_env(repo), + stdout=log_file, + stderr=subprocess.STDOUT, + preexec_fn=os.setsid, + text=True, + ) + try: + _wait_for_service(base_url, process) + except Exception: + log_file.close() + _stop_process(process) + raise + log_file.close() + return process, base_url + + +def _run_serving_benchmark(*, repo: Path, variant_dir: Path, base_url: str) -> dict[str, Any]: + metrics_path = variant_dir / "metrics.json" + response_dir = variant_dir / "responses" + command = [ + sys.executable, + str(REPO_ROOT / "benchmarks/diffusion/diffusion_benchmark_serving.py"), + "--base-url", + base_url, + "--endpoint", + "/v1/videos", + "--model", + MODEL_NAME, + "--dataset", + "random", + "--task", + "t2v", + "--num-prompts", + "1", + "--max-concurrency", + "1", + "--warmup-requests", + str(WARMUP_REQUESTS), + "--warmup-num-inference-steps", + str(WARMUP_NUM_INFERENCE_STEPS), + "--warmup-concurrency", + "1", + "--random-prompt", + PROMPT, + "--seed", + str(SEED), + "--disable-tqdm", + "--random-request-config", + _benchmark_config(), + "--output-file", + str(metrics_path), + "--save-response-dir", + str(response_dir), + ] + completed = _run(command, cwd=repo, env=_server_env(repo)) + (variant_dir / "benchmark.log").write_text(completed.stdout, encoding="utf-8") + metrics = _read_json(metrics_path) + saved_responses = metrics.get("saved_responses") or [] + if len(saved_responses) != 1: + raise RuntimeError(f"Expected one saved video response, got {saved_responses}") + return { + "repo": str(repo), + "metrics": metrics, + "video": saved_responses[0], + "log": str(variant_dir / "benchmark.log"), + "server_log": str(variant_dir / "server.log"), + } + + +def _run_variant(*, label: str, repo: Path, output_root: Path) -> dict[str, Any]: + variant_dir = output_root / label + variant_dir.mkdir(parents=True, exist_ok=True) + process, base_url = _start_server(repo=repo, variant_dir=variant_dir) + try: + return _run_serving_benchmark(repo=repo, variant_dir=variant_dir, base_url=base_url) + finally: + _stop_process(process) + + +@pytest.mark.benchmark +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_helios_distilled_pr_matches_main_acc_perf(accuracy_artifact_root: Path) -> None: + if not torch.cuda.is_available(): + pytest.skip("Helios-Distilled accuracy/perf comparison requires CUDA.") + + output_root = reset_artifact_dir(accuracy_artifact_root / "helios_distilled_acc_perf") + reference = _reference_repo(output_root) + output_json = output_root / "result.json" + + reference_result = _run_variant(label="reference", repo=reference, output_root=output_root) + candidate_result = _run_variant(label="candidate", repo=REPO_ROOT, output_root=output_root) + + with TemporaryDirectory() as tmp: + tmp_dir = Path(tmp) + reference_frames = _extract_frames(Path(reference_result["video"]), tmp_dir / "reference_frames") + candidate_frames = _extract_frames(Path(candidate_result["video"]), tmp_dir / "candidate_frames") + output_metrics = _video_metrics(reference_frames, candidate_frames) + + reference_time = float(reference_result["metrics"]["latency_mean"]) + candidate_time = float(candidate_result["metrics"]["latency_mean"]) + time_ratio = candidate_time / reference_time + result = { + "reference": reference_result, + "candidate": candidate_result, + "sampling": { + "model": MODEL_NAME, + "prompt": PROMPT, + "height": HEIGHT, + "width": WIDTH, + "num_frames": NUM_FRAMES, + "fps": FPS, + "seed": SEED, + "guidance_scale": GUIDANCE_SCALE, + "pyramid_num_inference_steps_list": list(PYRAMID_NUM_INFERENCE_STEPS), + "is_amplify_first_chunk": False, + }, + "performance": { + "reference_measured_time_s": reference_time, + "candidate_measured_time_s": candidate_time, + "candidate_over_reference_time_ratio": time_ratio, + }, + "output_metrics": output_metrics, + "thresholds": { + "min_psnr_db": MIN_PSNR_DB, + "max_mae": MAX_MAE, + "max_time_ratio": MAX_TIME_RATIO, + }, + } + output_json.write_text(json.dumps(result, indent=2, sort_keys=True, allow_nan=True), encoding="utf-8") + all_metrics = output_metrics["all_frames_metrics"] + print( + json.dumps( + { + "output_json": str(output_json), + "candidate_over_reference_time_ratio": time_ratio, + "all_frames_psnr_db": all_metrics["psnr_db"], + "all_frames_mae": all_metrics["mae"], + "all_frames_cosine_similarity": all_metrics["cosine_similarity"], + }, + indent=2, + allow_nan=True, + ) + ) + + assert all_metrics["psnr_db"] >= MIN_PSNR_DB + assert all_metrics["mae"] <= MAX_MAE + assert time_ratio <= MAX_TIME_RATIO + assert output_json.exists(), f"Expected Helios comparison result at {output_json}" diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py index c3a47ebaa..2d475861d 100644 --- a/vllm_omni/diffusion/models/helios/pipeline_helios.py +++ b/vllm_omni/diffusion/models/helios/pipeline_helios.py @@ -29,7 +29,6 @@ from vllm_omni.diffusion.models.helios.scheduling_helios import HeliosScheduler from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( @@ -504,7 +503,12 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip # Prepare frame indices if keep_first_frame: - indices = torch.arange(0, sum([1, *history_sizes, num_latent_frames_per_chunk])) + indices = torch.arange( + 0, + sum([1, *history_sizes, num_latent_frames_per_chunk]), + device=device, + dtype=torch.float32, + ) ( indices_prefix, indices_latents_history_long, @@ -514,7 +518,12 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip ) = indices.split([1, *history_sizes, num_latent_frames_per_chunk], dim=0) indices_latents_history_short = torch.cat([indices_prefix, indices_latents_history_1x], dim=0) else: - indices = torch.arange(0, sum([*history_sizes, num_latent_frames_per_chunk])) + indices = torch.arange( + 0, + sum([*history_sizes, num_latent_frames_per_chunk]), + device=device, + dtype=torch.float32, + ) ( indices_latents_history_long, indices_latents_history_mid, @@ -531,6 +540,7 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip attention_kwargs = {} vae_dtype = self.vae.dtype + latent_chunks_to_decode = [] # Chunked denoising loop for k in range(num_latent_chunk): @@ -645,16 +655,15 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip slice(-num_latent_frames_per_chunk, None), ) - current_latents = real_history_latents[index_slice].to(vae_dtype) / latents_std + latents_mean - current_video = self.vae.decode(current_latents, return_dict=False)[0] + if output_type != "latent": + current_latents = real_history_latents[index_slice].to(vae_dtype) / latents_std + latents_mean + latent_chunks_to_decode.append(current_latents) - if history_video is None: - history_video = current_video - else: - history_video = torch.cat([history_video, current_video], dim=2) + if output_type != "latent" and latent_chunks_to_decode: + batched_latents = torch.cat(latent_chunks_to_decode, dim=0) + batched_video = self.vae.decode(batched_latents, return_dict=False)[0] + history_video = torch.cat(list(batched_video.split(batch_size, dim=0)), dim=2) - if current_omni_platform.is_available(): - current_omni_platform.empty_cache() self._current_timestep = None if output_type == "latent": @@ -688,6 +697,9 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip ) -> torch.Tensor: """Single-stage denoising loop for one chunk.""" batch_size = latents.shape[0] + latents_history_short = latents_history_short.to(transformer_dtype) + latents_history_mid = latents_history_mid.to(transformer_dtype) + latents_history_long = latents_history_long.to(transformer_dtype) do_true_cfg = guidance_scale > 1.0 and negative_prompt_embeds is not None with self.progress_bar(total=len(timesteps)) as pbar: @@ -702,9 +714,9 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip "indices_latents_history_short": indices_latents_history_short, "indices_latents_history_mid": indices_latents_history_mid, "indices_latents_history_long": indices_latents_history_long, - "latents_history_short": latents_history_short.to(transformer_dtype), - "latents_history_mid": latents_history_mid.to(transformer_dtype), - "latents_history_long": latents_history_long.to(transformer_dtype), + "latents_history_short": latents_history_short, + "latents_history_mid": latents_history_mid, + "latents_history_long": latents_history_long, "attention_kwargs": attention_kwargs, "return_dict": False, } @@ -783,6 +795,9 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip ) -> torch.Tensor: """Pyramid multi-stage denoising for one chunk.""" batch_size, num_channel, num_frames_lat, height, width = latents.shape + latents_history_short = latents_history_short.to(transformer_dtype) + latents_history_mid = latents_history_mid.to(transformer_dtype) + latents_history_long = latents_history_long.to(transformer_dtype) # Downsample latents to the smallest pyramid level latents_flat = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames_lat, num_channel, height, width) @@ -853,9 +868,9 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip "indices_latents_history_short": indices_latents_history_short, "indices_latents_history_mid": indices_latents_history_mid, "indices_latents_history_long": indices_latents_history_long, - "latents_history_short": latents_history_short.to(transformer_dtype), - "latents_history_mid": latents_history_mid.to(transformer_dtype), - "latents_history_long": latents_history_long.to(transformer_dtype), + "latents_history_short": latents_history_short, + "latents_history_mid": latents_history_mid, + "latents_history_long": latents_history_long, "attention_kwargs": attention_kwargs, "return_dict": False, }