From e6dcbb6cb827192257661627f7e22fe6f65c856a Mon Sep 17 00:00:00 2001 From: Raghav Date: Sun, 10 May 2026 17:54:15 -0700 Subject: [PATCH] [feat] Add FLUX.1-dev pipeline parity test vs Diffusers FluxPipeline Compares FastVideo FluxPipeline decoded image against Diffusers FluxPipeline under identical prompt/seed/steps (4 steps, 256x256, seed=42, atol=5e-2). Requires official_weights/FLUX.1-dev and CUDA. Updates README and PORT_STATUS to reflect test exists, pending run. --- tests/local_tests/flux/PORT_STATUS.md | 8 +- tests/local_tests/flux/README.md | 13 ++ .../test_flux_dev_pipeline_parity.py | 158 ++++++++++++++++++ 3 files changed, 175 insertions(+), 4 deletions(-) create mode 100644 tests/local_tests/pipelines/test_flux_dev_pipeline_parity.py diff --git a/tests/local_tests/flux/PORT_STATUS.md b/tests/local_tests/flux/PORT_STATUS.md index ff919d220..e57fdc7d0 100644 --- a/tests/local_tests/flux/PORT_STATUS.md +++ b/tests/local_tests/flux/PORT_STATUS.md @@ -26,7 +26,7 @@ Strict-load evidence: pending (loader smoke test verifies class resolution only) | Test | Status | Notes | |---|---|---| | Pipeline smoke (2-step, 256×256) | pending local run | tests/local_tests/pipelines/test_flux_dev_pipeline_smoke.py | -| Pipeline parity vs Diffusers FluxPipeline | not yet written | blocker: needs weights + GPU time | +| Pipeline parity vs Diffusers FluxPipeline | pending local run | tests/local_tests/pipelines/test_flux_dev_pipeline_parity.py | ## Quality @@ -40,9 +40,9 @@ Strict-load evidence: pending (loader smoke test verifies class resolution only) 1. SSIM reference images not committed — test will FileNotFoundError before comparison. Resolution: run seed-ssim-references skill on A40/L40S after weights are available. -2. Pipeline parity test not written — smoke test only verifies non-null finite output. - Resolution: write tests/local_tests/pipelines/test_flux_dev_pipeline_parity.py - comparing latents or decoded image against Diffusers FluxPipeline. +2. Pipeline parity test written but not yet run — smoke test only verifies non-null finite output. + Resolution: run tests/local_tests/pipelines/test_flux_dev_pipeline_parity.py on GPU + with weights and record PASS output here. 3. DiT parity PASS not recorded — test exists but no local output committed. Resolution: run on GPU with weights, record pytest output here. diff --git a/tests/local_tests/flux/README.md b/tests/local_tests/flux/README.md index bc21831ad..ee57d78d8 100644 --- a/tests/local_tests/flux/README.md +++ b/tests/local_tests/flux/README.md @@ -33,6 +33,19 @@ pytest tests/local_tests/pipelines/test_flux_dev_pipeline_smoke.py -vs Status: requires weights — not run in CI. +## Pipeline Parity Test + +Compares FastVideo FluxPipeline decoded image output against Diffusers +FluxPipeline under identical prompt, seed, and inference parameters (4 +steps, 256×256, seed=42). Tolerance: atol=5e-2. + +```bash +FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \ +pytest tests/local_tests/pipelines/test_flux_dev_pipeline_parity.py -vs +``` + +Status: requires weights — pending local run (see PORT_STATUS.md). + ## DiT Parity Test Single-forward comparison of FastVideo `FluxTransformer2DModel` against diff --git a/tests/local_tests/pipelines/test_flux_dev_pipeline_parity.py b/tests/local_tests/pipelines/test_flux_dev_pipeline_parity.py new file mode 100644 index 000000000..63af03f29 --- /dev/null +++ b/tests/local_tests/pipelines/test_flux_dev_pipeline_parity.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end pipeline parity test for FLUX.1-dev. + +Compares FastVideo FluxPipeline decoded image output against Diffusers +FluxPipeline under identical prompt, seed, and inference parameters. + +Coverage scope: production_loader + implementation_subcomponent +Comparison target: decoded RGB image (CHW float32, values in [0, 1]) + +Requires ``official_weights/FLUX.1-dev`` (or set ``FLUX_DEV_ROOT``) and CUDA. + +Run from repo root:: + + FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \ + pytest tests/local_tests/pipelines/test_flux_dev_pipeline_parity.py -vs +""" +from __future__ import annotations + +import os +from pathlib import Path + +import pytest +import torch +from torch.testing import assert_close + +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29521") + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_FLUX_DEV_ROOT = os.environ.get( + "FLUX_DEV_ROOT", + str(_REPO_ROOT / "official_weights" / "FLUX.1-dev"), +) + +# Parity test parameters — small enough to run quickly, large enough to be meaningful. +_PROMPT = "a photo of a cat" +_HEIGHT = 256 +_WIDTH = 256 +_NUM_INFERENCE_STEPS = 4 +_SEED = 42 +_GUIDANCE_SCALE = 3.5 + +# Tolerance for full pipeline (text encode → denoise → VAE decode accumulates error). +# Tighter than 1e-2 is unrealistic across different attention backends and precision paths. +_ATOL = 5e-2 +_RTOL = 5e-2 + + +def _log_stats(label: str, t: torch.Tensor) -> None: + tf = t.float() + print( + f"[FLUX PARITY] {label}: shape={tuple(t.shape)} dtype={t.dtype} " + f"min={tf.min():.4f} max={tf.max():.4f} mean={tf.mean():.4f} std={tf.std():.4f}" + ) + + +def _requires_weights() -> bool: + return Path(_FLUX_DEV_ROOT, "model_index.json").is_file() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.skipif(not _requires_weights(), reason="requires official_weights/FLUX.1-dev (set FLUX_DEV_ROOT)") +def test_flux_dev_pipeline_image_parity() -> None: + """FastVideo FluxPipeline decoded image matches Diffusers FluxPipeline.""" + + device = torch.device("cuda") + dtype = torch.bfloat16 + + # ------------------------------------------------------------------ + # 1. Diffusers reference + # ------------------------------------------------------------------ + print("\n[FLUX PARITY] Running Diffusers reference pipeline...") + from diffusers import FluxPipeline as HFFluxPipeline + + hf_pipe = HFFluxPipeline.from_pretrained( + _FLUX_DEV_ROOT, + torch_dtype=dtype, + ).to(device) + hf_pipe.set_progress_bar_config(disable=True) + + generator = torch.Generator(device=device).manual_seed(_SEED) + hf_out = hf_pipe( + prompt=_PROMPT, + height=_HEIGHT, + width=_WIDTH, + num_inference_steps=_NUM_INFERENCE_STEPS, + guidance_scale=_GUIDANCE_SCALE, + generator=generator, + output_type="pt", + ) + # hf_out.images: [B, H, W, C] float in [0,1] + hf_image = hf_out.images[0].permute(2, 0, 1).float().cpu() # [C, H, W] + _log_stats("Diffusers output", hf_image) + + del hf_pipe + torch.cuda.empty_cache() + + # ------------------------------------------------------------------ + # 2. FastVideo pipeline + # ------------------------------------------------------------------ + print("[FLUX PARITY] Running FastVideo pipeline...") + from fastvideo.configs.pipelines.flux import FluxPipelineConfig + from fastvideo.distributed import ( + cleanup_dist_env_and_memory, + maybe_init_distributed_environment_and_model_parallel, + ) + from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType + from fastvideo.pipelines.basic.flux.flux_pipeline import FluxPipeline + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + torch.manual_seed(_SEED) + maybe_init_distributed_environment_and_model_parallel(1, 1) + try: + args = FastVideoArgs( + model_path=_FLUX_DEV_ROOT, + pipeline_config=FluxPipelineConfig(), + workload_type=WorkloadType.T2I, + hsdp_shard_dim=1, + pin_cpu_memory=False, + distributed_executor_backend="mp", + ) + pipeline = FluxPipeline(_FLUX_DEV_ROOT, args) + + batch = ForwardBatch( + data_type="image", + prompt=_PROMPT, + height=_HEIGHT, + width=_WIDTH, + seed=_SEED, + num_inference_steps=_NUM_INFERENCE_STEPS, + guidance_scale=_GUIDANCE_SCALE, + use_embedded_guidance=True, + true_cfg_scale=1.0, + num_videos_per_prompt=1, + save_video=False, + ) + fv_out = pipeline.forward(batch, args) + # fv_out.output: [B, C, H, W] float in [0, 1] + fv_image = fv_out.output[0].float().cpu() # [C, H, W] + _log_stats("FastVideo output", fv_image) + finally: + cleanup_dist_env_and_memory() + + # ------------------------------------------------------------------ + # 3. Compare + # ------------------------------------------------------------------ + print("[FLUX PARITY] Comparing outputs...") + assert hf_image.shape == fv_image.shape, ( + f"Shape mismatch: Diffusers {hf_image.shape} vs FastVideo {fv_image.shape}" + ) + + abs_diff = (hf_image - fv_image).abs() + max_diff = abs_diff.max().item() + mean_diff = abs_diff.mean().item() + print(f"[FLUX PARITY] max_diff={max_diff:.4f} mean_diff={mean_diff:.4f}") + + assert_close(hf_image, fv_image, atol=_ATOL, rtol=_RTOL) + print("[FLUX PARITY] PASSED")