wip

.buildkite/pipeline.yml

@@ -18,7 +18,8 @@ steps:
             - path: 
                 - "fastvideo/models/encoders/**"
                 - "fastvideo/models/loader/**"
-                - "fastvideo/tests/encoders/**"
+                - "tests/encoders/**"
+                - "tests/utils.py"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -31,7 +32,8 @@ steps:
             - path:
                 - "fastvideo/models/vaes/**"
                 - "fastvideo/models/loader/**"
-                - "fastvideo/tests/vaes/**"
+                - "tests/vaes/**"
+                - "tests/utils.py"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -44,7 +46,8 @@ steps:
             - path:
                 - "fastvideo/models/dits/**"
                 - "fastvideo/models/loader/**"
-                - "fastvideo/tests/transformers/**"
+                - "tests/transformers/**"
+                - "tests/utils.py"
                 - "fastvideo/layers/**"
                 - "fastvideo/attention/**"
                 - "pyproject.toml"
@@ -58,6 +61,8 @@ steps:
                   queue: "default"
             - path: 
                 - "fastvideo/**/*.py"
+                - "tests/ssim/**"
+                - "tests/utils.py"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -68,9 +73,10 @@ steps:
                 agents:
                   queue: "default"
             - path:
-                - "fastvideo/tests/lora/**"
+                - "tests/inference/lora/**"
                 - "fastvideo/models/loader/**"
-                - "fastvideo/tests/transformers/**"
+                - "tests/transformers/**"
+                - "tests/utils.py"
                 - "fastvideo/pipelines/**"
                 - "fastvideo/layers/lora/**"
                 - "pyproject.toml"
@@ -84,6 +90,7 @@ steps:
                   queue: "default"
             - path:
                 - "fastvideo/**"
+                - "tests/training/Vanilla/**"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -95,6 +102,7 @@ steps:
                   queue: "default"
             - path:
                 - "fastvideo/training/*distillation_pipeline.py"
+                - "tests/training/distill/**"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -106,6 +114,7 @@ steps:
                   queue: "default"
             - path:
                 - "fastvideo/**"
+                - "tests/training/lora/**"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -117,6 +126,7 @@ steps:
                   queue: "default"
             - path:
                 - "fastvideo/**"
+                - "tests/training/VSA/**"
                 - "csrc/attn/video_sparse_attn/**"
                 - "csrc/attn/video_sparse_attn/tk/**"
                 - "csrc/attn/video_sparse_attn/setup.py"
@@ -133,6 +143,7 @@ steps:
                   queue: "default"
             - path:
                 - "fastvideo/**"
+                - "tests/inference/STA/**"
                 - "csrc/attn/sliding_tile_attn/**"
                 - "csrc/attn/sliding_tile_attn/setup.py"
                 - "csrc/attn/sliding_tile_attn/config_sta.py"
@@ -190,6 +201,7 @@ steps:
             - path:
                 - "csrc/attn/vmoba_attn/vmoba/**"
                 - "fastvideo/attention/backends/vmoba.py"
+                - "tests/inference/vmoba/**"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -201,6 +213,8 @@ steps:
                   queue: "default"
             - path:
                 - "fastvideo/**"
+                - "tests/dataset/**"
+                - "tests/workflow/**"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:

.buildkite/scripts/pr_test.sh

@@ -50,7 +50,7 @@ else
     exit 1
 fi
 
-MODAL_TEST_FILE="fastvideo/tests/modal/pr_test.py"
+MODAL_TEST_FILE="tests/modal/pr_test.py"
 
 if [ -z "${TEST_TYPE:-}" ]; then
     log "Error: TEST_TYPE environment variable is not set"

.github/workflows/pr-test.yml

@@ -127,29 +127,35 @@ jobs:
             encoder-test:
               - 'fastvideo/models/encoders/**'
               - 'fastvideo/models/loader/**'
-              - 'fastvideo/tests/encoders/**'
+              - 'tests/encoders/**'
+              - 'tests/utils.py'
               - *common-paths
             vae-test:
               - 'fastvideo/models/vaes/**'
               - 'fastvideo/models/loader/**'
-              - 'fastvideo/tests/vaes/**'
+              - 'tests/vaes/**'
+              - 'tests/utils.py'
               - *common-paths
             transformer-test:
               - 'fastvideo/models/dits/**'
               - 'fastvideo/models/loader/**'
-              - 'fastvideo/tests/transformers/**'
+              - 'tests/transformers/**'
+              - 'tests/utils.py'
               - 'fastvideo/layers/**'
               - 'fastvideo/attention/**'
               - *common-paths
             training-test:
               - 'fastvideo/**'
+              - 'tests/training/Vanilla/**'
               - *common-paths
             training-test-VSA:
               - 'fastvideo/**'
+              - 'tests/training/VSA/**'
               - *common-paths
               - *vsa-kernel-paths
             inference-test-STA:
               - 'fastvideo/**'
+              - 'tests/inference/STA/**'
               - *common-paths
               - *sta-kernel-paths
             precision-test-STA:
@@ -160,6 +166,8 @@ jobs:
               - *vsa-kernel-paths
             unit-test:
               - 'fastvideo/**'
+              - 'tests/dataset/**'
+              - 'tests/workflow/**'
               - *common-paths
 
   encoder-test:
@@ -174,7 +182,7 @@ jobs:
       gpu_count: 1
       volume_size: 100
       image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/encoders -s"
+      test_command: "uv pip install -e .[test] && pytest ./tests/encoders -s"
       timeout_minutes: 30
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -192,7 +200,7 @@ jobs:
       gpu_count: 1
       volume_size: 100
       image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/vaes -s"
+      test_command: "uv pip install -e .[test] && pytest ./tests/vaes -s"
       timeout_minutes: 30
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -210,7 +218,7 @@ jobs:
       gpu_count: 1
       volume_size: 100
       image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/transformers -s"
+      test_command: "uv pip install -e .[test] && pytest ./tests/transformers -s"
       timeout_minutes: 30
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -236,7 +244,7 @@ jobs:
       volume_size: 200
       disk_size: 200
       image: "ghcr.io/${{ github.repository }}/fastvideo-dev:${{ matrix.python-version.tag }}"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/ssim -vs"
+      test_command: "uv pip install -e .[test] && pytest ./tests/ssim -vs"
       timeout_minutes: 60
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -255,7 +263,7 @@ jobs:
       volume_size: 100
       disk_size: 100
       image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/training/Vanilla -srP"
+      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./tests/training/Vanilla -srP"
       timeout_minutes: 30
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -275,7 +283,7 @@ jobs:
       volume_size: 100
       disk_size: 100
       image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/training/VSA -srP"
+      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./tests/training/VSA -srP"
       timeout_minutes: 30
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -295,7 +303,7 @@ jobs:
       volume_size: 100
       disk_size: 100
       image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/inference/STA -srP"
+      test_command: "uv pip install -e .[test] && pytest ./tests/inference/STA -srP"
       timeout_minutes: 30
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -369,7 +377,7 @@ jobs:
   #     volume_size: 100
   #     disk_size: 100
   #     image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-  #     test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/nightly/test_e2e_overfit_single_sample.py -vs"
+#     test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./tests/nightly/test_e2e_overfit_single_sample.py -vs"
   #     timeout_minutes: 30
   #   secrets:
   #     RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}

docs/source/inference/optimizations.md

@@ -1,32 +1,40 @@
 (inference-optimizations)=
+
 # Optimizations
 
 This page describes the various options for speeding up generation times in FastVideo.
 
 ## Table of Contents
+
 - Optimized Attention Backends
+
   - [Flash Attention](#optimizations-flash)
   - [Sliding Tile Attention](#optimizations-sta)
   - [Sage Attention](#optimizations-sage)
+  - [Sage Attention 3](#optimizations-sage3)
 
 - Caching Techniques
   - [TeaCache](#optimizations-teacache)
 
 (optimizations-backends)=
+
 ## Attention Backends
 
 ### Available Backends
+
 - Torch SDPA: `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA`
 - Flash Attention 2 and 3: `FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN`
 - Sliding Tile Attention: `FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN`
 - Video Sparse Attention: `FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN`
 - Sage Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN`
+- Sage Attention 3: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN_THREE`
 
 ### Configuring Backends
 
 There are two ways to configure the attention backend in FastVideo.
 
 #### 1. In Python
+
 In python, set the `FASTVIDEO_ATTENTION_BACKEND` environment variable before instantiating `VideoGenerator` like this:
 
 ```python
@@ -34,6 +42,7 @@ os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN"
 ```
 
 #### 2. In CLI
+
 You can also set the environment variable on the command line:
 
 ```bash
@@ -41,6 +50,7 @@ FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN python example.py
 ```
 
 (optimizations-flash)=
+
 ### Flash Attention
 
 **`FLASH_ATTN`**
@@ -57,7 +67,7 @@ And if using a Hopper+ GPU (ie H100), installing [Flash Attention 3](https://git
 git clone https://github.com/Dao-AILab/flash-attention.git && cd flash-attention
 
 cd hopper
-pip install ninja 
+pip install ninja
 python setup.py install
 ```
 
@@ -66,7 +76,9 @@ FastVideo will automatically detect and use `FA3` if it is installed when using
 :::
 
 (optimizations-sta)=
+
 ### Sliding Tile Attention
+
 **`SLIDING_TILE_ATTN`**
 
 ```bash
@@ -76,7 +88,9 @@ pip install st_attn==0.0.4
 Please see [this page](#sta-installation) for more installation instructions.
 
 (optimizations-vsa)=
+
 ### Video Sparse Attention
+
 **`VIDEO_SPARSE_ATTN`**
 
 ```bash
@@ -87,19 +101,45 @@ python setup_vsa.py install
 Please see [this page](#vsa-installation) for more installation instructions.
 
 (optimizations-sage)=
+
 ### Sage Attention
+
 **`SAGE_ATTN`**
 
 To use [SageAttention](https://github.com/thu-ml/SageAttention) 2.1.1, please compile from source:
 
 ```bash
 git clone https://github.com/thu-ml/SageAttention.git
-cd sageattention 
+cd sageattention
 python setup.py install  # or pip install -e .
 ```
 
+(optimizations-sage3)=
+
+### Sage Attention 3
+
+**`SAGE_ATTN_THREE`**
+
+[SageAttention 3](https://huggingface.co/jt-zhang/SageAttention3) is an advanced attention mechanism that leverages FP4 quantization and Blackwell GPU Tensor Cores for significant performance improvements.
+
+#### Hardware Requirements
+
+- RTX5090
+
+#### Installation
+
+Note that Sage Attention 3 requires `python>=3.13`, `torch>=2.8.0`, `CUDA >=12.8`. If you are using `uv` and using `torch==2.8.0` make sure that `sentencepiece==0.2.1` in the pyproject.toml file.
+
+To use Sage Attention 3 in FastVideo, first get access to the SageAttention3 code, then move `sageattn/` and `setup.py` to the directory `fastvideo/attention/backends`, then install from using:
+
+```bash
+python setup.py install
+```
+
 (optimizations-teacache)=
+
 ## Teacache
+
 TeaCache is an optimization technique supported in FastVideo that can significantly speed up video generation by skipping redundant calculations across diffusion steps. This guide explains how to enable and configure TeaCache for optimal performance in FastVideo.
 
 ### What is TeaCache?

fastvideo/attention/backends/sage_attn3.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+from fastvideo.attention.backends.sageattn.api import sageattn_blackwell
+
+from fastvideo.attention.backends.abstract import (AttentionBackend,
+                                                   AttentionImpl,
+                                                   AttentionMetadata,
+                                                   AttentionMetadataBuilder)
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class SageAttention3Backend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "SAGE_ATTN_THREE"
+
+    @staticmethod
+    def get_impl_cls() -> type["SageAttention3Impl"]:
+        return SageAttention3Impl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        raise NotImplementedError
+
+    # @staticmethod
+    # def get_metadata_cls() -> Type["AttentionMetadata"]:
+    #     return FlashAttentionMetadata
+
+
+class SageAttention3Impl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        output = sageattn_blackwell(query, key, value, is_causal=self.causal)
+        output = output.transpose(1, 2)
+        return output

fastvideo/configs/models/dits/base.py

@@ -15,13 +15,10 @@ class DiTArchConfig(ArchConfig):
     reverse_param_names_mapping: dict = field(default_factory=dict)
     lora_param_names_mapping: dict = field(default_factory=dict)
     _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
-        AttentionBackendEnum.SLIDING_TILE_ATTN,
-        AttentionBackendEnum.SAGE_ATTN,
-        AttentionBackendEnum.FLASH_ATTN,
-        AttentionBackendEnum.TORCH_SDPA,
-        AttentionBackendEnum.VIDEO_SPARSE_ATTN,
-        AttentionBackendEnum.VMOBA_ATTN,
-    )
+        AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN,
+        AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA,
+        AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN,
+        AttentionBackendEnum.SAGE_ATTN_THREE)
 
     hidden_size: int = 0
     num_attention_heads: int = 0

fastvideo/envs.py

@@ -175,6 +175,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # - "SLIDING_TILE_ATTN" : use Sliding Tile Attention
     # - "VIDEO_SPARSE_ATTN": use Video Sparse Attention
     # - "SAGE_ATTN": use Sage Attention
+    # - "SAGE_ATTN_THREE": use Sage Attention 3
     "FASTVIDEO_ATTENTION_BACKEND":
     lambda: os.getenv("FASTVIDEO_ATTENTION_BACKEND", None),
 

fastvideo/pipelines/stages/denoising.py

@@ -85,8 +85,9 @@ class DenoisingStage(PipelineStage):
                 AttentionBackendEnum.SLIDING_TILE_ATTN,
                 AttentionBackendEnum.VIDEO_SPARSE_ATTN,
                 AttentionBackendEnum.VMOBA_ATTN,
-                AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA
-            )  # hack
+                AttentionBackendEnum.FLASH_ATTN,
+                AttentionBackendEnum.TORCH_SDPA,
+                AttentionBackendEnum.SAGE_ATTN_THREE)  # hack
         )
 
     def forward(

fastvideo/platforms/cuda.py

@@ -115,6 +115,7 @@ class CudaPlatformBase(Platform):
 
         logger.info("Trying FASTVIDEO_ATTENTION_BACKEND=%s",
                     envs.FASTVIDEO_ATTENTION_BACKEND)
+        logger.info("Selected backend: %s", selected_backend)
         if selected_backend == AttentionBackendEnum.SLIDING_TILE_ATTN:
             try:
                 from st_attn import sliding_tile_attention  # noqa: F401
@@ -144,6 +145,20 @@ class CudaPlatformBase(Platform):
                 logger.info(
                     "Sage Attention backend is not installed. Fall back to Flash Attention."
                 )
+        elif selected_backend == AttentionBackendEnum.SAGE_ATTN_THREE:
+            try:
+                from fastvideo.attention.backends.sageattn.api import sageattn_blackwell  # noqa: F401
+
+                from fastvideo.attention.backends.sage_attn3 import (  # noqa: F401
+                    SageAttention3Backend)
+                logger.info("Using Sage Attention 3 backend.")
+
+                return "fastvideo.attention.backends.sage_attn3.SageAttention3Backend"
+            except ImportError as e:
+                logger.info(e)
+                logger.info(
+                    "Sage Attention 3 backend is not installed. Fall back to Flash Attention."
+                )
         elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
             try:
                 from vsa import block_sparse_attn  # noqa: F401

fastvideo/platforms/interface.py

@@ -18,6 +18,7 @@ class AttentionBackendEnum(enum.Enum):
     SLIDING_TILE_ATTN = enum.auto()
     TORCH_SDPA = enum.auto()
     SAGE_ATTN = enum.auto()
+    SAGE_ATTN_THREE = enum.auto()
     VIDEO_SPARSE_ATTN = enum.auto()
     VMOBA_ATTN = enum.auto()
     NO_ATTENTION = enum.auto()

tests/inference/lora/test_lora_inference_similarity.py

@@ -6,7 +6,7 @@ import pytest
 
 from fastvideo import VideoGenerator
 from fastvideo.logger import init_logger
-from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results
+from tests.utils import compute_video_ssim_torchvision, write_ssim_results
 from diffusers import DiffusionPipeline
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.pipelines import build_pipeline
@@ -192,4 +192,3 @@ def test_lora_inference_similarity(ATTENTION_BACKEND, model_id):
 
 
 
-

tests/modal/pr_test.py

@@ -64,35 +64,35 @@ def run_test(pytest_command: str):
 
 @app.function(gpu="L40S:1", image=image, timeout=900)
 def run_encoder_tests():
-    run_test("pytest ./fastvideo/tests/encoders -vs")
+    run_test("pytest ./tests/encoders -vs")
 
 @app.function(gpu="L40S:1", image=image, timeout=900)
 def run_vae_tests():
-    run_test("pytest ./fastvideo/tests/vaes -vs")
+    run_test("pytest ./tests/vaes -vs")
 
 @app.function(gpu="L40S:1", image=image, timeout=900)
 def run_transformer_tests():
-    run_test("pytest ./fastvideo/tests/transformers -vs")
+    run_test("pytest ./tests/transformers -vs")
 
 @app.function(gpu="L40S:2", image=image, timeout=2700)
 def run_ssim_tests():
-    run_test("pytest ./fastvideo/tests/ssim -vs")
+    run_test("pytest ./tests/ssim -vs")
 
 @app.function(gpu="L40S:4", image=image, timeout=900, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
 def run_training_tests():
-    run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/Vanilla -srP")
+    run_test("wandb login $WANDB_API_KEY && pytest ./tests/training/Vanilla -srP")
 
 @app.function(gpu="L40S:2", image=image, timeout=900, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
 def run_training_lora_tests():
-    run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP")
+    run_test("wandb login $WANDB_API_KEY && pytest ./tests/training/lora/test_lora_training.py -srP")
 
 @app.function(gpu="H100:2", image=image, timeout=900, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
 def run_training_tests_VSA():
-    run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/VSA -srP")
+    run_test("wandb login $WANDB_API_KEY && pytest ./tests/training/VSA -srP")
 
 @app.function(gpu="H100:2", image=image, timeout=900)
 def run_inference_tests_STA():
-    run_test("pytest ./fastvideo/tests/inference/STA -srP")
+    run_test("pytest ./tests/inference/STA -srP")
 
 @app.function(gpu="H100:1", image=image, timeout=900)
 def run_precision_tests_STA():
@@ -108,16 +108,16 @@ def run_precision_tests_vmoba():
 
 @app.function(gpu="L40S:1", image=image, timeout=900)
 def run_inference_tests_vmoba():
-    run_test('python fastvideo/tests/inference/vmoba/test_vmoba_inference.py')
+    run_test('python tests/inference/vmoba/test_vmoba_inference.py')
 
 @app.function(gpu="L40S:1", image=image, timeout=3600)
 def run_inference_lora_tests():
-    run_test("pytest ./fastvideo/tests/inference/lora/test_lora_inference_similarity.py -vs")
+    run_test("pytest ./tests/inference/lora/test_lora_inference_similarity.py -vs")
 
 @app.function(gpu="L40S:2", image=image, timeout=900)
 def run_distill_dmd_tests():
-    run_test("pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs")
+    run_test("pytest ./tests/training/distill/test_distill_dmd.py -vs")
 
 @app.function(gpu="L40S:1", image=image, timeout=900)
 def run_unit_test():
-    run_test("pytest ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ -vs")
+    run_test("pytest ./tests/dataset/ ./tests/workflow/ -vs")

tests/nightly/test_e2e_dmd_t2v_crush_smol.py

@@ -3,7 +3,7 @@ from pathlib import Path
 from huggingface_hub import snapshot_download
 import subprocess
 import sys
-from fastvideo.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision
+from tests.ssim.test_inference_similarity import compute_video_ssim_torchvision
 import shutil
 
 # Import the training pipeline
@@ -181,4 +181,4 @@ def test_e2e_overfit_single_sample():
 
 
 if __name__ == "__main__":
-    test_e2e_overfit_single_sample()
+    test_e2e_overfit_single_sample()

tests/nightly/test_e2e_i2v_overfit_single_sample.py

@@ -4,7 +4,7 @@ from huggingface_hub import snapshot_download
 import shutil
 import subprocess
 import sys
-from fastvideo.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision
+from tests.ssim.test_inference_similarity import compute_video_ssim_torchvision
 
 # Import the training pipeline
 sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
@@ -169,4 +169,4 @@ def test_e2e_overfit_single_sample():
 
 
 if __name__ == "__main__":
-    test_e2e_overfit_single_sample()
+    test_e2e_overfit_single_sample()

tests/nightly/test_e2e_overfit_single_sample.py

@@ -4,7 +4,7 @@ from huggingface_hub import snapshot_download
 import shutil
 import subprocess
 import sys
-from fastvideo.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision
+from tests.ssim.test_inference_similarity import compute_video_ssim_torchvision
 
 # Import the training pipeline
 sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
@@ -174,4 +174,4 @@ def test_e2e_overfit_single_sample():
 
 
 if __name__ == "__main__":
-    test_e2e_overfit_single_sample()
+    test_e2e_overfit_single_sample()

tests/old_tests/distributed_run.sh

@@ -10,10 +10,10 @@ torchrun --standalone --nnodes=1 --nproc_per_node=$num_gpus \
 num_gpus=2
 torchrun --standalone --nnodes=1 --nproc_per_node=$num_gpus \
     --master_port 29503  \
-    fastvideo/tests/test_hunyuanvideo_load.py --sequence_model_parallel_size $num_gpus 
+    tests/test_hunyuanvideo_load.py --sequence_model_parallel_size $num_gpus 
 
-torchrun --nnodes=1 --nproc_per_node=1 --master_port 29503  fastvideo/tests/test_llama_encoder.py
+torchrun --nnodes=1 --nproc_per_node=1 --master_port 29503  tests/test_llama_encoder.py
 
 
 export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
-torchrun --nnodes=1 --nproc_per_node=1 --master_port 29503  fastvideo/tests/test_clip_encoder.py
+torchrun --nnodes=1 --nproc_per_node=1 --master_port 29503  tests/test_clip_encoder.py

tests/ssim/README.md

@@ -2,7 +2,7 @@ The reference videos in the `*_reference_videos` directory are used as part of a
 
 `A40_reference_videos` are generated on A40s and so on.
 
-run `bash update_reference_videos.sh` from inside the `fastvideo/tests/ssim/` directory after running `test_inference_similarity.py` to update reference videos. Note: make sure to update the path to the corresponding device.
+run `bash update_reference_videos.sh` from inside the `tests/ssim/` directory after running `test_inference_similarity.py` to update reference videos. Note: make sure to update the path to the corresponding device.
 
 reference videos were generated on commit `4aeabbc629e0edf91477e80e795e7bb1823c71cb`
 causal videos were generated on commit b318063c0a4618f1d5d99ea82ca67a06aad0d19d

tests/ssim/test_causal_similarity.py

@@ -7,7 +7,7 @@ import pytest
 
 from fastvideo import VideoGenerator
 from fastvideo.logger import init_logger
-from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results
+from tests.utils import compute_video_ssim_torchvision, write_ssim_results
 from fastvideo.worker.multiproc_executor import MultiprocExecutor
 
 logger = init_logger(__name__)

tests/ssim/test_inference_similarity.py

@@ -7,7 +7,7 @@ import pytest
 
 from fastvideo import VideoGenerator
 from fastvideo.logger import init_logger
-from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results
+from tests.utils import compute_video_ssim_torchvision, write_ssim_results
 from fastvideo.worker.multiproc_executor import MultiprocExecutor
 
 logger = init_logger(__name__)

tests/training/VSA/test_training_loss_VSA.py

@@ -12,7 +12,8 @@ from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.utils import FlexibleArgumentParser
 
 wandb_name = "test_training_loss_VSA"
-reference_wandb_summary_file = "fastvideo/tests/training/VSA/reference_wandb_summary_VSA.json"
+_REFERENCE_DIR = Path(__file__).resolve().parent
+reference_wandb_summary_file = str(_REFERENCE_DIR / "reference_wandb_summary_VSA.json")
 
 NUM_NODES = "1"
 NUM_GPUS_PER_NODE = "2"

tests/training/Vanilla/test_training_loss.py

@@ -16,8 +16,9 @@ from fastvideo.utils import FlexibleArgumentParser
 from fastvideo.training.wan_training_pipeline import WanTrainingPipeline
 
 wandb_name = "test_training_loss"
-a40_reference_wandb_summary_file = "fastvideo/tests/training/Vanilla/a40_reference_wandb_summary.json"
-l40s_reference_wandb_summary_file = "fastvideo/tests/training/Vanilla/l40s_reference_wandb_summary.json"
+_REFERENCE_DIR = Path(__file__).resolve().parent
+a40_reference_wandb_summary_file = str(_REFERENCE_DIR / "a40_reference_wandb_summary.json")
+l40s_reference_wandb_summary_file = str(_REFERENCE_DIR / "l40s_reference_wandb_summary.json")
 
 NUM_NODES = "1"
 NUM_GPUS_PER_NODE = "2"

tests/training/lora/test_lora_training.py

@@ -9,7 +9,8 @@ import json
 from huggingface_hub import snapshot_download
 
 wandb_name = "test_lora_training"
-l40s_reference_wandb_summary_file = "fastvideo/tests/training/lora/l40s_reference_lora_wandb_summary.json"
+_REFERENCE_DIR = Path(__file__).resolve().parent
+l40s_reference_wandb_summary_file = str(_REFERENCE_DIR / "l40s_reference_lora_wandb_summary.json")
 
 NUM_NODES = "1"
 NUM_GPUS_PER_NODE = "2"
@@ -116,4 +117,4 @@ def test_lora_training():
 
 
 if __name__ == "__main__":
-    test_lora_training() 
+    test_lora_training()