[misc]: point DreamX-World at published FastVideo Diffusers repos

Advertise FastVideo/DreamX-World-5B-Cam-Diffusers and
FastVideo/DreamX-World-5B-Diffusers as the loadable ids in the registry
(raw GD-ML path detection is kept via the pattern detectors), default the
example and local tests to the hub ids so they no longer skip when local
converted dirs are absent, and drop the now-redundant pipeline class
override in the example (model_index.json carries DreamXWorldPipeline).

docker/Dockerfile

@@ -68,7 +68,7 @@ ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-p
 # cutlass-4.4 `cute.core.ThrMma` API, which crashes on the cutlass-dsl 4.5 that
 # flashinfer/quack pull in. After the wheel install we overlay this cutlass-4.5-safe
 # upstream cute (flash-attn-4) so the image runs FA4 instead of the FA2 fallback.
-ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083
+ARG FA4_CUTE_REF=940cd9680f3315f2f06b43ab5bea2c2cf2d96806
 
 # Provided automatically by BuildKit/buildx (e.g. "amd64" / "arm64") and used to
 # select the prebuilt flash-attn wheel. Empty under a plain `docker build` without
@@ -161,7 +161,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
         uv pip install flash-attn==${FLASH_ATTN_VERSION} --no-build-isolation; \
     fi
 
-# Overlay the CuTe-DSL-4.6-compatible upstream FA4 cute (FA4_CUTE_REF) over the
+# Overlay the cutlass-4.5-safe upstream FA4 cute (FA4_CUTE_REF) over the
 # wheel/source one so the image runs FA4, not the FA2 fallback. This pulls the FA4
 # runtime stack (cutlass-dsl, quack-kernels, apache-tvm-ffi, torch-c-dlpack-ext) --
 # the same deps the [dreamverse] extra already installs in CI; the installed torch

examples/inference/basic/basic_dreamx_world.py

@@ -1,3 +1,10 @@
+"""DreamX-World-5B-Cam camera-controlled video generation.
+
+Uses the pre-converted Diffusers checkpoint FastVideo/DreamX-World-5B-Cam-Diffusers.
+To convert the raw GD-ML/DreamX-World-5B-Cam checkpoint yourself, see
+scripts/checkpoint_conversion/dreamx_world_to_diffusers.py.
+"""
+
 import os
 
 from fastvideo import VideoGenerator
@@ -24,7 +31,6 @@ def main():
         vae_cpu_offload=True,
         text_encoder_cpu_offload=True,
         pin_cpu_memory=False,
-        override_pipeline_cls_name="DreamXWorldPipeline",
     )
 
     prompt = os.getenv(

fastvideo/distributed/parallel_state.py

@@ -23,7 +23,6 @@ If you only need to use the distributed environment without model parallelism,
  you can skip the model parallel initialization and destruction steps.
 """
 import contextlib
-import gc
 import os
 import pickle
 import weakref
@@ -1007,13 +1006,6 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
     if shutdown_ray:
         import ray  # Lazy import Ray
         ray.shutdown()
-    # Actually free GPU memory, as the name promises. FSDP-wrapped modules sit
-    # in reference cycles (module <-> FSDP state <-> hooks), so dropping the
-    # last user reference does not free their parameters until a gc pass runs.
-    # Without this, back-to-back model-load tests in one pytest session OOM.
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
 
 
 def get_same_node_ranks(pg: ProcessGroup | StatelessProcessGroup, source_rank: int = 0) -> list[int]:

fastvideo/registry.py

@@ -783,10 +783,12 @@ def _register_configs() -> None:
             "FastVideo/DreamX-World-5B-Cam-Diffusers",
         ],
         model_detectors=[
-            # Mutually exclusive with the AR detector below: Cam requires an
-            # explicit "cam" marker so hyphenated AR local paths (e.g.
-            # /ckpts/dreamx-world-5b-converted) don't first-match here —
-            # detector resolution is first-match in registration order.
+            # Pattern also catches the raw GD-ML/DreamX-World-5B-Cam id and
+            # local converted dirs. Mutually exclusive with the AR detector
+            # below: Cam requires an explicit "cam" marker so hyphenated AR
+            # local paths (e.g. /ckpts/dreamx-world-5b-converted) don't
+            # first-match here — detector resolution is first-match in
+            # registration order.
             lambda path:
             ("dreamx-world" in path.lower() and "cam" in path.lower()) or "dreamxworldpipeline" in path.lower()
         ],

fastvideo/tests/api/test_compat_translation.py

@@ -17,6 +17,46 @@ class TestLegacyTorchCompileKwargsTranslation:
     first-class :class:`CompileConfig` fields and anything unknown falls
     into ``extras``."""
 
+    def test_all_typed_keys_promoted(self) -> None:
+        config = legacy_from_pretrained_to_config(
+            "/models/ltx2",
+            {
+                "enable_torch_compile": True,
+                "torch_compile_kwargs": {
+                    "backend": "inductor",
+                    "fullgraph": True,
+                    "mode": "max-autotune-no-cudagraphs",
+                    "dynamic": False,
+                },
+            },
+        )
+        compile_config = config.engine.compile
+        assert compile_config.enabled is True
+        assert compile_config.backend == "inductor"
+        assert compile_config.fullgraph is True
+        assert compile_config.mode == "max-autotune-no-cudagraphs"
+        assert compile_config.dynamic is False
+        assert compile_config.extras == {}
+
+    def test_unknown_keys_land_in_extras(self) -> None:
+        config = legacy_from_pretrained_to_config(
+            "/models/ltx2",
+            {
+                "enable_torch_compile": True,
+                "torch_compile_kwargs": {
+                    "backend": "inductor",
+                    "options": {"triton.cudagraphs": False},
+                    "disable": False,
+                },
+            },
+        )
+        compile_config = config.engine.compile
+        assert compile_config.backend == "inductor"
+        assert compile_config.extras == {
+            "options": {"triton.cudagraphs": False},
+            "disable": False,
+        }
+
     def test_empty_kwargs_produces_empty_extras(self) -> None:
         config = legacy_from_pretrained_to_config(
             "/models/ltx2",
@@ -77,6 +117,13 @@ class TestLegacyLtx2VaeTilingTranslation:
     ``generator.pipeline.vae_tiling``; reverse direction emits the
     legacy name back to FastVideoArgs."""
 
+    def test_forward_routes_to_pipeline_vae_tiling(self) -> None:
+        config = legacy_from_pretrained_to_config(
+            "/models/ltx2",
+            {"ltx2_vae_tiling": False},
+        )
+        assert config.pipeline.vae_tiling is False
+
     def test_true_round_trips(self) -> None:
         config = legacy_from_pretrained_to_config(
             "/models/ltx2",

fastvideo/tests/contract/test_ci_test_collection.py

@@ -35,6 +35,7 @@ ALLOWLIST = {
     "hooks": "no lane yet — run manually",
     "layers": "no lane yet — torchrun FSDP dispatch tests, run manually",
     "nightly": "by design: nightly cadence, not per-PR",
+    "ops": "no lane yet — GPU op tests, run manually",
     "modal": "CI infrastructure itself, not a test suite",
 }
 

fastvideo/tests/modal/pr_test.py

@@ -187,8 +187,6 @@ def run_transformer_tests():
 
 
 @app.function(gpu="L40S:4",
-              cpu=16.0,
-              memory=65536,
               image=image,
               timeout=900,
               secrets=[
@@ -203,8 +201,6 @@ def run_training_tests():
 
 
 @app.function(gpu="L40S:2",
-              cpu=16.0,
-              memory=65536,
               image=image,
               timeout=900,
               secrets=[
@@ -280,7 +276,7 @@ def run_self_forcing_tests():
 @app.function(gpu="L40S:1", image=image, timeout=900)
 def run_unit_test():
     run_test(
-        "pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ ./fastvideo/tests/ops/ --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
+        "pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
     )
 
 
@@ -322,8 +318,6 @@ def run_dreamverse_app_tests():
 
 
 @app.function(gpu="L40S:1",
-              cpu=16.0,
-              memory=65536,
               image=image,
               timeout=1800,
               secrets=[
@@ -406,8 +400,6 @@ def run_lora_extraction_tests():
 
 
 @app.function(gpu="L40S:2",
-              cpu=16.0,
-              memory=65536,
               image=image,
               timeout=1800,
               secrets=[
@@ -435,8 +427,6 @@ def run_performance_tests():
         "export PERF_RUN_SOURCE='unknown'; "
         "export PERF_UPLOAD_POLICY='never'; "
         "fi; "
-        "(nvidia-smi --query-gpu=index,timestamp,clocks.sm,clocks.max.sm,power.draw,power.limit,temperature.gpu "
-        "--format=csv -l 10 > /tmp/gpu_telemetry.csv 2>/dev/null &); "
         "pytest ./fastvideo/tests/performance -vs; "
         "PYTEST_RC=$?; "
         "PERF_RC=0; "
@@ -445,8 +435,6 @@ def run_performance_tests():
         "PERF_RC=$?; "
         "fi; "
         "python ./fastvideo/tests/performance/dashboard.py || true; "
-        "echo '--- GPU telemetry (clocks.sm vs clocks.max.sm reveals capped hosts) ---'; "
-        "cat /tmp/gpu_telemetry.csv || true; "
         "FINAL_RC=$PYTEST_RC; "
         "if [ $FINAL_RC -eq 0 ]; then FINAL_RC=$PERF_RC; fi; "
         "exit $FINAL_RC")

fastvideo/tests/ops/quantization/test_absmax_fp8.py

@@ -34,9 +34,6 @@ class TestAbsMaxFP8LinearMethod(unittest.TestCase):
             method._convert_scale(scale)
 
     def test_create_weights_rejects_invalid_dtype(self):
-        # float64 is outside the supported set (bf16/fp16/fp32). The test
-        # originally passed float32, which create_weights has accepted since
-        # PR #981 — it never ran in CI, so the mismatch went unnoticed.
         method = AbsMaxFP8LinearMethod()
         layer = nn.Module()
         with self.assertRaisesRegex(AssertionError, "only supports"):
@@ -46,7 +43,7 @@ class TestAbsMaxFP8LinearMethod(unittest.TestCase):
                 output_partition_sizes=[3],
                 input_size=2,
                 output_size=3,
-                params_dtype=torch.float64,
+                params_dtype=torch.float32,
             )
 
     def test_absmax_fp8_parameter_weight_loader(self):

fastvideo/tests/ssim/test_dreamx_world_similarity.py

@@ -20,13 +20,22 @@ REQUIRED_GPUS = 1
 
 device_reference_folder = resolve_inference_device_reference_folder(logger)
 
+_LOCAL_CONVERTED_MODEL = Path("converted_weights/dreamx_world")
 _MODEL_PATH = os.getenv(
     "DREAMX_WORLD_SSIM_MODEL_PATH",
-    "FastVideo/DreamX-World-5B-Cam-Diffusers",
+    str(_LOCAL_CONVERTED_MODEL),
+)
+_LOCAL_AR_CANDIDATES = (
+    Path("/tmp/converted_dreamx_world_ar"),
+    Path("/root/data/dreamx_world_ar_converted"),
+)
+_DEFAULT_AR_MODEL_PATH = next(
+    (str(path) for path in _LOCAL_AR_CANDIDATES if path.exists()),
+    str(_LOCAL_AR_CANDIDATES[0]),
 )
 _AR_MODEL_PATH = os.getenv(
     "DREAMX_WORLD_AR_SSIM_MODEL_PATH",
-    "FastVideo/DreamX-World-5B-Diffusers",
+    _DEFAULT_AR_MODEL_PATH,
 )
 
 DREAMX_WORLD_PARAMS = {
@@ -128,6 +137,13 @@ def test_dreamx_world_inference_similarity(
     model_id: str,
     tmp_path: Path,
 ) -> None:
+    model_path = Path(str(DREAMX_WORLD_MODEL_TO_PARAMS[model_id]["model_path"]))
+    if not model_path.exists():
+        pytest.skip(
+            f"DreamX-World converted model path is missing: {model_path}. "
+            "Set DREAMX_WORLD_SSIM_MODEL_PATH to a FastVideo-loadable converted root."
+        )
+
     image_path = tmp_path / "dreamx_world_ssim_input.png"
     _write_deterministic_reference_image(image_path)
 
@@ -166,6 +182,13 @@ def test_dreamx_world_ar_inference_similarity(
     attention_backend_name: str,
     model_id: str,
 ) -> None:
+    model_path = Path(str(DREAMX_WORLD_AR_MODEL_TO_PARAMS[model_id]["model_path"]))
+    if not model_path.exists():
+        pytest.skip(
+            f"DreamX-World AR converted model path is missing: {model_path}. "
+            "Set DREAMX_WORLD_AR_SSIM_MODEL_PATH to a FastVideo-loadable converted root."
+        )
+
     run_text_to_video_similarity_test(
         logger=logger,
         script_dir=os.path.dirname(os.path.abspath(__file__)),

fastvideo/tests/transformers/test_hunyuanvideo.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import glob
+import json
+import os
+import pytest
+import torch
+
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size
+from fastvideo.logger import init_logger
+from fastvideo.models.loader.component_loader import TransformerLoader
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.utils import maybe_download_model
+from fastvideo.forward_context import set_forward_context
+from fastvideo.configs.models.dits import HunyuanVideoConfig
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29503"
+
+BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
+TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
+CONFIG_PATH = os.path.join(TRANSFORMER_PATH, "config.json")
+
+LOCAL_RANK = 0
+RANK = 0
+WORLD_SIZE = 1
+
+# Latent generated on commit c021e8a27cf437ac22827f2bc58b7f006561317f with 1 x L40S
+REFERENCE_LATENT = 93.23263549804688
+
+
+@pytest.mark.usefixtures("distributed_setup")
+@pytest.mark.skip(reason="Skipping test_hunyuanvideo_distributed")
+def test_hunyuanvideo_distributed():
+    logger.info(f"Initializing process: rank={RANK}, local_rank={LOCAL_RANK}, world_size={WORLD_SIZE}")
+
+    torch.cuda.set_device(f"cuda:{LOCAL_RANK}")
+
+    # Get tensor parallel info
+    sp_rank = get_sp_parallel_rank()
+    sp_world_size = get_sp_world_size()
+
+    logger.info(f"Process rank {RANK} initialized with SP rank {sp_rank} in SP world size {sp_world_size}")
+
+    config = json.load(open(CONFIG_PATH))
+    # remove   "_class_name": "HunyuanVideoTransformer3DModel",   "_diffusers_version": "0.32.0.dev0",
+    # TODO: write normalize config function
+    config.pop("_class_name")
+    config.pop("_diffusers_version")
+
+    precision_str = "bf16"
+    args = FastVideoArgs(
+        model_path=TRANSFORMER_PATH,
+        dit_cpu_offload=True,
+        pipeline_config=PipelineConfig(dit_config=HunyuanVideoConfig(), dit_precision=precision_str),
+    )
+    args.device = torch.device(f"cuda:{LOCAL_RANK}")
+
+    loader = TransformerLoader()
+    model = loader.load(TRANSFORMER_PATH, args)
+
+    model.eval()
+
+    # Create random inputs for testing
+    batch_size = 1
+    seq_len = 3
+    device = torch.device(f"cuda:{LOCAL_RANK}")
+
+    # Video latents [B, C, T, H, W]
+    hidden_states = torch.randn(batch_size, 16, 8, 16, 16, device=device, dtype=torch.bfloat16)
+    chunk_per_rank = hidden_states.shape[2] // sp_world_size
+    hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank : (sp_rank + 1) * chunk_per_rank]
+
+    # Text embeddings [B, L, D] (including global token)
+    encoder_hidden_states = torch.randn(batch_size, seq_len + 1, 4096, device=device, dtype=torch.bfloat16)
+
+    # Timestep
+    timestep = torch.tensor([500], device=device, dtype=torch.bfloat16)
+    forward_batch = ForwardBatch(
+        data_type="dummy",
+    )
+
+    # Disable gradients for inference
+    with torch.no_grad():
+        # Run inference on model
+        with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
+            with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch):
+                output = model(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    timestep=timestep,
+                )
+
+    latent = output.double().sum().item()
+
+    # Check if latents are similar
+    diff_output_latents = abs(REFERENCE_LATENT - latent)
+    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
+    assert diff_output_latents < 1e-4, f"Output latents differ significantly: max diff = {diff_output_latents}"

fastvideo/tests/transformers/test_matrixgame2.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import pytest
+import torch
+
+from fastvideo.configs.models.dits.matrixgame2 import MatrixGame2WanVideoConfig
+from fastvideo.configs.pipelines import PipelineConfig
+from fastvideo.distributed.parallel_state import (
+    get_sp_parallel_rank,
+    get_sp_world_size,
+)
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.models.loader.component_loader import TransformerLoader
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.utils import maybe_download_model
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29505"
+os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
+
+MODEL_PATH = maybe_download_model("FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers")
+TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
+REFERENCE_LATENT = -111835.58895772696
+
+
+@pytest.mark.skip(reason="Not reliably reproducible")
+@pytest.mark.usefixtures("distributed_setup")
+def test_matrixgame2_transformer():
+    transformer_path = TRANSFORMER_PATH
+
+    sp_rank = get_sp_parallel_rank()
+    sp_world_size = get_sp_world_size()
+
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    precision = torch.bfloat16
+    precision_str = "bf16"
+
+    args = FastVideoArgs(
+        model_path=transformer_path,
+        dit_cpu_offload=False,
+        use_fsdp_inference=False,
+        pipeline_config=PipelineConfig(
+            dit_config=MatrixGame2WanVideoConfig(), dit_precision=precision_str
+        ),
+    )
+    args.device = device
+
+    loader = TransformerLoader()
+    model = loader.load(transformer_path, args).to(device, dtype=precision)
+    model.eval()
+
+    total_params = sum(p.numel() for p in model.parameters())
+    weight_sum = sum(p.to(torch.float64).sum().item() for p in model.parameters())
+    weight_mean = weight_sum / total_params
+    logger.info("Total parameters: %s", total_params)
+    logger.info("Weight sum: %s", weight_sum)
+    logger.info("Weight mean: %s", weight_mean)
+
+    torch.manual_seed(42)
+
+    batch_size = 1
+    latent_frames = 15
+    latent_height = 44
+    latent_width = 80
+    image_dim = 1280
+    image_seq_len = 257
+
+    # Set num_frame_per_block to match actual latent_frames for action_model
+    model.num_frame_per_block = latent_frames
+    # Force full attention mode
+    model.local_attn_size = -1
+    # Clear cached block masks so they regenerate with correct settings
+    model.block_mask = None
+    model.block_mask_mouse = None
+    model.block_mask_keyboard = None
+
+    # Video latents [B, C, T, H, W]
+    x = torch.randn(batch_size, 16, latent_frames, latent_height, latent_width,
+                    device=device, dtype=precision)
+    cond_concat = torch.randn(batch_size, 20, latent_frames, latent_height, latent_width,
+                              device=device, dtype=precision)
+    hidden_states = torch.cat([x, cond_concat], dim=1)
+
+    if sp_world_size > 1:
+        chunk_per_rank = hidden_states.shape[2] // sp_world_size
+        hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank:(sp_rank + 1) * chunk_per_rank]
+
+    # Image embeddings
+    encoder_hidden_states_image = [
+        torch.randn(batch_size, image_seq_len, image_dim, device=device, dtype=precision)
+    ]
+
+    # Timestep
+    timestep = torch.full((batch_size, latent_frames), 500, device=device, dtype=precision)
+
+    # Action conditions
+    num_pixel_frames = (latent_frames - 1) * 4 + 1
+    mouse_cond = torch.randn(batch_size, num_pixel_frames, 2, device=device, dtype=precision)
+    keyboard_cond = torch.randn(batch_size, num_pixel_frames, 4, device=device, dtype=precision)
+
+    # Text embeddings (empty for i2v)
+    encoder_hidden_states = torch.zeros(batch_size, 0, model.config.hidden_size,
+                                        device=device, dtype=precision)
+
+    forward_batch = ForwardBatch(data_type="dummy")
+    model.num_frame_per_block = latent_frames
+
+    with torch.no_grad():
+        with torch.amp.autocast('cuda', dtype=precision):
+            with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch):
+                output = model(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_hidden_states_image=encoder_hidden_states_image,
+                    timestep=timestep,
+                    mouse_cond=mouse_cond,
+                    keyboard_cond=keyboard_cond,
+                )
+
+    latent = output.double().sum().item()
+
+    diff = abs(REFERENCE_LATENT - latent)
+    relative_diff = diff / abs(REFERENCE_LATENT)
+    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
+    logger.info(f"Absolute diff: {diff}, Relative diff: {relative_diff * 100:.4f}%")
+
+    # The 0.19% diff now
+    assert relative_diff < 0.005, f"Output latents differ significantly: relative diff = {relative_diff * 100:.4f}% (max allowed: 0.5%)"

fastvideo/tests/vaes/test_hunyuan_vae.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import os
+
+import pytest
+import torch
+from safetensors.torch import load_file
+
+from fastvideo.configs.pipelines import PipelineConfig
+from fastvideo.logger import init_logger
+
+# from fastvideo.models.vaes.hunyuanvae import (
+#     AutoencoderKLHunyuanVideo as MyHunyuanVAE)
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.models.loader.component_loader import VAELoader
+from fastvideo.configs.models.vaes import HunyuanVAEConfig
+from fastvideo.utils import maybe_download_model
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29503"
+
+BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
+VAE_PATH = os.path.join(MODEL_PATH, "vae")
+CONFIG_PATH = os.path.join(VAE_PATH, "config.json")
+
+# Latent generated on commit d71a4ebffc2034922fc379568b6a6aa722f3744c with 1 x A40
+# torch 2.7.1
+A40_REFERENCE_LATENT = -106.22467041015625
+
+# Latent generated on commit 2b54068960c41d42221e8b8719a374b499855029 with 1 x L40S
+L40S_REFERENCE_LATENT = -157.09130859375
+
+
+@pytest.mark.usefixtures("distributed_setup")
+@pytest.mark.skip(reason="Skipping test_hunyuan_vae")
+def test_hunyuan_vae():
+    device = torch.device("cuda:0")
+    precision = torch.bfloat16
+    precision_str = "bf16"
+    args = FastVideoArgs(
+        model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=HunyuanVAEConfig(), vae_precision=precision_str)
+    )
+    args.device = device
+    args.vae_cpu_offload = False
+
+    loader = VAELoader()
+    model = loader.load(VAE_PATH, args)
+
+    model.enable_tiling(
+        tile_sample_min_height=32,
+        tile_sample_min_width=32,
+        tile_sample_min_num_frames=8,
+        tile_sample_stride_height=16,
+        tile_sample_stride_width=16,
+        tile_sample_stride_num_frames=4,
+    )
+
+    batch_size = 1
+
+    # Video input [B, C, T, H, W]
+    input_tensor = torch.randn(batch_size, 3, 21, 64, 64, device=device, dtype=torch.bfloat16)
+
+    # Disable gradients for inference
+    with torch.no_grad():
+        latent = model.encode(input_tensor).mean.double().sum().item()
+
+    # Check if latents are similar
+    device_name = torch.cuda.get_device_name()
+    if "A40" in device_name:
+        REFERENCE_LATENT = A40_REFERENCE_LATENT
+    elif "L40S" in device_name:
+        REFERENCE_LATENT = L40S_REFERENCE_LATENT
+    else:
+        raise ValueError(f"Unknown device: {device_name}")
+
+    diff_encoded_latents = abs(REFERENCE_LATENT - latent)
+    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
+    assert diff_encoded_latents < 1e-4, f"Encoded latents differ significantly: max diff = {diff_encoded_latents}"

fastvideo/tests/vaes/test_official_wan_vae.py

@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import numpy as np
+import pytest
+import torch
+from diffusers import AutoencoderKLWan
+
+from fastvideo.configs.pipelines import PipelineConfig
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.loader.component_loader import VAELoader
+from fastvideo.configs.models.vaes import WanVAEConfig
+from fastvideo.utils import maybe_download_model
+
+try:
+    from wan.modules.vae2_2 import Wan2_2_VAE
+except ImportError:
+    Wan2_2_VAE = None
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29503"
+
+BASE_MODEL_PATH = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH)
+VAE_PATH = os.path.join(MODEL_PATH, "vae")
+
+
+@pytest.mark.skip(reason="disable test")
+@pytest.mark.usefixtures("distributed_setup")
+def test_official_wan_vae():
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    precision = torch.bfloat16
+    precision_str = "bf16"
+    args = FastVideoArgs(model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision=precision_str))
+    args.device = device
+    args.vae_cpu_offload = False
+
+    loader = VAELoader()
+    model2 = loader.load(VAE_PATH, args)
+    assert model2.use_feature_cache # Default to use the original WanVAE algorithm
+
+    model1 = Wan2_2_VAE(vae_pth="/mnt/weka/home/hao.zhang/wei/Wan2.2_VAE.pth", device=device, dtype=precision)
+
+    # Create identical inputs for both models
+    batch_size = 1
+
+    # Video input [B, C, T, H, W]
+    input_tensor = torch.randn(batch_size,
+                               3,
+                               81,
+                               32,
+                               32,
+                               device=device,
+                               dtype=precision)
+    # latent_tensor = torch.randn(batch_size,
+    #                             16,
+    #                             21,
+    #                             32,
+    #                             32,
+    #                             device=device,
+    #                             dtype=precision)
+
+    # Disable gradients for inference
+    with torch.no_grad():
+        # Test encoding
+        logger.info("Testing encoding...")
+        latent1 = model1.encode([input_tensor.squeeze(0)])[0].unsqueeze(0)
+        print("--------------------------------")
+        latent2 = model2.encode(input_tensor)
+        # Check if latents have the same shape
+        assert latent1.shape == latent2.mean.shape, f"Latent shapes don't match: {latent1.mean.shape} vs {latent2.mean.shape}"
+        # Check if latents are similar
+        max_diff_encode = torch.max(torch.abs(latent1 - latent2.mean))
+        mean_diff_encode = torch.mean(torch.abs(latent1 - latent2.mean))
+        logger.info("Maximum difference between encoded latents: %s",
+                    max_diff_encode.item())
+        logger.info("Mean difference between encoded latents: %s",
+                    mean_diff_encode.item())
+        assert max_diff_encode < 1e-5, f"Encoded latents differ significantly: max diff = {mean_diff_encode.item()}"
+        # Test decoding
+        logger.info("Testing decoding...")
+        # latent1_tensor = latent1.mode()
+        # mean1 = (torch.tensor(model1.config.latents_mean).view(
+        #     1, model1.config.z_dim, 1, 1, 1).to(input_tensor.device,
+        #                                         input_tensor.dtype))
+        # std1 = (1.0 / torch.tensor(model1.config.latents_std).view(
+        #     1, model1.config.z_dim, 1, 1, 1)).to(input_tensor.device,
+        #                                         input_tensor.dtype)
+        # latent1_tensor = latent1_tensor / std1 + mean1
+        # output1 = model1.decode(latent1_tensor).sample
+
+        output1 = model1.decode([latent1.squeeze(0)])
+        output1 = output1[0].unsqueeze(0)
+
+        mean2 = model2.config.arch_config.shift_factor.to(input_tensor.device, input_tensor.dtype)
+        std2 = model2.config.arch_config.scaling_factor.to(input_tensor.device, input_tensor.dtype)
+        latent2_tensor = latent2.mode()
+        latent2_tensor = latent2_tensor / std2 + mean2
+        output2 = model2.decode(latent2_tensor)
+        # Check if outputs have the same shape
+        assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
+
+        # Check if outputs are similar
+        max_diff_decode = torch.max(torch.abs(output1.float() - output2.float()))
+        mean_diff_decode = torch.mean(torch.abs(output1.float() - output2.float()))
+        logger.info("Maximum difference between decoded outputs: %s",
+                    max_diff_decode.item())
+        logger.info("Mean difference between decoded outputs: %s",
+                    mean_diff_decode.item())
+        assert max_diff_decode < 1e-5, f"Decoded outputs differ significantly: max diff = {max_diff_decode.item()}"

pyproject.toml

@@ -114,7 +114,7 @@ fastvideo-kernel = [
 imagebind = { git = "https://github.com/facebookresearch/ImageBind.git", rev = "53680b02d7e37b19b124fa37bae4b6c98c38f5be" }
 # FA4 cute, pinned to a cutlass-4.5-compatible revision. torch.compile support
 # comes from FastVideo's own custom_op wrappers.
-flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = "82d6441eec5d4dfec120153db2c0145ae855a083", subdirectory = "flash_attn/cute" }
+flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = "940cd9680f3315f2f06b43ab5bea2c2cf2d96806", subdirectory = "flash_attn/cute" }
 
 [project.optional-dependencies]
 

tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py

@@ -26,6 +26,7 @@ from fastvideo.models.loader.weight_utils import resolve_safetensors_files, safe
 REPO_ROOT = Path(__file__).resolve().parents[3]
 OFFICIAL_REF_DIR = Path(os.getenv("DREAMX_WORLD_OFFICIAL_REF_DIR", REPO_ROOT / "DreamX-World"))
 CONVERTED_AR_DIR = Path(os.getenv("DREAMX_WORLD_AR_CONVERTED_DIR", "/tmp/converted_dreamx_world_ar"))
+CONVERTED_AR_HF_REPO = "FastVideo/DreamX-World-5B-Diffusers"
 PARITY_SCOPE = "both"
 
 
@@ -161,7 +162,10 @@ def test_dreamx_world_ar_5b_config_matches_official_shape():
 def test_dreamx_world_ar_converted_5b_transformer_strict_loads():
     transformer_dir = CONVERTED_AR_DIR / "transformer"
     if not transformer_dir.exists():
-        pytest.skip(f"Converted AR transformer missing: {transformer_dir}")
+        # No local conversion: pull the published Diffusers transformer from the hub.
+        from huggingface_hub import snapshot_download
+
+        transformer_dir = Path(snapshot_download(CONVERTED_AR_HF_REPO, allow_patterns=["transformer/*"])) / "transformer"
     with torch.device("meta"):
         model = DreamXWorldARTransformer3DModel(make_dreamx_world_5b_ar_dit_config(), {})
     incompatible = load_model_from_full_model_state_dict(

tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py

@@ -38,6 +38,7 @@ OFFICIAL_REF_DIR = Path(os.getenv("DREAMX_WORLD_OFFICIAL_REF_DIR", REPO_ROOT / "
 LOCAL_WEIGHTS_DIR = Path(os.getenv("DREAMX_WORLD_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / "dreamx_world"))
 WAN_BASE_DIR = Path(os.getenv("DREAMX_WORLD_WAN_BASE_DIR", REPO_ROOT / "official_weights" / "Wan2.2-TI2V-5B"))
 CONVERTED_WEIGHTS_DIR = Path(os.getenv("DREAMX_WORLD_CONVERTED_WEIGHTS_DIR", REPO_ROOT / "converted_weights" / "dreamx_world"))
+CONVERTED_HF_REPO = "FastVideo/DreamX-World-5B-Cam-Diffusers"
 PARITY_SCOPE = "both"
 
 
@@ -130,7 +131,10 @@ def _load_fastvideo_transformer(device: torch.device, dtype: torch.dtype) -> tor
 def _load_fastvideo_transformer_strict(device: torch.device, dtype: torch.dtype) -> torch.nn.Module:
     transformer_dir = CONVERTED_WEIGHTS_DIR / "transformer"
     if not transformer_dir.exists():
-        pytest.skip(f"Converted DreamX transformer missing: {transformer_dir}")
+        # No local conversion: pull the published Diffusers transformer from the hub.
+        from huggingface_hub import snapshot_download
+
+        transformer_dir = Path(snapshot_download(CONVERTED_HF_REPO, allow_patterns=["transformer/*"])) / "transformer"
     safetensors_files = resolve_safetensors_files(str(transformer_dir))
     config = make_dreamx_world_5b_cam_dit_config()
     import fastvideo.models.dits.dreamx_world as fastvideo_dreamx

tests/local_tests/pipelines/test_dreamx_world_pipeline_parity.py

@@ -20,7 +20,8 @@ from torch.testing import assert_close
 
 os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
 
-MODEL_DIR = Path(os.getenv("DREAMX_WORLD_MODEL_DIR", "converted_weights/dreamx_world"))
+# Local converted dir or HF repo id; the loader downloads hub ids itself.
+MODEL_DIR = os.getenv("DREAMX_WORLD_MODEL_DIR", "FastVideo/DreamX-World-5B-Cam-Diffusers")
 
 
 pytestmark = pytest.mark.skipif(
@@ -68,9 +69,6 @@ def _close_generator(generator: Any) -> None:
 
 
 def test_dreamx_world_one_step_pipeline_latent_matches_manual_pass() -> None:
-    if not MODEL_DIR.exists():
-        pytest.fail(f"DreamX-World converted model directory is missing: {MODEL_DIR}")
-
     from fastvideo import VideoGenerator
     common_kwargs = dict(
         prompt="a quiet road through a futuristic city at sunrise",
@@ -88,7 +86,7 @@ def test_dreamx_world_one_step_pipeline_latent_matches_manual_pass() -> None:
     )
 
     generator = VideoGenerator.from_pretrained(
-        str(MODEL_DIR),
+        MODEL_DIR,
         num_gpus=1,
         use_fsdp_inference=False,
         dit_cpu_offload=False,

tests/local_tests/pipelines/test_dreamx_world_pipeline_smoke.py

@@ -13,7 +13,8 @@ from PIL import Image, ImageDraw
 
 os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
 
-MODEL_DIR = Path(os.getenv("DREAMX_WORLD_MODEL_DIR", "converted_weights/dreamx_world"))
+# Local converted dir or HF repo id; the loader downloads hub ids itself.
+MODEL_DIR = os.getenv("DREAMX_WORLD_MODEL_DIR", "FastVideo/DreamX-World-5B-Cam-Diffusers")
 
 
 
@@ -113,16 +114,13 @@ def test_dreamx_world_camera_stage_writes_y_camera() -> None:
 
 
 def test_dreamx_world_pipeline_load_generate_latent_smoke(tmp_path: Path) -> None:
-    if not MODEL_DIR.exists():
-        pytest.fail(f"DreamX-World converted model directory is missing: {MODEL_DIR}")
-
     from fastvideo import VideoGenerator
 
     image_path = tmp_path / "dreamx_world_smoke_input.png"
     _write_smoke_image(image_path)
 
     generator = VideoGenerator.from_pretrained(
-        str(MODEL_DIR),
+        MODEL_DIR,
         num_gpus=1,
         use_fsdp_inference=False,
         dit_cpu_offload=False,

tests/local_tests/test_nvfp4_fa4.py

@@ -10,6 +10,15 @@ import pytest
 
 os.environ["CUTE_DSL_ENABLE_TVM_FFI"] = "1"
 
+# Warm up flashinfer JIT before any FA4 imports
+from flashinfer.quantization import nvfp4_quantize, SfLayout
+_warmup = nvfp4_quantize(
+    torch.randn(4, 128, device="cuda", dtype=torch.bfloat16),
+    torch.ones(1, device="cuda", dtype=torch.float32),
+    sfLayout=SfLayout.layout_128x4, do_shuffle=False,
+)
+del _warmup
+
 
 def _make_impl(nheads, headdim, nvfp4=False):
     from fastvideo.attention.backends.flash_attn import FlashAttentionImpl
@@ -36,8 +45,7 @@ def _cuda_timer(fn, warmup=5, iters=20):
 
 
 @pytest.mark.skipif(
-    not torch.cuda.is_available()
-    or torch.cuda.get_device_capability() not in [(10, 0), (10, 3)],
+    torch.cuda.get_device_capability() not in [(10, 0), (10, 3)],
     reason="Requires Blackwell GPU (sm100a or sm103a)"
 )
 class TestNVFP4FA4:
@@ -47,17 +55,6 @@ class TestNVFP4FA4:
     MODEL_HEADDIM = 128
     MODEL_SEQLEN = 32760  # 480x832 video, 81 frames
 
-    @pytest.fixture(scope="class", autouse=True)
-    def _flashinfer_warmup(self):
-        # Warm up flashinfer JIT before any FA4 imports
-        from flashinfer.quantization import SfLayout, nvfp4_quantize
-        _warmup = nvfp4_quantize(
-            torch.randn(4, 128, device="cuda", dtype=torch.bfloat16),
-            torch.ones(1, device="cuda", dtype=torch.float32),
-            sfLayout=SfLayout.layout_128x4, do_shuffle=False,
-        )
-        del _warmup
-
     def test_quantize_fn_shapes(self):
         """_nvfp4_quantize_for_fa4 produces correct shapes and strides."""
         from fastvideo.attention.backends.flash_attn import _nvfp4_quantize_for_fa4