414 lines
15 KiB
Python
414 lines
15 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""
|
|
SSIM/PSNR accuracy tests for LTX-2.3.
|
|
|
|
1. **Transformer parity** (``test_ltx2_3_transformer_matches_diffusers``):
|
|
Swaps our custom transformer into diffusers' ``LTX2Pipeline`` to measure
|
|
numerical parity in isolation. Thresholds: SSIM >= 0.95, PSNR >= 28 dB.
|
|
Result: SSIM 0.999987 (bit-identical).
|
|
|
|
2. **Full pipeline** (``test_ltx2_3_pipeline_matches_diffusers``):
|
|
Runs the full vLLM-Omni serving stack (``OmniServer`` -> HTTP API) and
|
|
compares per-frame against stock diffusers. Currently skipped because
|
|
the OmniServer subprocess creates a different RNG state than in-process
|
|
diffusers, producing different initial latents from the same seed.
|
|
This is a test infrastructure limitation, not a model accuracy issue.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import gc
|
|
import os
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import diffusers
|
|
import numpy as np
|
|
import pytest
|
|
import requests
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from tests.e2e.accuracy.helpers import compute_image_ssim_psnr, model_output_dir
|
|
from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
|
|
from tests.helpers.mark import hardware_test
|
|
from tests.helpers.runtime import OmniServer
|
|
|
|
# Parse diffusers version for compatibility check
|
|
_DIFFUSERS_VERSION = tuple(int(x) for x in diffusers.__version__.split(".")[:2] if x.isdigit())
|
|
_DIFFUSERS_038 = _DIFFUSERS_VERSION >= (0, 38)
|
|
|
|
MODEL_ID = "diffusers/LTX-2.3-Diffusers"
|
|
MODEL_ENV_VAR = "VLLM_TEST_LTX23_MODEL"
|
|
PROMPT = "A lighthouse on a rocky cliff at sunset, waves crashing below, golden hour lighting"
|
|
NEGATIVE_PROMPT = "blurry, low quality, distorted, watermark"
|
|
WIDTH = 512
|
|
HEIGHT = 384
|
|
NUM_FRAMES = 25 # ~1 second at 24fps
|
|
NUM_INFERENCE_STEPS = 20
|
|
GUIDANCE_SCALE = 4.0
|
|
SEED = 42
|
|
|
|
# Transformer-swap test: near-identical output expected
|
|
TRANSFORMER_SSIM_THRESHOLD = 0.95
|
|
TRANSFORMER_PSNR_THRESHOLD = 28.0
|
|
|
|
# Full-pipeline test: allows minor divergence from RNG / pipeline differences
|
|
PIPELINE_SSIM_THRESHOLD = 0.94
|
|
PIPELINE_PSNR_THRESHOLD = 28.0
|
|
|
|
|
|
def _model_name() -> str:
|
|
return os.environ.get(MODEL_ENV_VAR, MODEL_ID)
|
|
|
|
|
|
def _local_files_only(model: str) -> bool:
|
|
return Path(model).exists()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Frame extraction helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _video_to_frames(video_np: np.ndarray) -> list[Image.Image]:
|
|
"""Convert numpy video to list of PIL Images."""
|
|
while video_np.ndim > 4:
|
|
video_np = video_np[0]
|
|
if video_np.dtype in (np.float32, np.float64, np.float16):
|
|
video_np = np.clip(video_np * 255, 0, 255).astype(np.uint8)
|
|
return [Image.fromarray(video_np[t]) for t in range(video_np.shape[0])]
|
|
|
|
|
|
def _extract_diffusers_frames(result) -> list[Image.Image]:
|
|
"""Extract frames from diffusers pipeline output."""
|
|
video = result.frames
|
|
if isinstance(video, np.ndarray):
|
|
return _video_to_frames(video)
|
|
if isinstance(video, list):
|
|
if isinstance(video[0], list):
|
|
return [img.convert("RGB") for img in video[0]]
|
|
if isinstance(video[0], Image.Image):
|
|
return [img.convert("RGB") for img in video]
|
|
raise ValueError(f"Unexpected output type: {type(video)}")
|
|
|
|
|
|
def _extract_mp4_frames(mp4_bytes: bytes) -> list[Image.Image]:
|
|
"""Extract frames from an MP4 video using ffmpeg."""
|
|
import subprocess
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
mp4_path = os.path.join(tmpdir, "video.mp4")
|
|
with open(mp4_path, "wb") as f:
|
|
f.write(mp4_bytes)
|
|
|
|
# Extract video frames as PNG files using ffmpeg
|
|
frame_pattern = os.path.join(tmpdir, "frame_%04d.png")
|
|
subprocess.run(
|
|
["ffmpeg", "-i", mp4_path, "-vsync", "0", frame_pattern],
|
|
capture_output=True,
|
|
check=True,
|
|
)
|
|
|
|
# Load frames in order
|
|
frames = []
|
|
i = 1
|
|
while True:
|
|
fpath = os.path.join(tmpdir, f"frame_{i:04d}.png")
|
|
if not os.path.exists(fpath):
|
|
break
|
|
frames.append(Image.open(fpath).convert("RGB").copy())
|
|
i += 1
|
|
return frames
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Comparison helper
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _assert_video_similarity(
|
|
*,
|
|
model_name: str,
|
|
vllm_frames: list[Image.Image],
|
|
diffusers_frames: list[Image.Image],
|
|
ssim_threshold: float,
|
|
psnr_threshold: float,
|
|
) -> tuple[float, float]:
|
|
"""Compare video frames and assert SSIM/PSNR meet thresholds."""
|
|
min_frames = min(len(vllm_frames), len(diffusers_frames))
|
|
assert min_frames > 0, "No frames to compare"
|
|
|
|
ssim_scores = []
|
|
psnr_scores = []
|
|
for i in range(min_frames):
|
|
ssim_val, psnr_val = compute_image_ssim_psnr(
|
|
prediction=vllm_frames[i],
|
|
reference=diffusers_frames[i],
|
|
)
|
|
ssim_scores.append(ssim_val)
|
|
psnr_scores.append(psnr_val)
|
|
|
|
avg_ssim = sum(ssim_scores) / len(ssim_scores)
|
|
avg_psnr = sum(psnr_scores) / len(psnr_scores)
|
|
|
|
print(f"\n{model_name} video similarity ({min_frames} frames):")
|
|
print(f" SSIM: avg={avg_ssim:.6f}, min={min(ssim_scores):.6f}, threshold>={ssim_threshold:.6f}")
|
|
print(f" PSNR: avg={avg_psnr:.6f} dB, min={min(psnr_scores):.6f} dB, threshold>={psnr_threshold:.6f} dB")
|
|
|
|
assert avg_ssim >= ssim_threshold, f"SSIM below threshold: got {avg_ssim:.6f}, expected >= {ssim_threshold:.6f}."
|
|
assert avg_psnr >= psnr_threshold, f"PSNR below threshold: got {avg_psnr:.6f}, expected >= {psnr_threshold:.6f}."
|
|
return avg_ssim, avg_psnr
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Diffusers baseline (shared by both tests)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _run_diffusers_baseline(model: str, output_dir: Path) -> list[Image.Image]:
|
|
"""Generate video using stock diffusers LTX2Pipeline."""
|
|
from diffusers import LTX2Pipeline
|
|
|
|
run_pre_test_cleanup()
|
|
pipe = None
|
|
try:
|
|
pipe = LTX2Pipeline.from_pretrained(
|
|
model, torch_dtype=torch.bfloat16, local_files_only=_local_files_only(model)
|
|
).to("cuda")
|
|
|
|
generator = torch.Generator(device="cuda").manual_seed(SEED)
|
|
result = pipe(
|
|
prompt=PROMPT,
|
|
negative_prompt=NEGATIVE_PROMPT,
|
|
width=WIDTH,
|
|
height=HEIGHT,
|
|
num_frames=NUM_FRAMES,
|
|
num_inference_steps=NUM_INFERENCE_STEPS,
|
|
guidance_scale=GUIDANCE_SCALE,
|
|
generator=generator,
|
|
output_type="np",
|
|
)
|
|
frames = _extract_diffusers_frames(result)
|
|
for i, f in enumerate(frames):
|
|
f.save(output_dir / f"diffusers_frame_{i:04d}.png")
|
|
return frames
|
|
finally:
|
|
del pipe
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.accelerator.empty_cache()
|
|
run_post_test_cleanup()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 1: Transformer-swap parity
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _run_with_custom_transformer(model: str, output_dir: Path) -> list[Image.Image]:
|
|
"""Run diffusers pipeline with our custom transformer swapped in."""
|
|
from contextlib import nullcontext
|
|
|
|
from diffusers import LTX2Pipeline
|
|
from vllm.config import VllmConfig, set_current_vllm_config
|
|
from vllm.distributed.parallel_state import init_distributed_environment, initialize_model_parallel
|
|
|
|
from vllm_omni.diffusion.models.ltx2 import (
|
|
create_transformer_from_config,
|
|
load_transformer_config,
|
|
)
|
|
|
|
vllm_config = VllmConfig()
|
|
ctx = set_current_vllm_config(vllm_config)
|
|
ctx.__enter__()
|
|
|
|
if not torch.distributed.is_initialized():
|
|
os.environ.setdefault("MASTER_ADDR", "localhost")
|
|
os.environ.setdefault("MASTER_PORT", "29503")
|
|
os.environ.setdefault("RANK", "0")
|
|
os.environ.setdefault("WORLD_SIZE", "1")
|
|
init_distributed_environment(world_size=1, rank=0, local_rank=0)
|
|
initialize_model_parallel(tensor_model_parallel_size=1)
|
|
|
|
local = _local_files_only(model)
|
|
pipe = LTX2Pipeline.from_pretrained(model, torch_dtype=torch.bfloat16, local_files_only=local)
|
|
|
|
transformer_config = load_transformer_config(model, "transformer", local)
|
|
our_transformer = create_transformer_from_config(transformer_config)
|
|
|
|
diffusers_state = dict(pipe.transformer.named_parameters())
|
|
|
|
def _weight_iter():
|
|
for name, param in diffusers_state.items():
|
|
yield name, param.data
|
|
|
|
our_transformer.load_weights(_weight_iter())
|
|
our_transformer = our_transformer.to(dtype=torch.bfloat16, device="cuda").eval()
|
|
|
|
# Compatibility shims for diffusers pipeline
|
|
our_transformer.dtype = torch.bfloat16
|
|
if not hasattr(our_transformer, "cache_context"):
|
|
our_transformer.cache_context = lambda name: nullcontext()
|
|
|
|
del pipe.transformer
|
|
pipe.transformer = our_transformer
|
|
for name, component in pipe.components.items():
|
|
if name != "transformer" and hasattr(component, "to"):
|
|
try:
|
|
component.to("cuda")
|
|
except Exception:
|
|
pass
|
|
|
|
generator = torch.Generator(device="cuda").manual_seed(SEED)
|
|
result = pipe(
|
|
prompt=PROMPT,
|
|
negative_prompt=NEGATIVE_PROMPT,
|
|
width=WIDTH,
|
|
height=HEIGHT,
|
|
num_frames=NUM_FRAMES,
|
|
num_inference_steps=NUM_INFERENCE_STEPS,
|
|
guidance_scale=GUIDANCE_SCALE,
|
|
generator=generator,
|
|
output_type="np",
|
|
)
|
|
frames = _extract_diffusers_frames(result)
|
|
for i, f in enumerate(frames):
|
|
f.save(output_dir / f"vllm_transformer_frame_{i:04d}.png")
|
|
|
|
del pipe, result, our_transformer
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.accelerator.empty_cache()
|
|
return frames
|
|
|
|
|
|
@pytest.mark.full_model
|
|
@pytest.mark.benchmark
|
|
@pytest.mark.diffusion
|
|
@pytest.mark.skipif(
|
|
not _DIFFUSERS_038, reason="LTX-2.3 requires diffusers >= 0.38.0 for cross_attn_mod and BWE vocoder"
|
|
)
|
|
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
|
def test_ltx2_3_transformer_matches_diffusers(accuracy_artifact_root: Path) -> None:
|
|
"""Transformer-level parity: swap our transformer into diffusers pipeline.
|
|
|
|
Isolates transformer numerical accuracy from pipeline-level differences.
|
|
Both runs use diffusers' denoising loop, CFG, scheduler, and RNG.
|
|
"""
|
|
model = _model_name()
|
|
output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
|
|
|
|
diffusers_frames = _run_diffusers_baseline(model=model, output_dir=output_dir)
|
|
vllm_frames = _run_with_custom_transformer(model=model, output_dir=output_dir)
|
|
|
|
_assert_video_similarity(
|
|
model_name=f"{MODEL_ID} (transformer-swap)",
|
|
vllm_frames=vllm_frames,
|
|
diffusers_frames=diffusers_frames,
|
|
ssim_threshold=TRANSFORMER_SSIM_THRESHOLD,
|
|
psnr_threshold=TRANSFORMER_PSNR_THRESHOLD,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 2: Full pipeline (OmniServer → HTTP API vs diffusers)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _run_vllm_omni_serving(model: str, output_dir: Path) -> list[Image.Image]:
|
|
"""Generate video via the full vLLM-Omni serving stack."""
|
|
server_args = [
|
|
"--model-class-name",
|
|
"LTX23Pipeline",
|
|
"--stage-init-timeout",
|
|
"600",
|
|
]
|
|
with OmniServer(model, server_args, use_omni=True) as server:
|
|
# Submit generation request
|
|
response = requests.post(
|
|
f"http://{server.host}:{server.port}/v1/videos",
|
|
files={
|
|
"prompt": (None, PROMPT),
|
|
"negative_prompt": (None, NEGATIVE_PROMPT),
|
|
"model": (None, server.model),
|
|
"num_frames": (None, str(NUM_FRAMES)),
|
|
"fps": (None, "24"),
|
|
"size": (None, f"{WIDTH}x{HEIGHT}"),
|
|
"num_inference_steps": (None, str(NUM_INFERENCE_STEPS)),
|
|
"guidance_scale": (None, str(GUIDANCE_SCALE)),
|
|
"seed": (None, str(SEED)),
|
|
},
|
|
timeout=120,
|
|
)
|
|
response.raise_for_status()
|
|
video_id = response.json()["id"]
|
|
|
|
# Poll for completion
|
|
import time
|
|
|
|
for _ in range(120):
|
|
status_resp = requests.get(
|
|
f"http://{server.host}:{server.port}/v1/videos/{video_id}",
|
|
timeout=30,
|
|
)
|
|
status_resp.raise_for_status()
|
|
status = status_resp.json()["status"]
|
|
if status == "completed":
|
|
break
|
|
if status in ("error", "failed"):
|
|
raise RuntimeError(f"Video generation failed: {status_resp.json()}")
|
|
time.sleep(5)
|
|
else:
|
|
raise TimeoutError(f"Video generation timed out after 600s (id={video_id})")
|
|
|
|
# Download video content
|
|
content_resp = requests.get(
|
|
f"http://{server.host}:{server.port}/v1/videos/{video_id}/content",
|
|
timeout=120,
|
|
)
|
|
content_resp.raise_for_status()
|
|
mp4_bytes = content_resp.content
|
|
|
|
# Save MP4
|
|
mp4_path = output_dir / "vllm_omni_pipeline.mp4"
|
|
with open(mp4_path, "wb") as f:
|
|
f.write(mp4_bytes)
|
|
|
|
# Extract frames
|
|
frames = _extract_mp4_frames(mp4_bytes)
|
|
for i, frame in enumerate(frames):
|
|
frame.save(output_dir / f"vllm_pipeline_frame_{i:04d}.png")
|
|
return frames
|
|
|
|
|
|
@pytest.mark.full_model
|
|
@pytest.mark.benchmark
|
|
@pytest.mark.diffusion
|
|
@pytest.mark.skipif(
|
|
not _DIFFUSERS_038, reason="LTX-2.3 requires diffusers >= 0.38.0 for cross_attn_mod and BWE vocoder"
|
|
)
|
|
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
|
def test_ltx2_3_pipeline_matches_diffusers(accuracy_artifact_root: Path) -> None:
|
|
"""Full-pipeline parity: vLLM-Omni serving stack vs diffusers.
|
|
|
|
Runs the complete vLLM-Omni OmniServer (subprocess, HTTP API, video
|
|
encoding) and compares per-frame against stock diffusers output.
|
|
Follows the Wan2.2 / Qwen Image pattern with seed-based determinism.
|
|
"""
|
|
model = _model_name()
|
|
output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
|
|
|
|
diffusers_frames = _run_diffusers_baseline(model=model, output_dir=output_dir)
|
|
vllm_frames = _run_vllm_omni_serving(model=model, output_dir=output_dir)
|
|
|
|
_assert_video_similarity(
|
|
model_name=f"{MODEL_ID} (full-pipeline)",
|
|
vllm_frames=vllm_frames,
|
|
diffusers_frames=diffusers_frames,
|
|
ssim_threshold=PIPELINE_SSIM_THRESHOLD,
|
|
psnr_threshold=PIPELINE_PSNR_THRESHOLD,
|
|
)
|