From 60f61c810190f4e5bb8c29716fc2d25caa62ea39 Mon Sep 17 00:00:00 2001
From: William Lin <SolitaryThinker@users.noreply.github.com>
Date: Fri, 8 Aug 2025 18:45:03 -0700
Subject: [PATCH] [bugfix] fix pyproject install and VSA precision test (#726)

---
 .buildkite/pipeline.yml                    |  1 +
 csrc/attn/README.md                        |  2 +-
 fastvideo/pipelines/pipeline_batch_info.py |  8 +++++---
 fastvideo/tests/modal/pr_test.py           |  4 ++--
 pyproject.toml                             | 22 +++++++++++++++++++++-
 5 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 01f4ee505..9f3850b84 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -163,6 +163,7 @@ steps:
             - path:
                 - "csrc/attn/vsa/**"
                 - "csrc/attn/tk/**"
+                - "csrc/attn/tests/test_vsa.py"
                 - "csrc/attn/setup_vsa.py"
                 - "csrc/attn/config_vsa.py"
                 - "csrc/attn/vsa.cpp"
diff --git a/csrc/attn/README.md b/csrc/attn/README.md
index b7d226546..30af4bf34 100644
--- a/csrc/attn/README.md
+++ b/csrc/attn/README.md
@@ -84,7 +84,7 @@ out = sliding_tile_attention(q, k, v, window_size, 0, False)
 ### Test
 ```bash
 python tests/test_sta.py # test STA
-python tests/test_block_sparse.py # test VSA
+python tests/test_vsa.py # test VSA
 ```
 ### Benchmark
 ```bash
diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py
index b5acc76f6..ba79b95d0 100644
--- a/fastvideo/pipelines/pipeline_batch_info.py
+++ b/fastvideo/pipelines/pipeline_batch_info.py
@@ -9,11 +9,13 @@ in a functional manner, reducing the need for explicit parameter passing.
 
 import pprint
 from dataclasses import asdict, dataclass, field
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import PIL.Image
 import torch
-from torchcodec.decoders import VideoDecoder
+
+if TYPE_CHECKING:
+    from torchcodec.decoders import VideoDecoder
 
 from fastvideo.attention import AttentionMetadata
 from fastvideo.configs.sample.teacache import TeaCacheParams, WanTeaCacheParams
@@ -198,5 +200,5 @@ class TrainingBatch:
 
 @dataclass
 class PreprocessBatch(ForwardBatch):
-    video_loader: list[VideoDecoder] = field(default_factory=list)
+    video_loader: list["VideoDecoder"] = field(default_factory=list)
     video_file_name: list[str] = field(default_factory=list)
diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py
index 05e81e971..6dea8bf1d 100644
--- a/fastvideo/tests/modal/pr_test.py
+++ b/fastvideo/tests/modal/pr_test.py
@@ -52,7 +52,7 @@ def run_test(pytest_command: str):
     git clone {git_repo} /FastVideo &&
     cd /FastVideo &&
     {checkout_command} &&
-    uv pip install -e .[test] --index-strategy unsafe-best-match &&
+    uv pip install -e .[test] &&
     {pytest_command}
     """
     
@@ -100,7 +100,7 @@ def run_precision_tests_STA():
 
 @app.function(gpu="H100:1", image=image, timeout=900)
 def run_precision_tests_VSA():
-    run_test("python csrc/attn/tests/test_block_sparse.py")
+    run_test("python csrc/attn/tests/test_vsa.py")
 
 @app.function(gpu="L40S:1", image=image, timeout=3600)
 def run_inference_lora_tests():
diff --git a/pyproject.toml b/pyproject.toml
index 9e7f58a81..8c3051941 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -53,7 +53,27 @@ dependencies = [
 ]
 
 [tool.uv]
-extra-index-url = ["https://download.pytorch.org/whl/cu128"]
+prerelease = "allow"
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
+  { index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
+]
+torchvision = [
+  { index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
+  { index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
+]
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cu128"
+url = "https://download.pytorch.org/whl/cu128"
+explicit = true
 
 [project.optional-dependencies]