zmq worker

.github/workflows/docs.yml

@@ -8,14 +8,12 @@ on:
       - main
     paths:
       - "docs/**/*.md"
-      - "fastvideo/v1/examples/**/*.py"
   pull_request:
     branches: 
       - main
     types: [opened, ready_for_review, synchronize, reopened]
     paths:
       - "docs/**/*.md"
-      - "fastvideo/v1/examples/**/*.py"
 
   # Allows you to run this workflow manually from the Actions tab
   workflow_dispatch:

.github/workflows/pr-test.yml

@@ -116,7 +116,7 @@ jobs:
           --gpu-count 1
           --volume-size 100
           --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
-          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s"
+          --test-command "pip install -e . && pytest ./fastvideo/v1/tests/encoders -s"
 
       - name: Terminate RunPod Instances
         if: ${{ always() }}
@@ -164,7 +164,7 @@ jobs:
           --gpu-count 1
           --volume-size 100
           --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
-          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
+          --test-command "pip install -e . && pytest ./fastvideo/v1/tests/vaes -s"
 
       - name: Terminate RunPod Instances
         if: ${{ always() }}
@@ -212,7 +212,7 @@ jobs:
           --gpu-count 1
           --volume-size 100
           --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
-          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
+          --test-command "pip install -e . && pytest ./fastvideo/v1/tests/transformers -s"
 
       - name: Terminate RunPod Instances
         if: ${{ always() }}
@@ -261,7 +261,7 @@ jobs:
           --disk-size 200
           --volume-size 200
           --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
-          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/ssim -vs"
+          --test-command "pip install -e . && pytest ./fastvideo/v1/tests/ssim -vs"
 
       - name: Terminate RunPod Instances
         if: ${{ always() }}

.gitignore

@@ -55,7 +55,4 @@ docs/source/getting_started/examples/
 *.pkl
 
 # Reference videos
-!fastvideo/v1/tests/ssim/reference_videos/**/*.mp4
-
-# Static images
-!docs/source/_static/images/**/*.png
+!fastvideo/v1/tests/ssim/reference_videos/**/*.mp4

.pre-commit-config.yaml

@@ -19,7 +19,6 @@ exclude: |
     fastvideo/sample/.*|
     fastvideo/train\.py|
     fastvideo/utils/.*|
-    fastvideo/v1/examples/.*|
     .github/workflows/fastvideo-publish.yml|
     .github/workflows/sta-publish.yml
   )
@@ -41,10 +40,10 @@ repos:
   - id: codespell
     additional_dependencies: ['tomli']
     args: ['--toml', 'pyproject.toml']
-- repo: https://github.com/PyCQA/isort
-  rev: 6.0.1
-  hooks:
-  - id: isort
+# - repo: https://github.com/PyCQA/isort
+#   rev: 0a0b7a830386ba6a31c2ec8316849ae4d1b8240d # 6.0.0
+#   hooks:
+#   - id: isort
 - repo: https://github.com/jackdewinter/pymarkdown
   rev: v0.9.29
   hooks:

Dockerfile

@@ -2,7 +2,7 @@ FROM nvidia/cuda:12.4.1-devel-ubuntu20.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 
-WORKDIR /FastVideo
+WORKDIR /app
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
     wget \
@@ -34,15 +34,4 @@ RUN conda run -n fastvideo-dev pip install --no-cache-dir --upgrade pip && \
 
 COPY . .
 
-RUN conda run -n fastvideo-dev pip install --no-cache-dir -e .[dev]
-
-# Remove authentication headers
-RUN git config --unset-all http.https://github.com/.extraheader || true
-
-# Set up automatic conda environment activation for all shells
-RUN echo 'source /opt/conda/etc/profile.d/conda.sh' >> /root/.bashrc && \
-    echo 'conda activate fastvideo-dev' >> /root/.bashrc && \
-    # Ensure .bashrc is sourced for SSH login shells
-    echo 'if [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
-
 EXPOSE 22

README.md

@@ -29,7 +29,7 @@ Dev in progress and highly experimental.
 - ```2024/12/17```: `FastVideo` v1.0 is released.
 
 ## 🔧 Installation from source
-The code is tested on Python 3.10-3.12, CUDA 12.4 and H100.
+The code is tested on Python 3.10.0, CUDA 12.4 and H100.
 
 ```
 # Clone FastVideo
@@ -47,7 +47,7 @@ To try Sliding Tile Attention (optional), please follow the instruction in [csrc
 You can also install the Sliding Tile Attention package using
 
 ```
-pip install st_attn==0.0.4
+pip install st_attn==0.0.3
 ```
 
 ## 🚀 Inference

csrc/sliding_tile_attention/setup.py

@@ -9,7 +9,7 @@ target = target.lower()
 
 # Package metadata
 PACKAGE_NAME = "st_attn"
-VERSION = "0.0.4"
+VERSION = "0.0.3"
 AUTHOR = "Hao AI Lab"
 DESCRIPTION = "Sliding Tile Atteniton Kernel Used in FastVideo"
 URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/sliding_tile_attention"

csrc/sliding_tile_attention/st_attn.cpp

@@ -10,7 +10,7 @@
 
 #ifdef TK_COMPILE_ATTN
 extern torch::Tensor sta_forward(
-    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_w_size, int kernel_h_size, int text_length, bool process_text, bool has_text, int kernel_aspect_ratio_flag
+    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_w_size, int kernel_h_size, int text_length, bool process_text, bool has_text
 ); 
 #endif
 

csrc/sliding_tile_attention/st_attn/__init__.py

@@ -4,13 +4,8 @@ import torch
 from st_attn_cuda import sta_fwd
 
 
-def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_text=True, img_latent_shape='30*48*80'):
+def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_text=True):
     seq_length = q_all.shape[2]
-    img_latent_shape_mapping = {
-        '30x48x80':1,
-        '36x48x48':2,
-        '18x48x80':3, 
-    }
     if has_text:
         assert q_all.shape[
             2] >= 115200, "STA currently only supports video with latent size (30, 48, 80), which is 117 frames x 768 x 1280 pixels"
@@ -22,14 +17,8 @@ def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_te
             k_all = torch.cat([k_all, k_all[:, :, -pad_size:]], dim=2)
             v_all = torch.cat([v_all, v_all[:, :, -pad_size:]], dim=2)
     else:
-        if img_latent_shape == '36x48x48': # Stepvideo 204x768x68
-            assert q_all.shape[2] == 82944
-        elif img_latent_shape == '18x48x80': # Wan 69x768x1280
-            assert q_all.shape[2] == 69120
-        else:
-            raise ValueError(f"Unsupported {img_latent_shape}, current shape is {q_all.shape}, only support '36x48x48' for Stepvideo and '18x48x80' for Wan")
+        assert q_all.shape[2] == 82944
 
-    kernel_aspect_ratio_flag = img_latent_shape_mapping[img_latent_shape]
     hidden_states = torch.empty_like(q_all)
     # This for loop is ugly. but it is actually quite efficient. The sequence dimension alone can already oversubscribe SMs
     for head_index, (t_kernel, h_kernel, w_kernel) in enumerate(window_size):
@@ -40,7 +29,7 @@ def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_te
                                                                                       head_index:head_index + 1],
                                               hidden_states[batch:batch + 1, head_index:head_index + 1])
 
-            _ = sta_fwd(q_head, k_head, v_head, o_head, t_kernel, h_kernel, w_kernel, text_length, False, has_text, kernel_aspect_ratio_flag)
+            _ = sta_fwd(q_head, k_head, v_head, o_head, t_kernel, h_kernel, w_kernel, text_length, False, has_text)
     if has_text:
-        _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, True, True, kernel_aspect_ratio_flag)
+        _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, True, True)
     return hidden_states[:, :, :seq_length]

csrc/sliding_tile_attention/st_attn/st_attn_h100.cu

@@ -359,7 +359,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
 #include <iostream>
 
 torch::Tensor 
-sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_h_size, int kernel_w_size, int text_length, bool process_text, bool has_text, int kernel_aspect_ratio_flag)
+sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_h_size, int kernel_w_size, int text_length, bool process_text, bool has_text)
 {
     CHECK_INPUT(q);
     CHECK_INPUT(k);
@@ -558,267 +558,123 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
 
         } else {
             dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
-            if (kernel_aspect_ratio_flag == 2){
-                if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-                }  else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 6) {
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 1, 3, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false,1, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-                } else if (kernel_t_size == 6 && kernel_h_size == 3 && kernel_w_size == 3) {
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            }  else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 6) {
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 1, 1, 3, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false,1, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-                } else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 6){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            } else if (kernel_t_size == 6 && kernel_h_size == 3 && kernel_w_size == 3) {
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-                }else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 3){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            } else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 6){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-                } else if (kernel_t_size ==6 && kernel_h_size == 3 && kernel_w_size == 6){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            }else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 3){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-                } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 6){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 1){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                }  else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 6){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                }  else if (kernel_t_size == 1 && kernel_h_size == 1 && kernel_w_size == 6){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 1){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else {
-                    // print error
-                    std::cout << "Invalid kernel size" << std::endl;
-                    //print kernel size
-                    std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl;
-                }
-            }
-            else if (kernel_aspect_ratio_flag == 3) {
-                if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
+            } else if (kernel_t_size ==6 && kernel_h_size == 3 && kernel_w_size == 6){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 1, 1, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 1, 1, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-
-                } else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 5) {
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 1, 2, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false,1, 1, 2, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-
-                } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 5){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 2, 2, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 2, 2, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-
-                } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 1){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 3, 0, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 3, 0, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-
-                } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 7){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 2, 3, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 2, 3, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 9){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 2, 4, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 2, 4, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 10){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 3, 5, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 3, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 3){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 3, 1, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 3, 1, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 3 && kernel_h_size == 1 && kernel_w_size == 1){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 0, 0, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 0, 0, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 10){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 3, 5, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 0, 3, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 10){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 2, 5, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 0, 2, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 7){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 3, 3, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 0, 3, 3, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 7){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 2, 3, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 0, 2, 3, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 9){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 2, 4, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 0, 2, 4, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 3 && kernel_h_size == 1 && kernel_w_size == 10){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 0, 5, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false,1, 0, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 10){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 1, 5, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false,1, 1, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 3 && kernel_w_size == 10){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 1, 5, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false,0, 1, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 5){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 3, 2, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false,0, 3, 2, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else {
-                    // print error
-                    std::cout << "Invalid kernel size" << std::endl;
-                    //print kernel size
-                    std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl;
-                }
-            }
-
-            else {
-                std::cout << "Unsupported kernel_aspect_ratio_flag: " << kernel_aspect_ratio_flag << std::endl;
+            } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 6){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 1){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            }  else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 6){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            }  else if (kernel_t_size == 1 && kernel_h_size == 1 && kernel_w_size == 6){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 1){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            } else {
+                // print error
+                std::cout << "Invalid kernel size" << std::endl;
+                //print kernel size
+                std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl;
             }
 
         }

docs/source/developer_guide/overview.md

@@ -50,89 +50,3 @@ pre-commit run --all-files
 # Unit tests
 pytest tests/
 ```
-
----
-## 🐳 Using the FastVideo Docker Image
-
-If you prefer a containerized development environment or want to avoid managing dependencies manually, you can use our prebuilt Docker image:
-
-**Image:** [`ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest`](https://ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev)
-
-### Starting the container
-
-```bash
-docker run --gpus all -it ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest
-```
-
-This will:
-
-- Start the container with GPU access  
-- Drop you into a shell with the `fastvideo-dev` Conda environment preconfigured
-
-### Using the container
-
-```bash
-# Conda environment should already be active
-# FastVideo package installed in editable mode
-
-# Pull the latest changes from remote
-cd /FastVideo
-git pull
-
-# Run linters and tests
-pre-commit run --all-files
-pytest tests/
-```
-
----
-
-## 📦 Developing FastVideo on RunPod
-
-You can easily use the FastVideo Docker image as a custom container on [RunPod](https://www.runpod.io) for development or experimentation.
-
-### Creating a new pod
-
-Choose a GPU that supports CUDA 12.4
-
-![RunPod CUDA selection](../_static/images/runpod_cuda.png)
-
-When creating your pod template, use this image:
-
-```
-ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest
-```
-
-Paste Container Start Command to support SSH ([RunPod Docs](https://docs.runpod.io/pods/configuration/use-ssh)):
-
-```bash
-bash -c "apt update;DEBIAN_FRONTEND=noninteractive apt-get install openssh-server -y;mkdir -p ~/.ssh;cd $_;chmod 700 ~/.ssh;echo \"$PUBLIC_KEY\" >> authorized_keys;chmod 700 authorized_keys;service ssh start;sleep infinity"
-```
-
-![RunPod template configuration](../_static/images/runpod_template.png)
-
-After deploying, the pod will take a few minutes to pull the image and start the SSH service.
-
-![RunPod ssh](../_static/images/runpod_ssh.png)
-
-### Working with the pod
-
-After SSH'ing into your pod, you'll find the `fastvideo-dev` Conda environment already activated.
-
-To pull in the latest changes from the GitHub repo:
-
-```bash
-cd /FastVideo
-git pull
-```
-
-`If you have a persistent volume and want to keep your code changes, you can move /FastVideo to /workspace/FastVideo, or simply clone the repository there.`
-
-Run your development workflows as usual:
-
-```bash
-# Run linters
-pre-commit run --all-files
-
-# Run tests
-pytest tests/
-```

docs/source/generate_examples.py

@@ -9,7 +9,7 @@ from typing import Optional
 
 ROOT_DIR = Path(__file__).parent.parent.parent.resolve()
 ROOT_DIR_RELATIVE = '../../../..'
-EXAMPLE_DIR = ROOT_DIR / "fastvideo/v1/examples"
+EXAMPLE_DIR = ROOT_DIR / "examples"
 EXAMPLE_DOC_DIR = ROOT_DIR / "docs/source/getting_started/examples"
 
 
@@ -208,7 +208,6 @@ def generate_examples():
     glob_patterns = ["*.py", "*.md", "*.sh"]
     # Find categorised examples
     for category in category_indices:
-        print(category)
         category_dir = EXAMPLE_DIR / category
         globs = [category_dir.glob(pattern) for pattern in glob_patterns]
         for path in itertools.chain(*globs):
@@ -229,7 +228,6 @@ def generate_examples():
 
     # Generate the example documentation
     for example in sorted(examples, key=lambda e: e.path.stem):
-        print(example)
         doc_path = EXAMPLE_DOC_DIR / f"{example.path.stem}.md"
         with open(doc_path, "w+") as f:
             f.write(example.generate())

docs/source/getting_started/installation.md

@@ -2,20 +2,12 @@
 
 # 🔧 Installation
 
-FastVideo currently only supports Linux and NVIDIA CUDA GPUs.
+FastVideo currently only supports Linux and CUDA GPUs. The code is tested on Python 3.10.0 and CUDA 12.4, primarily with NVIDIA H100 GPUs.
 
-FastVideo has been tested on the following GPUs, but it should work on any GPUs that supports CUDA 12.4+, please create an issue if you discover any issues:
-- RTX 4090
-- A40
-- L40S
-- A100
-- H100
+## Prerequisites
 
-## Requirements
-
-- OS: Linux
-- Python: 3.10-3.12
-- CUDA 12.4+ (Untested on CUDA < 12.4)
+- CUDA 12.4 installed and supported
+- Linux operating system
 
 ## Installation Options
 
@@ -27,9 +19,7 @@ pip install fastvideo
 
 ### Option 2: Installation from Source
 
-We recommend using a Python environment such as Conda.
-
-#### 1. [Optional] Install Miniconda (if not already installed)
+#### 1. Install Miniconda (if not already installed)
 
 ```bash
 wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
@@ -37,7 +27,7 @@ bash Miniconda3-latest-Linux-x86_64.sh
 source ~/.bashrc
 ```
 
-#### 2. [Optional] Create and activate a Conda environment for FastVideo
+#### 2. Create and activate a Conda environment for FastVideo
 
 ```bash
 conda create -n fastvideo python=3.10 -y
@@ -66,7 +56,7 @@ pip install -e .
 pip install flash-attn==2.7.0.post2 --no-build-isolation
 ```
 
-### Sliding Tile Attention (STA) (Requires CUDA 12.4+ and H100)
+### Sliding Tile Attention (STA)
 
 To try Sliding Tile Attention (optional), please follow the instructions in [csrc/sliding_tile_attention/README.md](#sta-installation) to install STA.
 

docs/source/index.md

@@ -69,7 +69,6 @@ sliding_tile_attention/demo
 :caption: Inference
 :maxdepth: 1
 
-inference/wanvideo
 inference/stepvideo
 inference/hunyuanvideo
 inference/fasthunyuan

fastvideo/v1/attention/backends/sliding_tile_attn.py

@@ -20,39 +20,20 @@ logger = init_logger(__name__)
 
 
 # TODO(will-refactor): move this to a utils file
-def dict_to_3d_list(mask_strategy) -> List[List[List[Optional[torch.Tensor]]]]:
-    indices = [tuple(map(int, key.split('_'))) for key in mask_strategy]
-
-    max_timesteps_idx = max(
-        timesteps_idx for timesteps_idx, layer_idx, head_idx in indices) + 1
-    max_layer_idx = max(layer_idx
-                        for timesteps_idx, layer_idx, head_idx in indices) + 1
-    max_head_idx = max(head_idx
-                       for timesteps_idx, layer_idx, head_idx in indices) + 1
-
-    result = [[[None for _ in range(max_head_idx)]
-               for _ in range(max_layer_idx)] for _ in range(max_timesteps_idx)]
-
+def dict_to_3d_list(mask_strategy,
+                    t_max=50,
+                    l_max=60,
+                    h_max=24) -> List[List[List[Optional[torch.Tensor]]]]:
+    result = [[[None for _ in range(h_max)] for _ in range(l_max)]
+              for _ in range(t_max)]
+    if mask_strategy is None:
+        return result
     for key, value in mask_strategy.items():
-        timesteps_idx, layer_idx, head_idx = map(int, key.split('_'))
-        result[timesteps_idx][layer_idx][head_idx] = value
-
+        t, layer, h = map(int, key.split('_'))
+        result[t][layer][h] = value
     return result
 
 
-class RangeDict(dict):
-
-    def __getitem__(self, item):
-        for key in self.keys():
-            if isinstance(key, tuple):
-                low, high = key
-                if low <= item <= high:
-                    return super().__getitem__(key)
-            elif key == item:
-                return super().__getitem__(key)
-        raise KeyError(f"seq_len {item} not supported for STA")
-
-
 class SlidingTileAttentionBackend(AttentionBackend):
 
     accept_output_buffer: bool = True
@@ -122,73 +103,52 @@ class SlidingTileAttentionImpl(AttentionImpl):
 
         with open(config_file) as f:
             mask_strategy = json.load(f)
-        mask_strategy = dict_to_3d_list(mask_strategy)
 
+        mask_strategy = dict_to_3d_list(mask_strategy)
         self.prefix = prefix
         self.mask_strategy = mask_strategy
         sp_group = get_sp_group()
         self.sp_size = sp_group.world_size
-        # STA config
-        self.STA_base_tile_size = [6, 8, 8]
-        self.img_latent_shape_mapping = RangeDict({
-            (115200, 115456): '30x48x80',
-            82944: '36x48x48',
-            69120: '18x48x80',
-        })
-        self.full_window_mapping = {
-            '30x48x80': [5, 6, 10],
-            '36x48x48': [6, 6, 6],
-            '18x48x80': [3, 6, 10]
-        }
 
     def tile(self, x: torch.Tensor) -> torch.Tensor:
         x = rearrange(x,
                       "b (sp t h w) head d -> b (t sp h w) head d",
                       sp=self.sp_size,
-                      t=self.img_latent_shape_int[0] // self.sp_size,
-                      h=self.img_latent_shape_int[1],
-                      w=self.img_latent_shape_int[2])
+                      t=30 // self.sp_size,
+                      h=48,
+                      w=80)
         return rearrange(
             x,
             "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
-            n_t=self.full_window_size[0],
-            n_h=self.full_window_size[1],
-            n_w=self.full_window_size[2],
-            ts_t=self.STA_base_tile_size[0],
-            ts_h=self.STA_base_tile_size[1],
-            ts_w=self.STA_base_tile_size[2])
+            n_t=5,
+            n_h=6,
+            n_w=10,
+            ts_t=6,
+            ts_h=8,
+            ts_w=8)
 
     def untile(self, x: torch.Tensor) -> torch.Tensor:
         x = rearrange(
             x,
             "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
-            n_t=self.full_window_size[0],
-            n_h=self.full_window_size[1],
-            n_w=self.full_window_size[2],
-            ts_t=self.STA_base_tile_size[0],
-            ts_h=self.STA_base_tile_size[1],
-            ts_w=self.STA_base_tile_size[2])
+            n_t=5,
+            n_h=6,
+            n_w=10,
+            ts_t=6,
+            ts_h=8,
+            ts_w=8)
         return rearrange(x,
                          "b (t sp h w) head d -> b (sp t h w) head d",
                          sp=self.sp_size,
-                         t=self.img_latent_shape_int[0] // self.sp_size,
-                         h=self.img_latent_shape_int[1],
-                         w=self.img_latent_shape_int[2])
+                         t=30 // self.sp_size,
+                         h=48,
+                         w=80)
 
     def preprocess_qkv(
         self,
         qkv: torch.Tensor,
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
-        img_sequence_length = qkv.shape[1]
-        self.img_latent_shape_str = self.img_latent_shape_mapping[
-            img_sequence_length]
-        self.full_window_size = self.full_window_mapping[
-            self.img_latent_shape_str]
-        self.img_latent_shape_int = list(
-            map(int, self.img_latent_shape_str.split('x')))
-        self.img_seq_length = self.img_latent_shape_int[
-            0] * self.img_latent_shape_int[1] * self.img_latent_shape_int[2]
         return self.tile(qkv)
 
     def postprocess_output(
@@ -213,15 +173,11 @@ class SlidingTileAttentionImpl(AttentionImpl):
         timestep = attn_metadata.current_timestep
         # pattern:'.double_blocks.0.attn.impl' or '.single_blocks.0.attn.impl'
         layer_idx = int(self.prefix.split('.')[-3])
-
         # TODO: remove hardcode
-
-        text_length = q.shape[1] - self.img_seq_length
-        has_text = text_length > 0
-
-        query = q.transpose(1, 2).contiguous()
-        key = k.transpose(1, 2).contiguous()
-        value = v.transpose(1, 2).contiguous()
+        text_length = q.shape[1] - (30 * 48 * 80)
+        query = q.transpose(1, 2)
+        key = k.transpose(1, 2)
+        value = v.transpose(1, 2)
 
         head_num = query.size(1)
         sp_group = get_sp_group()
@@ -231,11 +187,7 @@ class SlidingTileAttentionImpl(AttentionImpl):
             self.mask_strategy[timestep][layer_idx][head_idx + start_head]
             for head_idx in range(head_num)
         ]
-        # if has_text is False:
-        #     from IPython import embed
-        #     embed()
-        hidden_states = sliding_tile_attention(
-            query, key, value, windows, text_length, has_text,
-            self.img_latent_shape_str).transpose(1, 2)
+        hidden_states = sliding_tile_attention(query, key, value, windows,
+                                               text_length).transpose(1, 2)
 
         return hidden_states

fastvideo/v1/attention/layer.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Optional, Tuple
+from typing import List, Optional
 
 import torch
 import torch.nn as nn
@@ -25,8 +25,7 @@ class DistributedAttention(nn.Module):
                  num_kv_heads: Optional[int] = None,
                  softmax_scale: Optional[float] = None,
                  causal: bool = False,
-                 supported_attention_backends: Optional[Tuple[_Backend,
-                                                              ...]] = None,
+                 supported_attention_backends: Optional[List[_Backend]] = None,
                  prefix: str = "",
                  **extra_impl_args) -> None:
         super().__init__()
@@ -147,8 +146,7 @@ class LocalAttention(nn.Module):
                  num_kv_heads: Optional[int] = None,
                  softmax_scale: Optional[float] = None,
                  causal: bool = False,
-                 supported_attention_backends: Optional[Tuple[_Backend,
-                                                              ...]] = None,
+                 supported_attention_backends: Optional[List[_Backend]] = None,
                  **extra_impl_args) -> None:
         super().__init__()
         if softmax_scale is None:

fastvideo/v1/attention/selector.py

@@ -3,8 +3,7 @@
 
 import os
 from contextlib import contextmanager
-from functools import cache
-from typing import Generator, Optional, Tuple, Type, cast
+from typing import Generator, List, Optional, Type, cast
 
 import torch
 
@@ -82,17 +81,7 @@ def get_global_forced_attn_backend() -> Optional[_Backend]:
 def get_attn_backend(
     head_size: int,
     dtype: torch.dtype,
-    supported_attention_backends: Optional[Tuple[_Backend, ...]] = None,
-) -> Type[AttentionBackend]:
-    return _cached_get_attn_backend(head_size, dtype,
-                                    supported_attention_backends)
-
-
-@cache
-def _cached_get_attn_backend(
-    head_size: int,
-    dtype: torch.dtype,
-    supported_attention_backends: Optional[Tuple[_Backend, ...]] = None,
+    supported_attention_backends: Optional[List[_Backend]] = None,
 ) -> Type[AttentionBackend]:
     # Check whether a particular choice of backend was
     # previously forced.

fastvideo/v1/configs/__init__.py

@@ -0,0 +1,9 @@
+from fastvideo.v1.configs.hunyuan import HunyuanConfig, FastHunyuanConfig
+from fastvideo.v1.configs.wan import WanT2V480PConfig, WanI2V480PConfig
+from fastvideo.v1.configs.base import BaseConfig, SlidingTileAttnConfig
+from fastvideo.v1.configs.registry import get_pipeline_config_cls_for_name
+
+__all__ = [
+    "HunyuanConfig", "FastHunyuanConfig", "BaseConfig", "SlidingTileAttnConfig",
+    "WanT2V480PConfig", "WanI2V480PConfig", "get_pipeline_config_cls_for_name"
+]

fastvideo/v1/configs/base.py

@@ -0,0 +1,67 @@
+from dataclasses import dataclass, field
+from typing import Optional, Dict, Any
+
+
+@dataclass
+class BaseConfig:
+    """Base configuration for all pipeline architectures."""
+
+    # Video parameters
+    height: int = 720
+    width: int = 1280
+    num_frames: int = 125
+    fps: int = 24
+
+    # Video generation parameters
+    num_inference_steps: int = 50
+    guidance_scale: float = 1.0
+    seed: int = 1024
+    guidance_rescale: float = 0.0
+    embedded_cfg_scale: float = 6.0
+    flow_shift: Optional[float] = None
+    use_cpu_offload: bool = False
+    disable_autocast: bool = False
+
+    # Model configuration
+    precision: str = "bf16"
+
+    # VAE configuration
+    vae_precision: str = "fp16"
+    vae_tiling: bool = True
+    vae_sp: bool = True
+    vae_scale_factor: Optional[int] = None
+
+    # DiT configuration
+    num_channels_latents: Optional[int] = None
+
+    # Image encoder configuration
+    image_encoder_precision: str = "fp32"
+
+    # Text encoder configuration
+    text_encoder_precision: str = "fp16"
+    text_len: int = -1
+    hidden_state_skip_layer: int = 0
+
+    # STA (Spatial-Temporal Attention) parameters
+    mask_strategy_file_path: Optional[str] = None
+    enable_torch_compile: bool = False
+
+    neg_prompt: Optional[str] = None
+
+
+@dataclass
+class SlidingTileAttnConfig(BaseConfig):
+    """Configuration for sliding tile attention."""
+
+    # Override any BaseConfig defaults as needed
+    # Add sliding tile specific parameters
+    window_size: int = 16
+    stride: int = 8
+
+    # You can provide custom defaults for inherited fields
+    height: int = 576
+    width: int = 1024
+
+    # Additional configuration specific to sliding tile attention
+    pad_to_square: bool = False
+    use_overlap_optimization: bool = True

fastvideo/v1/configs/hunyuan.py

@@ -1,16 +1,12 @@
 from dataclasses import dataclass
-from fastvideo.v1.configs.pipelines.base import BaseConfig
+from fastvideo.v1.configs.base import BaseConfig
 
-from fastvideo.v1.configs.models import VAEConfig
-from fastvideo.v1.configs.models.vaes import HunyuanVAEConfig
 
 @dataclass
 class HunyuanConfig(BaseConfig):
     """Base configuration for HunYuan pipeline architecture."""
 
     # HunyuanConfig-specific parameters with defaults
-    # VAE
-    vae_config: VAEConfig = HunyuanVAEConfig()
     # Denoising stage
     embedded_cfg_scale: int = 6
     flow_shift: int = 7
@@ -30,10 +26,6 @@ class HunyuanConfig(BaseConfig):
     text_encoder_precision_2: str = "fp16"
     text_len_2: int = 77
 
-    def __post_init__(self):
-        self.vae_config.load_encoder = False
-        self.vae_config.load_decoder = True
-
 
 @dataclass
 class FastHunyuanConfig(HunyuanConfig):

fastvideo/v1/configs/models/__init__.py

@@ -1,7 +0,0 @@
-from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
-from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
-
-__all__ = [
-    "ArchConfig", "ModelConfig",
-    "VAEArchConfig", "VAEConfig"
-]

fastvideo/v1/configs/models/base.py

@@ -1,47 +0,0 @@
-from dataclasses import dataclass, fields
-from typing import Dict, Any
-
-# 1. ArchConfig contains all fields from diffuser's/transformer's config.json (i.e. all fields related to the architecture of the model)
-# 2. ArchConfig should be inherited & overriden by each model arch_config
-# 3. Any field in ArchConfig is fixed upon initialization, and should be hidden away from users
-@dataclass
-class ArchConfig:
-    pass
-
-@dataclass
-class ModelConfig:
-    # Every model config parameter can be categorized into either ArchConfig or everything else
-    # Diffuser/Transformer parameters
-    arch_config: ArchConfig = ArchConfig()
-
-    # FastVideo-specific parameters here
-    # i.e. STA, quantization, teacache
-
-    # This should be used only when loading from transformers/diffusers
-    def update_model_arch(
-        self, 
-        source_model_dict: Dict[str, Any]
-    ) -> None:
-        arch_config = self.arch_config
-        valid_fields = {f.name for f in fields(arch_config)}
-        
-        for key, value in source_model_dict.items():
-            if key in valid_fields:
-                setattr(arch_config, key, value)
-            else:
-                raise AttributeError(f"{type(arch_config).__name__} has no field '{key}'")
-
-    def update_model_config(
-        self,
-        source_model_dict: Dict[str, Any]    
-    ) -> None:
-        assert "arch_config" not in source_model_dict.keys(), "Source model config shouldn't contain arch_config."
-        
-        valid_fields = {f.name for f in fields(self)}
-
-        for key, value in source_model_dict.items():
-            if key in valid_fields:
-                setattr(self, key, value)
-            else:
-                print(f"{type(self).__name__} does not contain field '{key}'!")
-                raise AttributeError(f"Invalid field: {key}")

fastvideo/v1/configs/models/vaes/__init__.py

@@ -1,7 +0,0 @@
-from fastvideo.v1.configs.models.vaes.hunyuanvae import HunyuanVAEConfig, HunyuanVAEArchConfig
-from fastvideo.v1.configs.models.vaes.wanvae import WanVAEConfig, WanVAEArchConfig
-
-__all__ = [
-    "HunyuanVAEConfig", "HunyuanVAEArchConfig", 
-    "WanVAEConfig", "WanVAEArchConfig"
-]

fastvideo/v1/configs/models/vaes/base.py

@@ -1,36 +0,0 @@
-from dataclasses import dataclass
-from typing import Union
-
-import torch
-
-from fastvideo.v1.configs.models import ArchConfig, ModelConfig
-
-@dataclass
-class VAEArchConfig(ArchConfig):
-    scaling_factor: Union[float, torch.tensor] = 0
-
-    temporal_compression_ratio: int = 4
-    spatial_compression_ratio: int = 8
-
-@dataclass
-class VAEConfig(ModelConfig):
-    arch_config: VAEArchConfig = VAEArchConfig()
-
-    # FastVideoVAE-specific parameters
-    load_encoder: bool = True
-    load_decoder: bool = True
-
-    tile_sample_min_height: int = 256
-    tile_sample_min_width: int = 256
-    tile_sample_min_num_frames: int = 16
-    tile_sample_stride_height: int = 192
-    tile_sample_stride_width: int = 192
-    tile_sample_stride_num_frames: int = 12
-    blend_num_frames: int = 0
-
-    use_tiling: bool = True
-    use_temporal_tiling: bool = True
-    use_parallel_tiling: bool = True
-
-    def __post_init__(self):
-        self.blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames

fastvideo/v1/configs/models/vaes/hunyuanvae.py

@@ -1,37 +0,0 @@
-from dataclasses import dataclass
-from typing import Tuple
-
-from fastvideo.v1.configs.models.vaes.base import VAEConfig, VAEArchConfig
-
-@dataclass
-class HunyuanVAEArchConfig(VAEArchConfig):
-    in_channels: int = 3
-    out_channels: int = 3
-    latent_channels: int = 16
-    down_block_types: Tuple[str, ...] = (
-        "HunyuanVideoDownBlock3D",
-        "HunyuanVideoDownBlock3D",
-        "HunyuanVideoDownBlock3D",
-        "HunyuanVideoDownBlock3D",
-    )
-    up_block_types: Tuple[str, ...] = (
-        "HunyuanVideoUpBlock3D",
-        "HunyuanVideoUpBlock3D",
-        "HunyuanVideoUpBlock3D",
-        "HunyuanVideoUpBlock3D",
-    )
-    block_out_channels: Tuple[int, ...] = (128, 256, 512, 512)
-    layers_per_block: int = 2
-    act_fn: str = "silu"
-    norm_num_groups: int = 32
-    scaling_factor: float = 0.476986
-    spatial_compression_ratio: int = 8
-    temporal_compression_ratio: int = 4
-    mid_block_add_attention: bool = True
-
-    def __post_init__(self):
-        self.spatial_compression_ratio: int = 2**(len(self.block_out_channels)-1)
-
-@dataclass
-class HunyuanVAEConfig(VAEConfig):
-    arch_config: VAEArchConfig = HunyuanVAEArchConfig()

fastvideo/v1/configs/models/vaes/wanvae.py

@@ -1,72 +0,0 @@
-from dataclasses import dataclass
-from typing import Tuple
-
-import torch
-
-from fastvideo.v1.configs.models.vaes.base import VAEConfig, VAEArchConfig
-
-@dataclass
-class WanVAEArchConfig(VAEArchConfig):
-    base_dim: int = 96
-    z_dim: int = 16
-    dim_mult: Tuple[int, ...] = (1, 2, 4, 4)
-    num_res_blocks: int = 2
-    attn_scales: Tuple[float, ...] = ()
-    temperal_downsample: Tuple[bool, ...] = (False, True, True)
-    dropout: float = 0.0
-    latents_mean: Tuple[float, ...] = (
-        -0.7571,
-        -0.7089,
-        -0.9113,
-        0.1075,
-        -0.1745,
-        0.9653,
-        -0.1517,
-        1.5508,
-        0.4134,
-        -0.0715,
-        0.5517,
-        -0.3632,
-        -0.1922,
-        -0.9497,
-        0.2503,
-        -0.2921,
-    )
-    latents_std: Tuple[float, ...] = (
-        2.8184,
-        1.4541,
-        2.3275,
-        2.6558,
-        1.2196,
-        1.7708,
-        2.6052,
-        2.0743,
-        3.2687,
-        2.1526,
-        2.8652,
-        1.5579,
-        1.6382,
-        1.1253,
-        2.8251,
-        1.9160,
-    )
-    temporal_compression_ratio = 4
-    spatial_compression_ratio = 8
-    
-    def __post_init__(self):
-        self.scaling_factor: torch.tensor = 1.0 / torch.tensor(self.latents_std).view(
-            1, self.z_dim, 1, 1, 1)
-        self.shift_factor: torch.tensor = torch.tensor(self.latents_mean).view(
-            1, self.z_dim, 1, 1, 1)
-
-@dataclass
-class WanVAEConfig(VAEConfig):
-    arch_config: VAEArchConfig = WanVAEArchConfig()
-    use_feature_cache: bool = True
-
-    use_tiling: bool = False
-    use_temporal_tiling: bool = False
-    use_parallel_tiling: bool = False
-
-    def __post_init__(self):
-        self.blend_num_frames = (self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames) * 2

fastvideo/v1/configs/pipelines/__init__.py

@@ -1,9 +0,0 @@
-from fastvideo.v1.configs.pipelines.hunyuan import HunyuanConfig, FastHunyuanConfig
-from fastvideo.v1.configs.pipelines.wan import WanT2V480PConfig, WanI2V480PConfig
-from fastvideo.v1.configs.pipelines.base import BaseConfig, SlidingTileAttnConfig
-from fastvideo.v1.configs.pipelines.registry import get_pipeline_config_cls_for_name
-
-__all__ = [
-    "HunyuanConfig", "FastHunyuanConfig", "BaseConfig", "SlidingTileAttnConfig",
-    "WanT2V480PConfig", "WanI2V480PConfig", "get_pipeline_config_cls_for_name"
-]

fastvideo/v1/configs/pipelines/base.py

@@ -1,103 +0,0 @@
-from dataclasses import dataclass, asdict, fields
-from typing import Optional, Dict, Any
-import json
-
-from fastvideo.v1.configs.models import ModelConfig, VAEConfig
-from fastvideo.v1.utils import shallow_asdict
-
-@dataclass
-class BaseConfig:
-    """Base configuration for all pipeline architectures."""
-
-    # Video parameters
-    height: int = 720
-    width: int = 1280
-    num_frames: int = 125
-    fps: int = 24
-
-    # Video generation parameters
-    num_inference_steps: int = 50
-    guidance_scale: float = 1.0
-    seed: int = 1024
-    guidance_rescale: float = 0.0
-    embedded_cfg_scale: float = 6.0
-    flow_shift: Optional[float] = None
-    use_cpu_offload: bool = False
-    disable_autocast: bool = False
-
-    # Model configuration
-    precision: str = "bf16"
-
-    # VAE configuration
-    vae_precision: str = "fp16"
-    vae_tiling: bool = True
-    vae_sp: bool = True
-    # vae_scale_factor: Optional[int] = None # Deprecated
-    vae_config: VAEConfig = VAEConfig()
-
-    # DiT configuration
-    num_channels_latents: Optional[int] = None # Deprecated
-
-    # Image encoder configuration
-    image_encoder_precision: str = "fp32"
-
-    # Text encoder configuration
-    text_encoder_precision: str = "fp16"
-    text_len: int = -1 # Deprecated
-    hidden_state_skip_layer: int = 0 # Deprecated
-
-    # STA (Spatial-Temporal Attention) parameters
-    mask_strategy_file_path: Optional[str] = None
-    enable_torch_compile: bool = False
-
-    neg_prompt: Optional[str] = None
-
-    def dump_to_json(self, file_path: str):
-        output_dict = shallow_asdict(self)
-        for key, value in output_dict.items():
-            if isinstance(value, ModelConfig):
-                model_dict = asdict(value)
-                # Model Arch Config should be hidden away from the users
-                model_dict.pop("arch_config")
-                output_dict[key] = model_dict
-        
-        with open(file_path, "w") as f:
-            json.dump(output_dict, f, indent=2)
-
-    def load_from_json(self, file_path: str):
-        with open(file_path, "r") as f:
-            input_pipeline_dict = json.load(f)
-        self.update_pipeline_config(input_pipeline_dict)
-
-    def update_pipeline_config(
-        self,
-        source_pipeline_dict: Dict[str, Any]    
-    ) -> None:
-        for f in fields(self):
-            key = f.name
-            if key in source_pipeline_dict:
-                current_value = getattr(self, key)
-                new_value = source_pipeline_dict[key]
-
-                # If it's a nested ModelConfig, update it recursively
-                if isinstance(current_value, ModelConfig):
-                    current_value.update_model_config(new_value)
-                else:
-                    setattr(self, key, new_value)
-
-@dataclass
-class SlidingTileAttnConfig(BaseConfig):
-    """Configuration for sliding tile attention."""
-
-    # Override any BaseConfig defaults as needed
-    # Add sliding tile specific parameters
-    window_size: int = 16
-    stride: int = 8
-
-    # You can provide custom defaults for inherited fields
-    height: int = 576
-    width: int = 1024
-
-    # Additional configuration specific to sliding tile attention
-    pad_to_square: bool = False
-    use_overlap_optimization: bool = True

fastvideo/v1/configs/registry.py

@@ -1,15 +1,14 @@
 """Registry for pipeline weight-specific configurations."""
 
 import os
-from typing import Callable, Dict, Optional, Type
+from typing import Dict, Type, Optional, Callable
 
-from fastvideo.v1.configs.pipelines.base import BaseConfig
-from fastvideo.v1.configs.pipelines.hunyuan import HunyuanConfig, FastHunyuanConfig
-from fastvideo.v1.configs.pipelines.wan import WanT2V480PConfig, WanI2V480PConfig
+from fastvideo.v1.configs.base import BaseConfig
+from fastvideo.v1.configs.hunyuan import HunyuanConfig, FastHunyuanConfig
+from fastvideo.v1.configs.wan import WanT2V480PConfig, WanI2V480PConfig
 
+from fastvideo.v1.utils import maybe_download_model_index, verify_model_config_and_directory
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.utils import (maybe_download_model_index,
-                                verify_model_config_and_directory)
 
 logger = init_logger(__name__)
 
@@ -42,7 +41,7 @@ PIPELINE_FALLBACK_CONFIG: Dict[str, Type[BaseConfig]] = {
 
 
 def get_pipeline_config_cls_for_name(
-        pipeline_name_or_path: str) -> Optional[type[BaseConfig]]:
+        pipeline_name_or_path: str) -> Optional[BaseConfig]:
     """Get the appropriate config class for specific pretrained weights."""
 
     if os.path.exists(pipeline_name_or_path):

fastvideo/v1/configs/wan.py

@@ -1,19 +1,12 @@
 from dataclasses import dataclass
-from fastvideo.v1.configs.pipelines.base import BaseConfig
+from fastvideo.v1.configs.base import BaseConfig
 
-from fastvideo.v1.configs.models import VAEConfig
-from fastvideo.v1.configs.models.vaes import WanVAEConfig
 
 @dataclass
 class WanT2V480PConfig(BaseConfig):
     """Base configuration for Wan T2V 1.3B pipeline architecture."""
 
     # WanConfig-specific parameters with defaults
-    # VAE
-    vae_config: VAEConfig = WanVAEConfig()
-    vae_tiling: bool = False
-    vae_sp: bool = False
-
     # Video parameters
     height: int = 480
     width: int = 832
@@ -37,9 +30,6 @@ class WanT2V480PConfig(BaseConfig):
 
     # WanConfig-specific added parameters
 
-    def __post_init__(self):
-        self.vae_config.load_encoder = False
-        self.vae_config.load_decoder = True
 
 @dataclass
 class WanI2V480PConfig(WanT2V480PConfig):
@@ -52,7 +42,3 @@ class WanI2V480PConfig(WanT2V480PConfig):
 
     # Precision for each component
     image_encoder_precision: str = "fp32"
-
-    def __post_init__(self):
-        self.vae_config.load_encoder = True
-        self.vae_config.load_decoder = True

fastvideo/v1/configs/wan_1.3B_t2v_pipeline.json

@@ -1,41 +0,0 @@
-{
-  "height": 480,
-  "width": 832,
-  "num_frames": 81,
-  "fps": 16,
-  "num_inference_steps": 50,
-  "guidance_scale": 3.0,
-  "seed": 1024,
-  "guidance_rescale": 0.0,
-  "embedded_cfg_scale": 6.0,
-  "flow_shift": 3,
-  "use_cpu_offload": true,
-  "disable_autocast": false,
-  "precision": "bf16",
-  "vae_precision": "fp16",
-  "vae_tiling": false,
-  "vae_sp": false,
-  "vae_config": {
-    "load_encoder": false,
-    "load_decoder": true,
-    "tile_sample_min_height": 256,
-    "tile_sample_min_width": 256,
-    "tile_sample_min_num_frames": 16,
-    "tile_sample_stride_height": 192,
-    "tile_sample_stride_width": 192,
-    "tile_sample_stride_num_frames": 12,
-    "blend_num_frames": 8,
-    "use_tiling": false,
-    "use_temporal_tiling": false,
-    "use_parallel_tiling": false,
-    "use_feature_cache": true
-  },
-  "num_channels_latents": null,
-  "image_encoder_precision": "fp32",
-  "text_encoder_precision": "fp32",
-  "text_len": 512,
-  "hidden_state_skip_layer": 0,
-  "mask_strategy_file_path": null,
-  "enable_torch_compile": false,
-  "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
-}

fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json

@@ -1,41 +0,0 @@
-{
-  "height": 480,
-  "width": 832,
-  "num_frames": 81,
-  "fps": 16,
-  "num_inference_steps": 40,
-  "guidance_scale": 5.0,
-  "seed": 1024,
-  "guidance_rescale": 0.0,
-  "embedded_cfg_scale": 6.0,
-  "flow_shift": 3,
-  "use_cpu_offload": true,
-  "disable_autocast": false,
-  "precision": "bf16",
-  "vae_precision": "fp16",
-  "vae_tiling": false,
-  "vae_sp": false,
-  "vae_config": {
-    "load_encoder": true,
-    "load_decoder": true,
-    "tile_sample_min_height": 256,
-    "tile_sample_min_width": 256,
-    "tile_sample_min_num_frames": 16,
-    "tile_sample_stride_height": 192,
-    "tile_sample_stride_width": 192,
-    "tile_sample_stride_num_frames": 12,
-    "blend_num_frames": 8,
-    "use_tiling": false,
-    "use_temporal_tiling": false,
-    "use_parallel_tiling": false,
-    "use_feature_cache": true
-  },
-  "num_channels_latents": null,
-  "image_encoder_precision": "fp32",
-  "text_encoder_precision": "fp32",
-  "text_len": 512,
-  "hidden_state_skip_layer": 0,
-  "mask_strategy_file_path": null,
-  "enable_torch_compile": false,
-  "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
-}

fastvideo/v1/distributed/__init__.py

@@ -2,10 +2,11 @@
 
 from fastvideo.v1.distributed.communication_op import *
 from fastvideo.v1.distributed.parallel_state import (
-    cleanup_dist_env_and_memory, get_sequence_model_parallel_rank,
-    get_sequence_model_parallel_world_size, get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size, get_world_group,
-    init_distributed_environment, initialize_model_parallel)
+    init_distributed_environment,
+    initialize_model_parallel,
+    get_sequence_model_parallel_rank,
+    get_sequence_model_parallel_world_size,
+)
 from fastvideo.v1.distributed.utils import *
 
 __all__ = [
@@ -13,8 +14,4 @@ __all__ = [
     "initialize_model_parallel",
     "get_sequence_model_parallel_rank",
     "get_sequence_model_parallel_world_size",
-    "get_tensor_model_parallel_rank",
-    "get_tensor_model_parallel_world_size",
-    "cleanup_dist_env_and_memory",
-    "get_world_group",
 ]

fastvideo/v1/entrypoints/cli/cli_types.py

@@ -9,7 +9,8 @@ from fastvideo.v1.utils import FlexibleArgumentParser
 class CLISubcommand:
     """Base class for CLI subcommands"""
 
-    name: str
+    def __init__(self) -> None:
+        self.name = ""
 
     def cmd(self, args: argparse.Namespace) -> None:
         """Execute the command with the given arguments"""

fastvideo/v1/entrypoints/cli/generate.py

@@ -73,6 +73,10 @@ class GenerateSubcommand(CLISubcommand):
             required=False,
             help="Read CLI options from a config YAML file.")
 
+        generate_parser.add_argument("--num-gpus",
+                                     type=int,
+                                     default=1,
+                                     help="Number of GPUs to use")
         generate_parser.add_argument("--master-port",
                                      type=int,
                                      default=None,

fastvideo/v1/entrypoints/video_generator.py

@@ -8,8 +8,8 @@ diffusion models.
 
 import os
 import time
-from dataclasses import asdict
 from typing import Any, Callable, Dict, List, Optional, Union
+from dataclasses import asdict
 
 import imageio
 import numpy as np
@@ -17,11 +17,12 @@ import torch
 import torchvision
 from einops import rearrange
 
-from fastvideo.v1.configs.pipelines import get_pipeline_config_cls_for_name, BaseConfig
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines import ForwardBatch
-from fastvideo.v1.utils import align_to, shallow_asdict
+from fastvideo.v1.pipelines import (ForwardBatch)
+from fastvideo.v1.configs import get_pipeline_config_cls_for_name
+
+from fastvideo.v1.utils import align_to
 from fastvideo.v1.worker.executor import Executor
 
 logger = init_logger(__name__)
@@ -52,7 +53,6 @@ class VideoGenerator:
                         model_path: str,
                         device: Optional[str] = None,
                         torch_dtype: Optional[torch.dtype] = None,
-                        pipeline_config: Optional[Union[str | BaseConfig]] = None,
                         **kwargs) -> "VideoGenerator":
         """
         Create a video generator from a pretrained model.
@@ -65,36 +65,28 @@ class VideoGenerator:
                 
         Returns:
             The created video generator
-
-        Priority level: Default pipeline config < User's pipeline config < User's kwargs
         """
 
-        config = None
-        # 1. If users provide a pipeline config, it will override the default pipeline config
-        if isinstance(pipeline_config, BaseConfig):
-            config = pipeline_config
-        else:
-            config_cls = get_pipeline_config_cls_for_name(model_path)
-            if config_cls is not None:
-                config = config_cls()
-                if isinstance(pipeline_config, str):
-                    config.load_from_json(pipeline_config)
+        config_cls = get_pipeline_config_cls_for_name(model_path)
+        config = config_cls()
 
-        # 2. If users also provide some kwargs, it will override the pipeline config.
-        # The user kwargs shouldn't contain model config parameters!
         if config is None:
-            logger.warning("No config found for model %s, using default config",
-                           model_path)
-            config_args = kwargs
+            logger.warning(f"No config found for model {model_path}, using default config")
+            config_args = {}
         else:
-            config_args = shallow_asdict(config)
-            config_args.update(kwargs)
+            config_args = asdict(config)
+
+        # override config_args with kwargs
+        config_args.update(kwargs)
 
         fastvideo_args = FastVideoArgs(
             model_path=model_path,
             device_str=device or "cuda" if torch.cuda.is_available() else "cpu",
             **config_args)
-        fastvideo_args.check_fastvideo_args()
+
+
+        if torch_dtype is not None:
+            fastvideo_args.dtype = torch_dtype
 
         return cls.from_fastvideo_args(fastvideo_args)
 
@@ -124,7 +116,6 @@ class VideoGenerator:
     def generate_video(
         self,
         prompt: str,
-        image_path: Optional[str] = None,
         negative_prompt: Optional[str] = None,
         output_path: Optional[str] = None,
         save_video: bool = True,
@@ -162,11 +153,9 @@ class VideoGenerator:
             Either the output dictionary or the list of frames depending on return_frames
         """
         # Create a copy of inference args to avoid modifying the original
-        fastvideo_args = self.fastvideo_args
+        fastvideo_args = self.fastvideo_args.copy()
 
         # Override parameters if provided
-        if image_path is not None:
-            fastvideo_args.image_path = image_path
         if negative_prompt is not None:
             fastvideo_args.neg_prompt = negative_prompt
         if num_inference_steps is not None:
@@ -184,6 +173,10 @@ class VideoGenerator:
         if seed is not None:
             fastvideo_args.seed = seed
 
+        # Store callback info
+        fastvideo_args.callback = callback
+        fastvideo_args.callback_steps = callback_steps
+
         # Validate inputs
         if not isinstance(prompt, str):
             raise TypeError(
@@ -236,7 +229,6 @@ class VideoGenerator:
         device = torch.device(fastvideo_args.device_str)
         batch = ForwardBatch(
             prompt=prompt,
-            image_path=fastvideo_args.image_path,
             negative_prompt=fastvideo_args.neg_prompt,
             num_videos_per_prompt=fastvideo_args.num_videos,
             height=fastvideo_args.height,
@@ -253,11 +245,13 @@ class VideoGenerator:
 
         # Run inference
         start_time = time.time()
-        output_batch = self.executor.execute_forward(batch, fastvideo_args)
-        samples = output_batch
+        samples = self.pipeline.forward(
+            batch=batch,
+            fastvideo_args=fastvideo_args,
+        ).output
 
         gen_time = time.time() - start_time
-        logger.info("Generated successfully in %.2f seconds", gen_time)
+        logger.info(f"Generated successfully in {gen_time:.2f} seconds")
 
         # Process outputs
         videos = rearrange(samples, "b c t h w -> t b c h w")
@@ -274,7 +268,7 @@ class VideoGenerator:
                 os.makedirs(os.path.dirname(save_path), exist_ok=True)
                 video_path = os.path.join(save_path, f"{prompt[:100]}.mp4")
                 imageio.mimsave(video_path, frames, fps=fastvideo_args.fps)
-                logger.info("Saved video to %s", video_path)
+                logger.info(f"Saved video to {video_path}")
             else:
                 logger.warning("No output path provided, video not saved")
 
@@ -288,3 +282,121 @@ class VideoGenerator:
                 (target_height, target_width, fastvideo_args.num_frames),
                 "generation_time": gen_time
             }
+
+    def batch_generate(self,
+                       prompts: List[str],
+                       output_path: Optional[str] = None,
+                       **kwargs) -> List[Dict[str, Any]]:
+        """
+        Generate videos for a batch of prompts.
+        
+        Args:
+            prompts: List of prompts to generate videos for
+            output_path: Path to save the videos (overrides the one in fastvideo_args)
+            **kwargs: Additional parameters to pass to generate_video
+            
+        Returns:
+            List of output dictionaries from each generation
+        """
+        if output_path:
+            self.fastvideo_args.output_path = output_path
+
+        results = []
+        for prompt in prompts:
+            result = self.generate_video(prompt=prompt, **kwargs)
+            results.append(result)
+
+        return results
+
+    def image_to_video(self,
+                       image: Union[str, torch.Tensor, np.ndarray],
+                       prompt: Optional[str] = None,
+                       strength: float = 0.8,
+                       **kwargs) -> Dict[str, Any]:
+        """
+        Generate a video from an initial image.
+        
+        Args:
+            image: Input image (path, tensor, or numpy array)
+            prompt: Text prompt to guide the generation
+            strength: How much to transform the original image (0-1)
+            **kwargs: Additional parameters to pass to generate_video
+            
+        Returns:
+            Output dictionary from the generation
+        """
+        # Load image if path is provided
+        if isinstance(image, str):
+            # Implementation would depend on your image loading utilities
+            # This is a placeholder for the concept
+            image_tensor = self._load_image(image)
+        elif isinstance(image, np.ndarray):
+            # Convert numpy array to tensor
+            image_tensor = torch.from_numpy(image).permute(2, 0, 1) / 255.0
+        else:
+            image_tensor = image
+
+        # Add image to inference args
+        fastvideo_args = self.fastvideo_args.copy()
+        fastvideo_args.init_image = image_tensor
+        fastvideo_args.strength = strength
+
+        # Generate video
+        return self.generate_video(prompt=prompt or "",
+                                   fastvideo_args=fastvideo_args,
+                                   **kwargs)
+
+    def to(self, device: Union[str, torch.device]) -> "VideoGenerator":
+        """
+        Move the model to the specified device.
+        
+        Args:
+            device: The device to move the model to
+            
+        Returns:
+            Self for chaining
+        """
+        device_str = str(device)
+        self.fastvideo_args.device_str = device_str
+        self.fastvideo_args.device = torch.device(device_str)
+
+        # Move pipeline components to device
+        self.pipeline.to(device)
+
+        return self
+
+    def _load_image(self, image_path: str) -> torch.Tensor:
+        """
+        Load an image from a path and convert to tensor.
+        
+        Args:
+            image_path: Path to the image
+            
+        Returns:
+            Tensor representation of the image
+        """
+        # Placeholder implementation - would need to be implemented
+        # based on your image loading utilities
+        import PIL.Image
+        from torchvision import transforms
+
+        image = PIL.Image.open(image_path).convert("RGB")
+        transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+        ])
+        return transform(image).unsqueeze(0)
+
+
+def load_prompts_from_file(prompt_path: str) -> List[str]:
+    """
+    Load prompts from a file.
+    
+    Args:
+        prompt_path: Path to the file containing prompts
+        
+    Returns:
+        List of prompts
+    """
+    with open(prompt_path) as f:
+        return [line.strip() for line in f.readlines()]

fastvideo/v1/examples/basic.py

@@ -1,25 +1,20 @@
 from fastvideo import VideoGenerator
 
 def main():
-    # FastVideo will automatically use the optimal default arguments for the
-    # model. 
-    # If a local path is provided, FastVideo will make a best effort
-    # attempt to identify the optimal arguments.
+
+    # This will automatically handle distributed setup if num_gpus > 1
     generator = VideoGenerator.from_pretrained(
         "FastVideo/FastHunyuan-Diffusers",
-        # if num_gpus > 1, FastVideo will automatically handle distributed setup
-        num_gpus=4,
+        num_gpus=2,
+        distributed_executor_backend="mp",
     )
 
     # Generate videos with the same simple API, regardless of GPU count
     prompt = "A beautiful woman in a red dress walking down a street"
     video = generator.generate_video(prompt)
 
-    # Generate another video with a different prompt, without reloading the
-    # model!
     prompt2 = "A beautiful woman in a blue dress walking down a street"
     video2 = generator.generate_video(prompt2)
 
-
-if __name__ == "__main__":
+if __name__ == '__main__':
     main()

fastvideo/v1/examples/inference/basic/README.md

@@ -1,28 +0,0 @@
-# Basic Video Generation Tutorial
-The `VideoGenerator` class provides the primary Python interface for doing offline video generation, which is interacting with a diffusion pipeline without using a separate inference api server.
-
-
-## Usage
-The first script in this example shows the most basic usage of FastVideo. If you are new to Python and FastVideo, you should start here.
-
-```bash
-python fastvideo/v1/examples/inference/basic/basic.py
-```
-
-# Basic Walkthrough
-
-All you need to generate videos using multi-gpus from state-of-the-art diffusion pipelines is the following few lines! 
-
-```python
-from fastvideo import VideoGenerator
-
-generator = VideoGenerator.from_pretrained(
-    "FastVideo/FastHunyuan-Diffusers",
-    num_gpus=2,
-)
-
-prompt = "A beautiful woman in a red dress walking down a street"
-video = generator.generate_video(prompt)
-```
-
-More to come! These examples and APIs are still under construction!

fastvideo/v1/examples/inference/gradio/README.md

@@ -1,81 +0,0 @@
-# FastVideo VideoGenerator Gradio Demo
-
-This is a Gradio-based web interface for generating videos using the FastVideo framework. The demo allows users to create videos from text prompts with various customization options.
-
-## Overview
-
-The demo uses the FastVideo framework to generate videos based on text prompts. It provides a simple web interface built with Gradio that allows users to:
-
-- Enter text prompts to generate videos
-- Customize video parameters (dimensions, number of frames, etc.)
-- Use negative prompts to guide the generation process
-- Set or randomize seeds for reproducibility
-
----
-
-## Requirements
-
-- Linux-based OS
-- Python 3.10
-- Cuda 12.4
-- FastVideo
-
-## Installation
-
-```bash
-pip install fastvideo
-```
-
-## Usage
-
-Run the demo with:
-
-```bash
-python fastvideo/v1/examples/inference/gradio/gradio_demo.py
-```
-
-This will start a web server at `http://0.0.0.0:7860` where you can access the interface.
-
----
-
-## Model Initialization
-
-```python
-args = FastVideoArgs(model_path="FastVideo/FastHunyuan-Diffusers", num_gpus=2)
-
-generator = VideoGenerator.from_pretrained(
-    model_path=args.model_path,
-    num_gpus=args.num_gpus
-)
-```
-
-This demo initializes a `VideoGenerator` with the minimum required arguments for inference. Users can seamlessly adjust inference options between generations, including prompts, resolution, video length, or even the number of inference steps, *without ever needing to reload the model*.
-
-## Video Generation
-
-The core functionality is in the `generate_video` function, which:
-1. Processes user inputs
-2. Uses the FastVideo VideoGenerator from earlier to run inference (`generator.generate_video()`)
-3. Returns an output path that Gradio uses to display the generated video
-
-## Gradio Interface
-
-The interface is built with several components:
-- A text input for the prompt
-- A video display for the result
-- Inference options in a collapsible accordion:
-  - Height and width sliders
-  - Number of frames slider
-  - Guidance scale slider
-  - Inference steps slider
-  - Negative prompt options
-  - Seed controls
-
-### Inference Options
-
-- **Height/Width**: Control the resolution of the generated video
-- **Number of Frames**: Set how many frames to generate
-- **Guidance Scale**: Control how closely the generation follows the prompt
-- **Inference Steps**: More steps can improve quality but take longer
-- **Negative Prompt**: Specify what you don't want to see in the video
-- **Seed**: Control randomness for reproducible results

fastvideo/v1/examples/inference/gradio/gradio_demo.py

@@ -1,141 +0,0 @@
-import os
-import gradio as gr
-import torch
-
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo import VideoGenerator
-
-
-if __name__ == "__main__":
-    args = FastVideoArgs(model_path="FastVideo/FastHunyuan-Diffusers", num_gpus=2)
-
-    generator = VideoGenerator.from_pretrained(
-        model_path=args.model_path,
-        num_gpus=args.num_gpus
-    )
-
-    def generate_video(
-        prompt,
-        negative_prompt,
-        use_negative_prompt,
-        seed,
-        guidance_scale,
-        num_frames,
-        height,
-        width,
-        num_inference_steps,
-        randomize_seed=False,
-    ):
-        if randomize_seed:
-            seed = torch.randint(0, 1000000, (1, )).item()
-
-        if not use_negative_prompt:
-            negative_prompt = None
-
-        generator.generate_video(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_inference_steps=num_inference_steps,
-            num_frames=num_frames,
-            height=height,
-            width=width,
-            guidance_scale=guidance_scale,
-            seed=seed
-        )
-
-        output_path = os.path.join(args.output_path, f"{prompt[:100]}.mp4")
-
-        return output_path, seed
-
-    examples = [
-    "A hand enters the frame, pulling a sheet of plastic wrap over three balls of dough placed on a wooden surface. The plastic wrap is stretched to cover the dough more securely. The hand adjusts the wrap, ensuring that it is tight and smooth over the dough. The scene focuses on the hand’s movements as it secures the edges of the plastic wrap. No new objects appear, and the camera remains stationary, focusing on the action of covering the dough.",
-    "A vintage train snakes through the mountains, its plume of white steam rising dramatically against the jagged peaks. The cars glint in the late afternoon sun, their deep crimson and gold accents lending a touch of elegance. The tracks carve a precarious path along the cliffside, revealing glimpses of a roaring river far below. Inside, passengers peer out the large windows, their faces lit with awe as the landscape unfolds.",
-    "A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background. Strings of fairy lights hang above, casting a warm, golden glow over the scene. Groups of people gather around high tables, their laughter blending with the soft rhythm of live jazz. The aroma of freshly mixed cocktails and charred appetizers wafts through the air, mingling with the cool night breeze.",
-    ]
-
-    with gr.Blocks() as demo:
-        gr.Markdown("# FastVideo Inference Demo")
-
-        with gr.Group():
-            with gr.Row():
-                prompt = gr.Text(
-                    label="Prompt",
-                    show_label=False,
-                    max_lines=1,
-                    placeholder="Enter your prompt",
-                    container=False,
-                )
-                run_button = gr.Button("Run", scale=0)
-            result = gr.Video(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced options", open=False):
-            with gr.Group():
-                with gr.Row():
-                    height = gr.Slider(
-                        label="Height",
-                        minimum=256,
-                        maximum=1024,
-                        step=32,
-                        value=args.height,
-                    )
-                    width = gr.Slider(label="Width", minimum=256, maximum=1024, step=32, value=args.width)
-
-                with gr.Row():
-                    num_frames = gr.Slider(
-                        label="Number of Frames",
-                        minimum=21,
-                        maximum=163,
-                        value=45,
-                    )
-                    guidance_scale = gr.Slider(
-                        label="Guidance Scale",
-                        minimum=1,
-                        maximum=12,
-                        value=args.guidance_scale,
-                    )
-                    num_inference_steps = gr.Slider(
-                        label="Inference Steps",
-                        minimum=4,
-                        maximum=100,
-                        value=6,
-                    )
-
-                with gr.Row():
-                    use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False)
-                negative_prompt = gr.Text(
-                    label="Negative prompt",
-                    max_lines=1,
-                    placeholder="Enter a negative prompt",
-                    visible=False,
-                )
-
-                seed = gr.Slider(label="Seed", minimum=0, maximum=1000000, step=1, value=args.seed)
-                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-                seed_output = gr.Number(label="Used Seed")
-
-        gr.Examples(examples=examples, inputs=prompt)
-
-        use_negative_prompt.change(
-            fn=lambda x: gr.update(visible=x),
-            inputs=use_negative_prompt,
-            outputs=negative_prompt,
-        )
-
-        run_button.click(
-            fn=generate_video,
-            inputs=[
-                prompt,
-                negative_prompt,
-                use_negative_prompt,
-                seed,
-                guidance_scale,
-                num_frames,
-                height,
-                width,
-                num_inference_steps,
-                randomize_seed,
-            ],
-            outputs=[result, seed_output],
-        )
-
-    demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860)

fastvideo/v1/fastvideo_args.py

@@ -7,10 +7,8 @@ import dataclasses
 from contextlib import contextmanager
 from typing import List, Optional
 
-from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import FlexibleArgumentParser
-
-from fastvideo.v1.configs.models import VAEConfig
+from fastvideo.v1.logger import init_logger
 
 logger = init_logger(__name__)
 
@@ -52,10 +50,9 @@ class FastVideoArgs:
 
     # VAE configuration
     vae_precision: str = "fp16"
-    vae_tiling: bool = True # Might change in between forward passes
-    vae_sp: bool = False # Might change in between forward passes
-    # vae_scale_factor: Optional[int] = None # Deprecated
-    vae_config: VAEConfig = VAEConfig()
+    vae_tiling: bool = True
+    vae_sp: bool = False
+    vae_scale_factor: Optional[int] = None
 
     # DiT configuration
     num_channels_latents: Optional[int] = None
@@ -124,12 +121,19 @@ class FastVideoArgs:
             type=str,
             help="Directory containing StepVideo model",
         )
+        parser.add_argument(
+            "--distributed-executor-backend",
+            type=str,
+            default=FastVideoArgs.distributed_executor_backend,
+            choices=["torch"],
+            help="Backend for distributed execution",
+        )
 
         # distributed_executor_backend
         parser.add_argument(
             "--distributed-executor-backend",
             type=str,
-            choices=["mp"],
+            choices=["mp", "torch"],
             default=FastVideoArgs.distributed_executor_backend,
             help="The distributed executor backend to use",
         )
@@ -441,9 +445,6 @@ class FastVideoArgs:
         if self.sp_size is None:
             self.sp_size = self.num_gpus
 
-        if self.num_gpus < max(self.tp_size, self.sp_size):
-            self.num_gpus = max(self.tp_size, self.sp_size)
-
         if self.tp_size != self.sp_size:
             raise ValueError(
                 f"tp_size ({self.tp_size}) must be equal to sp_size ({self.sp_size})"

fastvideo/v1/logger.py

@@ -20,13 +20,6 @@ FASTVIDEO_LOGGING_CONFIG_PATH = envs.FASTVIDEO_LOGGING_CONFIG_PATH
 FASTVIDEO_LOGGING_LEVEL = envs.FASTVIDEO_LOGGING_LEVEL
 FASTVIDEO_LOGGING_PREFIX = envs.FASTVIDEO_LOGGING_PREFIX
 
-RED = '\033[91m'
-GREEN = '\033[92m'
-RESET = '\033[0;0m'
-
-_warned_local_main_process = False
-_warned_main_process = False
-
 _FORMAT = (f"{FASTVIDEO_LOGGING_PREFIX}%(levelname)s %(asctime)s "
            "[%(filename)s:%(lineno)d] %(message)s")
 _DATE_FORMAT = "%m-%d %H:%M:%S"
@@ -75,68 +68,6 @@ def _print_warning_once(logger: Logger, msg: str) -> None:
     logger.warning(msg, stacklevel=2)
 
 
-# TODO(will): add env variable to control this process-aware logging behavior
-def _info(logger: Logger,
-          msg: object,
-          *args: Any,
-          main_process_only: bool = False,
-          local_main_process_only: bool = True,
-          **kwargs: Any) -> None:
-    """Process-aware INFO level logging function.
-    
-    This function controls logging behavior based on the process rank, allowing for 
-    selective logging from specific processes in a distributed environment.
-    
-    Args:
-        logger: The logger instance to use for logging
-        msg: The message format string to log
-        *args: Format string arguments
-        main_process_only: If True, only log if this is the global main process (RANK=0)
-        local_main_process_only: If True, only log if this is the local main process (LOCAL_RANK=0)
-        **kwargs: Additional keyword arguments to pass to the logger.log method
-            - stacklevel: Defaults to 2 to show the original caller's location
-    
-    Note:
-        - When both main_process_only and local_main_process_only are True, 
-          the message will be logged only if both conditions are met
-        - When both are False, the message will be logged from all processes
-        - By default, only logs from processes with LOCAL_RANK=0
-    """
-    try:
-        local_rank = int(os.environ["LOCAL_RANK"])
-        rank = int(os.environ["RANK"])
-    except Exception:
-        local_rank = 0
-        rank = 0
-
-    is_main_process = rank == 0
-    is_local_main_process = local_rank == 0
-
-    if (main_process_only and is_main_process) or (local_main_process_only
-                                                   and is_local_main_process):
-        logger.log(logging.INFO, msg, *args, **kwargs)
-
-    global _warned_local_main_process, _warned_main_process
-
-    if not _warned_local_main_process and local_main_process_only:
-        logger.warning(
-            '%s is_local_main_process is set to True, logging only from the local main process.%s',
-            GREEN,
-            RESET,
-        )
-        _warned_local_main_process = True
-    if not _warned_main_process and main_process_only:
-        logger.warning(
-            '%s is_main_process_only is set to True, logging only from the main process.%s',
-            GREEN,
-            RESET,
-        )
-        _warned_main_process = True
-
-    if not main_process_only and not local_main_process_only:
-        logger.log(logging.INFO, msg, *args, **kwargs)
-
-
 class _FastvideoLogger(Logger):
     """
     Note:
@@ -160,20 +91,6 @@ class _FastvideoLogger(Logger):
         """
         _print_warning_once(self, msg)
 
-    def info(  # type: ignore[override]
-            self,
-            msg: object,
-            *args: Any,
-            main_process_only: bool = False,
-            local_main_process_only: bool = True,
-            **kwargs: Any) -> None:
-        _info(self,
-              msg,
-              *args,
-              main_process_only=main_process_only,
-              local_main_process_only=local_main_process_only,
-              **kwargs)
-
 
 def _configure_fastvideo_root_logger() -> None:
     logging_config = dict[str, Any]()
@@ -211,6 +128,7 @@ def _configure_fastvideo_root_logger() -> None:
         dictConfig(logging_config)
 
 
+# TODO: add rank_zero_only log
 def init_logger(name: str) -> _FastvideoLogger:
     """The main purpose of this function is to ensure that loggers are
     retrieved in such a way that we can be sure the root fastvideo logger has
@@ -221,12 +139,10 @@ def init_logger(name: str) -> _FastvideoLogger:
     methods_to_patch = {
         "info_once": _print_info_once,
         "warning_once": _print_warning_once,
-        "info": _info,
     }
 
     for method_name, method in methods_to_patch.items():
-        setattr(logger, method_name,
-                MethodType(method, logger))  # type: ignore[arg-type]
+        setattr(logger, method_name, MethodType(method, logger))
 
     return cast(_FastvideoLogger, logger)
 

fastvideo/v1/models/dits/base.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from abc import ABC, abstractmethod
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Union
 
 import torch
 from torch import nn
@@ -15,9 +15,7 @@ class BaseDiT(nn.Module, ABC):
     _param_names_mapping: dict
     hidden_size: int
     num_attention_heads: int
-    # always supports torch_sdpa
-    _supported_attention_backends: Tuple[_Backend,
-                                         ...] = (_Backend.TORCH_SDPA, )
+    _supported_attention_backends: List[_Backend] = []
 
     def __init_subclass__(cls) -> None:
         required_class_attrs = [
@@ -57,5 +55,5 @@ class BaseDiT(nn.Module, ABC):
                 )
 
     @property
-    def supported_attention_backends(self) -> Tuple[_Backend, ...]:
+    def supported_attention_backends(self) -> List[_Backend]:
         return self._supported_attention_backends

fastvideo/v1/models/dits/hunyuanvideo.py

@@ -92,7 +92,7 @@ class MMDoubleStreamBlock(nn.Module):
         num_attention_heads: int,
         mlp_ratio: float,
         dtype: Optional[torch.dtype] = None,
-        supported_attention_backends: Optional[Tuple[_Backend, ...]] = None,
+        supported_attention_backends: Optional[List[_Backend]] = None,
         prefix: str = "",
     ):
         super().__init__()
@@ -299,7 +299,7 @@ class MMSingleStreamBlock(nn.Module):
         num_attention_heads: int,
         mlp_ratio: float = 4.0,
         dtype: Optional[torch.dtype] = None,
-        supported_attention_backends: Optional[Tuple[_Backend, ...]] = None,
+        supported_attention_backends: Optional[List[_Backend]] = None,
         prefix: str = "",
     ):
         super().__init__()
@@ -436,8 +436,9 @@ class HunyuanVideoTransformer3DModel(BaseDiT):
         lambda n, m: "single" in n and str.isdigit(n.split(".")[-1]),
         lambda n, m: "refiner" in n and str.isdigit(n.split(".")[-1]),
     ]
-    _supported_attention_backends = (_Backend.SLIDING_TILE_ATTN,
-                                     _Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
+    _supported_attention_backends = [
+        _Backend.SLIDING_TILE_ATTN, _Backend.FLASH_ATTN, _Backend.TORCH_SDPA
+    ]
     _param_names_mapping = {
         # 1. context_embedder.time_text_embed submodules (specific rules, applied first):
         r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$":
@@ -571,7 +572,7 @@ class HunyuanVideoTransformer3DModel(BaseDiT):
         pooled_projection_dim: int = 768,
         rope_theta: int = 256,
         qk_norm: str = "rms_norm",  #TODO(PY)
-        prefix="Hunyuan",
+        prefix="",
     ):
         super().__init__()
         hidden_size = attention_head_dim * num_attention_heads
@@ -895,8 +896,9 @@ class IndividualTokenRefinerBlock(nn.Module):
             num_heads=num_attention_heads,
             head_size=hidden_size // num_attention_heads,
             # TODO: remove hardcode; remove STA
-            supported_attention_backends=(_Backend.FLASH_ATTN,
-                                          _Backend.TORCH_SDPA),
+            supported_attention_backends=[
+                _Backend.FLASH_ATTN, _Backend.TORCH_SDPA
+            ],
         )
 
     def forward(self, x, c):

fastvideo/v1/models/dits/wanvideo.py

@@ -114,14 +114,14 @@ class WanSelfAttention(nn.Module):
         self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
 
         # Scaled dot product attention
-        self.attn = LocalAttention(
-            num_heads=num_heads,
-            head_size=self.head_dim,
-            dropout_rate=0,
-            softmax_scale=None,
-            causal=False,
-            supported_attention_backends=(_Backend.FLASH_ATTN,
-                                          _Backend.TORCH_SDPA))
+        self.attn = LocalAttention(num_heads=num_heads,
+                                   head_size=self.head_dim,
+                                   dropout_rate=0,
+                                   softmax_scale=None,
+                                   causal=False,
+                                   supported_attention_backends=[
+                                       _Backend.FLASH_ATTN, _Backend.TORCH_SDPA
+                                   ])
 
     def forward(self, x: torch.Tensor, context: torch.Tensor,
                 context_lens: int):
@@ -163,14 +163,13 @@ class WanT2VCrossAttention(WanSelfAttention):
 class WanI2VCrossAttention(WanSelfAttention):
 
     def __init__(
-        self,
-        dim: int,
-        num_heads: int,
-        window_size=(-1, -1),
-        qk_norm=True,
-        eps=1e-6,
-        supported_attention_backends: Optional[Tuple[_Backend, ...]] = None
-    ) -> None:
+            self,
+            dim: int,
+            num_heads: int,
+            window_size=(-1, -1),
+            qk_norm=True,
+            eps=1e-6,
+            supported_attention_backends: Optional[List[str]] = None) -> None:
         super().__init__(dim, num_heads, window_size, qk_norm, eps,
                          supported_attention_backends)
 
@@ -211,17 +210,17 @@ class WanI2VCrossAttention(WanSelfAttention):
 
 class WanTransformerBlock(nn.Module):
 
-    def __init__(self,
-                 dim: int,
-                 ffn_dim: int,
-                 num_heads: int,
-                 qk_norm: str = "rms_norm_across_heads",
-                 cross_attn_norm: bool = False,
-                 eps: float = 1e-6,
-                 added_kv_proj_dim: Optional[int] = None,
-                 supported_attention_backends: Optional[Tuple[_Backend,
-                                                              ...]] = None,
-                 prefix: str = ""):
+    def __init__(
+        self,
+        dim: int,
+        ffn_dim: int,
+        num_heads: int,
+        qk_norm: str = "rms_norm_across_heads",
+        cross_attn_norm: bool = False,
+        eps: float = 1e-6,
+        added_kv_proj_dim: Optional[int] = None,
+        supported_attention_backends: Optional[List[_Backend]] = None,
+    ):
         super().__init__()
 
         # 1. Self-attention
@@ -234,8 +233,7 @@ class WanTransformerBlock(nn.Module):
             num_heads=num_heads,
             head_size=dim // num_heads,
             causal=False,
-            supported_attention_backends=supported_attention_backends,
-            prefix=f"{prefix}.attn1")
+            supported_attention_backends=supported_attention_backends)
         self.hidden_dim = dim
         self.num_attention_heads = num_heads
         dim_head = dim // num_heads
@@ -353,8 +351,7 @@ class WanTransformer3DModel(BaseDiT):
     _fsdp_shard_conditions = [
         lambda n, m: "blocks" in n and str.isdigit(n.split(".")[-1]),
     ]
-    _supported_attention_backends = (_Backend.SLIDING_TILE_ATTN,
-                                     _Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
+    _supported_attention_backends = [_Backend.FLASH_ATTN, _Backend.TORCH_SDPA]
     _param_names_mapping = {
         r"^patch_embedding\.(.*)$":
         r"patch_embedding.proj.\1",
@@ -394,24 +391,25 @@ class WanTransformer3DModel(BaseDiT):
         r"blocks.\1.self_attn_residual_norm.norm.\2",
     }
 
-    def __init__(self,
-                 patch_size: Tuple[int, int, int] = (1, 2, 2),
-                 text_len=512,
-                 num_attention_heads: int = 40,
-                 attention_head_dim: int = 128,
-                 in_channels: int = 16,
-                 out_channels: int = 16,
-                 text_dim: int = 4096,
-                 freq_dim: int = 256,
-                 ffn_dim: int = 13824,
-                 num_layers: int = 40,
-                 cross_attn_norm: bool = True,
-                 qk_norm: str = "rms_norm_across_heads",
-                 eps: float = 1e-6,
-                 image_dim: Optional[int] = None,
-                 added_kv_proj_dim: Optional[int] = None,
-                 rope_max_seq_len: int = 1024,
-                 prefix="Wan") -> None:
+    def __init__(
+        self,
+        patch_size: Tuple[int, int, int] = (1, 2, 2),
+        text_len=512,
+        num_attention_heads: int = 40,
+        attention_head_dim: int = 128,
+        in_channels: int = 16,
+        out_channels: int = 16,
+        text_dim: int = 4096,
+        freq_dim: int = 256,
+        ffn_dim: int = 13824,
+        num_layers: int = 40,
+        cross_attn_norm: bool = True,
+        qk_norm: str = "rms_norm_across_heads",
+        eps: float = 1e-6,
+        image_dim: Optional[int] = None,
+        added_kv_proj_dim: Optional[int] = None,
+        rope_max_seq_len: int = 1024,
+    ) -> None:
         super().__init__()
 
         inner_dim = num_attention_heads * attention_head_dim
@@ -438,16 +436,11 @@ class WanTransformer3DModel(BaseDiT):
 
         # 3. Transformer blocks
         self.blocks = nn.ModuleList([
-            WanTransformerBlock(inner_dim,
-                                ffn_dim,
-                                num_attention_heads,
-                                qk_norm,
-                                cross_attn_norm,
-                                eps,
+            WanTransformerBlock(inner_dim, ffn_dim, num_attention_heads,
+                                qk_norm, cross_attn_norm, eps,
                                 added_kv_proj_dim,
-                                self._supported_attention_backends,
-                                prefix=f"{prefix}.blocks.{i}")
-            for i in range(num_layers)
+                                self._supported_attention_backends)
+            for _ in range(num_layers)
         ])
 
         # 4. Output norm & projection

fastvideo/v1/models/encoders/base.py

@@ -1,4 +1,4 @@
-from typing import Tuple
+from typing import List
 
 from torch import nn
 
@@ -6,8 +6,7 @@ from fastvideo.v1.platforms import _Backend
 
 
 class BaseEncoder(nn.Module):
-    _supported_attention_backends: Tuple[_Backend,
-                                         ...] = (_Backend.TORCH_SDPA, )
+    _supported_attention_backends: List[_Backend] = []
 
     def __init__(self, *args, **kwargs) -> None:
         super().__init__()
@@ -20,5 +19,5 @@ class BaseEncoder(nn.Module):
         pass
 
     @property
-    def supported_attention_backends(self) -> Tuple[_Backend, ...]:
+    def supported_attention_backends(self) -> List[_Backend]:
         return self._supported_attention_backends

fastvideo/v1/models/encoders/clip.py

@@ -469,7 +469,7 @@ class CLIPTextTransformer(nn.Module):
 
 
 class CLIPTextModel(BaseEncoder):
-    _supported_attention_backends = (_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
+    _supported_attention_backends = [_Backend.FLASH_ATTN, _Backend.TORCH_SDPA]
 
     def __init__(
         self,
@@ -619,7 +619,7 @@ class CLIPVisionModel(BaseEncoder, SupportsQuant):
     config_class = CLIPVisionConfig
     main_input_name = "pixel_values"
     packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
-    _supported_attention_backends = (_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
+    _supported_attention_backends = [_Backend.FLASH_ATTN, _Backend.TORCH_SDPA]
 
     def __init__(
         self,

fastvideo/v1/models/encoders/llama.py

@@ -280,7 +280,7 @@ class LlamaDecoderLayer(nn.Module):
 
 
 class LlamaModel(BaseEncoder):
-    _supported_attention_backends = (_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
+    _supported_attention_backends = [_Backend.FLASH_ATTN, _Backend.TORCH_SDPA]
 
     def __init__(self,
                  config: LlamaConfig,

fastvideo/v1/models/loader/component_loader.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import dataclasses
-from dataclasses import asdict
 import glob
 import os
 import time
@@ -311,11 +310,9 @@ class VAELoader(ComponentLoader):
         assert class_name is not None, "Model config does not contain a _class_name attribute. Only diffusers format is supported."
         config.pop("_diffusers_version")
 
-        vae_config = fastvideo_args.vae_config
-        vae_config.update_model_arch(config)
         vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
 
-        vae = vae_cls(vae_config).to(fastvideo_args.device)
+        vae = vae_cls(**config).to(fastvideo_args.device)
 
         # Find all safetensors files
         safetensors_list = glob.glob(
@@ -325,7 +322,7 @@ class VAELoader(ComponentLoader):
             safetensors_list
         ) == 1, f"Found {len(safetensors_list)} safetensors files in {model_path}"
         loaded = safetensors_load_file(safetensors_list[0])
-        vae.load_state_dict(loaded, strict=False) # We might only load encoder or decoder
+        vae.load_state_dict(loaded)
         dtype = PRECISION_TO_TYPE[fastvideo_args.vae_precision]
         vae = vae.eval().to(dtype)
 
@@ -346,10 +343,6 @@ class TransformerLoader(ComponentLoader):
                 "Only diffusers format is supported.")
         model_config.pop("_diffusers_version")
 
-        # Config from Diffusers supercedes fastvideo's model config
-        # dit_config = fastvideo_args.dit_config
-        # model_config.update(dit_config)
-
         model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
 
         # Find all safetensors files

fastvideo/v1/models/vaes/common.py

@@ -11,7 +11,6 @@ from diffusers.utils.torch_utils import randn_tensor
 
 from fastvideo.v1.distributed import (get_sequence_model_parallel_rank,
                                       get_sequence_model_parallel_world_size)
-from fastvideo.v1.configs.models import VAEConfig
 
 
 class ParallelTiledVAE(ABC):
@@ -21,36 +20,29 @@ class ParallelTiledVAE(ABC):
     tile_sample_stride_height: int
     tile_sample_stride_width: int
     tile_sample_stride_num_frames: int
-    blend_num_frames: int
     use_tiling: bool
     use_temporal_tiling: bool
     use_parallel_tiling: bool
+    temporal_compression_ratio: int
+    spatial_compression_ratio: int
+    scaling_factor: Union[float, torch.tensor]
 
-    def __init__(self, config: VAEConfig, **kwargs) -> None:
-        self.config = config
-        self.arch_config = config.arch_config
-        self.tile_sample_min_height = config.tile_sample_min_height
-        self.tile_sample_min_width = config.tile_sample_min_width
-        self.tile_sample_min_num_frames = config.tile_sample_min_num_frames
-        self.tile_sample_stride_height = config.tile_sample_stride_height
-        self.tile_sample_stride_width = config.tile_sample_stride_width
-        self.tile_sample_stride_num_frames = config.tile_sample_stride_num_frames
-        self.blend_num_frames = config.blend_num_frames
-        self.use_tiling = config.use_tiling
-        self.use_temporal_tiling = config.use_temporal_tiling
-        self.use_parallel_tiling = config.use_parallel_tiling
+    def __init__(self, *args, **kwargs) -> None:
+        # Check if subclass has defined all required properties
+        required_attributes = [
+            'tile_sample_min_height', 'tile_sample_min_width',
+            'tile_sample_min_num_frames', 'tile_sample_stride_height',
+            'tile_sample_stride_width', 'tile_sample_stride_num_frames',
+            'spatial_compression_ratio', 'temporal_compression_ratio',
+            'use_tiling', 'use_temporal_tiling', 'use_parallel_tiling',
+            'scaling_factor'
+        ]
 
-    @property
-    def temporal_compression_ratio(self) -> int:
-        return self.arch_config.temporal_compression_ratio
-    
-    @property
-    def spatial_compression_ratio(self) -> int:
-        return self.arch_config.spatial_compression_ratio
-    
-    @property
-    def scaling_factor(self) -> Union[float, torch.tensor]:
-        return self.arch_config.scaling_factor
+        for attr in required_attributes:
+            if not hasattr(self, attr):
+                raise AttributeError(
+                    f"Subclasses of ParallelVAE must define '{attr}' property")
+        self.blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
 
     @abstractmethod
     def _encode(self, *args, **kwargs) -> torch.Tensor:
@@ -416,10 +408,6 @@ class ParallelTiledVAE(ABC):
         tile_sample_stride_height: Optional[int] = None,
         tile_sample_stride_width: Optional[int] = None,
         tile_sample_stride_num_frames: Optional[int] = None,
-        blend_num_frames: Optional[int] = None,
-        use_tiling: Optional[bool] = None,
-        use_temporal_tiling: Optional[bool] = None,
-        use_parallel_tiling: Optional[bool] = None,
     ) -> None:
         r"""
         Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
@@ -451,13 +439,7 @@ class ParallelTiledVAE(ABC):
         self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height
         self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width
         self.tile_sample_stride_num_frames = tile_sample_stride_num_frames or self.tile_sample_stride_num_frames
-        if blend_num_frames is not None:
-            self.blend_num_frames = blend_num_frames
-        else:
-            self.blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
-        self.use_tiling = use_tiling or self.use_tiling
-        self.use_temporal_tiling = use_temporal_tiling or self.use_temporal_tiling
-        self.use_parallel_tiling = use_parallel_tiling or self.use_parallel_tiling
+        self.blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
 
     def disable_tiling(self) -> None:
         r"""

fastvideo/v1/models/vaes/hunyuanvae.py

@@ -15,7 +15,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple, Union, cast
+from typing import Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -24,8 +24,8 @@ import torch.nn.functional as F
 import torch.utils.checkpoint
 
 from fastvideo.v1.layers.activation import get_act_fn
+from fastvideo.v1.models.utils import auto_attributes
 from fastvideo.v1.models.vaes.common import ParallelTiledVAE
-from fastvideo.v1.configs.models.vaes import HunyuanVAEConfig, HunyuanVAEArchConfig
 
 
 def prepare_causal_attention_mask(
@@ -773,53 +773,95 @@ class AutoencoderKLHunyuanVideo(nn.Module, ParallelTiledVAE):
 
     _supports_gradient_checkpointing = True
 
+    @auto_attributes
     def __init__(
         self,
-        config: HunyuanVAEConfig,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        latent_channels: int = 16,
+        down_block_types: Tuple[str, ...] = (
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+        ),
+        up_block_types: Tuple[str, ...] = (
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+        ),
+        block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+        layers_per_block: int = 2,
+        act_fn: str = "silu",
+        norm_num_groups: int = 32,
+        scaling_factor: float = 0.476986,
+        spatial_compression_ratio: int = 8,
+        temporal_compression_ratio: int = 4,
+        mid_block_add_attention: bool = True,
+        load_encoder: bool = True,
+        load_decoder: bool = True,
     ) -> None:
         super().__init__()
-        ParallelTiledVAE.__init__(self, config)
-        arch_config: HunyuanVAEArchConfig = cast(HunyuanVAEArchConfig, config.arch_config)
 
         # TODO(will): only pass in config. We do this by manually defining a
         # config for hunyuan vae
-        self.block_out_channels = arch_config.block_out_channels
-        
-        if config.load_encoder:
+        self.block_out_channels = block_out_channels
+
+        if load_encoder:
             self.encoder = HunyuanVideoEncoder3D(
-                in_channels=arch_config.in_channels,
-                out_channels=arch_config.latent_channels,
-                down_block_types=arch_config.down_block_types,
-                block_out_channels=arch_config.block_out_channels,
-                layers_per_block=arch_config.layers_per_block,
-                norm_num_groups=arch_config.norm_num_groups,
-                act_fn=arch_config.act_fn,
+                in_channels=in_channels,
+                out_channels=latent_channels,
+                down_block_types=down_block_types,
+                block_out_channels=block_out_channels,
+                layers_per_block=layers_per_block,
+                norm_num_groups=norm_num_groups,
+                act_fn=act_fn,
                 double_z=True,
-                mid_block_add_attention=arch_config.mid_block_add_attention,
-                temporal_compression_ratio=arch_config.temporal_compression_ratio,
-                spatial_compression_ratio=arch_config.spatial_compression_ratio,
+                mid_block_add_attention=mid_block_add_attention,
+                temporal_compression_ratio=temporal_compression_ratio,
+                spatial_compression_ratio=spatial_compression_ratio,
             )
-            self.quant_conv = nn.Conv3d(2 * arch_config.latent_channels,
-                                        2 * arch_config.latent_channels,
+            self.quant_conv = nn.Conv3d(2 * latent_channels,
+                                        2 * latent_channels,
                                         kernel_size=1)
 
-        if config.load_decoder:
+        if load_decoder:
             self.decoder = HunyuanVideoDecoder3D(
-                in_channels=arch_config.latent_channels,
-                out_channels=arch_config.out_channels,
-                up_block_types=arch_config.up_block_types,
-                block_out_channels=arch_config.block_out_channels,
-                layers_per_block=arch_config.layers_per_block,
-                norm_num_groups=arch_config.norm_num_groups,
-                act_fn=arch_config.act_fn,
-                time_compression_ratio=arch_config.temporal_compression_ratio,
-                spatial_compression_ratio=arch_config.spatial_compression_ratio,
-                mid_block_add_attention=arch_config.mid_block_add_attention,
+                in_channels=latent_channels,
+                out_channels=out_channels,
+                up_block_types=up_block_types,
+                block_out_channels=block_out_channels,
+                layers_per_block=layers_per_block,
+                norm_num_groups=norm_num_groups,
+                act_fn=act_fn,
+                time_compression_ratio=temporal_compression_ratio,
+                spatial_compression_ratio=spatial_compression_ratio,
+                mid_block_add_attention=mid_block_add_attention,
             )
-            self.post_quant_conv = nn.Conv3d(arch_config.latent_channels,
-                                             arch_config.latent_channels,
+            self.post_quant_conv = nn.Conv3d(latent_channels,
+                                             latent_channels,
                                              kernel_size=1)
 
+        # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent
+        # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the
+        # intermediate tiles together, the memory requirement can be lowered.
+        self.use_tiling = True
+        self.use_temporal_tiling = True
+        self.use_parallel_tiling = True
+        self.scaling_factor = scaling_factor
+
+        # The minimal tile height and width for spatial tiling to be used
+        self.tile_sample_min_height = 256
+        self.tile_sample_min_width = 256
+        self.tile_sample_min_num_frames = 16
+
+        # The minimal distance between two spatial tiles
+        self.tile_sample_stride_height = 192
+        self.tile_sample_stride_width = 192
+        self.tile_sample_stride_num_frames = 12
+        ParallelTiledVAE.__init__(self)
+
     def _encode(self, x: torch.Tensor) -> torch.Tensor:
         x = self.encoder(x)
         enc = self.quant_conv(x)

fastvideo/v1/models/vaes/wanvae.py

@@ -14,39 +14,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import contextvars
-from contextlib import contextmanager
-from typing import Optional, Tuple, Union, cast
+from typing import Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import torch.utils.checkpoint
 
 from fastvideo.v1.layers.activation import get_act_fn
-from fastvideo.v1.models.vaes.common import ParallelTiledVAE, DiagonalGaussianDistribution
-from fastvideo.v1.configs.models.vaes import WanVAEConfig, WanVAEArchConfig
+from fastvideo.v1.models.utils import auto_attributes
+from fastvideo.v1.models.vaes.common import ParallelTiledVAE
 
 CACHE_T = 2
 
-is_first_frame = contextvars.ContextVar("is_first_frame", default=False)
-feat_cache = contextvars.ContextVar("feat_cache", default=None)
-feat_idx = contextvars.ContextVar("feat_idx", default=0)
-
-
-@contextmanager
-def forward_context(first_frame_arg=False,
-                    feat_cache_arg=None,
-                    feat_idx_arg=None):
-    is_first_frame_token = is_first_frame.set(first_frame_arg)
-    feat_cache_token = feat_cache.set(feat_cache_arg)
-    feat_idx_token = feat_idx.set(feat_idx_arg)
-    try:
-        yield
-    finally:
-        is_first_frame.reset(is_first_frame_token)
-        feat_cache.reset(feat_cache_token)
-        feat_idx.reset(feat_idx_token)
-
 
 class WanCausalConv3d(nn.Conv3d):
     r"""
@@ -80,17 +60,12 @@ class WanCausalConv3d(nn.Conv3d):
         )
         self.padding: Tuple[int, int, int]
         # Set up causal padding
-        self._padding: Tuple[int, ...] = (self.padding[2], self.padding[2],
-                                          self.padding[1], self.padding[1],
-                                          2 * self.padding[0], 0)
+        self._padding = (self.padding[2], self.padding[2], self.padding[1],
+                         self.padding[1], 2 * self.padding[0], 0)
         self.padding = (0, 0, 0)
 
-    def forward(self, x, cache_x=None):
+    def forward(self, x):
         padding = list(self._padding)
-        if cache_x is not None and self._padding[4] > 0:
-            cache_x = cache_x.to(x.device)
-            x = torch.cat([cache_x, x], dim=2)
-            padding[4] -= cache_x.shape[2]
         x = F.pad(x, padding)
         return super().forward(x)
 
@@ -182,82 +157,28 @@ class WanResample(nn.Module):
             self.time_conv = WanCausalConv3d(dim,
                                              dim, (3, 1, 1),
                                              stride=(2, 1, 1),
-                                             padding=(0, 0, 0))
+                                             padding=(1, 0, 0))
 
         else:
             self.resample = nn.Identity()
 
-    def forward(self, x):
+    def forward(self, x, first_frame=False):
         b, c, t, h, w = x.size()
-        first_frame = is_first_frame.get()
         if first_frame:
             assert t == 1
-        _feat_cache = feat_cache.get()
-        _feat_idx = feat_idx.get()
-        if self.mode == "upsample3d":
-            if _feat_cache is not None:
-                idx = _feat_idx
-                if _feat_cache[idx] is None:
-                    _feat_cache[idx] = "Rep"
-                    _feat_idx += 1
-                else:
-                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
-                    if cache_x.shape[2] < 2 and _feat_cache[
-                            idx] is not None and _feat_cache[idx] != "Rep":
-                        # cache last frame of last two chunk
-                        cache_x = torch.cat([
-                            _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
-                                cache_x.device), cache_x
-                        ],
-                                            dim=2)
-                    if cache_x.shape[2] < 2 and _feat_cache[
-                            idx] is not None and _feat_cache[idx] == "Rep":
-                        cache_x = torch.cat([
-                            torch.zeros_like(cache_x).to(cache_x.device),
-                            cache_x
-                        ],
-                                            dim=2)
-                    if _feat_cache[idx] == "Rep":
-                        x = self.time_conv(x)
-                    else:
-                        x = self.time_conv(x, _feat_cache[idx])
-                    _feat_cache[idx] = cache_x
-                    _feat_idx += 1
-
-                    x = x.reshape(b, 2, c, t, h, w)
-                    x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]),
-                                    3)
-                    x = x.reshape(b, c, t * 2, h, w)
-                feat_cache.set(_feat_cache)
-                feat_idx.set(_feat_idx)
-            elif not first_frame and hasattr(self, "time_conv"):
-                x = self.time_conv(x)
-                x = x.reshape(b, 2, c, t, h, w)
-                x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
-                x = x.reshape(b, c, t * 2, h, w)
+        if self.mode == "upsample3d" and not first_frame and hasattr(
+                self, "time_conv"):
+            x = self.time_conv(x)
+            x = x.reshape(b, 2, c, t, h, w)
+            x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
+            x = x.reshape(b, c, t * 2, h, w)
         t = x.shape[2]
         x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
         x = self.resample(x)
         x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4)
-
-        _feat_cache = feat_cache.get()
-        _feat_idx = feat_idx.get()
-        if self.mode == "downsample3d":
-            if _feat_cache is not None:
-                idx = _feat_idx
-                if _feat_cache[idx] is None:
-                    _feat_cache[idx] = x.clone()
-                    _feat_idx += 1
-                else:
-                    cache_x = x[:, :, -1:, :, :].clone()
-                    x = self.time_conv(
-                        torch.cat([_feat_cache[idx][:, :, -1:, :, :], x], 2))
-                    _feat_cache[idx] = cache_x
-                    _feat_idx += 1
-                feat_cache.set(_feat_cache)
-                feat_idx.set(_feat_idx)
-            elif not first_frame and hasattr(self, "time_conv"):
-                x = self.time_conv(x)
+        if self.mode == "downsample3d" and not first_frame and hasattr(
+                self, "time_conv"):
+            x = self.time_conv(x)
         return x
 
 
@@ -301,25 +222,7 @@ class WanResidualBlock(nn.Module):
         x = self.norm1(x)
         x = self.nonlinearity(x)
 
-        _feat_cache = feat_cache.get()
-        _feat_idx = feat_idx.get()
-        if _feat_cache is not None:
-            idx = _feat_idx
-            cache_x = x[:, :, -CACHE_T:, :, :].clone()
-            if cache_x.shape[2] < 2 and _feat_cache[idx] is not None:
-                cache_x = torch.cat([
-                    _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
-                        cache_x.device), cache_x
-                ],
-                                    dim=2)
-
-            x = self.conv1(x, _feat_cache[idx])
-            _feat_cache[idx] = cache_x
-            _feat_idx += 1
-            feat_cache.set(_feat_cache)
-            feat_idx.set(_feat_idx)
-        else:
-            x = self.conv1(x)
+        x = self.conv1(x)
 
         # Second normalization and activation
         x = self.norm2(x)
@@ -328,25 +231,7 @@ class WanResidualBlock(nn.Module):
         # Dropout
         x = self.dropout(x)
 
-        _feat_cache = feat_cache.get()
-        _feat_idx = feat_idx.get()
-        if _feat_cache is not None:
-            idx = _feat_idx
-            cache_x = x[:, :, -CACHE_T:, :, :].clone()
-            if cache_x.shape[2] < 2 and _feat_cache[idx] is not None:
-                cache_x = torch.cat([
-                    _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
-                        cache_x.device), cache_x
-                ],
-                                    dim=2)
-
-            x = self.conv2(x, _feat_cache[idx])
-            _feat_cache[idx] = cache_x
-            _feat_idx += 1
-            feat_cache.set(_feat_cache)
-            feat_idx.set(_feat_idx)
-        else:
-            x = self.conv2(x)
+        x = self.conv2(x)
 
         # Add residual connection
         return x + h
@@ -515,30 +400,15 @@ class WanEncoder3d(nn.Module):
 
         self.gradient_checkpointing = False
 
-    def forward(self, x):
-        _feat_cache = feat_cache.get()
-        _feat_idx = feat_idx.get()
-        if _feat_cache is not None:
-            idx = _feat_idx
-            cache_x = x[:, :, -CACHE_T:, :, :].clone()
-            if cache_x.shape[2] < 2 and _feat_cache[idx] is not None:
-                # cache last frame of last two chunk
-                cache_x = torch.cat([
-                    _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
-                        cache_x.device), cache_x
-                ],
-                                    dim=2)
-            x = self.conv_in(x, _feat_cache[idx])
-            _feat_cache[idx] = cache_x
-            _feat_idx += 1
-            feat_cache.set(_feat_cache)
-            feat_idx.set(_feat_idx)
-        else:
-            x = self.conv_in(x)
+    def forward(self, x, first_frame=False):
+        x = self.conv_in(x)
 
         ## downsamples
         for layer in self.down_blocks:
-            x = layer(x)
+            if isinstance(layer, WanResample):
+                x = layer(x, first_frame=first_frame)
+            else:
+                x = layer(x)
 
         ## middle
         x = self.mid_block(x)
@@ -546,26 +416,7 @@ class WanEncoder3d(nn.Module):
         ## head
         x = self.norm_out(x)
         x = self.nonlinearity(x)
-
-        _feat_cache = feat_cache.get()
-        _feat_idx = feat_idx.get()
-        if _feat_cache is not None:
-            idx = _feat_idx
-            cache_x = x[:, :, -CACHE_T:, :, :].clone()
-            if cache_x.shape[2] < 2 and _feat_cache[idx] is not None:
-                # cache last frame of last two chunk
-                cache_x = torch.cat([
-                    _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
-                        cache_x.device), cache_x
-                ],
-                                    dim=2)
-            x = self.conv_out(x, _feat_cache[idx])
-            _feat_cache[idx] = cache_x
-            _feat_idx += 1
-            feat_cache.set(_feat_cache)
-            feat_idx.set(_feat_idx)
-        else:
-            x = self.conv_out(x)
+        x = self.conv_out(x)
         return x
 
 
@@ -614,7 +465,7 @@ class WanUpBlock(nn.Module):
 
         self.gradient_checkpointing = False
 
-    def forward(self, x):
+    def forward(self, x, first_frame=False):
         """
         Forward pass through the upsampling block.
 
@@ -630,7 +481,7 @@ class WanUpBlock(nn.Module):
             x = resnet(x)
 
         if self.upsamplers is not None:
-            x = self.upsamplers[0](x)
+            x = self.upsamplers[0](x, first_frame=first_frame)
         return x
 
 
@@ -718,57 +569,21 @@ class WanDecoder3d(nn.Module):
 
         self.gradient_checkpointing = False
 
-    def forward(self, x):
+    def forward(self, x, first_frame=False):
         ## conv1
-        _feat_cache = feat_cache.get()
-        _feat_idx = feat_idx.get()
-        if _feat_cache is not None:
-            idx = _feat_idx
-            cache_x = x[:, :, -CACHE_T:, :, :].clone()
-            if cache_x.shape[2] < 2 and _feat_cache[idx] is not None:
-                # cache last frame of last two chunk
-                cache_x = torch.cat([
-                    _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
-                        cache_x.device), cache_x
-                ],
-                                    dim=2)
-            x = self.conv_in(x, _feat_cache[idx])
-            _feat_cache[idx] = cache_x
-            _feat_idx += 1
-            feat_cache.set(_feat_cache)
-            feat_idx.set(_feat_idx)
-        else:
-            x = self.conv_in(x)
+        x = self.conv_in(x)
 
         ## middle
         x = self.mid_block(x)
 
         ## upsamples
         for up_block in self.up_blocks:
-            x = up_block(x)
+            x = up_block(x, first_frame=first_frame)
 
         ## head
         x = self.norm_out(x)
         x = self.nonlinearity(x)
-        _feat_cache = feat_cache.get()
-        _feat_idx = feat_idx.get()
-        if _feat_cache is not None:
-            idx = _feat_idx
-            cache_x = x[:, :, -CACHE_T:, :, :].clone()
-            if cache_x.shape[2] < 2 and _feat_cache[idx] is not None:
-                # cache last frame of last two chunk
-                cache_x = torch.cat([
-                    _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
-                        cache_x.device), cache_x
-                ],
-                                    dim=2)
-            x = self.conv_out(x, _feat_cache[idx])
-            _feat_cache[idx] = cache_x
-            _feat_idx += 1
-            feat_cache.set(_feat_cache)
-            feat_idx.set(_feat_idx)
-        else:
-            x = self.conv_out(x)
+        x = self.conv_out(x)
         return x
 
 
@@ -780,89 +595,96 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
 
     _supports_gradient_checkpointing = False
 
+    @auto_attributes
     def __init__(self,
-                 config: WanVAEConfig,
-                ) -> None:
+                 base_dim: int = 96,
+                 z_dim: int = 16,
+                 dim_mult: Tuple[int, ...] = (1, 2, 4, 4),
+                 num_res_blocks: int = 2,
+                 attn_scales: Tuple[float, ...] = (),
+                 temperal_downsample: Tuple[bool, ...] = (False, True, True),
+                 dropout: float = 0.0,
+                 latents_mean: Tuple[float, ...] = (
+                     -0.7571,
+                     -0.7089,
+                     -0.9113,
+                     0.1075,
+                     -0.1745,
+                     0.9653,
+                     -0.1517,
+                     1.5508,
+                     0.4134,
+                     -0.0715,
+                     0.5517,
+                     -0.3632,
+                     -0.1922,
+                     -0.9497,
+                     0.2503,
+                     -0.2921,
+                 ),
+                 latents_std: Tuple[float, ...] = (
+                     2.8184,
+                     1.4541,
+                     2.3275,
+                     2.6558,
+                     1.2196,
+                     1.7708,
+                     2.6052,
+                     2.0743,
+                     3.2687,
+                     2.1526,
+                     2.8652,
+                     1.5579,
+                     1.6382,
+                     1.1253,
+                     2.8251,
+                     1.9160,
+                 ),
+                 load_encoder: bool = True,
+                 load_decoder: bool = True) -> None:
         super().__init__()
-        ParallelTiledVAE.__init__(self, config)
 
-        self.arch_config = cast(WanVAEArchConfig, self.arch_config)
-        
-        self.z_dim = self.arch_config.z_dim
-        self.temperal_downsample = list(self.arch_config.temperal_downsample)
-        self.temperal_upsample = list(self.arch_config.temperal_downsample)[::-1]
-        self.latents_mean = list(self.arch_config.latents_mean)
-        self.latents_std = list(self.arch_config.latents_std)
-        self.shift_factor = self.arch_config.shift_factor
+        self.z_dim = z_dim
+        self.temperal_downsample = list(temperal_downsample)
+        self.temperal_upsample = list(temperal_downsample)[::-1]
+        self.latents_mean = list(latents_mean)
+        self.latents_std = list(latents_std)
+        self.scaling_factor = 1.0 / torch.tensor(self.config.latents_std).view(
+            1, self.config.z_dim, 1, 1, 1)
+        self.shift_factor = torch.tensor(self.config.latents_mean).view(
+            1, self.config.z_dim, 1, 1, 1)
 
-        if config.load_encoder:
-            self.encoder = WanEncoder3d(self.arch_config.base_dim, self.z_dim * 2, self.arch_config.dim_mult,
-                                        self.arch_config.num_res_blocks, self.arch_config.attn_scales,
-                                        self.temperal_downsample, self.arch_config.dropout)
-        self.quant_conv = WanCausalConv3d(self.z_dim * 2, self.z_dim * 2, 1)
-        self.post_quant_conv = WanCausalConv3d(self.z_dim, self.z_dim, 1)
+        if load_encoder:
+            self.encoder = WanEncoder3d(base_dim, z_dim * 2, dim_mult,
+                                        num_res_blocks, attn_scales,
+                                        self.temperal_downsample, dropout)
+        self.quant_conv = WanCausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.post_quant_conv = WanCausalConv3d(z_dim, z_dim, 1)
 
-        if config.load_decoder:
-            self.decoder = WanDecoder3d(self.arch_config.base_dim, self.z_dim, self.arch_config.dim_mult,
-                                        self.arch_config.num_res_blocks, self.arch_config.attn_scales,
-                                        self.temperal_upsample, self.arch_config.dropout)
+        if load_decoder:
+            self.decoder = WanDecoder3d(base_dim, z_dim, dim_mult,
+                                        num_res_blocks, attn_scales,
+                                        self.temperal_upsample, dropout)
 
-        self.use_feature_cache = config.use_feature_cache
+        self.use_tiling = True
+        self.use_temporal_tiling = False
+        self.use_parallel_tiling = False
+        self.spatial_compression_ratio = 8
+        self.temporal_compression_ratio = 4
 
-    def clear_cache(self) -> None:
+        # The minimal tile height and width for spatial tiling to be used
+        self.tile_sample_min_height = 256
+        self.tile_sample_min_width = 256
+        self.tile_sample_min_num_frames = 16
 
-        def _count_conv3d(model) -> int:
-            count = 0
-            for m in model.modules():
-                if isinstance(m, WanCausalConv3d):
-                    count += 1
-            return count
-
-        if self.config.load_decoder:
-            self._conv_num = _count_conv3d(self.decoder)
-            self._conv_idx = 0
-            self._feat_map = [None] * self._conv_num
-        # cache encode
-        if self.config.load_encoder:
-            self._enc_conv_num = _count_conv3d(self.encoder)
-            self._enc_conv_idx = 0
-            self._enc_feat_map = [None] * self._enc_conv_num
-
-    def encode(self, x: torch.Tensor) -> torch.Tensor:
-        if self.use_feature_cache:
-            self.clear_cache()
-            with forward_context(feat_cache_arg=self._enc_feat_map,
-                                 feat_idx_arg=self._enc_conv_idx):
-                t = x.shape[2]
-                iter_ = 1 + (t - 1) // 4
-                for i in range(iter_):
-                    feat_idx.set(0)
-                    if i == 0:
-                        out = self.encoder(x[:, :, :1, :, :])
-                    else:
-                        out_ = self.encoder(x[:, :,
-                                              1 + 4 * (i - 1):1 + 4 * i, :, :])
-                        out = torch.cat([out, out_], 2)
-            enc = self.quant_conv(out)
-            mu, logvar = enc[:, :self.z_dim, :, :, :], enc[:,
-                                                           self.z_dim:, :, :, :]
-            enc = torch.cat([mu, logvar], dim=1)
-            enc = DiagonalGaussianDistribution(enc)
-            self.clear_cache()
-        else:
-            for block in self.encoder.down_blocks:
-                if isinstance(block,
-                              WanResample) and block.mode == "downsample3d":
-                    _padding = list(block.time_conv._padding)
-                    _padding[4] = 2
-                    block.time_conv._padding = tuple(_padding)
-            enc = ParallelTiledVAE.encode(self, x)
-
-        return enc
+        # The minimal distance between two spatial tiles
+        self.tile_sample_stride_height = 192
+        self.tile_sample_stride_width = 192
+        self.tile_sample_stride_num_frames = 12
+        ParallelTiledVAE.__init__(self)
 
     def _encode(self, x: torch.Tensor, first_frame=False) -> torch.Tensor:
-        with forward_context(first_frame_arg=first_frame):
-            out = self.encoder(x)
+        out = self.encoder(x, first_frame=first_frame)
         enc = self.quant_conv(out)
         mu, logvar = enc[:, :self.z_dim, :, :, :], enc[:, self.z_dim:, :, :, :]
         enc = torch.cat([mu, logvar], dim=1)
@@ -886,32 +708,9 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
         enc = torch.cat([first_frame, enc], dim=2)
         return enc
 
-    def decode(self, z: torch.Tensor) -> torch.Tensor:
-        if self.use_feature_cache:
-            self.clear_cache()
-            iter_ = z.shape[2]
-            x = self.post_quant_conv(z)
-            with forward_context(feat_cache_arg=self._feat_map,
-                                 feat_idx_arg=self._conv_idx):
-                for i in range(iter_):
-                    feat_idx.set(0)
-                    if i == 0:
-                        out = self.decoder(x[:, :, i:i + 1, :, :])
-                    else:
-                        out_ = self.decoder(x[:, :, i:i + 1, :, :])
-                        out = torch.cat([out, out_], 2)
-
-            out = torch.clamp(out, min=-1.0, max=1.0)
-            self.clear_cache()
-        else:
-            out = ParallelTiledVAE.decode(self, z)
-
-        return out
-
     def _decode(self, z: torch.Tensor, first_frame=False) -> torch.Tensor:
         x = self.post_quant_conv(z)
-        with forward_context(first_frame_arg=first_frame):
-            out = self.decoder(x)
+        out = self.decoder(x, first_frame=first_frame)
 
         out = torch.clamp(out, min=-1.0, max=1.0)
 

fastvideo/v1/pipelines/hunyuan/hunyuan_pipeline.py

@@ -6,6 +6,8 @@ This module contains an implementation of the Hunyuan video diffusion pipeline
 using the modular pipeline architecture.
 """
 
+from diffusers.image_processor import VaeImageProcessor
+
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
@@ -69,6 +71,14 @@ class HunyuanVideoPipeline(ComposedPipelineBase):
         """
         Initialize the pipeline.
         """
+        vae_scale_factor = 2**(len(self.get_module("vae").block_out_channels) -
+                               1)
+        fastvideo_args.vae_scale_factor = vae_scale_factor
+
+        self.image_processor = VaeImageProcessor(
+            vae_scale_factor=vae_scale_factor)
+        self.add_module("image_processor", self.image_processor)
+
         num_channels_latents = self.get_module("transformer").in_channels
         fastvideo_args.num_channels_latents = num_channels_latents
 

fastvideo/v1/pipelines/stages/clip_image_encoding.py

@@ -7,8 +7,8 @@ This module contains implementations of image encoding stages for diffusion pipe
 
 import torch
 
-from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
+from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.vision_utils import load_image
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch

fastvideo/v1/pipelines/stages/clip_text_encoding.py

@@ -7,8 +7,8 @@ This module contains implementations of prompt encoding stages for diffusion pip
 
 import torch
 
-from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
+from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage

fastvideo/v1/pipelines/stages/decoding.py

@@ -10,7 +10,6 @@ from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
 from fastvideo.v1.utils import PRECISION_TO_TYPE
-from fastvideo.v1.models.vaes.common import ParallelTiledVAE
 
 logger = init_logger(__name__)
 
@@ -23,8 +22,8 @@ class DecodingStage(PipelineStage):
     output format (e.g., pixel values).
     """
 
-    def __init__(self, vae: ParallelTiledVAE) -> None:
-        self.vae: ParallelTiledVAE = vae
+    def __init__(self, vae) -> None:
+        self.vae = vae
 
     def forward(
         self,

fastvideo/v1/pipelines/stages/denoising.py

@@ -13,12 +13,11 @@ from tqdm.auto import tqdm
 
 from fastvideo.v1.attention import get_attn_backend
 from fastvideo.v1.distributed import (get_sequence_model_parallel_rank,
-                                      get_sequence_model_parallel_world_size,
-                                      get_world_group)
+                                      get_sequence_model_parallel_world_size)
 from fastvideo.v1.distributed.communication_op import (
     sequence_model_parallel_all_gather)
-from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
+from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
@@ -179,9 +178,10 @@ class DenoisingStage(PipelineStage):
                     self.attn_backend = get_attn_backend(
                         head_size=attn_head_size,
                         dtype=torch.float16,  # TODO(will): hack
-                        supported_attention_backends=(
+                        supported_attention_backends=[
                             _Backend.SLIDING_TILE_ATTN, _Backend.FLASH_ATTN,
-                            _Backend.TORCH_SDPA)  # hack
+                            _Backend.TORCH_SDPA
+                        ]  # hack
                     )
                     if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
                         self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
@@ -304,11 +304,7 @@ class DenoisingStage(PipelineStage):
         Returns:
             A tqdm progress bar.
         """
-        local_rank = get_world_group().local_rank
-        if local_rank == 0:
-            return tqdm(iterable=iterable, total=total)
-        else:
-            return tqdm(iterable=iterable, total=total, disable=True)
+        return tqdm(iterable=iterable, total=total)
 
     def rescale_noise_cfg(self,
                           noise_cfg,

fastvideo/v1/pipelines/stages/encoding.py

@@ -15,7 +15,6 @@ from fastvideo.v1.models.vision_utils import (get_default_height_width,
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
 from fastvideo.v1.utils import PRECISION_TO_TYPE
-from fastvideo.v1.models.vaes.common import ParallelTiledVAE
 
 logger = init_logger(__name__)
 
@@ -28,8 +27,8 @@ class EncodingStage(PipelineStage):
     input format (e.g., latents).
     """
 
-    def __init__(self, vae: ParallelTiledVAE) -> None:
-        self.vae: ParallelTiledVAE = vae
+    def __init__(self, vae) -> None:
+        self.vae = vae
 
     def forward(
         self,

fastvideo/v1/pipelines/stages/latent_preparation.py

@@ -6,6 +6,7 @@ from diffusers.utils.torch_utils import randn_tensor
 
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.vaes.common import ParallelTiledVAE
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
 
@@ -20,9 +21,10 @@ class LatentPreparationStage(PipelineStage):
     denoised during the diffusion process.
     """
 
-    def __init__(self, scheduler) -> None:
+    def __init__(self, scheduler, vae=None) -> None:
         super().__init__()
         self.scheduler = scheduler
+        self.vae = vae
 
     def forward(
         self,
@@ -42,7 +44,7 @@ class LatentPreparationStage(PipelineStage):
 
         # Adjust video length based on VAE version if needed
         if hasattr(self, 'adjust_video_length'):
-            batch = self.adjust_video_length(batch, fastvideo_args)
+            batch = self.adjust_video_length(self.vae, batch, fastvideo_args)
         # Determine batch size
         if isinstance(batch.prompt, list):
             batch_size = len(batch.prompt)
@@ -68,14 +70,15 @@ class LatentPreparationStage(PipelineStage):
             raise ValueError("Height and width must be provided")
 
         assert fastvideo_args.num_channels_latents is not None
+        assert fastvideo_args.vae_scale_factor is not None
 
         # Calculate latent shape
         shape = (
             batch_size,
             fastvideo_args.num_channels_latents,
             num_frames,
-            height // fastvideo_args.vae_config.arch_config.spatial_compression_ratio,
-            width // fastvideo_args.vae_config.arch_config.spatial_compression_ratio,
+            height // fastvideo_args.vae_scale_factor,
+            width // fastvideo_args.vae_scale_factor,
         )
 
         # Validate generator if it's a list
@@ -103,7 +106,7 @@ class LatentPreparationStage(PipelineStage):
 
         return batch
 
-    def adjust_video_length(self, batch: ForwardBatch,
+    def adjust_video_length(self, vae: ParallelTiledVAE, batch: ForwardBatch,
                             fastvideo_args: FastVideoArgs) -> ForwardBatch:
         """
         Adjust video length based on VAE version.
@@ -116,7 +119,7 @@ class LatentPreparationStage(PipelineStage):
             The batch with adjusted video length.
         """
         video_length = batch.num_frames
-        temporal_scale_factor = fastvideo_args.vae_config.arch_config.temporal_compression_ratio
+        temporal_scale_factor = vae.temporal_compression_ratio if vae is not None else 4
         # TODO
         batch.num_frames = (video_length - 1) // temporal_scale_factor + 1
         return batch

fastvideo/v1/pipelines/stages/llama_encoding.py

@@ -9,8 +9,8 @@ from typing import TypedDict
 
 import torch
 
-from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
+from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage

fastvideo/v1/pipelines/stages/t5_encoding.py

@@ -6,8 +6,8 @@ This module contains implementations of prompt encoding stages for diffusion pip
 """
 import torch
 
-from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
+from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage

fastvideo/v1/pipelines/wan/wan_i2v_pipeline.py

@@ -9,13 +9,10 @@ using the modular pipeline architecture.
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
-
-# isort: off
 from fastvideo.v1.pipelines.stages import (
     CLIPImageEncodingStage, ConditioningStage, DecodingStage, DenoisingStage,
     EncodingStage, InputValidationStage, LatentPreparationStage,
     T5EncodingStage, TimestepPreparationStage)
-# isort: on
 
 logger = init_logger(__name__)
 
@@ -54,7 +51,8 @@ class WanImageToVideoPipeline(ComposedPipelineBase):
 
         self.add_stage(stage_name="latent_preparation_stage",
                        stage=LatentPreparationStage(
-                           scheduler=self.get_module("scheduler")))
+                           scheduler=self.get_module("scheduler"),
+                           vae=self.get_module("vae")))
 
         self.add_stage(stage_name="image_latent_preparation_stage",
                        stage=EncodingStage(vae=self.get_module("vae")))
@@ -71,6 +69,9 @@ class WanImageToVideoPipeline(ComposedPipelineBase):
         """
         Initialize the pipeline.
         """
+        vae_scale_factor = self.get_module("vae").spatial_compression_ratio
+        fastvideo_args.vae_scale_factor = vae_scale_factor
+
         num_channels_latents = self.get_module("transformer").out_channels
         fastvideo_args.num_channels_latents = num_channels_latents
 

fastvideo/v1/pipelines/wan/wan_pipeline.py

@@ -47,7 +47,8 @@ class WanPipeline(ComposedPipelineBase):
 
         self.add_stage(stage_name="latent_preparation_stage",
                        stage=LatentPreparationStage(
-                           scheduler=self.get_module("scheduler")))
+                           scheduler=self.get_module("scheduler"),
+                           vae=self.get_module("vae")))
 
         self.add_stage(stage_name="denoising_stage",
                        stage=DenoisingStage(
@@ -61,6 +62,9 @@ class WanPipeline(ComposedPipelineBase):
         """
         Initialize the pipeline.
         """
+        vae_scale_factor = self.get_module("vae").spatial_compression_ratio
+        fastvideo_args.vae_scale_factor = vae_scale_factor
+
         num_channels_latents = self.get_module("transformer").in_channels
         fastvideo_args.num_channels_latents = num_channels_latents
 

fastvideo/v1/tests/vaes/test_hunyuan_vae.py

@@ -8,11 +8,8 @@ import torch
 from safetensors.torch import load_file
 
 from fastvideo.v1.logger import init_logger
-# from fastvideo.v1.models.vaes.hunyuanvae import (
-#     AutoencoderKLHunyuanVideo as MyHunyuanVAE)
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.models.loader.component_loader import VAELoader
-from fastvideo.v1.configs.models.vaes import HunyuanVAEConfig
+from fastvideo.v1.models.vaes.hunyuanvae import (
+    AutoencoderKLHunyuanVideo as MyHunyuanVAE)
 from fastvideo.v1.utils import maybe_download_model
 
 logger = init_logger(__name__)
@@ -34,14 +31,21 @@ REFERENCE_LATENT = -105.51324462890625
 @pytest.mark.usefixtures("distributed_setup")
 def test_hunyuan_vae():
     device = torch.device("cuda:0")
-    precision = torch.bfloat16
-    precision_str = "bf16"
-    args = FastVideoArgs(model_path=VAE_PATH, vae_precision=precision_str)
-    args.device = device
-    args.vae_config = HunyuanVAEConfig()
+    # Initialize the two model implementations
+    config = json.load(open(CONFIG_PATH))
+    config.pop("_class_name")
+    config.pop("_diffusers_version")
+    model = MyHunyuanVAE(**config).to(torch.bfloat16)
 
-    loader = VAELoader()
-    model = loader.load(VAE_PATH, "", args)
+    loaded = load_file(os.path.join(VAE_PATH,
+                                    "diffusion_pytorch_model.safetensors"))
+    model.load_state_dict(loaded)
+
+    # Set model to eval mode
+    model.eval()
+
+    # Move to GPU
+    model = model.to(device)
 
     model.enable_tiling(tile_sample_min_height=32,
                          tile_sample_min_width=32,

fastvideo/v1/tests/vaes/test_wan_vae.py

@@ -9,7 +9,6 @@ from diffusers import AutoencoderKLWan
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.loader.component_loader import VAELoader
-from fastvideo.v1.configs.models.vaes import WanVAEConfig
 from fastvideo.v1.utils import maybe_download_model
 
 logger = init_logger(__name__)
@@ -31,11 +30,9 @@ def test_wan_vae():
     precision_str = "bf16"
     args = FastVideoArgs(model_path=VAE_PATH, vae_precision=precision_str)
     args.device = device
-    args.vae_config = WanVAEConfig()
 
     loader = VAELoader()
     model2 = loader.load(VAE_PATH, "", args)
-    assert model2.use_feature_cache # Default to use the original WanVAE algorithm
 
     model1 = AutoencoderKLWan.from_pretrained(
         VAE_PATH, torch_dtype=precision).to(device).eval()
@@ -51,48 +48,43 @@ def test_wan_vae():
                                32,
                                device=device,
                                dtype=precision)
-    # latent_tensor = torch.randn(batch_size,
-    #                             16,
-    #                             21,
-    #                             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).latent_dist
+        latent1 = model1.encode(input_tensor).latent_dist.mean
         print("--------------------------------")
-        latent2 = model2.encode(input_tensor)
+        latent2 = model2.encode(input_tensor).mean
         # Check if latents have the same shape
-        assert latent1.mean.shape == latent2.mean.shape, f"Latent shapes don't match: {latent1.mean.shape} vs {latent2.mean.shape}"
+        assert latent1.shape == latent2.shape, f"Latent shapes don't match: {latent1.shape} vs {latent2.shape}"
+        assert latent1.shape == latent2.shape, f"Latent shapes don't match: {latent1.shape} vs {latent2.shape}"
         # Check if latents are similar
-        max_diff_encode = torch.max(torch.abs(latent1.mean - latent2.mean))
-        mean_diff_encode = torch.mean(torch.abs(latent1.mean - latent2.mean))
+        max_diff_encode = torch.max(torch.abs(latent1 - latent2))
+        mean_diff_encode = torch.mean(torch.abs(latent1 - latent2))
         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()}"
+        assert mean_diff_encode < 5e-1, f"Encoded latents differ significantly: mean 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
-
-        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)
+        latents_mean = (torch.tensor(model1.config.latents_mean).view(
+            1, model1.config.z_dim, 1, 1, 1).to(latent_tensor.device,
+                                                latent_tensor.dtype))
+        latents_std = 1.0 / torch.tensor(model1.config.latents_std).view(
+            1, model1.config.z_dim, 1, 1, 1).to(latent_tensor.device,
+                                                latent_tensor.dtype)
+        latent_tensor = latent_tensor / latents_std + latents_mean
+        output2 = model2.decode(latent_tensor)
+        output1 = model1.decode(latent_tensor).sample
         # Check if outputs have the same shape
         assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
 
@@ -103,4 +95,4 @@ def test_wan_vae():
                     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 = {mean_diff_decode.item()}"
+        assert mean_diff_decode < 1e-1, f"Decoded outputs differ significantly: mean diff = {mean_diff_decode.item()}"

fastvideo/v1/utils.py

@@ -2,23 +2,24 @@
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/utils.py
 
 import argparse
-import ctypes
 import hashlib
 import importlib
+import ctypes
+import signal
 import inspect
 import json
 import math
+import traceback
 import os
-import signal
 import sys
 import tempfile
-import traceback
-from dataclasses import asdict, fields, is_dataclass
-from functools import partial, wraps
-from typing import (Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar,
-                    Union, cast)
-
+from functools import wraps, partial
+from typing import Any, Dict, List, Optional, Type, TypeVar, Union, cast, Callable, Tuple
+from dataclasses import asdict, fields
 import cloudpickle
+import zmq
+import psutil
+
 import filelock
 import torch
 import yaml
@@ -477,7 +478,6 @@ def maybe_download_model_index(model_name_or_path: str) -> Dict[str, Any]:
         The parsed model_index.json as a dictionary
     """
     import tempfile
-
     from huggingface_hub import hf_hub_download
 
     # If it's a local path, verify it directly
@@ -550,29 +550,70 @@ def run_method(obj: Any, method: Union[str, bytes, Callable], args: tuple[Any],
         func = partial(method, obj)  # type: ignore
     return func(*args, **kwargs)
 
-def shallow_asdict(obj):
-    if not is_dataclass(obj):
-        raise TypeError("Expected dataclass instance")
-    return {f.name: getattr(obj, f.name) for f in fields(obj)}
 
-def kill_itself_when_parent_died() -> None:
+def diff_keys(a, b):
+    return [k for k in asdict(a) if asdict(a)[k] != asdict(b)[k]]
+
+
+def update_in_place(target, source, ignore_fields=()):
+    for f in fields(target):
+        if hasattr(source, f.name) and f.name not in list(ignore_fields):
+            setattr(target, f.name, getattr(source, f.name))
+
+
+def get_zmq_socket(
+    context: zmq.Context, socket_type: zmq.SocketType, endpoint: str, bind: bool
+):
+    mem = psutil.virtual_memory()
+    total_mem = mem.total / 1024**3
+    available_mem = mem.available / 1024**3
+    if total_mem > 32 and available_mem > 16:
+        buf_size = int(0.5 * 1024**3)
+    else:
+        buf_size = -1
+
+    socket = context.socket(socket_type)
+
+    def set_send_opt():
+        socket.setsockopt(zmq.SNDHWM, 0)
+        socket.setsockopt(zmq.SNDBUF, buf_size)
+
+    def set_recv_opt():
+        socket.setsockopt(zmq.RCVHWM, 0)
+        socket.setsockopt(zmq.RCVBUF, buf_size)
+
+    if socket_type == zmq.PUSH:
+        set_send_opt()
+    elif socket_type == zmq.PULL:
+        set_recv_opt()
+    elif socket_type == zmq.DEALER:
+        set_send_opt()
+        set_recv_opt()
+    else:
+        raise ValueError(f"Unsupported socket type: {socket_type}")
+
+    if bind:
+        socket.bind(endpoint)
+    else:
+        socket.connect(endpoint)
+
+    return socket
+
+def kill_itself_when_parent_died():
     # if sys.platform == "linux":
-    # sigkill this process when parent worker manager dies
+        # sigkill this process when parent worker manager dies
     PR_SET_PDEATHSIG = 1
     libc = ctypes.CDLL("libc.so.6")
     libc.prctl(PR_SET_PDEATHSIG, signal.SIGKILL)
     # else:
     #     logger.warning("kill_itself_when_parent_died is only supported in linux.")
 
-
-def get_exception_traceback() -> str:
+def get_exception_traceback():
     etype, value, tb = sys.exc_info()
     err_str = "".join(traceback.format_exception(etype, value, tb))
     return err_str
 
-
 class TypeBasedDispatcher:
-
     def __init__(self, mapping: List[Tuple[Type, Callable]]):
         self._mapping = mapping
 
@@ -580,4 +621,4 @@ class TypeBasedDispatcher:
         for ty, fn in self._mapping:
             if isinstance(obj, ty):
                 return fn(obj)
-        raise ValueError(f"Invalid object: {obj}")
+        raise ValueError(f"Invalid object: {obj}")

fastvideo/v1/worker/executor.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
-from typing import (Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union,
-                    cast)
+from typing import Callable, List, Tuple, Any, Optional, Union, TypeVar, Dict
 
+import torch
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.pipelines import ForwardBatch
 from fastvideo.v1.utils import init_logger
@@ -10,7 +10,6 @@ logger = init_logger(__name__)
 
 _R = TypeVar("_R")
 
-
 class Executor(ABC):
 
     def __init__(self, fastvideo_args: FastVideoArgs):
@@ -24,28 +23,25 @@ class Executor(ABC):
 
     @classmethod
     def get_class(cls, fastvideo_args: FastVideoArgs) -> type["Executor"]:
-        if fastvideo_args.distributed_executor_backend == "mp":
+        if fastvideo_args.distributed_executor_backend == "torch":
+            raise ValueError("Torch is not supported yet")
+            from fastvideo.v1.worker.torchrun_executor import TorchRunExecutor
+            return TorchRunExecutor
+        elif fastvideo_args.distributed_executor_backend == "mp":
             from fastvideo.v1.worker.multiproc_executor import MultiprocExecutor
-            return cast(type["Executor"], MultiprocExecutor)
+            return MultiprocExecutor
         else:
             raise ValueError(
                 f"Unsupported distributed executor backend: {fastvideo_args.distributed_executor_backend}"
             )
 
-    def execute_forward(
+    def execute_model(
         self,
         forward_batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
     ) -> ForwardBatch:
-        outputs: List[Dict[str,
-                           Any]] = self.collective_rpc("execute_forward",
-                                                       kwargs={
-                                                           "forward_batch":
-                                                           forward_batch,
-                                                           "fastvideo_args":
-                                                           fastvideo_args
-                                                       })
-        return cast(ForwardBatch, outputs[0]["output_batch"])
+        output = self.collective_rpc("execute_forward",
+                                     args=(forward_batch, ))
+        return output
 
     @abstractmethod
     def collective_rpc(self,
@@ -75,4 +71,4 @@ class Executor(ABC):
             It is recommended to use this API to only pass control messages,
             and set up data-plane communication to pass data.
         """
-        raise NotImplementedError
+        raise NotImplementedError

fastvideo/v1/worker/gpu_worker.py

@@ -1,58 +1,128 @@
-import contextlib
-import faulthandler
-import gc
-import multiprocessing as mp
-import os
-import signal
-import sys
-from typing import Any, Dict, Optional, TextIO, cast
+from abc import ABC
 
-import psutil
 import torch
 
-from fastvideo.v1.distributed import (cleanup_dist_env_and_memory,
-                                      init_distributed_environment,
-                                      initialize_model_parallel)
-from fastvideo.v1.fastvideo_args import FastVideoArgs
+from typing import Optional, List
+import multiprocessing as mp
+import setproctitle
+import psutil
+import faulthandler
+import signal
+import traceback
+
+from fastvideo.v1.fastvideo_args import FastVideoArgs, set_current_fastvideo_args
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines import ForwardBatch, build_pipeline
-from fastvideo.v1.utils import (get_exception_traceback,
-                                kill_itself_when_parent_died)
+from fastvideo.v1.pipelines import ForwardBatch
+from fastvideo.v1.utils import run_method, update_environment_variables
+from fastvideo.v1.distributed import (
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+import zmq
+from fastvideo.v1.utils import get_zmq_socket, kill_itself_when_parent_died, get_exception_traceback, TypeBasedDispatcher
+from fastvideo.v1.worker.io_struct import RpcReqInput, RpcReqOutput, GenerateRequest
+
 
 logger = init_logger(__name__)
 
-# ANSI color codes
-CYAN = '\033[1;36m'
-RESET = '\033[0;0m'
+# class WorkerWrapper5
+#     def __init__(self, fastvideo_args: FastVideoArgs):
+#         self.fastvideo_args = fastvideo_args
+    
+#     def update_environment_variables(self, envs_list: List[Dict[str, str]]) -> None:
+#         envs = envs_list[self.rpc_rank]
+#         key = 'CUDA_VISIBLE_DEVICES'
+#         if key in envs and key in os.environ:
+#             # overwriting CUDA_VISIBLE_DEVICES is desired behavior
+#             # suppress the warning in `update_environment_variables`
+#             del os.environ[key]
+#         update_environment_variables(envs)
+    
+#     def init_worker(self, all_kwargs: List[Dict[str, Any]]) -> None:
+#         kwargs = all_kwargs[self.rpc_rank]
+#         self.fastvideo_args = kwargs.get("fastvideo_args", None)
+#         assert self.fastvideo_args is not None, (
+#             "fastvideo_args is required to initialize the worker")
+
+#         with set_current_fastvideo_args(self.fastvideo_args):
+#             self.worker = Worker(**kwargs)
+#             assert self.worker is not None
+
+#     def init_device(self):
+#         self.worker.init_device()
+
+#     def execute_method(self, method: Union[str, bytes], *args, **kwargs):
+#         try:
+#             # method resolution order:
+#             # if a method is defined in this class, it will be called directly.
+#             # otherwise, since we define `__getattr__` and redirect attribute
+#             # query to `self.worker`, the method will be called on the worker.
+#             return run_method(self, method, args, kwargs)
+#         except Exception as e:
+#             # if the driver worker also execute methods,
+#             # exceptions in the rest worker may cause deadlock in rpc like ray
+#             # see https://github.com/vllm-project/vllm/issues/3455
+#             # print the error and inform the user to solve the error
+#             msg = (f"Error executing method {method!r}. "
+#                    "This might cause deadlock in distributed execution.")
+#             logger.exception(msg)
+#             raise e
+        
+#     def __getattr__(self, attr):
+#         return getattr(self.worker, attr)
 
 
-class Worker:
+class Worker(ABC):
 
     def __init__(self, fastvideo_args: FastVideoArgs, local_rank: int,
-                 rank: int, pipe):
+                 rank: int, is_driver_worker: bool = False):
         self.fastvideo_args = fastvideo_args
         self.local_rank = local_rank
         self.rank = rank
-        # TODO(will): don't hardcode this
+        # TODO: don't hardcode this
         self.distributed_init_method = "env://"
-        self.pipe = pipe
+        self.is_driver_worker = is_driver_worker
 
-        self.init_device()
+        context = zmq.Context(2)
+        self.recv_from_rpc = get_zmq_socket(context, zmq.DEALER, "ipc://fastvideo_rpc_broadcast", False)
+        # self.send_to_rpc = get_zmq_socket(context, zmq.PUSH, "inproc://fastvideo_rpc_broadcast", True)
 
         # Init request dispatcher
-        # TODO(will): add request dispatcher: use TypeBasedDispatcher from
-        # utils.py
-        # self._request_dispatcher = TypeBasedDispatcher(
-        #     [
-        # (RpcReqInput, self.handle_rpc_request),
-        # (GenerateRequest, self.handle_generate_request),
-        # (ExpertDistributionReq, self.expert_distribution_handle),
-        #     ]
-        # )
+        self._request_dispatcher = TypeBasedDispatcher(
+            [
+                # (TokenizedGenerateReqInput, self.handle_generate_request),
+                # (TokenizedEmbeddingReqInput, self.handle_embedding_request),
+                # (FlushCacheReq, self.flush_cache_wrapped),
+                # (AbortReq, self.abort_request),
+                # (OpenSessionReqInput, self.open_session),
+                # (CloseSessionReqInput, self.close_session),
+                # (UpdateWeightFromDiskReqInput, self.update_weights_from_disk),
+                # (InitWeightsUpdateGroupReqInput, self.init_weights_update_group),
+                # (
+                #     UpdateWeightsFromDistributedReqInput,
+                #     self.update_weights_from_distributed,
+                # ),
+                # (UpdateWeightsFromTensorReqInput, self.update_weights_from_tensor),
+                # (GetWeightsByNameReqInput, self.get_weights_by_name),
+                # (ReleaseMemoryOccupationReqInput, self.release_memory_occupation),
+                # (ResumeMemoryOccupationReqInput, self.resume_memory_occupation),
+                # (ProfileReq, self.profile),
+                # (GetInternalStateReq, self.get_internal_state),
+                # (SetInternalStateReq, self.set_internal_state),
+                (RpcReqInput, self.handle_rpc_request),
+                (GenerateRequest, self.handle_generate_request),
+                # (ExpertDistributionReq, self.expert_distribution_handle),
+            ]
+        )
+    
+    def handle_rpc_request(self, req: RpcReqInput) -> RpcReqOutput:
+        pass
+    
+    def handle_generate_request(self, req: GenerateRequest) -> None:
+        logger.info(f"Worker {self.rank} received generate request")
+        pass
 
-    def init_device(self) -> None:
-        """Initialize the device for the worker."""
-        assert self.fastvideo_args.device_str is not None
+    def init_device(self):
         if self.fastvideo_args.device_str.startswith("cuda"):
             # torch.distributed.all_reduce does not free the input tensor until
             # the synchronization point. This causes the memory usage to grow
@@ -67,87 +137,47 @@ class Worker:
             self.device = torch.device(f"cuda:{self.local_rank}")
             torch.cuda.set_device(self.device)
 
-            # _check_if_gpu_supports_dtype(self.model_config.dtype)
+            _check_if_gpu_supports_dtype(self.model_config.dtype)
             gc.collect()
             torch.cuda.empty_cache()
             self.init_gpu_memory = torch.cuda.mem_get_info()[0]
         else:
-            raise ValueError(
-                f"Unsupported device: {self.fastvideo_args.device_str}")
-
-        os.environ["MASTER_ADDR"] = "localhost"
-        os.environ["MASTER_PORT"] = "29503"
-        os.environ["LOCAL_RANK"] = str(self.local_rank)
-        os.environ["RANK"] = str(self.rank)
+            raise ValueError(f"Unsupported device: {self.fastvideo_args.device_str}")
 
         # Initialize the distributed environment.
         init_worker_distributed_environment(self.fastvideo_args, self.rank,
                                             self.distributed_init_method,
                                             self.local_rank)
-
+        
         self.pipeline = build_pipeline(self.fastvideo_args)
 
-    def execute_forward(self, forward_batch: ForwardBatch,
-                        fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        self.fastvideo_args.num_inference_steps = fastvideo_args.num_inference_steps
-        output_batch = self.pipeline.forward(forward_batch, self.fastvideo_args)
-        return cast(ForwardBatch, output_batch)
+    def execute_forward(self, forward_batch: ForwardBatch) -> torch.Tensor:
+        return self.pipeline.forward(forward_batch)
+    
+    def recv_requests(self) -> List[ForwardBatch]:
+        pass
 
-    def shutdown(self) -> Dict[str, Any]:
-        """Gracefully shut down the worker process"""
-        logger.info("Worker %d shutting down...",
-                    self.rank,
-                    local_main_process_only=False)
-        # Clean up resources
-        if hasattr(self, 'pipeline') and self.pipeline is not None:
-            # Clean up pipeline resources if needed
-            pass
-        # Release CUDA resources
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-        # Destroy the distributed environment
-        cleanup_dist_env_and_memory(shutdown_ray=False)
-
-        logger.info("Worker %d shutdown complete",
-                    self.rank,
-                    local_main_process_only=False)
-        return {"status": "shutdown_complete"}
+    def send_response(self, output: torch.Tensor) -> None:
+        pass
 
     def event_loop(self) -> None:
         """Event loop for the worker."""
-        logger.info("Worker %d starting event loop...",
-                    self.rank,
-                    local_main_process_only=False)
+        logger.info(f"Worker {self.rank} starting event loop...")
+        recv_rpc = self.recv_from_rpc.recv_pyobj()
+        assert recv_rpc == "wait_until_ready"
+        logger.info(f"Worker {self.rank} received wait_until_ready")
+        self.recv_from_rpc.send_pyobj(f"ready{self.rank}")
+        logger.info(f"Worker {self.rank} sent ready")
+        logger.info(f"Worker {self.rank} started event loop")
         while True:
-            recv_rpc = self.pipe.recv()
-            method_name = recv_rpc.get('method')
-
-            # Handle shutdown request
-            if method_name == 'shutdown':
-                response = self.shutdown()
-                with contextlib.suppress(Exception):
-                    self.pipe.send(response)
-                break  # Exit the loop
-
-            # Handle regular RPC calls
-            if method_name == 'execute_forward':
-                forward_batch = recv_rpc['kwargs']['forward_batch']
-                fastvideo_args = recv_rpc['kwargs']['fastvideo_args']
-                output_batch = self.execute_forward(forward_batch,
-                                                    fastvideo_args)
-                self.pipe.send({"output_batch": output_batch.output.cpu()})
-            else:
-                # Handle other methods dynamically if needed
-                args = recv_rpc.get('args', ())
-                kwargs = recv_rpc.get('kwargs', {})
-                if hasattr(self, method_name):
-                    method = getattr(self, method_name)
-                    result = method(*args, **kwargs)
-                    self.pipe.send(result)
-                else:
-                    self.pipe.send({"error": f"Unknown method: {method_name}"})
-
+            logger.info(f"Worker {self.rank} waiting for RPC")
+            recv_rpc = self.recv_from_rpc.recv_pyobj()
+            logger.info(f"Received RPC: {recv_rpc}")
+            # assert isinstance(recv_rpc, RpcReqInput)
+            inputs = self._request_dispatcher(recv_rpc)
+            output = self.execute_forward(inputs)
+            self.send_response(output)
+        
 
 def init_worker_distributed_environment(
     fastvideo_args: FastVideoArgs,
@@ -175,64 +205,26 @@ def init_worker_distributed_environment(
 
 
 def run_worker_process(fastvideo_args: FastVideoArgs, local_rank: int,
-                       rank: int, pipe):
-    # Add process-specific prefix to stdout and stderr
-    process_name = mp.current_process().name
-    pid = os.getpid()
-    _add_prefix(sys.stdout, process_name, pid)
-    _add_prefix(sys.stderr, process_name, pid)
-
-    # Config the process
-    kill_itself_when_parent_died()
-    faulthandler.enable()
-    parent_process = psutil.Process().parent()
-
-    logger.info("Worker %d initializing...",
-                rank,
-                local_main_process_only=False)
+                 rank: int, pipe_writer: mp.Pipe):
+    print(f"Worker {rank} starting...")
     try:
-        worker = Worker(fastvideo_args, local_rank, rank, pipe)
-        logger.info("Worker %d sending ready", rank)
-        pipe.send({
-            "status": "ready",
-            "local_rank": local_rank,
-        })
-        try:
-            worker.event_loop()
-        except KeyboardInterrupt:
-            logger.info(
-                "Worker %d received KeyboardInterrupt, shutting down...",
-                rank,
-                local_main_process_only=False)
-            worker.shutdown()
+        # Config the process
+        kill_itself_when_parent_died()
+        prefix = f"{local_rank}"
+        setproctitle.setproctitle(f"fastvideo::gpu_worker{prefix.replace(' ', '_')}")
+        faulthandler.enable()
+        parent_process = psutil.Process().parent()
+
+        logger.info(f"Worker {rank} initializing...")
+        worker = Worker(fastvideo_args, local_rank, rank, False)
+        pipe_writer.send(
+            {
+                "status": "ready",
+                "local_rank": local_rank,
+            }
+        )
+        worker.event_loop()
     except Exception:
         traceback = get_exception_traceback()
-        logger.error("Worker %d hit an exception: %s", rank, traceback)
+        logger.error(f"Worker {rank} hit an exception: {traceback}")
         parent_process.send_signal(signal.SIGQUIT)
-
-
-def _add_prefix(file: TextIO, worker_name: str, pid: int) -> None:
-    """Prepend each output line with process-specific prefix"""
-
-    prefix = f"{CYAN}({worker_name} pid={pid}){RESET} "
-    file_write = file.write
-
-    def write_with_prefix(s: str):
-        if not s:
-            return
-        if file.start_new_line:  # type: ignore[attr-defined]
-            file_write(prefix)
-        idx = 0
-        while (next_idx := s.find('\n', idx)) != -1:
-            next_idx += 1
-            file_write(s[idx:next_idx])
-            if next_idx == len(s):
-                file.start_new_line = True  # type: ignore[attr-defined]
-                return
-            file_write(prefix)
-            idx = next_idx
-        file_write(s[idx:])
-        file.start_new_line = False  # type: ignore[attr-defined]
-
-    file.start_new_line = True  # type: ignore[attr-defined]
-    file.write = write_with_prefix  # type: ignore[method-assign]

fastvideo/v1/worker/io_struct.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+from typing import Optional, Dict
+
+@dataclass
+class RpcReqInput:
+    method: str
+    parameters: Optional[Dict] = None
+
+
+@dataclass
+class RpcReqOutput:
+    success: bool
+    message: str
+
+@dataclass
+class GenerateRequest:
+    prompt: str
+    seed: int

fastvideo/v1/worker/multiproc_executor.py

@@ -1,15 +1,18 @@
-import atexit
-import contextlib
+import signal
+import psutil
 import multiprocessing as mp
 import time
+import pickle
+import cloudpickle
+import zmq
+from typing import List, Callable, Any, Optional, Union
 from multiprocessing.process import BaseProcess
-from typing import Any, Callable, List, Optional, Union, cast
 
+from fastvideo.v1.worker.executor import Executor
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.v1.worker.executor import Executor
 from fastvideo.v1.worker.gpu_worker import run_worker_process
+from fastvideo.v1.utils import get_zmq_socket
 
 logger = init_logger(__name__)
 
@@ -17,137 +20,96 @@ logger = init_logger(__name__)
 class MultiprocExecutor(Executor):
 
     def _init_executor(self) -> None:
-        self.world_size = self.fastvideo_args.num_gpus
-        self.shutting_down = False
+        # The child processes will send SIGUSR1 when unrecoverable
+        # errors happen.
+        def sigusr1_handler(signum, frame):
+            logger.fatal(
+                "MulitprocExecutor got fatal signal from worker processes, "
+                "shutting down. See stack trace above for root cause issue.")
+            # Propagate error up to parent process.
+            parent_process = psutil.Process().parent()
+            parent_process.send_signal(signal.SIGUSR1)
+            self.shutdown()
 
-        # this will force the use of the `spawn` multiprocessing start if cuda
-        # is initialized
+        signal.signal(signal.SIGUSR1, sigusr1_handler)
+
+        self.world_size = self.fastvideo_args.num_gpus
+
+        # Set mp start method
         mp.set_start_method("spawn", force=True)
 
         self.workers: List[BaseProcess] = []
-        self.worker_pipes = []
-
-        # Create pipes and start workers
+        self.worker_pipe_readers: List[mp.Pipe] = []
         for rank in range(self.world_size):
-            executor_pipe, worker_pipe = mp.Pipe(duplex=True)
-            self.worker_pipes.append(executor_pipe)
+            reader, writer = mp.Pipe(duplex=False)
+            self.worker_pipe_readers.append(reader)
 
             worker = mp.Process(target=run_worker_process,
-                                name=f"FVWorkerProc-{rank}",
-                                kwargs=dict(fastvideo_args=self.fastvideo_args,
-                                            local_rank=rank,
-                                            rank=rank,
-                                            pipe=worker_pipe))
+                                     args=(self.fastvideo_args, rank, rank,
+                                            writer))
             worker.start()
             self.workers.append(worker)
-        logger.info("Workers: %s", self.workers)
+        logger.info(f"Workers: {self.workers}")
+        for reader in self.worker_pipe_readers:
+            data = reader.recv()
+            assert data["status"] == "ready"
+            logger.info(f"Worker {data['local_rank']} ready")
 
-        # Wait for all workers to be ready
-        for idx, pipe in enumerate(self.worker_pipes):
-            data = pipe.recv()
-            if data["status"] != "ready" or data["local_rank"] != idx:
-                raise RuntimeError(f"Worker {idx} failed to start")
-        logger.info("%d workers ready", self.world_size)
+        context = zmq.Context(2)
+        self.send_to_rpc = get_zmq_socket(
+            context, zmq.DEALER, "ipc://fastvideo_rpc_broadcast", True
+        )
+        logger.info("Sending wait_until_ready to workers")
+        self.send_to_rpc.send_pyobj("wait_until_ready")
+        self.send_to_rpc.send_pyobj("wait_until_ready")
+        logger.info("Sent wait_until_ready to workers")
+        recv_req = self.send_to_rpc.recv_pyobj(zmq.BLOCKY)
+        print(f"Received wait_until_ready from workers: {recv_req}")
+        recv_req = self.send_to_rpc.recv_pyobj(zmq.BLOCKY)
+        print(f"Received wait_until_ready from workers: {recv_req}")
 
-        # Register shutdown on exit
-        atexit.register(self.shutdown)
 
-    def execute_forward(self, forward_batch: ForwardBatch,
-                        fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        responses = self.collective_rpc("execute_forward",
-                                        kwargs={
-                                            "forward_batch": forward_batch,
-                                            "fastvideo_args": fastvideo_args
-                                        })
-        return cast(ForwardBatch, responses[0]["output_batch"])
+
+
+        # self.collective_rpc("wait_until_ready")
 
     def collective_rpc(self,
                        method: Union[str, Callable],
                        timeout: Optional[float] = None,
                        args: tuple = (),
                        kwargs: Optional[dict] = None) -> list[Any]:
+        start_time = time.monotonic()
         kwargs = kwargs or {}
 
+        # NOTE: If the args are heterogeneous, then we pack them into a list,
+        # and unpack them in the method of every worker, because every worker
+        # knows their own rank.
         try:
-            for pipe in self.worker_pipes:
-                pipe.send({"method": method, "args": args, "kwargs": kwargs})
+            if isinstance(method, str):
+                send_method = method
+            else:
+                send_method = cloudpickle.dumps(
+                    method, protocol=pickle.HIGHEST_PROTOCOL)
+            self.rpc_broadcast_mq.enqueue((send_method, args, kwargs))
+
+            responses = [None] * self.world_size
+            for w in self.workers:
+                dequeue_timeout = timeout - (time.monotonic() - start_time
+                                             ) if timeout is not None else None
+                status, result = w.worker_response_mq.dequeue(
+                    timeout=dequeue_timeout)
+
+                if status != WorkerProc.ResponseStatus.SUCCESS:
+                    if isinstance(result, Exception):
+                        raise result
+                    else:
+                        raise RuntimeError("Worker failed")
+
+                responses[w.rank] = result
 
-            responses = []
-            for pipe in self.worker_pipes:
-                response = pipe.recv()
-                responses.append(response)
             return responses
         except TimeoutError as e:
             raise TimeoutError(f"RPC call to {method} timed out.") from e
         except Exception as e:
             # Re-raise any other exceptions
-            raise e
-
-    def shutdown(self) -> None:
-        """Properly shut down the executor and its workers"""
-        if hasattr(self, 'shutting_down') and self.shutting_down:
-            return  # Prevent multiple shutdown calls
-
-        logger.info("Shutting down MultiprocExecutor...")
-        self.shutting_down = True
-
-        # First try gentle termination
-        try:
-            # Send termination message to all workers
-            for pipe in self.worker_pipes:
-                with contextlib.suppress(Exception):
-                    pipe.send({"method": "shutdown", "args": (), "kwargs": {}})
-
-            # Give workers some time to exit gracefully
-            start_time = time.time()
-            while time.time() - start_time < 5.0:  # 5 seconds timeout
-                if all(not worker.is_alive() for worker in self.workers):
-                    break
-                time.sleep(0.1)
-
-            # Force terminate any remaining workers
-            for worker in self.workers:
-                if worker.is_alive():
-                    worker.terminate()
-
-            # Final timeout for terminate
-            start_time = time.time()
-            while time.time() - start_time < 2.0:  # 2 seconds timeout
-                if all(not worker.is_alive() for worker in self.workers):
-                    break
-                time.sleep(0.1)
-
-            # Kill if still alive
-            for worker in self.workers:
-                if worker.is_alive():
-                    worker.kill()
-                worker.join(timeout=1.0)
-
-        except Exception as e:
-            logger.error("Error during shutdown: %s", e)
-            # Last resort, try to kill all workers
-            for worker in self.workers:
-                with contextlib.suppress(Exception):
-                    if worker.is_alive():
-                        worker.kill()
-
-        # Clean up pipes
-        for pipe in self.worker_pipes:
-            with contextlib.suppress(Exception):
-                pipe.close()
-
-        self.workers = []
-        self.worker_pipes = []
-        logger.info("MultiprocExecutor shutdown complete")
-
-    def __del__(self):
-        """Ensure cleanup on garbage collection"""
-        self.shutdown()
-
-    def __enter__(self):
-        """Support for context manager protocol"""
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        """Ensure cleanup when exiting context"""
-        self.shutdown()
+            raise e

fastvideo/v1/worker/worker_wrappers.py

@@ -0,0 +1,7 @@
+class WorkerWrapper:
+
+    def __init__(self, worker: Worker):
+        self.worker = worker
+
+    def __call__(self, *args, **kwargs):
+        return self.worker(*args, **kwargs)

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fastvideo"
-version = "0.0.2"
+version = "0.0.1.dev1"
 description = "FastVideo"
 readme = "README.md"
 requires-python = ">=3.8"

scripts/inference/v1_inference_fasthunyuan.sh

@@ -1,24 +0,0 @@
-#!/bin/bash
-
-num_gpus=4
-export MODEL_BASE=FastVideo/FastHunyuan-Diffusers
-export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
-# Note that the tp_size and sp_size should be the same and equal to the number
-# of GPUs. They are used for different parallel groups. sp_size is used for
-# dit model and tp_size is used for encoder models.
-torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
-    fastvideo/v1/sample/v1_fastvideo_inference.py \
-    --sp_size $num_gpus \
-    --tp_size $num_gpus \
-    --height 720 \
-    --width 1280 \
-    --num_frames 125 \
-    --num_inference_steps 6 \
-    --guidance_scale 1 \
-    --embedded_cfg_scale 6 \
-    --flow_shift 17 \
-    --prompt_path ./assets/prompt.txt \
-    --seed 1024 \
-    --output_path outputs_video/ \
-    --model_path $MODEL_BASE \
-    --vae-sp

scripts/inference/v1_inference_wan_STA.sh

@@ -1,27 +0,0 @@
-#!/bin/bash
-
-num_gpus=2
-export FASTVIDEO_ATTENTION_CONFIG=assets/mask_strategy_wan.json
-export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
-export MODEL_BASE=Wan-AI/Wan2.1-T2V-14B-Diffusers
-# Note that the tp_size and sp_size should be the same and equal to the number
-# of GPUs. They are used for different parallel groups. sp_size is used for
-# dit model and tp_size is used for encoder models.
-torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
-    fastvideo/v1/sample/v1_fastvideo_inference.py \
-    --sp_size $num_gpus \
-    --tp_size $num_gpus \
-    --height 768 \
-    --width 1280 \
-    --num_frames 69 \
-    --num_inference_steps 50 \
-    --fps 16 \
-    --guidance_scale 5.0 \
-    --prompt_path ./assets/prompt.txt \
-    --neg_prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
-    --seed 12345 \
-    --output_path outputs_video/ \
-    --model_path $MODEL_BASE \
-    --vae-sp \
-    --text-encoder-precision "fp32" \
-    --use-cpu-offload