[fix] Fix latent handoff and FastWan-QAD example path

examples/inference/optimizations/FastWan_QAD_TAEHV.py

@@ -18,12 +18,11 @@
     - TAEHV weights ``taew2_1.pth`` (https://github.com/madebyollin/taehv)
 
 Usage:
-    python fp4_linear_taehv_wan2_1_1_3b.py                    # FP4 + TAEHV + compile
-    python fp4_linear_taehv_wan2_1_1_3b.py --no-taehv         # FP4 + full Wan VAE
-    python fp4_linear_taehv_wan2_1_1_3b.py --no-compile       # eager
-    python fp4_linear_taehv_wan2_1_1_3b.py --baseline         # dense bf16 reference
-    python fp4_linear_taehv_wan2_1_1_3b.py --distilled_model ''  # base Wan2.1 weights
-    python fp4_linear_taehv_wan2_1_1_3b.py --warmups 5 --benchmark-runs 20  # timing stats (default)
+    python FastWan_QAD_TAEHV.py                    # FP4 + TAEHV + compile
+    python FastWan_QAD_TAEHV.py --no-taehv         # FP4 + full Wan VAE
+    python FastWan_QAD_TAEHV.py --no-compile       # eager
+    python FastWan_QAD_TAEHV.py --baseline         # dense bf16 reference
+    python FastWan_QAD_TAEHV.py --warmups 5 --benchmark-runs 20  # timing stats (default)
 """
 
 import argparse
@@ -42,13 +41,12 @@
 
 OUTPUT_PATH = "video_samples"
 
-# Distilled, quantization-aware (QAD) transformer for Wan2.1-1.3B (3 steps,
-# guidance 1.0). Loaded on top of the base Wan2.1 pipeline; pass
-# ``--distilled_model ''`` to run the base weights instead.
-DEFAULT_DISTILLED_MODEL = "FastVideo/FastWan-QAD-1.3B"
-DISTILLED_WEIGHTS_FILE = (
-    "generator_inference_transformer/diffusion_pytorch_model.safetensors"
-)
+# FastWan-QAD-1.3B is already a complete Diffusers pipeline. Keep the
+# distilled overlay disabled by default; users can still pass a local
+# safetensors file or compatible repo via ``--distilled_model``.
+DEFAULT_MODEL = "FastVideo/FastWan-QAD-1.3B"
+DEFAULT_DISTILLED_MODEL = ""
+DISTILLED_WEIGHTS_FILE = "transformer/diffusion_pytorch_model.safetensors"
 
 # TAEHV checkpoint for Wan2.1. Clone https://github.com/madebyollin/taehv to get
 # ``taew2_1.pth`` (Wan 2.1 / Wan 2.2-14B / Qwen-Image all use this VAE).
@@ -183,7 +181,7 @@ def build_generator(args: argparse.Namespace) -> VideoGenerator:
 def main() -> None:
     parser = argparse.ArgumentParser(
         description="FP4 linear Wan2.1-1.3B with TAEHV decoding benchmark")
-    parser.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    parser.add_argument("--model", default=DEFAULT_MODEL,
                         help="Model path or HuggingFace ID")
     parser.add_argument("--baseline", action="store_true",
                         help="Run dense bf16 instead of FP4 linear")
@@ -198,7 +196,7 @@ def main() -> None:
     parser.add_argument("--distilled_model", default=DEFAULT_DISTILLED_MODEL,
                         help="HuggingFace ID (or local path) of a distilled "
                              "transformer checkpoint to load on top of --model. "
-                             "Pass '' to use the base --model weights instead.")
+                             "Pass '' to use --model weights directly.")
     parser.add_argument("--num_gpus", type=int, default=1)
     parser.add_argument("--infer_steps", type=int, default=3)
     parser.add_argument("--guidance_scale", type=float, default=1.0)

fastvideo/tests/modal/pr_test.py

@@ -276,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/ ./fastvideo/tests/ops/ ./fastvideo/tests/worker/ --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
     )
 
 

fastvideo/tests/worker/test_multiproc_executor.py

@@ -0,0 +1,24 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from fastvideo.worker.multiproc_executor import _prepare_worker_output_for_parent
+
+
+def test_prepare_worker_output_keeps_non_latent_tensor_on_device():
+    result = torch.empty((1, 2))
+    args = SimpleNamespace(output_type="pil")
+
+    assert _prepare_worker_output_for_parent(result, args) is result
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
+def test_prepare_worker_output_moves_latent_cuda_tensor_to_cpu():
+    result = torch.empty((1, 16, 1, 8, 8), device="cuda")
+    args = SimpleNamespace(output_type="latent")
+
+    prepared = _prepare_worker_output_for_parent(result, args)
+
+    assert prepared.device.type == "cpu"
+    assert prepared.shape == result.shape

fastvideo/worker/multiproc_executor.py

@@ -73,6 +73,13 @@ class StreamingResult:
     error: Exception | None = None
 
 
+def _prepare_worker_output_for_parent(result: Any, fastvideo_args: FastVideoArgs | None) -> Any:
+    if (fastvideo_args is not None and getattr(fastvideo_args, "output_type", None) == "latent"
+            and isinstance(result, torch.Tensor) and result.is_cuda):
+        return result.detach().cpu()
+    return result
+
+
 class MultiprocExecutor(Executor):
 
     def _init_executor(self) -> None:
@@ -693,8 +700,13 @@ def worker_busy_loop(self) -> None:
                         logging_info = None
                         if envs.FASTVIDEO_STAGE_LOGGING:
                             logging_info = output_batch.logging_info
-                        # result tensor shared by CUDA IPC to avoid serialization overhead
                         result = output_batch.output
+                        # Latent outputs are small and are commonly consumed by
+                        # the parent process (for example, external TAEHV
+                        # decoding). Keep them off CUDA IPC so consumer/WSL2
+                        # drivers do not need to import a worker-owned CUDA
+                        # memory handle.
+                        result = _prepare_worker_output_for_parent(result, fastvideo_args)
                         extra = output_batch.extra or {}
                         extra["peak_memory_mb"] = (torch.cuda.max_memory_allocated() / (1024 * 1024))
                         self.pipe.send({