nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge branch 'main' of github.com:hao-ai-lab/FastVideo

Browse Source
This commit is contained in:
“BrianChen1129”
2025-08-11 03:44:32 +00:00
parent 690362e0b6 8cae5e4708
commit 741dfb2a23
95 changed files with 3399 additions and 619 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/pipeline.yml
+1
View File
@@ -163,6 +163,7 @@ steps:
- path:
- "csrc/attn/vsa/**"
- "csrc/attn/tk/**"
- "csrc/attn/tests/test_vsa.py"
- "csrc/attn/setup_vsa.py"
- "csrc/attn/config_vsa.py"
- "csrc/attn/vsa.cpp"
.github/workflows/vsa-publish.yml
+257
View File
@@ -0,0 +1,257 @@
name: Publish Video Sparse Attention Kernel to PyPI on Version Change
on:
push:
branches:
- main
paths:
- "csrc/attn/setup_vsa.py"
workflow_dispatch:
jobs:
check-version-change:
runs-on: ubuntu-latest
outputs:
version-changed: ${{ steps.check-version.outputs.changed }}
new-version: ${{ steps.check-version.outputs.new-version }}
steps:
- name: Checkout code
uses: actions/checkout@v4
with:
fetch-depth: 2
- name: Check if version changed
id: check-version
run: |
cd csrc/attn
# Get current commit's version
NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup_vsa.py)
echo "New version: $NEW_VERSION"
# Get previous version from git history
OLD_VERSION=$(git show HEAD~1:./setup_vsa.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
echo "Old version: $OLD_VERSION"
if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
echo "Version changed from $OLD_VERSION to $NEW_VERSION"
echo "changed=true" >> $GITHUB_OUTPUT
echo "new-version=$NEW_VERSION" >> $GITHUB_OUTPUT
else
echo "Version did not change"
echo "changed=false" >> $GITHUB_OUTPUT
fi
build_wheels:
name: Build Wheel
needs: check-version-change
if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
runs-on: ${{ matrix.os }}
strategy:
fail-fast: false
matrix:
# Using ubuntu-20.04 instead of 22.04 for more compatibility (glibc). Ideally we'd use the
# manylinux docker image, but I haven't figured out how to install CUDA on manylinux.
os: [ubuntu-22.04]
python-version: ['3.10', '3.11', '3.12', '3.13']
# For version reference https://pytorch.org/get-started/previous-versions/
torch-cuda:
- torch-version: '2.5.1'
cuda-version: '12.4.1'
torch-cuda-short: 'cu124'
- torch-version: '2.6.0'
cuda-version: '12.6.3'
torch-cuda-short: 'cu126'
- torch-version: '2.7.1'
cuda-version: '12.8.0'
torch-cuda-short: 'cu128'
steps:
- name: Free up disk space
run: |
echo "Initial disk space:"
df -h
# Remove large directories
sudo rm -rf /usr/share/dotnet
sudo rm -rf /usr/local/lib/android
sudo rm -rf /opt/ghc
sudo rm -rf /usr/local/share/boost
sudo rm -rf /usr/share/swift
sudo rm -rf /usr/local/lib/node_modules
sudo rm -rf /usr/local/share/powershell
sudo rm -rf /usr/share/rust
sudo rm -rf /usr/local/.ghcup
# Remove cached files
sudo rm -rf /var/lib/apt/lists/*
sudo rm -rf /var/cache/apt/archives/*
echo "Disk space after cleanup:"
df -h
- name: Checkout
uses: actions/checkout@v4
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python-version }}
- name: Install CUDA ${{ matrix.torch-cuda.cuda-version }}
uses: Jimver/cuda-toolkit@v0.2.21
id: cuda-toolkit
with:
cuda: ${{ matrix.torch-cuda.cuda-version }}
linux-local-args: '["--toolkit"]'
method: 'network'
- name: Install dependencies (GCC, Clang, CUDA Paths, Git)
run: |
sudo apt update
sudo apt install -y git patchelf gcc-11 g++-11 clang-11
sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
# Allow Git to Access Safe Directory
git config --global --add safe.directory /__w/FastVideo/FastVideo
# Set CUDA environment variables
export CUDA_HOME=/usr/local/cuda-${{ matrix.torch-cuda.cuda-version }}
export PATH=${CUDA_HOME}/bin:${PATH}
export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
# Verify installation
gcc --version
g++ --version
clang-11 --version
nvcc --version
- name: Install PyTorch ${{ matrix.torch-cuda.torch-version }}+cu${{ matrix.torch-cuda.cuda-version }}
run: |
pip install --upgrade pip
# With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error
# AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable
pip install typing-extensions==4.12.2
# We want to figure out the CUDA version to download pytorch
# e.g. we can have system CUDA version being 11.7 but if torch==1.12 then we need to download the wheel from cu116
# see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix
pip install --no-cache-dir torch==${{ matrix.torch-cuda.torch-version }} --index-url https://download.pytorch.org/whl/${{matrix.torch-cuda.torch-cuda-short}}
nvcc --version
python --version
python -c "import torch; print('PyTorch:', torch.__version__)"
python -c "import torch; print('CUDA:', torch.version.cuda)"
python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
- name: Build wheel
run: |
export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH
# We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA version is 11.7 and pytorch cuda version is 11.6
# https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/utils/cpp_extension.py#L810
# However this still fails so I'm using a newer version of setuptools
pip install setuptools
pip install ninja packaging wheel
cd csrc/attn # Move into the correct folder
git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
python setup_vsa.py bdist_wheel --dist-dir=dist
- name: Rename wheel file
run: |
cd csrc/attn
CUDA_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.torch-version }} | cut -d. -f1,2)
# Get the correct version format
tmpname=cu${CUDA_SHORT_VERSION}torch${TORCH_SHORT_VERSION}
wheel_name=$(ls dist/*whl | xargs -n 1 basename | sed "s/-/+$tmpname-/2")
# Rename with version information
ls dist/*whl |xargs -I {} mv {} dist/${wheel_name}
echo "wheel_name=${wheel_name}" >> $GITHUB_ENV
- name: Upload wheel artifact
uses: actions/upload-artifact@v4
with:
name: ${{ env.wheel_name }}
path: csrc/attn/dist/*.whl
retention-days: 90
publish_package:
name: Publish package
needs: [build_wheels, check-version-change]
if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
runs-on: ubuntu-22.04
permissions:
id-token: write # Needed for OIDC Trusted Publishing
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: '3.10'
- name: Install CUDA 12.4.1
uses: Jimver/cuda-toolkit@v0.2.21
id: cuda-toolkit
with:
cuda: 12.4.1
linux-local-args: '["--toolkit"]'
method: 'network'
sub-packages: '["nvcc"]'
- name: Install dependencies (GCC, Clang, CUDA Paths, Git)
run: |
sudo apt update
sudo apt install -y git patchelf gcc-11 g++-11 clang-11
sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
# Allow Git to Access Safe Directory
git config --global --add safe.directory /__w/FastVideo/FastVideo
# Set CUDA environment variables
export CUDA_HOME=/usr/local/cuda-12.4.1
export PATH=${CUDA_HOME}/bin:${PATH}
export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
# Verify installation
gcc --version
g++ --version
clang-11 --version
nvcc --version
- name: Install PyTorch 2.5.1+cu12.4.1
run: |
pip install --upgrade pip
# With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error
# AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable
pip install typing-extensions==4.12.2
# We want to figure out the CUDA version to download pytorch
# e.g. we can have system CUDA version being 11.7 but if torch==1.12 then we need to download the wheel from cu116
# see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix
export TORCH_CUDA_VERSION=124
pip install --no-cache-dir torch==2.5.1 --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}
nvcc --version
python --version
python -c "import torch; print('PyTorch:', torch.__version__)"
python -c "import torch; print('CUDA:', torch.version.cuda)"
python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
- name: Build source distribution
run: |
export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH
# We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA version is 11.7 and pytorch cuda version is 11.6
# https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/utils/cpp_extension.py#L810
# However this still fails so I'm using a newer version of setuptools
pip install setuptools
pip install ninja packaging wheel
cd csrc/attn # Move into the correct folder
git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
python setup_vsa.py sdist --dist-dir=dist
- name: Publish release distributions to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
with:
packages-dir: csrc/attn/dist/
README.md
+5 -2
View File
@@ -7,7 +7,7 @@
FastVideo features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.
<p align="center">
| <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastWan2.1-T2V-1.3B-Diffusers" target="_blank"><b>FastWan2.1</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastWan2.2-TI2V-5B-Diffusers" target="_blank"><b>FastWan2.2</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> | 🟣💬 <a href="https://ibb.co/qqPzbrw" target="_blank"> <b> WeChat </b> </a> |
| 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/collections/FastVideo/fastwan-6886a305d9799c8cd1496408" target="_blank"><b>FastWan</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> | 🟣💬 <a href="https://ibb.co/qqPzbrw" target="_blank"> <b> WeChat </b> </a> |
</p>
<div align="center">
@@ -64,12 +64,15 @@ See below for recipes and datasets:
## Inference
### Generating Your First Video
Here's a minimal example to generate a video using the default settings. Create a file called `example.py` with the following code:
Here's a minimal example to generate a video using the default settings. Make sure VSA kernels are [installed](https://hao-ai-lab.github.io/FastVideo/video_sparse_attention/installation.html). Create a file called `example.py` with the following code:
```python
import os
from fastvideo import VideoGenerator
def main():
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
# Create a video generator with a pre-trained model
generator = VideoGenerator.from_pretrained(
"FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
csrc/attn/README.md
+1 -1
View File
@@ -84,7 +84,7 @@ out = sliding_tile_attention(q, k, v, window_size, 0, False)
### Test
```bash
python tests/test_sta.py # test STA
python tests/test_block_sparse.py # test VSA
python tests/test_vsa.py # test VSA
```
### Benchmark
```bash
csrc/attn/benchmarks/bench_sta.py
-2
View File
@@ -86,8 +86,6 @@ def benchmark_attention(configurations):
# print(f"Average TFLOPS: {tflops_bwd}")
# print("=" * 60)
torch.cuda.empty_cache()
return results
csrc/attn/setup_vsa.py
+10 -13
View File
@@ -51,19 +51,16 @@ for k in kernels:
source_files.append(sources[k]['source_files'][target])
cpp_flags.append(f'-DTK_COMPILE_{k.replace(" ", "_").upper()}')
ext_modules = []
import torch
major, minor = torch.cuda.get_device_capability(0)
if major == 9 and minor == 0:# check if H100
ext_modules = [
CUDAExtension('vsa_cuda',
sources=source_files,
extra_compile_args={
'cxx': cpp_flags,
'nvcc': cuda_flags
},
libraries=['cuda'])
]
ext_modules = [
CUDAExtension('vsa_cuda',
sources=source_files,
extra_compile_args={
'cxx': cpp_flags,
'nvcc': cuda_flags
},
libraries=['cuda'])
]
csrc/attn/vsa/__init__.py
+4 -3
View File
@@ -1,12 +1,13 @@
import torch
from typing import Tuple
block_sparse_attn=None
try:
import torch
major, minor = torch.cuda.get_device_capability(0)
if major == 9 and minor == 0:# check if H100
from vsa_cuda import block_sparse_fwd, block_sparse_bwd
from vsa.block_sparse_wrapper import block_sparse_attn_SM90
block_sparse_attn = block_sparse_attn_SM90
except ImportError:
else:
from vsa.block_sparse_wrapper import block_sparse_attn_triton
block_sparse_fwd = None
block_sparse_bwd = None
csrc/attn/vsa/block_sparse_wrapper.py
+93 -90
View File
@@ -11,95 +11,6 @@ from typing import Tuple, Optional
@torch.library.custom_op("vsa::block_sparse_attn_SM90", mutates_args=(), device_types="cuda")
def block_sparse_attn_SM90(
q_padded: torch.Tensor,
k_padded: torch.Tensor,
v_padded: torch.Tensor,
block_map: torch.Tensor,
variable_block_sizes: torch.Tensor,
)-> Tuple[torch.Tensor, torch.Tensor]:
q_padded = q_padded.contiguous()
k_padded = k_padded.contiguous()
v_padded = v_padded.contiguous()
q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map)
variable_block_sizes = variable_block_sizes.int()
o_padded, lse_padded = block_sparse_fwd(q_padded, k_padded, v_padded, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes)
return o_padded, lse_padded
@torch.library.register_fake("vsa::block_sparse_attn_SM90")
def _block_sparse_attn_SM90_fake(
q_padded: torch.Tensor,
k_padded: torch.Tensor,
v_padded: torch.Tensor,
block_map: torch.Tensor,
variable_block_sizes: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
q_padded, k_padded, v_padded = [x.contiguous() for x in (q_padded, k_padded, v_padded)]
B, H, S, D = q_padded.shape
o_padded = torch.empty_like(q_padded)
lse_padded = torch.empty((B, H, S, 1), device=q_padded.device, dtype=torch.float32)
return o_padded, lse_padded
@torch.library.custom_op("vsa::block_sparse_attn_backward_SM90", mutates_args=(), device_types="cuda")
def block_sparse_attn_backward_SM90(
grad_output_padded: torch.Tensor,
q_padded: torch.Tensor,
k_padded: torch.Tensor,
v_padded: torch.Tensor,
o_padded: torch.Tensor,
lse_padded: torch.Tensor,
block_map: torch.Tensor,
variable_block_sizes: torch.Tensor,
)-> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
grad_output_padded = grad_output_padded.contiguous()
k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(block_map.transpose(-1, -2))
grad_q_padded, grad_k_padded, grad_v_padded = block_sparse_bwd(
q_padded, k_padded, v_padded, o_padded, lse_padded, grad_output_padded, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes
)
grad_q_padded = grad_q_padded.to(grad_output_padded.dtype)
grad_k_padded = grad_k_padded.to(grad_output_padded.dtype)
grad_v_padded = grad_v_padded.to(grad_output_padded.dtype)
return grad_q_padded, grad_k_padded, grad_v_padded
@torch.library.register_fake("vsa::block_sparse_attn_backward_SM90")
def _block_sparse_attn_backward_SM90_fake(
grad_output_padded: torch.Tensor,
q_padded: torch.Tensor,
k_padded: torch.Tensor,
v_padded: torch.Tensor,
o_padded: torch.Tensor,
lse_padded: torch.Tensor,
block_map: torch.Tensor,
variable_block_sizes: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
torch._check(grad_output_padded.dtype == torch.bfloat16)
torch._check(lse_padded.dtype == torch.float32)
grad_output_padded = grad_output_padded.contiguous()
dq = torch.empty_like(grad_output_padded)
dk = torch.empty_like(grad_output_padded)
dv = torch.empty_like(grad_output_padded)
return dq, dk, dv
def backward_SM90(ctx, grad_output1, grad_output2):
q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes= ctx.saved_tensors
dq, dk, dv = block_sparse_attn_backward_SM90(grad_output1, q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes)
return dq, dk, dv, None, None
def setup_context_SM90(ctx, inputs, output):
q_padded, k_padded, v_padded, block_map, variable_block_sizes = inputs
o_padded, lse_padded = output
ctx.save_for_backward(q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes)
block_sparse_attn_SM90.register_autograd(backward_SM90, setup_context=setup_context_SM90)
@torch.library.custom_op("vsa::block_sparse_attn_triton", mutates_args=(), device_types="cuda")
def block_sparse_attn_triton(
q: torch.Tensor,
@@ -179,4 +90,96 @@ def setup_context_triton(ctx, inputs, output):
o_padded, M = output
ctx.save_for_backward(q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes)
block_sparse_attn_triton.register_autograd(backward_triton, setup_context=setup_context_triton)
block_sparse_attn_triton.register_autograd(backward_triton, setup_context=setup_context_triton)
major, minor = torch.cuda.get_device_capability(0)
if major == 9 and minor == 0:# check if H100
@torch.library.custom_op("vsa::block_sparse_attn_SM90", mutates_args=(), device_types="cuda")
def block_sparse_attn_SM90(
q_padded: torch.Tensor,
k_padded: torch.Tensor,
v_padded: torch.Tensor,
block_map: torch.Tensor,
variable_block_sizes: torch.Tensor,
)-> Tuple[torch.Tensor, torch.Tensor]:
q_padded = q_padded.contiguous()
k_padded = k_padded.contiguous()
v_padded = v_padded.contiguous()
q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map)
variable_block_sizes = variable_block_sizes.int()
o_padded, lse_padded = block_sparse_fwd(q_padded, k_padded, v_padded, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes)
return o_padded, lse_padded
@torch.library.register_fake("vsa::block_sparse_attn_SM90")
def _block_sparse_attn_SM90_fake(
q_padded: torch.Tensor,
k_padded: torch.Tensor,
v_padded: torch.Tensor,
block_map: torch.Tensor,
variable_block_sizes: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
q_padded, k_padded, v_padded = [x.contiguous() for x in (q_padded, k_padded, v_padded)]
B, H, S, D = q_padded.shape
o_padded = torch.empty_like(q_padded)
lse_padded = torch.empty((B, H, S, 1), device=q_padded.device, dtype=torch.float32)
return o_padded, lse_padded
@torch.library.custom_op("vsa::block_sparse_attn_backward_SM90", mutates_args=(), device_types="cuda")
def block_sparse_attn_backward_SM90(
grad_output_padded: torch.Tensor,
q_padded: torch.Tensor,
k_padded: torch.Tensor,
v_padded: torch.Tensor,
o_padded: torch.Tensor,
lse_padded: torch.Tensor,
block_map: torch.Tensor,
variable_block_sizes: torch.Tensor,
)-> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
grad_output_padded = grad_output_padded.contiguous()
k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(block_map.transpose(-1, -2))
grad_q_padded, grad_k_padded, grad_v_padded = block_sparse_bwd(
q_padded, k_padded, v_padded, o_padded, lse_padded, grad_output_padded, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes
)
grad_q_padded = grad_q_padded.to(grad_output_padded.dtype)
grad_k_padded = grad_k_padded.to(grad_output_padded.dtype)
grad_v_padded = grad_v_padded.to(grad_output_padded.dtype)
return grad_q_padded, grad_k_padded, grad_v_padded
@torch.library.register_fake("vsa::block_sparse_attn_backward_SM90")
def _block_sparse_attn_backward_SM90_fake(
grad_output_padded: torch.Tensor,
q_padded: torch.Tensor,
k_padded: torch.Tensor,
v_padded: torch.Tensor,
o_padded: torch.Tensor,
lse_padded: torch.Tensor,
block_map: torch.Tensor,
variable_block_sizes: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
torch._check(grad_output_padded.dtype == torch.bfloat16)
torch._check(lse_padded.dtype == torch.float32)
grad_output_padded = grad_output_padded.contiguous()
dq = torch.empty_like(grad_output_padded)
dk = torch.empty_like(grad_output_padded)
dv = torch.empty_like(grad_output_padded)
return dq, dk, dv
def backward_SM90(ctx, grad_output1, grad_output2):
q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes= ctx.saved_tensors
dq, dk, dv = block_sparse_attn_backward_SM90(grad_output1, q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes)
return dq, dk, dv, None, None
def setup_context_SM90(ctx, inputs, output):
q_padded, k_padded, v_padded, block_map, variable_block_sizes = inputs
o_padded, lse_padded = output
ctx.save_for_backward(q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes)
block_sparse_attn_SM90.register_autograd(backward_SM90, setup_context=setup_context_SM90)
docs/source/distillation/dmd.md
+1 -1
View File
@@ -6,7 +6,7 @@ We introduce a new finetuning strategy - **Sparse-distill**, which jointly integ
We provide two distilled models:
- **[FastWan2.1-T2V-1.3B-Diffusers](https://huggingface.co/FastVideo/FastWan2.1-T2V-1.3B-Diffusers)**: 3-step inference, up to **20 FPS** on H100 GPU
- **[FastWan2.1-T2V-1.3B-Diffusers](https://huggingface.co/FastVideo/FastWan2.1-T2V-1.3B-Diffusers)**: 3-step inference, up to **16 FPS** on H100 GPU
- **[FastWan2.1-T2V-14B-480P-Diffusers](https://huggingface.co/FastVideo/FastWan2.1-T2V-14B-480P-Diffusers)**: 3-step inference, up to **60x speed up** at 480P, **90x speed up** at 720P for denoising loop
- **[FastWan2.2-TI2V-5B-FullAttn-Diffusers](https://huggingface.co/FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers)**: 3-step inference, up to **50x speed up** at 720P for denoising loop
docs/source/index.md
+1 -1
View File
@@ -9,7 +9,7 @@
:::{raw} html
<p style="text-align:center">
<strong>FastVideo is An unified inference and post-training framework for accelerated video generation.
<strong>FastVideo is a unified inference and post-training framework for accelerated video generation.
</strong>
</p>
docs/source/sliding_tile_attention/installation.md
+6
View File
@@ -53,3 +53,9 @@ out = sliding_tile_attention(q, k, v, window_size, text_length)
out = sliding_tile_attention(q, k, v, window_size, 0, False)
```
# 🚀Inference
```bash
bash scripts/inference/v1_inference_wan_STA.sh
```
docs/source/video_sparse_attention/installation.md
+6
View File
@@ -59,3 +59,9 @@ from vsa import video_sparse_attn
output = video_sparse_attn(q, k, v, variable_block_sizes, topk, block_size, compress_attn_weight)
```
# 🚀Inference
```bash
bash scripts/inference/v1_inference_wan_VSA.sh
```
examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.slurm
+1 -1
View File
@@ -86,7 +86,7 @@ validation_args=(
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 200
--validation_sampling_steps "3"
--validation_guidance_scale "1.0" # not used for dmd inference
--validation_guidance_scale "6.0" # not used for dmd inference
)
# Optimizer arguments
examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slurm
+1 -1
View File
@@ -86,7 +86,7 @@ validation_args=(
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 200
--validation_sampling_steps "3"
--validation_guidance_scale "1.0" # not used for dmd inference
--validation_guidance_scale "6.0" # not used for dmd inference
)
# Optimizer arguments
examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.slurm
+1 -1
View File
@@ -86,7 +86,7 @@ validation_args=(
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 200
--validation_sampling_steps "3"
--validation_guidance_scale "1.0" # not used for dmd inference
--validation_guidance_scale "6.0" # not used for dmd inference
)
# Optimizer arguments
examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/README.md
+13
View File
@@ -0,0 +1,13 @@
# Wan2.2-5B Distill Example
These are end-to-end example scripts for distilling Wan2.2 TI2V 5B model DMD+VSA methods.
### 0. Make sure you have installed VSA
```bash
cd csrc/attn
git submodule update --init --recursive
python setup_vsa.py install
```
### Data-free Distillation
When `--simulate_generator_forward` is enabled, distillation becomes data-free by simulating intermediate steps through forward inference of the generator. This helps avoid training–inference mismatch. See Section 4.5 of [DMD2](https://arxiv.org/pdf/2405.14867) for details.
examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B.sh
+1 -1
View File
@@ -87,7 +87,7 @@ validation_args=(
--validation_dataset_file "$VALIDATION_DIR"
--validation_steps 200
--validation_sampling_steps "3"
--validation_guidance_scale "1.0" # not used for dmd inference
--validation_guidance_scale "6.0" # not used for dmd inference
)
# Optimizer arguments
examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B_VSA.sh
+1 -1
View File
@@ -87,7 +87,7 @@ validation_args=(
--validation_dataset_file "$VALIDATION_DIR"
--validation_steps 200
--validation_sampling_steps "3"
--validation_guidance_scale "1.0" # not used for dmd inference
--validation_guidance_scale "6.0" # not used for dmd inference
)
# Optimizer arguments
examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/README.md
+11 -1
View File
@@ -9,4 +9,14 @@ git submodule update --init --recursive
python setup_vsa.py install
```
### TODO
### 1. Download dataset:
```bash
bash examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/download_dataset.sh
```
### 2. Configure and run distillation:
#### For DMD-only distillation:
```bash
bash examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/
```
examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh
+1 -1
View File
@@ -63,7 +63,7 @@ validation_args=(
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 200
--validation_sampling_steps "3"
--validation_guidance_scale "1.0" # not used for dmd inference
--validation_guidance_scale "6.0" # not used for dmd inference
)
# Optimizer arguments
examples/inference/basic/basic.py
+1 -2
View File
@@ -16,8 +16,7 @@ def main():
dit_cpu_offload=False,
vae_cpu_offload=False,
text_encoder_cpu_offload=True,
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
pin_cpu_memory=False,
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
# image_encoder_cpu_offload=False,
)
examples/inference/basic/basic_dmd.py
+1 -1
View File
@@ -17,7 +17,7 @@ def main():
use_fsdp_inference=True,
# Adjust these offload parameters if you have < 32GB of VRAM
text_encoder_cpu_offload=True,
pin_cpu_memory=False,
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
dit_cpu_offload=False,
vae_cpu_offload=False,
VSA_sparsity=0.8,
examples/inference/basic/basic_wan2_2.py
+48
View File
@@ -0,0 +1,48 @@
from fastvideo import VideoGenerator
# from fastvideo.configs.sample import SamplingParam
OUTPUT_PATH = "video_samples_wan2_2_14B_t2v"
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.
generator = VideoGenerator.from_pretrained(
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
# FastVideo will automatically handle distributed setup
num_gpus=1,
use_fsdp_inference=True,
dit_cpu_offload=True, # DiT need to be offloaded for MoE
vae_cpu_offload=False,
text_encoder_cpu_offload=True,
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
pin_cpu_memory=True,
# image_encoder_cpu_offload=False,
)
# sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
# sampling_param.num_frames = 45
# sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
# Generate videos with the same simple API, regardless of GPU count
prompt = (
"A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
"wide with interest. The playful yet serene atmosphere is complemented by soft "
"natural light filtering through the petals. Mid-shot, warm and cheerful tones."
)
_ = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, height=720, width=1280, num_frames=81)
# video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
# Generate another video with a different prompt, without reloading the
# model!
prompt2 = (
"A majestic lion strides across the golden savanna, its powerful frame "
"glistening under the warm afternoon sun. The tall grass ripples gently in "
"the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
"embodying the raw energy of the wild. Low angle, steady tracking shot, "
"cinematic.")
_ = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, height=720, width=1280, num_frames=81)
if __name__ == "__main__":
main()
examples/inference/basic/basic_wan2_2_i2v.py
+30
View File
@@ -0,0 +1,30 @@
from fastvideo import VideoGenerator
# from fastvideo.configs.sample import SamplingParam
OUTPUT_PATH = "video_samples_wan2_2_14B_i2v"
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.
generator = VideoGenerator.from_pretrained(
"Wan-AI/Wan2.2-I2V-A14B-Diffusers",
# FastVideo will automatically handle distributed setup
num_gpus=1,
use_fsdp_inference=True,
dit_cpu_offload=True, # DiT need to be offloaded for MoE
vae_cpu_offload=False,
text_encoder_cpu_offload=True,
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
pin_cpu_memory=True,
# image_encoder_cpu_offload=False,
)
prompt = "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside."
image_path = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG"
video = generator.generate_video(prompt, image_path=image_path, output_path=OUTPUT_PATH, save_video=True, height=832, width=480, num_frames=81)
if __name__ == "__main__":
main()
examples/inference/gradio/README.md
-59
View File
@@ -1,59 +0,0 @@
# FastVideo 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
---
## Usage
Run the demo with:
```bash
python 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
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
examples/inference/gradio/gradio_demo.py
-169
View File
@@ -1,169 +0,0 @@
import argparse
import os
from copy import deepcopy
import gradio as gr
import torch
from fastvideo import VideoGenerator
from fastvideo.configs.sample.base import SamplingParam
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="FastVideo Gradio Demo")
parser.add_argument("--model_path",
type=str,
default="FastVideo/FastHunyuan-diffusers",
help="Path to the model")
parser.add_argument("--num_gpus",
type=int,
default=1,
help="Number of GPUs to use")
parser.add_argument("--output_path",
type=str,
default="outputs",
help="Path to save generated videos")
parsed_args = parser.parse_args()
# args = FastVideoArgs(model_path="FastVideo/FastHunyuan-Diffusers", num_gpus=2)
generator = VideoGenerator.from_pretrained(
model_path=parsed_args.model_path, num_gpus=parsed_args.num_gpus)
default_params = SamplingParam.from_pretrained(parsed_args.model_path)
def generate_video(
prompt,
negative_prompt,
use_negative_prompt,
seed,
guidance_scale,
num_frames,
height,
width,
num_inference_steps,
randomize_seed=False,
):
params = deepcopy(default_params)
params.prompt = prompt
params.negative_prompt = negative_prompt
params.seed = seed
params.guidance_scale = guidance_scale
params.num_frames = num_frames
params.height = height
params.width = width
params.num_inference_steps = num_inference_steps
if randomize_seed:
params.seed = torch.randint(0, 1000000, (1, )).item()
if not use_negative_prompt:
params.negative_prompt = None
generator.generate_video(prompt=prompt, sampling_param=params)
output_path = os.path.join(parsed_args.output_path,
f"{params.prompt[:100]}.mp4")
return output_path, params.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=default_params.height,
)
width = gr.Slider(label="Width",
minimum=256,
maximum=1024,
step=32,
value=default_params.width)
with gr.Row():
num_frames = gr.Slider(
label="Number of Frames",
minimum=21,
maximum=163,
value=default_params.num_frames,
)
guidance_scale = gr.Slider(
label="Guidance Scale",
minimum=1,
maximum=12,
value=default_params.guidance_scale,
)
num_inference_steps = gr.Slider(
label="Inference Steps",
minimum=4,
maximum=100,
value=default_params.num_inference_steps,
)
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=default_params.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=default_params.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)
examples/inference/gradio/gradio_frontend.py
+730
View File
@@ -0,0 +1,730 @@
import argparse
import os
import requests
import base64
import time
import gradio as gr
from fastvideo.configs.sample.base import SamplingParam
MODEL_PATH_MAPPING = {
"FastWan2.1-T2V-1.3B": "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
"FastWan2.2-TI2V-5B-FullAttn": "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
}
class RayServeClient:
def __init__(self, backend_url: str):
self.backend_url = backend_url
self.session = requests.Session()
def check_health(self) -> bool:
try:
response = self.session.get(f"{self.backend_url}/health", timeout=5)
return response.status_code == 200
except requests.exceptions.RequestException:
return False
def generate_video(self, request_data: dict) -> dict:
start_time = time.time()
try:
headers = {"Content-Type": "application/json"}
response = self.session.post(
f"{self.backend_url}/generate_video",
json=request_data,
headers=headers,
timeout=300
)
round_trip_time = time.time() - start_time
if response.status_code == 200:
result = response.json()
backend_total = result.get("total_time", 0)
network_time = round_trip_time - backend_total
result["network_time"] = network_time
return result
else:
return {"success": False, "error_message": f"HTTP {response.status_code}: {response.text}"}
except requests.exceptions.RequestException as e:
return {"success": False, "error_message": f"Request failed: {str(e)}"}
def save_video_from_base64(video_data: str, output_dir: str, prompt: str) -> str:
if not video_data:
return None
try:
if video_data.startswith('data:video/'):
video_data = video_data.split(',')[1]
video_bytes = base64.b64decode(video_data)
safe_prompt = prompt[:50].replace(' ', '_').replace('/', '_').replace('\\', '_')
video_filename = f"{safe_prompt}.mp4"
video_path = os.path.join(output_dir, video_filename)
os.makedirs(output_dir, exist_ok=True)
with open(video_path, 'wb') as f:
f.write(video_bytes)
return video_path
except Exception as e:
print(f"Failed to save video: {e}")
return None
def create_timing_display(inference_time, encoding_time, network_time, total_time, stage_execution_times, num_frames):
dit_denoising_time = f"{stage_execution_times[5]:.2f}s" if len(stage_execution_times) > 5 else "N/A"
timing_html = f"""
<div style="margin: 10px 0;">
<h3 style="text-align: center; margin-bottom: 10px;">⏱️ Timing Breakdown</h3>
<div style="display: grid; grid-template-columns: repeat(5, 1fr); gap: 10px; margin-bottom: 10px;">
<div class="timing-card timing-card-highlight">
<div style="font-size: 20px;">🚀</div>
<div style="font-weight: bold; margin: 3px 0; font-size: 14px;">DiT Denoising</div>
<div style="font-size: 16px; color: #ffa200; font-weight: bold;">{dit_denoising_time}</div>
</div>
<div class="timing-card">
<div style="font-size: 20px;">🧠</div>
<div style="font-weight: bold; margin: 3px 0; font-size: 14px;">E2E (w. vae/text encoder)</div>
<div style="font-size: 16px; color: #2563eb;">{inference_time:.2f}s</div>
</div>
<div class="timing-card">
<div style="font-size: 20px;">🎬</div>
<div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Video Encoding</div>
<div style="font-size: 16px; color: #dc2626;">{encoding_time:.2f}s</div>
</div>
<div class="timing-card">
<div style="font-size: 20px;">🌐</div>
<div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Network Transfer</div>
<div style="font-size: 16px; color: #059669;">{network_time:.2f}s</div>
</div>
<div class="timing-card">
<div style="font-size: 20px;">📊</div>
<div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Total Processing</div>
<div style="font-size: 18px; color: #0277bd;">{total_time:.2f}s</div>
</div>
</div>"""
if inference_time > 0:
fps = num_frames / inference_time
timing_html += f"""
<div class="performance-card" style="margin-top: 15px;">
<span style="font-weight: bold;">Generation Speed: </span>
<span style="font-size: 18px; color: #6366f1; font-weight: bold;">{fps:.1f} frames/second</span>
</div>"""
return timing_html + "</div>"
def load_example_prompts():
def contains_chinese(text):
return any('\u4e00' <= char <= '\u9fff' for char in text)
def load_from_file(filepath):
prompts, labels = [], []
try:
with open(filepath, "r", encoding='utf-8') as f:
for line in f:
line = line.strip()
if line and not contains_chinese(line):
label = line[:100] + "..." if len(line) > 100 else line
labels.append(label)
prompts.append(line)
except Exception as e:
print(f"Warning: Could not read {filepath}: {e}")
return prompts, labels
examples, example_labels = load_from_file("prompts/prompts_final.txt")
if not examples:
examples = ["A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background."]
example_labels = ["Crowded rooftop bar at night"]
return examples, example_labels
def create_gradio_interface(backend_url: str, default_params: dict[str, SamplingParam]):
client = RayServeClient(backend_url)
def generate_video(
prompt, negative_prompt, use_negative_prompt, seed, guidance_scale,
num_frames, height, width, randomize_seed, model_selection, progress
):
if not client.check_health():
return None, f"Backend is not available. Please check if Ray Serve is running at {backend_url}", ""
# Validate dimensions
max_pixels = 720 * 1280
if height * width > max_pixels:
return None, f"Video dimensions too large. Maximum: 720x1280 pixels", ""
if progress:
progress(0.1, desc="Checking backend health...")
request_data = {
"prompt": prompt,
"negative_prompt": negative_prompt,
"use_negative_prompt": use_negative_prompt,
"seed": seed,
"guidance_scale": guidance_scale,
"num_frames": num_frames,
"height": height,
"width": width,
"randomize_seed": randomize_seed,
"return_frames": False,
"image_path": None,
"model_path": MODEL_PATH_MAPPING.get(model_selection, "FastVideo/FastWan2.1-T2V-1.3B-Diffusers")
}
if progress:
progress(0.4, desc="Generating video...")
response = client.generate_video(request_data)
if progress:
progress(0.8, desc="Processing response...")
if response.get("success", False):
video_data = response.get("video_data", "")
used_seed = response.get("seed", seed)
inference_time = response.get("inference_time", 0.0)
encoding_time = response.get("encoding_time", 0.0)
total_time = response.get("total_time", 0.0)
network_time = response.get("network_time", 0.0)
stage_execution_times = response.get("stage_execution_times", [])
timing_details = create_timing_display(
inference_time, encoding_time, network_time, total_time,
stage_execution_times, num_frames
)
if video_data:
if progress:
progress(0.9, desc="Saving video...")
video_path = save_video_from_base64(video_data, "outputs", prompt)
if progress:
progress(1.0, desc="Generation complete!")
if video_path and os.path.exists(video_path):
return video_path, used_seed, timing_details
else:
return None, "Failed to save video", ""
else:
return None, "No video data received from backend", ""
else:
error_msg = response.get("error_message", "Unknown error occurred")
return None, f"Generation failed: {error_msg}", ""
examples, example_labels = load_example_prompts()
theme = gr.themes.Base().set(
button_primary_background_fill="#2563eb",
button_primary_background_fill_hover="#1d4ed8",
button_primary_text_color="white",
slider_color="#2563eb",
checkbox_background_color_selected="#2563eb",
)
def get_default_values(model_name):
model_path = MODEL_PATH_MAPPING.get(model_name)
if model_path and model_path in default_params:
params = default_params[model_path]
return {
'height': params.height,
'width': params.width,
'num_frames': params.num_frames,
'guidance_scale': params.guidance_scale,
'seed': params.seed,
}
return {
'height': 448,
'width': 832,
'num_frames': 61,
'guidance_scale': 3.0,
'seed': 1024,
}
initial_values = get_default_values("FastWan2.1-T2V-1.3B")
with gr.Blocks(title="FastWan", theme=theme) as demo:
gr.Image("fastvideo-logos/main/svg/full.svg", show_label=False, container=False, height=80)
gr.HTML("""
<div style="text-align: center; margin-bottom: 10px;">
<p style="font-size: 18px;"> Make Video Generation Go Blurrrrrrr </p>
<p style="font-size: 18px;"> <a href="https://github.com/hao-ai-lab/FastVideo/tree/main" target="_blank">Code</a> | <a href="https://hao-ai-lab.github.io/blogs/fastvideo_post_training/" target="_blank">Blog</a> | <a href="https://hao-ai-lab.github.io/FastVideo/" target="_blank">Docs</a> </p>
</div>
""")
with gr.Accordion("🎥 What Is FastVideo?", open=False):
gr.HTML("""
<div style="padding: 20px; line-height: 1.6;">
<p style="font-size: 16px; margin-bottom: 15px;">
FastVideo is an inference and post-training framework for diffusion models. It features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.
</p>
</div>
""")
with gr.Row():
model_selection = gr.Dropdown(
choices=list(MODEL_PATH_MAPPING.keys()),
value="FastWan2.1-T2V-1.3B",
label="Select Model",
interactive=True
)
with gr.Row():
example_dropdown = gr.Dropdown(
choices=example_labels,
label="Example Prompts",
value=None,
interactive=True,
allow_custom_value=False
)
with gr.Row():
with gr.Column(scale=6):
prompt = gr.Text(
label="Prompt",
show_label=False,
max_lines=3,
placeholder="Describe your scene...",
container=False,
lines=3,
autofocus=True,
)
with gr.Column(scale=1, min_width=120, elem_classes="center-button"):
run_button = gr.Button("Run", variant="primary", size="lg")
with gr.Row():
with gr.Column():
error_output = gr.Text(label="Error", visible=False)
timing_display = gr.Markdown(label="Timing Breakdown", visible=False)
with gr.Row(equal_height=True, elem_classes="main-content-row"):
with gr.Column(scale=1, elem_classes="advanced-options-column"):
with gr.Group():
gr.HTML("<div style='margin: 0 0 15px 0; text-align: center; font-size: 16px;'>Advanced Options</div>")
with gr.Row():
height = gr.Number(
label="Height",
value=initial_values['height'],
interactive=False,
container=True
)
width = gr.Number(
label="Width",
value=initial_values['width'],
interactive=False,
container=True
)
with gr.Row():
num_frames = gr.Number(
label="Number of Frames",
value=initial_values['num_frames'],
interactive=False,
container=True
)
guidance_scale = gr.Slider(
label="Guidance Scale",
minimum=1,
maximum=12,
value=initial_values['guidance_scale'],
)
with gr.Row():
use_negative_prompt = gr.Checkbox(
label="Use negative prompt", value=False)
negative_prompt = gr.Text(
label="Negative prompt",
max_lines=3,
lines=3,
placeholder="Enter a negative prompt",
visible=False,
)
seed = gr.Slider(
label="Seed",
minimum=0,
maximum=1000000,
step=1,
value=initial_values['seed'],
)
randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
seed_output = gr.Number(label="Used Seed")
with gr.Column(scale=1, elem_classes="video-column"):
result = gr.Video(
label="Generated Video",
show_label=True,
height=466,
width=600,
container=True,
elem_classes="video-component"
)
gr.HTML("""
<style>
.center-button {
display: flex !important;
justify-content: center !important;
height: 100% !important;
padding-top: 1.4em !important;
}
.gradio-container {
max-width: 1200px !important;
margin: 0 auto !important;
}
.main {
max-width: 1200px !important;
margin: 0 auto !important;
}
.gr-form, .gr-box, .gr-group {
max-width: 1200px !important;
}
.gr-video {
max-width: 500px !important;
margin: 0 auto !important;
}
.main-content-row {
display: flex !important;
align-items: flex-start !important;
min-height: 500px !important;
gap: 20px !important;
}
.advanced-options-column,
.video-column {
display: flex !important;
flex-direction: column !important;
flex: 1 !important;
min-height: 400px !important;
align-items: stretch !important;
}
.video-column > * {
margin-top: 0 !important;
}
.video-column .gr-video,
.video-component {
margin-top: 0 !important;
padding-top: 0 !important;
}
.video-column .gr-video .gr-form {
margin-top: 0 !important;
}
.advanced-options-column .gr-group,
.video-column .gr-video {
margin-top: 0 !important;
vertical-align: top !important;
}
.advanced-options-column > *:last-child,
.video-column > *:last-child {
flex-grow: 0 !important;
}
@media (max-width: 1400px) {
.main-content-row {
min-height: 600px !important;
}
.advanced-options-column,
.video-column {
min-height: 600px !important;
}
}
@media (max-width: 1200px) {
.main-content-row {
flex-direction: column !important;
align-items: stretch !important;
}
.advanced-options-column,
.video-column {
min-height: auto !important;
width: 100% !important;
}
}
.timing-card {
background: var(--background-fill-secondary) !important;
border: 1px solid var(--border-color-primary) !important;
color: var(--body-text-color) !important;
padding: 10px;
border-radius: 8px;
text-align: center;
min-height: 80px;
display: flex;
flex-direction: column;
justify-content: center;
}
.timing-card-highlight {
background: var(--background-fill-primary) !important;
border: 2px solid var(--color-accent) !important;
}
.performance-card {
background: var(--background-fill-secondary) !important;
border: 1px solid var(--border-color-primary) !important;
color: var(--body-text-color) !important;
padding: 10px;
border-radius: 6px;
text-align: center;
}
.gr-number input[readonly] {
background-color: var(--background-fill-secondary) !important;
border: 1px solid var(--border-color-primary) !important;
color: var(--body-text-color-subdued) !important;
cursor: default !important;
text-align: center !important;
font-weight: 500 !important;
}
</style>
""")
def on_example_select(example_label):
if example_label and example_label in example_labels:
index = example_labels.index(example_label)
return examples[index]
return ""
example_dropdown.change(
fn=on_example_select,
inputs=example_dropdown,
outputs=prompt,
)
gr.HTML("""
<div style="text-align: center; margin-top: 10px; margin-bottom: 15px;">
<p style="font-size: 16px; margin: 0;">The compute for this demo is generously provided by <a href="https://www.gmicloud.ai/" target="_blank">GMI Cloud</a>. Note that this demo is meant to showcase FastWan's quality and that under a large number of requests, generation speed may be affected. We are also rate-limiting users to 3 requests per minute.</p>
</div>
""")
use_negative_prompt.change(
fn=lambda x: gr.update(visible=x),
inputs=use_negative_prompt,
outputs=negative_prompt,
)
def on_model_selection_change(selected_model):
if not selected_model:
selected_model = "FastWan2.1-T2V-1.3B"
model_path = MODEL_PATH_MAPPING.get(selected_model)
if model_path and model_path in default_params:
params = default_params[model_path]
return (
gr.update(value=params.height),
gr.update(value=params.width),
gr.update(value=params.num_frames),
gr.update(value=params.guidance_scale),
gr.update(value=params.seed),
)
return (
gr.update(value=448),
gr.update(value=832),
gr.update(value=61),
gr.update(value=3.0),
gr.update(value=1024),
)
model_selection.change(
fn=on_model_selection_change,
inputs=model_selection,
outputs=[height, width, num_frames, guidance_scale, seed],
)
def handle_generation(*args, progress=None, request: gr.Request = None):
model_selection, prompt, negative_prompt, use_negative_prompt, seed, guidance_scale, num_frames, height, width, randomize_seed = args
result_path, seed_or_error, timing_details = generate_video(
prompt, negative_prompt, use_negative_prompt, seed, guidance_scale,
num_frames, height, width, randomize_seed, model_selection, progress
)
if result_path and os.path.exists(result_path):
return (
result_path,
seed_or_error,
gr.update(visible=False),
gr.update(visible=True, value=timing_details),
)
else:
return (
None,
seed_or_error,
gr.update(visible=True, value=seed_or_error),
gr.update(visible=False),
)
run_button.click(
fn=handle_generation,
inputs=[
model_selection,
prompt,
negative_prompt,
use_negative_prompt,
seed,
guidance_scale,
num_frames,
height,
width,
randomize_seed,
],
outputs=[result, seed_output, error_output, timing_display],
concurrency_limit=20,
)
return demo
def main():
parser = argparse.ArgumentParser(description="FastVideo Gradio Frontend")
parser.add_argument("--backend_url", type=str, default="http://localhost:8000",
help="URL of the Ray Serve backend")
parser.add_argument("--t2v_model_paths", type=str,
default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.1-T2V-14B-Diffusers",
help="Comma separated list of paths to the T2V model(s)")
parser.add_argument("--host", type=str, default="0.0.0.0",
help="Host to bind to")
parser.add_argument("--port", type=int, default=7860,
help="Port to bind to")
args = parser.parse_args()
default_params = {}
model_paths = args.t2v_model_paths.split(",")
for model_path in model_paths:
default_params[model_path] = SamplingParam.from_pretrained(model_path)
demo = create_gradio_interface(args.backend_url, default_params)
print(f"Starting Gradio frontend at http://{args.host}:{args.port}")
print(f"Backend URL: {args.backend_url}")
print(f"T2V Models: {args.t2v_model_paths}")
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import HTMLResponse, FileResponse
import uvicorn
app = FastAPI()
@app.get("/logo.png")
def get_logo():
return FileResponse(
"fastvideo-logos/main/svg/full.svg",
media_type="image/svg+xml",
headers={
"Cache-Control": "public, max-age=3600",
"Access-Control-Allow-Origin": "*"
}
)
@app.get("/favicon.ico")
def get_favicon():
favicon_path = "fastvideo-logos/main/svg/icon-simple.svg"
if os.path.exists(favicon_path):
return FileResponse(
favicon_path,
media_type="image/svg+xml",
headers={
"Cache-Control": "public, max-age=3600",
"Access-Control-Allow-Origin": "*"
}
)
else:
raise HTTPException(status_code=404, detail="Favicon not found")
@app.get("/", response_class=HTMLResponse)
def index(request: Request):
base_url = str(request.base_url).rstrip('/')
return f"""
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
<title>FastWan</title>
<meta name="title" content="FastWan">
<meta name="description" content="Make video generation go blurrrrrrr">
<meta name="keywords" content="FastVideo, video generation, AI, machine learning, FastWan">
<meta property="og:type" content="website">
<meta property="og:url" content="{base_url}/">
<meta property="og:title" content="FastWan">
<meta property="og:description" content="Make video generation go blurrrrrrr">
<meta property="og:image" content="{base_url}/logo.png">
<meta property="og:image:width" content="1200">
<meta property="og:image:height" content="630">
<meta property="og:site_name" content="FastWan">
<meta property="twitter:card" content="summary_large_image">
<meta property="twitter:url" content="{base_url}/">
<meta property="twitter:title" content="FastWan">
<meta property="twitter:description" content="Make video generation go blurrrrrrr">
<meta property="twitter:image" content="{base_url}/logo.png">
<link rel="icon" type="image/png" sizes="32x32" href="/favicon.ico">
<link rel="icon" type="image/png" sizes="16x16" href="/favicon.ico">
<link rel="apple-touch-icon" href="/favicon.ico">
<style>
body, html {{
margin: 0;
padding: 0;
height: 100%;
overflow: hidden;
}}
iframe {{
width: 100%;
height: 100vh;
border: none;
}}
</style>
</head>
<body>
<iframe src="/gradio" width="100%" height="100%" style="border: none;"></iframe>
</body>
</html>
"""
app = gr.mount_gradio_app(
app,
demo,
path="/gradio",
allowed_paths=[os.path.abspath("outputs"), os.path.abspath("fastvideo-logos")]
)
uvicorn.run(app, host=args.host, port=args.port)
if __name__ == "__main__":
main()
examples/inference/gradio/ray_serve_backend.py
+379
View File
@@ -0,0 +1,379 @@
import time
import os
import torch
import base64
import io
from copy import deepcopy
from typing import Dict, Any, Optional, List
import signal
import sys
import ray
from ray import serve
from fastapi import FastAPI, Request, Response
from pydantic import BaseModel
import numpy as np
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded
import imageio
from ray.serve.handle import DeploymentHandle
from prometheus_client import Counter, Histogram, generate_latest
NUM_GPUS = 16
DEFAULT_FPS = 16
SEED_RANGE_MAX = 1_000_000
SUPPORTED_MODELS = [
"FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
"FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
]
MODEL_CONFIGS = {
"1.3B": {
"num_cpus": 2,
"text_encoder_cpu_offload": False,
"dit_cpu_offload": False,
"vae_cpu_offload": False,
"VSA_sparsity": 0.8,
},
"14B": {
"num_cpus": 16,
"text_encoder_cpu_offload": True,
"dit_cpu_offload": True,
"vae_cpu_offload": False,
"VSA_sparsity": 0.9,
}
}
class VideoGenerationRequest(BaseModel):
prompt: str
negative_prompt: Optional[str] = None
use_negative_prompt: bool = False
seed: int = 42
guidance_scale: float = 7.5
num_frames: int = 21
height: int = 448
width: int = 832
randomize_seed: bool = False
return_frames: bool = False
model_path: Optional[str] = None
class VideoGenerationResponse(BaseModel):
video_data: Optional[str] = None
seed: int
success: bool
error_message: Optional[str] = None
generation_time: Optional[float] = None
model_load_time: Optional[float] = None
inference_time: Optional[float] = None
encoding_time: Optional[float] = None
total_time: Optional[float] = None
stage_names: Optional[List[str]] = None
stage_execution_times: Optional[List[float]] = None
def encode_video_to_base64(frames: List[np.ndarray], fps: int = DEFAULT_FPS) -> str:
if not frames:
return ""
try:
buffer = io.BytesIO()
imageio.mimsave(buffer, frames, fps=fps, format="mp4")
buffer.seek(0)
video_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
return f"data:video/mp4;base64,{video_base64}"
except Exception as e:
print(f"Warning: Failed to encode video: {e}")
return ""
def setup_model_environment(model_path: str) -> None:
if "fullattn" in model_path.lower():
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"
else:
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
os.environ["FASTVIDEO_STAGE_LOGGING"] = "1"
def process_generation_result(result: Any) -> tuple[List[np.ndarray], float, List[str], List[float]]:
frames = result if isinstance(result, list) else result.get("frames", [])
generation_time = result.get("generation_time", 0.0) if isinstance(result, dict) else 0.0
logging_info = result.get("logging_info", None)
if logging_info:
stage_names = logging_info.get_execution_order()
stage_execution_times = [
logging_info.get_stage_info(stage_name).get("execution_time", 0.0)
for stage_name in stage_names
]
else:
stage_names = []
stage_execution_times = []
return frames, generation_time, stage_names, stage_execution_times
def prepare_sampling_params(video_request: VideoGenerationRequest, default_params: Any) -> Any:
params = deepcopy(default_params)
params.prompt = video_request.prompt
if video_request.use_negative_prompt:
params.negative_prompt = video_request.negative_prompt
params.seed = (video_request.seed if not video_request.randomize_seed
else torch.randint(0, SEED_RANGE_MAX, (1,)).item())
params.randomize_seed = video_request.randomize_seed
params.guidance_scale = video_request.guidance_scale
params.num_frames = video_request.num_frames
params.height = video_request.height
params.width = video_request.width
params.save_video = False
params.return_frames = False
return params
class BaseModelDeployment:
def __init__(self, model_path: str, output_path: str = "outputs"):
self.model_path = model_path
self.output_path = output_path
self.generator = None
self.default_params = None
os.makedirs(self.output_path, exist_ok=True)
setup_model_environment(self.model_path)
def _initialize_generator(self, config: Dict[str, Any]) -> None:
from fastvideo.entrypoints.video_generator import VideoGenerator
from fastvideo.configs.sample.base import SamplingParam
print(f"Initializing model: {self.model_path}")
self.generator = VideoGenerator.from_pretrained(
model_path=self.model_path,
num_gpus=1,
use_fsdp_inference=True,
text_encoder_cpu_offload=config["text_encoder_cpu_offload"],
dit_cpu_offload=config["dit_cpu_offload"],
vae_cpu_offload=config["vae_cpu_offload"],
VSA_sparsity=config["VSA_sparsity"],
enable_stage_verification=False,
)
self.default_params = SamplingParam.from_pretrained(self.model_path)
def generate_video(self, video_request: VideoGenerationRequest) -> VideoGenerationResponse:
total_start_time = time.time()
params = prepare_sampling_params(video_request, self.default_params)
inference_start_time = time.time()
result = self.generator.generate_video(
prompt=video_request.prompt,
sampling_param=params,
save_video=False,
return_frames=False,
)
inference_time = time.time() - inference_start_time
frames, generation_time, stage_names, stage_execution_times = process_generation_result(result)
encoding_start_time = time.time()
video_data = encode_video_to_base64(frames, fps=DEFAULT_FPS)
encoding_time = time.time() - encoding_start_time
total_time = time.time() - total_start_time
return VideoGenerationResponse(
video_data=video_data,
seed=params.seed,
success=True,
generation_time=generation_time,
inference_time=inference_time,
encoding_time=encoding_time,
total_time=total_time,
stage_names=stage_names,
stage_execution_times=stage_execution_times,
)
@serve.deployment(
ray_actor_options={"num_cpus": 2, "num_gpus": 1, "runtime_env": {"conda": "fv"}},
)
class T2VModelDeployment(BaseModelDeployment):
def __init__(self, t2v_model_path: str, output_path: str = "outputs"):
super().__init__(t2v_model_path, output_path)
self._initialize_generator(MODEL_CONFIGS["1.3B"])
print("✅ T2V model initialized successfully")
@serve.deployment(
ray_actor_options={"num_cpus": 16, "num_gpus": 1, "runtime_env": {"conda": "fv"}},
)
class T2V14BModelDeployment(BaseModelDeployment):
def __init__(self, t2v_14b_model_path: str, output_path: str = "outputs"):
super().__init__(t2v_14b_model_path, output_path)
# Override environment for 14B model
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
self._initialize_generator(MODEL_CONFIGS["14B"])
print("✅ T2V 14B model initialized successfully")
app = FastAPI()
limiter = Limiter(key_func=get_remote_address)
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
@serve.deployment(num_replicas=50, ray_actor_options={"num_cpus": 2})
@serve.ingress(app)
class FastVideoAPI:
def __init__(self, t2v_deployments: Dict[str, DeploymentHandle]):
self.t2v_deployments = t2v_deployments
# Initialize Prometheus metrics
self.request_count = Counter('fastvideo_requests_total', 'Total FastVideo requests', ['model_type', 'status'])
self.request_duration = Histogram('fastvideo_request_duration_seconds', 'FastVideo request duration', ['model_type'])
self.video_generation_time = Histogram('fastvideo_video_generation_seconds', 'Video generation time', ['model_type'])
def _get_model_name(self, model_path: Optional[str]) -> str:
return model_path.split('/')[-1] if model_path else "unknown"
def _record_metrics(self, model_name: str, status: str, duration: float, response: Optional[VideoGenerationResponse] = None) -> None:
self.request_count.labels(model_type=model_name, status=status).inc()
self.request_duration.labels(model_type=model_name).observe(duration)
if response and hasattr(response, 'generation_time') and response.generation_time:
self.video_generation_time.labels(model_type=model_name).observe(response.generation_time)
@app.post("/generate_video", response_model=VideoGenerationResponse)
@limiter.limit("10/minute")
async def generate_video(self, request: Request, video_request: VideoGenerationRequest) -> VideoGenerationResponse:
"""Route the request to the appropriate model deployment based on model_path."""
start_time = time.time()
model_name = self._get_model_name(video_request.model_path)
try:
if video_request.model_path not in self.t2v_deployments:
raise ValueError(f"Model {video_request.model_path} not found")
response_ref = self.t2v_deployments[video_request.model_path].generate_video.remote(video_request)
response = await response_ref
self._record_metrics(model_name, "success", time.time() - start_time, response)
return response
except Exception as e:
self._record_metrics(model_name, "error", time.time() - start_time)
return VideoGenerationResponse(
video_data=None,
seed=video_request.seed,
success=False,
error_message=str(e),
generation_time=0,
inference_time=0,
encoding_time=0,
total_time=0,
)
@app.get("/health")
@limiter.limit("10/minute")
async def health_check(self, request: Request) -> Dict[str, str]:
return {"status": "healthy"}
@app.get("/metrics")
async def metrics(self) -> Response:
return Response(generate_latest(), media_type="text/plain")
def validate_configuration(model_paths: List[str], replicas: List[int]) -> None:
assert len(model_paths) == len(replicas), "Number of models and replicas must match"
assert sum(replicas) <= NUM_GPUS, f"Total replicas ({sum(replicas)}) must be <= {NUM_GPUS}"
for model, replica_count in zip(model_paths, replicas):
assert model in SUPPORTED_MODELS, f"Model {model} not supported"
assert replica_count > 0, f"Replicas must be greater than 0"
def start_ray_serve(
*,
t2v_model_paths: str,
t2v_model_replicas: str,
output_path: str = "outputs",
host: str = "0.0.0.0",
port: int = 8000,
) -> None:
if not ray.is_initialized():
ray.init()
model_paths = t2v_model_paths.split(",")
replicas = [int(r) for r in t2v_model_replicas.split(",")]
validate_configuration(model_paths, replicas)
t2v_deps = {}
for model_path, replica_count in zip(model_paths, replicas):
t2v_dep = T2VModelDeployment.options(num_replicas=replica_count).bind(model_path, output_path)
t2v_deps[model_path] = t2v_dep
api = FastVideoAPI.bind(t2v_deps)
serve.run(api, route_prefix="/", name="fast_video")
print(f"Ray Serve backend started at http://{host}:{port}")
for model_path, replica_count in zip(model_paths, replicas):
print(f"T2V Model: {model_path} | Replicas: {replica_count}")
print(f"Health check: http://{host}:{port}/health")
print(f"Video generation endpoint: http://{host}:{port}/generate_video")
def setup_signal_handlers() -> None:
signal.signal(signal.SIGINT, lambda *_: sys.exit(0))
signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="FastVideo Ray Serve Backend")
parser.add_argument("--t2v_model_paths",
type=str,
default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
help="Comma separated list of paths to the T2V model(s)")
parser.add_argument("--t2v_model_replicas",
type=str,
default="4,4",
help="Comma separated list of number of replicas for the T2V model(s)")
parser.add_argument("--output_path",
type=str,
default="outputs",
help="Path to save generated videos")
parser.add_argument("--host",
type=str,
default="0.0.0.0",
help="Host to bind to")
parser.add_argument("--port",
type=int,
default=8000,
help="Port to bind to")
args = parser.parse_args()
model_paths = args.t2v_model_paths.split(",")
replicas = [int(r) for r in args.t2v_model_replicas.split(",")]
validate_configuration(model_paths, replicas)
start_ray_serve(
t2v_model_paths=args.t2v_model_paths,
t2v_model_replicas=args.t2v_model_replicas,
output_path=args.output_path,
host=args.host,
port=args.port,
)
setup_signal_handlers()
print("✅ FastVideo backend is running. Press Ctrl-C to stop.")
while True:
time.sleep(3600)
examples/inference/gradio/start.sh
+3
View File
@@ -0,0 +1,3 @@
python examples/inference/gradio/start_ray_serve_app.py \
--t2v_model_paths "FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers" \
--t2v_model_replicas "4,4"
examples/inference/gradio/start_ray_serve_app.py
+257
View File
@@ -0,0 +1,257 @@
"""
Startup script for FastVideo with Ray Serve backend and Gradio frontend.
This script starts both the backend and frontend services.
"""
import argparse
import os
import subprocess
import sys
import time
import threading
import signal
import requests
from pathlib import Path
from typing import Dict, Any, Optional
DEFAULT_BACKEND_HOST = "0.0.0.0"
DEFAULT_BACKEND_PORT = 8000
DEFAULT_FRONTEND_HOST = "0.0.0.0"
DEFAULT_FRONTEND_PORT = 7860
DEFAULT_OUTPUT_PATH = "outputs"
DEFAULT_T2V_MODELS = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers"
DEFAULT_T2V_REPLICAS = "4,4"
HEALTH_CHECK_TIMEOUT = 5
HEALTH_CHECK_MAX_RETRIES = 100
HEALTH_CHECK_INTERVAL = 2
PROCESS_SHUTDOWN_TIMEOUT = 5
PROCESS_MONITOR_INTERVAL = 1
PROJECT_ROOT = Path(__file__).parent.parent.parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
class ServiceManager:
def __init__(self, args: argparse.Namespace):
self.args = args
self.backend_process: Optional[subprocess.Popen] = None
self.frontend_process: Optional[subprocess.Popen] = None
self.backend_url = f"http://{args.backend_host}:{args.backend_port}"
def check_backend_health(self, max_retries: int = HEALTH_CHECK_MAX_RETRIES) -> bool:
health_url = f"{self.backend_url}/health"
for attempt in range(max_retries):
try:
response = requests.get(health_url, timeout=HEALTH_CHECK_TIMEOUT)
if response.status_code == 200:
print(f"✅ Backend is healthy at {self.backend_url}")
return True
except requests.exceptions.RequestException:
pass
if attempt < max_retries - 1:
print(f"⏳ Waiting for backend to start... ({attempt + 1}/{max_retries})")
time.sleep(HEALTH_CHECK_INTERVAL)
print(f"❌ Backend failed to start within {max_retries * HEALTH_CHECK_INTERVAL} seconds")
return False
def _create_monitor_thread(self, process: subprocess.Popen, service_name: str) -> threading.Thread:
def monitor():
if process.stdout:
for line in process.stdout:
print(f"[{service_name}] {line.rstrip()}")
thread = threading.Thread(target=monitor, daemon=True)
thread.start()
return thread
def _start_service(self, script_name: str, args_dict: Dict[str, Any], service_name: str) -> subprocess.Popen:
script_path = Path(__file__).parent / script_name
cmd = [sys.executable, str(script_path)]
for key, value in args_dict.items():
cmd.extend([f"--{key}", str(value)])
print(f"🚀 Starting {service_name}...")
print(f"Command: {' '.join(cmd)}")
process = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
universal_newlines=True,
bufsize=1
)
self._create_monitor_thread(process, service_name.upper())
return process
def start_backend(self) -> subprocess.Popen:
backend_args = {
"t2v_model_paths": self.args.t2v_model_paths,
"t2v_model_replicas": self.args.t2v_model_replicas,
"output_path": self.args.output_path,
"host": self.args.backend_host,
"port": self.args.backend_port
}
self.backend_process = self._start_service("ray_serve_backend.py", backend_args, "backend")
return self.backend_process
def start_frontend(self) -> subprocess.Popen:
frontend_args = {
"backend_url": self.backend_url,
"t2v_model_paths": self.args.t2v_model_paths,
"host": self.args.frontend_host,
"port": self.args.frontend_port
}
self.frontend_process = self._start_service("gradio_frontend.py", frontend_args, "frontend")
return self.frontend_process
def shutdown_services(self) -> None:
print("\n🛑 Shutting down services...")
processes = []
if self.frontend_process:
self.frontend_process.terminate()
processes.append(("frontend", self.frontend_process))
if self.backend_process:
self.backend_process.terminate()
processes.append(("backend", self.backend_process))
for name, process in processes:
try:
process.wait(timeout=PROCESS_SHUTDOWN_TIMEOUT)
print(f"✅ {name.capitalize()} stopped gracefully")
except subprocess.TimeoutExpired:
print(f"⚠️ Force killing {name} process...")
process.kill()
print("✅ All services stopped")
def monitor_processes(self) -> None:
if not self.backend_process or not self.frontend_process:
print("❌ Processes not properly initialized")
return
try:
while True:
if self.frontend_process.poll() is not None:
print("❌ Frontend process died unexpectedly")
break
if self.backend_process.poll() is not None:
print("❌ Backend process died unexpectedly")
break
time.sleep(PROCESS_MONITOR_INTERVAL)
except KeyboardInterrupt:
pass
self.shutdown_services()
def setup_signal_handlers(service_manager: ServiceManager) -> None:
def signal_handler(signum: int, frame: Any) -> None:
service_manager.shutdown_services()
sys.exit(0)
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
def print_startup_info(args: argparse.Namespace) -> None:
print("🎬 FastVideo Ray Serve App")
print("=" * 50)
print(f"T2V Models: {args.t2v_model_paths}")
print(f"T2V Model Replicas: {args.t2v_model_replicas}")
print(f"Output: {args.output_path}")
print(f"Backend: http://{args.backend_host}:{args.backend_port}")
print(f"Frontend: http://{args.frontend_host}:{args.frontend_port}")
print("=" * 50)
def parse_arguments() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="FastVideo Ray Serve App")
parser.add_argument("--t2v_model_paths",
type=str,
default=DEFAULT_T2V_MODELS,
help="Comma separated list of paths to the T2V model(s)")
parser.add_argument("--t2v_model_replicas",
type=str,
default=DEFAULT_T2V_REPLICAS,
help="Comma separated list of number of replicas for the T2V model(s)")
parser.add_argument("--output_path",
type=str,
default=DEFAULT_OUTPUT_PATH,
help="Path to save generated videos")
parser.add_argument("--backend_host",
type=str,
default=DEFAULT_BACKEND_HOST,
help="Backend host to bind to")
parser.add_argument("--backend_port",
type=int,
default=DEFAULT_BACKEND_PORT,
help="Backend port to bind to")
parser.add_argument("--frontend_host",
type=str,
default=DEFAULT_FRONTEND_HOST,
help="Frontend host to bind to")
parser.add_argument("--frontend_port",
type=int,
default=DEFAULT_FRONTEND_PORT,
help="Frontend port to bind to")
parser.add_argument("--skip_backend_check",
action="store_true",
help="Skip backend health check")
return parser.parse_args()
def main() -> None:
args = parse_arguments()
os.makedirs(args.output_path, exist_ok=True)
print_startup_info(args)
service_manager = ServiceManager(args)
setup_signal_handlers(service_manager)
try:
service_manager.start_backend()
if not args.skip_backend_check:
if not service_manager.check_backend_health():
print("❌ Backend failed to start. Terminating...")
service_manager.shutdown_services()
sys.exit(1)
service_manager.start_frontend()
print("\n🎉 Both services are starting up!")
print(f"📺 Frontend will be available at: http://{args.frontend_host}:{args.frontend_port}")
print(f"🔧 Backend API will be available at: http://{args.backend_host}:{args.backend_port}")
print("\nPress Ctrl+C to stop both services...")
service_manager.monitor_processes()
except Exception as e:
print(f"❌ Unexpected error: {e}")
service_manager.shutdown_services()
sys.exit(1)
if __name__ == "__main__":
main()
examples/inference/lora/wan_lora_inference.py
+2 -2
View File
@@ -10,7 +10,7 @@ def main():
dit_cpu_offload=False,
vae_cpu_offload=False,
text_encoder_cpu_offload=True,
pin_cpu_memory=False,
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
lora_path="benjamin-paine/steamboat-willie-1.3b",
lora_nickname="steamboat"
)
@@ -18,7 +18,7 @@ def main():
"height": 480,
"width": 832,
"num_frames": 81,
"guidance_scale": 5.0,
"guidance_scale": 6.0,
"num_inference_steps": 32,
"seed": 42,
}
examples/inference/lora/wan_lora_inference_from_ckpt.py
+12 -6
View File
@@ -11,22 +11,22 @@ def main():
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
num_gpus=1,
dit_cpu_offload=False,
vae_cpu_offload=False,
vae_cpu_offload=True,
text_encoder_cpu_offload=True,
pin_cpu_memory=False,
lora_path="checkpoints/wan_t2v_finetune_lora/checkpoint-1250/transformer",
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
lora_path="checkpoints/wan_t2v_finetune_lora/checkpoint-1000/transformer",
lora_nickname="crush_smol"
)
kwargs = {
"height": 480,
"width": 832,
"num_frames": 77,
"guidance_scale": 5.0,
"guidance_scale": 6.0,
"num_inference_steps": 50,
"seed": 42,
}
# Generate video with LoRA style
prompt = "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table."
prompt = "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press."
video = generator.generate_video(
prompt,
@@ -34,6 +34,12 @@ def main():
save_video=True,
**kwargs
)
prompt = "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press."
video = generator.generate_video(
prompt,
output_path=OUTPUT_PATH,
save_video=True,
**kwargs
)
if __name__ == "__main__":
main()
examples/inference/optimizations/attention_example.py
+1 -1
View File
@@ -14,7 +14,7 @@ def main():
dit_cpu_offload=False,
vae_cpu_offload=False,
text_encoder_cpu_offload=True,
pin_cpu_memory=False,
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
)
load_time = time.perf_counter() - start_time
print(f"Model loading time: {load_time:.2f} seconds")
examples/inference/optimizations/teacache_example.py
+1 -1
View File
@@ -12,7 +12,7 @@ def main():
dit_cpu_offload=False,
vae_cpu_offload=False,
text_encoder_cpu_offload=True,
pin_cpu_memory=False,
pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
)
load_time = time.perf_counter() - start_time
print(f"Model loading time: {load_time:.2f} seconds")
examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh
+1 -1
View File
@@ -54,7 +54,7 @@ validation_args=(
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 100
--validation_sampling_steps "40"
--validation_guidance_scale "1.0"
--validation_guidance_scale "6.0"
)
# Optimizer arguments
examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm
+1 -1
View File
@@ -88,7 +88,7 @@ validation_args=(
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 100
--validation_sampling_steps "40"
--validation_guidance_scale "1.0"
--validation_guidance_scale "6.0"
)
# Optimizer arguments
examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.sh
+1 -1
View File
@@ -54,7 +54,7 @@ validation_args=(
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 100
--validation_sampling_steps "40"
--validation_guidance_scale "1.0"
--validation_guidance_scale "6.0"
)
# Optimizer arguments
examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.slurm
+1 -1
View File
@@ -88,7 +88,7 @@ validation_args=(
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 100
--validation_sampling_steps "40"
--validation_guidance_scale "1.0"
--validation_guidance_scale "6.0"
)
# Optimizer arguments
examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v_lora.sh
+1 -1
View File
@@ -54,7 +54,7 @@ validation_args=(
--validation_preprocessed_path "$VALIDATION_DIR"
--validation_steps 100
--validation_sampling_steps "40"
--validation_guidance_scale "1.0"
--validation_guidance_scale "6.0"
)
# Optimizer arguments
examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh
+3 -2
View File
@@ -52,9 +52,9 @@ dataset_args=(
validation_args=(
--log_validation
--validation_dataset_file $VALIDATION_DATASET_FILE
--validation_steps 50
--validation_steps 200
--validation_sampling_steps "50"
--validation_guidance_scale "1.0"
--validation_guidance_scale "6.0"
)
# Optimizer arguments
@@ -77,6 +77,7 @@ miscellaneous_args=(
--num_euler_timesteps 50
--ema_start_step 0
--enable_gradient_checkpointing_type "full"
# --resume_from_checkpoint "checkpoints/wan_t2v_finetune/checkpoint-2500"
)
torchrun \
examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.slurm
+2 -2
View File
@@ -85,14 +85,14 @@ validation_args=(
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 100
--validation_sampling_steps "50"
--validation_guidance_scale "1.0"
--validation_guidance_scale "6.0"
)
# Optimizer arguments
optimizer_args=(
--learning_rate 5e-5
--mixed_precision "bf16"
--checkpointing_steps 500
--checkpointing_steps 400
--weight_decay 1e-4
--max_grad_norm 1.0
)
examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_lora.sh
+3 -3
View File
@@ -52,16 +52,16 @@ dataset_args=(
validation_args=(
--log_validation
--validation_dataset_file $VALIDATION_DATASET_FILE
--validation_steps 50
--validation_steps 200
--validation_sampling_steps "50"
--validation_guidance_scale "1.0"
--validation_guidance_scale "6.0"
)
# Optimizer arguments
optimizer_args=(
--learning_rate 5e-5
--mixed_precision "bf16"
--checkpointing_steps 500
--checkpointing_steps 400
--weight_decay 1e-4
--max_grad_norm 1.0
)
examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh
+24
View File
@@ -0,0 +1,24 @@
#!/bin/bash
GPU_NUM=1 # 2,4,8
MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
MODEL_TYPE="wan"
DATASET_PATH="data/crush-smol-test/"
OUTPUT_DIR="data/crush-smol_processed_t2v/"
torchrun --nproc_per_node=$GPU_NUM \
fastvideo/pipelines/preprocess/v1_preprocessing_new.py \
--model_path $MODEL_PATH \
--mode preprocess \
--workload_type t2v \
--preprocess.dataset_path $DATASET_PATH \
--preprocess.dataset_output_dir $OUTPUT_DIR \
--preprocess.preprocess_video_batch_size 4 \
--preprocess.dataloader_num_workers 0 \
--preprocess.max_height 480 \
--preprocess.max_width 832 \
--preprocess.num_frames 77 \
--preprocess.train_fps 16 \
--preprocess.samples_per_file 8 \
--preprocess.flush_frequency 8 \
--preprocess.video_length_tolerance_range 5
fastvideo/configs/models/dits/wanvideo.py
+1 -1
View File
@@ -96,7 +96,7 @@ class WanVideoArchConfig(DiTArchConfig):
super().__post_init__()
self.out_channels = self.out_channels or self.in_channels
self.hidden_size = self.num_attention_heads * self.attention_head_dim
self.num_channels_latents = self.in_channels if self.added_kv_proj_dim is None else self.out_channels
self.num_channels_latents = self.out_channels
@dataclass
fastvideo/configs/pipelines/registry.py
+3 -1
View File
@@ -31,7 +31,9 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
"FastVideo/FastWan2.2-TI2V-5B-Diffusers": FastWan2_2_TI2V_5B_Config,
"FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig,
"FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers": WanT2V720PConfig,
"Wan-AI/Wan2.2-TI2V-5B-Diffusers": WanT2V720PConfig
"Wan-AI/Wan2.2-TI2V-5B-Diffusers": WanT2V720PConfig,
"Wan-AI/Wan2.2-T2V-A14B-Diffusers": WanT2V480PConfig,
"Wan-AI/Wan2.2-I2V-A14B-Diffusers": WanI2V480PConfig,
# Add other specific weight variants
}
fastvideo/configs/sample/base.py
+1 -1
View File
@@ -20,7 +20,7 @@ class SamplingParam:
# Text inputs
prompt: str | list[str] | None = None
negative_prompt: str | None = None
negative_prompt: str = "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"
prompt_path: str | None = None
output_path: str = "outputs/"
output_video_name: str | None = None
fastvideo/configs/sample/registry.py
+15 -10
View File
@@ -6,12 +6,19 @@ from typing import Any
from fastvideo.configs.sample.hunyuan import (FastHunyuanSamplingParam,
HunyuanSamplingParam)
from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam
from fastvideo.configs.sample.wan import (FastWanT2V480PConfig,
Wan2_2_TI2V_5B_SamplingParam,
WanI2V_14B_480P_SamplingParam,
WanI2V_14B_720P_SamplingParam,
WanT2V_1_3B_SamplingParam,
WanT2V_14B_SamplingParam)
# isort: off
from fastvideo.configs.sample.wan import (
FastWanT2V480PConfig,
Wan2_2_I2V_A14B_SamplingParam,
Wan2_2_T2V_A14B_SamplingParam,
Wan2_2_TI2V_5B_SamplingParam,
WanI2V_14B_480P_SamplingParam,
WanI2V_14B_720P_SamplingParam,
WanT2V_1_3B_SamplingParam,
WanT2V_14B_SamplingParam,
)
# isort: on
from fastvideo.logger import init_logger
from fastvideo.utils import (maybe_download_model_index,
verify_model_config_and_directory)
@@ -29,10 +36,8 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
"FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig,
"Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
"FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
# "Wan-AI/Wan2.2-T2V-A14B-Diffusers":
# Wan2_2_T2V_A14B_SamplingParam,
# "Wan-AI/Wan2.2-I2V-A14B-Diffusers":
# Wan2_2_I2V_A14B_SamplingParam,
"Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam,
"Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam,
# Add other specific weight variants
}
fastvideo/configs/sample/wan.py
+8 -2
View File
@@ -129,9 +129,15 @@ class Wan2_2_TI2V_5B_SamplingParam(Wan2_2_Base_SamplingParam):
@dataclass
class Wan2_2_T2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam):
pass
guidance_scale: float = 4.0
guidance_scale_2: float = 3.0
num_inference_steps: int = 40
fps: int = 16
@dataclass
class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam):
pass
guidance_scale: float = 3.5
guidance_scale_2: float = 3.5
num_inference_steps: int = 40
fps: int = 16
fastvideo/entrypoints/video_generator.py
+9 -2
View File
@@ -9,6 +9,7 @@ diffusion models.
import math
import os
import time
from copy import deepcopy
from typing import Any
import imageio
@@ -202,6 +203,8 @@ class VideoGenerator:
if sampling_param is None:
sampling_param = SamplingParam.from_pretrained(
fastvideo_args.model_path)
else:
sampling_param = deepcopy(sampling_param)
kwargs["prompt"] = prompt
sampling_param.update(kwargs)
@@ -275,6 +278,7 @@ class VideoGenerator:
width: {target_width}
video_length: {sampling_param.num_frames}
prompt: {prompt}
image_path: {sampling_param.image_path}
neg_prompt: {sampling_param.negative_prompt}
seed: {sampling_param.seed}
infer_steps: {sampling_param.num_inference_steps}
@@ -304,7 +308,8 @@ class VideoGenerator:
# Run inference
start_time = time.perf_counter()
output_batch = self.executor.execute_forward(batch, fastvideo_args)
samples = output_batch
samples = output_batch.output
logging_info = output_batch.logging_info
gen_time = time.perf_counter() - start_time
logger.info("Generated successfully in %.2f seconds", gen_time)
@@ -334,9 +339,11 @@ class VideoGenerator:
else:
return {
"samples": samples,
"frames": frames,
"prompts": prompt,
"size": (target_height, target_width, batch.num_frames),
"generation_time": gen_time
"generation_time": gen_time,
"logging_info": logging_info,
}
def set_lora_adapter(self,
fastvideo/fastvideo_args.py
+3
View File
@@ -158,6 +158,9 @@ class FastVideoArgs:
# # DMD parameters
# dmd_denoising_steps: List[int] | None = field(default=None)
# MoE parameters used by Wan2.2
boundary_ratio: float | None = None
@property
def training_mode(self) -> bool:
return not self.inference_mode
fastvideo/layers/lora/linear.py
+3 -3
View File
@@ -101,15 +101,15 @@ class BaseLayerWithLoRA(nn.Module):
B: torch.Tensor,
training_mode: bool = False,
lora_path: str | None = None) -> None:
self.lora_A = A # share storage with weights in the pipeline
self.lora_B = B
self.lora_A = torch.nn.Parameter(
A) # share storage with weights in the pipeline
self.lora_B = torch.nn.Parameter(B)
self.disable_lora = False
if not training_mode:
self.merge_lora_weights()
self.lora_path = lora_path
@torch.no_grad()
# @torch.compile()
def merge_lora_weights(self) -> None:
if self.disable_lora:
return
fastvideo/models/dits/wanvideo.py
-2
View File
@@ -672,11 +672,9 @@ class WanTransformer3DModel(CachableDiT):
for block in self.blocks:
hidden_states = block(hidden_states, encoder_hidden_states,
timestep_proj, freqs_cis)
# if teacache is enabled, we need to cache the original hidden states
if enable_teacache:
self.maybe_cache_states(hidden_states, original_hidden_states)
# 5. Output norm, projection & unpatchify
shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2,
dim=1)
fastvideo/models/loader/component_loader.py
+2 -2
View File
@@ -72,8 +72,7 @@ class ComponentLoader(ABC):
module_loaders = {
"scheduler": (SchedulerLoader, "diffusers"),
"transformer": (TransformerLoader, "diffusers"),
"real_score_transformer": (TransformerLoader, "diffusers"),
"fake_score_transformer": (TransformerLoader, "diffusers"),
"transformer_2": (TransformerLoader, "diffusers"),
"vae": (VAELoader, "diffusers"),
"text_encoder": (TextEncoderLoader, "transformers"),
"text_encoder_2": (TextEncoderLoader, "transformers"),
@@ -452,6 +451,7 @@ class TransformerLoader(ComponentLoader):
hsdp_replicate_dim=fastvideo_args.hsdp_replicate_dim,
hsdp_shard_dim=fastvideo_args.hsdp_shard_dim,
cpu_offload=fastvideo_args.dit_cpu_offload,
pin_cpu_memory=fastvideo_args.pin_cpu_memory,
fsdp_inference=fastvideo_args.use_fsdp_inference,
# TODO(will): make these configurable
param_dtype=torch.bfloat16,
fastvideo/models/schedulers/base.py
+2 -1
View File
@@ -8,10 +8,11 @@ import torch
class BaseScheduler(ABC):
timesteps: torch.Tensor
order: int
num_train_timesteps: int
def __init__(self, *args, **kwargs) -> None:
# Check if subclass has defined all required properties
required_attributes = ['timesteps', 'order']
required_attributes = ['timesteps', 'order', 'num_train_timesteps']
for attr in required_attributes:
if not hasattr(self, attr):
fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py
+3
View File
@@ -134,6 +134,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
self.timesteps = sigmas * num_train_timesteps
self.num_train_timesteps = num_train_timesteps
self._step_index: int | None = None
self._begin_index: int | None = None
@@ -145,6 +146,8 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
self.sigma_min = self.sigmas[-1].item()
self.sigma_max = self.sigmas[0].item()
BaseScheduler.__init__(self)
@property
def shift(self) -> float:
"""
fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py
+3
View File
@@ -117,6 +117,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
self.sigmas = sigmas
self.timesteps = sigmas * num_train_timesteps
self.num_train_timesteps = num_train_timesteps
self.model_outputs = [None] * solver_order
self.timestep_list: list[Any | None] = [None] * solver_order
@@ -132,6 +133,8 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
self.sigma_min = self.sigmas[-1].item()
self.sigma_max = self.sigmas[0].item()
BaseScheduler.__init__(self)
@property
def step_index(self):
"""
fastvideo/models/schedulers/scheduling_unipc_multistep.py
+1
View File
@@ -273,6 +273,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
self.num_inference_steps = None
timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy()
self.timesteps = torch.from_numpy(timesteps)
self.num_train_timesteps = num_train_timesteps
self.model_outputs = [None] * solver_order
self.timestep_list = [None] * solver_order
self.lower_order_nums = 0
fastvideo/pipelines/basic/wan/wan_i2v_pipeline.py
+14 -12
View File
@@ -12,12 +12,10 @@ from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
from fastvideo.pipelines.lora_pipeline import LoRAPipeline
# isort: off
from fastvideo.pipelines.stages import (ImageEncodingStage, ConditioningStage,
DecodingStage, DenoisingStage,
EncodingStage, InputValidationStage,
LatentPreparationStage,
TextEncodingStage,
TimestepPreparationStage)
from fastvideo.pipelines.stages import (
ImageEncodingStage, ConditioningStage, DecodingStage, DenoisingStage,
ImageVAEEncodingStage, InputValidationStage, LatentPreparationStage,
TextEncodingStage, TimestepPreparationStage)
# isort: on
from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (
FlowUniPCMultistepScheduler)
@@ -48,11 +46,14 @@ class WanImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase):
tokenizers=[self.get_module("tokenizer")],
))
self.add_stage(stage_name="image_encoding_stage",
stage=ImageEncodingStage(
image_encoder=self.get_module("image_encoder"),
image_processor=self.get_module("image_processor"),
))
if (self.get_module("image_encoder") is not None
and self.get_module("image_processor") is not None):
self.add_stage(
stage_name="image_encoding_stage",
stage=ImageEncodingStage(
image_encoder=self.get_module("image_encoder"),
image_processor=self.get_module("image_processor"),
))
self.add_stage(stage_name="conditioning_stage",
stage=ConditioningStage())
@@ -67,11 +68,12 @@ class WanImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase):
transformer=self.get_module("transformer")))
self.add_stage(stage_name="image_latent_preparation_stage",
stage=EncodingStage(vae=self.get_module("vae")))
stage=ImageVAEEncodingStage(vae=self.get_module("vae")))
self.add_stage(stage_name="denoising_stage",
stage=DenoisingStage(
transformer=self.get_module("transformer"),
transformer_2=self.get_module("transformer_2"),
scheduler=self.get_module("scheduler")))
self.add_stage(stage_name="decoding_stage",
fastvideo/pipelines/basic/wan/wan_pipeline.py
+1
View File
@@ -61,6 +61,7 @@ class WanPipeline(LoRAPipeline, ComposedPipelineBase):
self.add_stage(stage_name="denoising_stage",
stage=DenoisingStage(
transformer=self.get_module("transformer"),
transformer_2=self.get_module("transformer_2", None),
scheduler=self.get_module("scheduler"),
pipeline=self))
fastvideo/pipelines/composed_pipeline_base.py
+35 -12
View File
@@ -37,6 +37,7 @@ class ComposedPipelineBase(ABC):
is_video_pipeline: bool = False # To be overridden by video pipelines
_required_config_modules: list[str] = []
_extra_config_module_map: dict[str, str] = {}
training_args: TrainingArgs | None = None
fastvideo_args: FastVideoArgs | TrainingArgs | None = None
modules: dict[str, torch.nn.Module] = {}
@@ -148,9 +149,7 @@ class ComposedPipelineBase(ABC):
# model is loaded with the correct precision. Subsequently we will
# use FSDP2's MixedPrecisionPolicy to set the precision for the
# fwd, bwd, and other operations' precision.
# fastvideo_args.precision = fastvideo_args.master_weight_type
assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training'
# assert fastvideo_args.precision == 'fp32', 'only fp32 is supported for training'
logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args)
@@ -239,6 +238,18 @@ class ComposedPipelineBase(ABC):
model_index.pop("_class_name")
model_index.pop("_diffusers_version")
# @TODO(Wei): Temporary hack
if "boundary_ratio" in model_index and model_index[
"boundary_ratio"] is not None:
logger.info(
"MoE pipeline detected. Adding transformer_2 to self.required_config_modules..."
)
self.required_config_modules.append("transformer_2")
if fastvideo_args.boundary_ratio is None:
logger.info(
"MoE pipeline detected. Setting boundary ratio to %s",
model_index["boundary_ratio"])
fastvideo_args.boundary_ratio = model_index["boundary_ratio"]
model_index.pop("boundary_ratio", None)
model_index.pop("expand_timesteps", None)
@@ -248,12 +259,20 @@ class ComposedPipelineBase(ABC):
) > 1, "model_index.json must contain at least one pipeline module"
for module_name in self.required_config_modules:
if module_name not in model_index:
if module_name not in model_index and module_name in self._extra_config_module_map:
extra_module_value = self._extra_config_module_map[module_name]
logger.warning(
"model_index.json does not contain a %s module, adding %s to model_index",
module_name, module_name)
if 'transformer' in module_name:
model_index[module_name] = model_index['transformer']
"model_index.json does not contain a %s module, but found {%s: %s} in _extra_config_module_map, adding to model_index.",
module_name, module_name, extra_module_value)
if extra_module_value in model_index:
logger.info("Using module %s for %s", extra_module_value,
module_name)
model_index[module_name] = model_index[extra_module_value]
continue
else:
raise ValueError(
f"Required module key: {module_name} value: {model_index.get(module_name)} was not found in loaded modules {model_index.keys()}"
)
# all the component models used by the pipeline
required_modules = self.required_config_modules
@@ -263,6 +282,7 @@ class ComposedPipelineBase(ABC):
for module_name, (transformers_or_diffusers,
architecture) in model_index.items():
if transformers_or_diffusers is None:
self.required_config_modules.remove(module_name)
continue
if module_name not in required_modules:
logger.info("Skipping module %s", module_name)
@@ -271,14 +291,17 @@ class ComposedPipelineBase(ABC):
logger.info("Using module %s already provided", module_name)
modules[module_name] = loaded_modules[module_name]
continue
if 'transformer' in module_name:
loading_module_name = module_name.split("_")[-1]
# we load the module from the extra config module map if it exists
if module_name in self._extra_config_module_map:
load_module_name = self._extra_config_module_map[module_name]
else:
loading_module_name = module_name
load_module_name = module_name
component_model_path = os.path.join(self.model_path,
loading_module_name)
load_module_name)
module = PipelineComponentLoader.load_module(
module_name=module_name,
module_name=load_module_name,
component_model_path=component_model_path,
transformers_or_diffusers=transformers_or_diffusers,
fastvideo_args=fastvideo_args,
fastvideo/pipelines/pipeline_batch_info.py
+61 -8
View File
@@ -9,15 +9,55 @@ in a functional manner, reducing the need for explicit parameter passing.
import pprint
from dataclasses import asdict, dataclass, field
from typing import Any
from typing import TYPE_CHECKING, Any
import PIL.Image
import torch
if TYPE_CHECKING:
from torchcodec.decoders import VideoDecoder
import time
from collections import OrderedDict
from fastvideo.attention import AttentionMetadata
from fastvideo.configs.sample.teacache import TeaCacheParams, WanTeaCacheParams
class PipelineLoggingInfo:
"""Simple approach using OrderedDict to track stage metrics."""
def __init__(self):
# OrderedDict preserves insertion order and allows easy access
self.stages: OrderedDict[str, dict[str, Any]] = OrderedDict()
def add_stage_execution_time(self, stage_name: str, execution_time: float):
"""Add execution time for a stage."""
if stage_name not in self.stages:
self.stages[stage_name] = {}
self.stages[stage_name]['execution_time'] = execution_time
self.stages[stage_name]['timestamp'] = time.time()
def add_stage_metric(self, stage_name: str, metric_name: str, value: Any):
"""Add any metric for a stage."""
if stage_name not in self.stages:
self.stages[stage_name] = {}
self.stages[stage_name][metric_name] = value
def get_stage_info(self, stage_name: str) -> dict[str, Any]:
"""Get all info for a specific stage."""
return self.stages.get(stage_name, {})
def get_execution_order(self) -> list[str]:
"""Get stages in execution order."""
return list(self.stages.keys())
def get_total_execution_time(self) -> float:
"""Get total pipeline execution time."""
return sum(
stage.get('execution_time', 0) for stage in self.stages.values())
@dataclass
class ForwardBatch:
"""
@@ -75,15 +115,15 @@ class ForwardBatch:
image_latent: torch.Tensor | None = None
# Latent dimensions
height_latents: int | None = None
width_latents: int | None = None
num_frames: int = 1 # Default for image models
height_latents: list[int] | int | None = None
width_latents: list[int] | int | None = None
num_frames: list[int] | int = 1 # Default for image models
num_frames_round_down: bool = False # Whether to round down num_frames if it's not divisible by num_gpus
# Original dimensions (before VAE scaling)
height: int | None = None
width: int | None = None
fps: int | None = None
height: list[int] | int | None = None
width: list[int] | int | None = None
fps: list[int] | int | None = None
# Timesteps
timesteps: torch.Tensor | None = None
@@ -93,6 +133,7 @@ class ForwardBatch:
# Scheduler parameters
num_inference_steps: int = 50
guidance_scale: float = 1.0
guidance_scale_2: float | None = None
guidance_rescale: float = 0.0
eta: float = 0.0
sigmas: list[float] | None = None
@@ -128,6 +169,10 @@ class ForwardBatch:
# VSA parameters
VSA_sparsity: float = 0.0
# Logging info
logging_info: PipelineLoggingInfo = field(
default_factory=PipelineLoggingInfo)
def __post_init__(self):
"""Initialize dependent fields after dataclass initialization."""
@@ -136,6 +181,8 @@ class ForwardBatch:
self.do_classifier_free_guidance = True
if self.negative_prompt_embeds is None:
self.negative_prompt_embeds = []
if self.guidance_scale_2 is None:
self.guidance_scale_2 = self.guidance_scale
def __str__(self):
return pprint.pformat(asdict(self), indent=2, width=120)
@@ -189,4 +236,10 @@ class TrainingBatch:
fake_score_loss: float = 0.0
dmd_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
fake_score_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
fake_score_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
@dataclass
class PreprocessBatch(ForwardBatch):
video_loader: list["VideoDecoder"] = field(default_factory=list)
video_file_name: list[str] = field(default_factory=list)
fastvideo/pipelines/pipeline_registry.py
+3 -3
View File
@@ -65,7 +65,7 @@ class _PipelineRegistry:
arch = _PIPELINE_NAME_TO_ARCHITECTURE_NAME[pipeline_name_in_config]
return set(self.pipelines[pipeline_type.value][arch].keys())
def _load_preprocessing_pipeline_cls(
def _load_preprocess_pipeline_cls(
self, workload_type: WorkloadType,
arch: str) -> type[ComposedPipelineBase] | None:
if workload_type == WorkloadType.I2V:
@@ -90,7 +90,7 @@ class _PipelineRegistry:
return None
if pipeline_type == PipelineType.PREPROCESS:
return self._load_preprocessing_pipeline_cls(workload_type, arch)
return self._load_preprocess_pipeline_cls(workload_type, arch)
elif pipeline_type == PipelineType.BASIC:
return self.pipelines[
pipeline_type.value][arch][pipeline_name_in_config]
@@ -131,7 +131,7 @@ def import_pipeline_classes(
Import pipeline classes based on the pipeline type and workload type.
Args:
pipeline_types: The pipeline types to load (basic, preprocessing, training).
pipeline_types: The pipeline types to load (basic, preprocess, training).
If None, loads all types.
Returns:
fastvideo/pipelines/preprocess/__init__.py
View File
fastvideo/pipelines/preprocess/preprocess_stages.py
+73
View File
@@ -0,0 +1,73 @@
from collections.abc import Callable
from typing import cast
import numpy as np
import torch
from einops import rearrange
from torchvision import transforms
from fastvideo.dataset.transform import (CenterCropResizeVideo,
TemporalRandomCrop)
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.pipelines.pipeline_batch_info import (ForwardBatch,
PreprocessBatch)
from fastvideo.pipelines.stages.base import PipelineStage
class VideoTransformStage(PipelineStage):
"""
Crop a video in temporal dimension.
"""
def __init__(self, train_fps: int, num_frames: int, max_height: int,
max_width: int, do_temporal_sample: bool) -> None:
self.train_fps = train_fps
self.num_frames = num_frames
if do_temporal_sample:
self.temporal_sample_fn: Callable | None = TemporalRandomCrop(
num_frames)
else:
self.temporal_sample_fn = None
self.video_transform = transforms.Compose([
CenterCropResizeVideo((max_height, max_width)),
])
def forward(self, batch: ForwardBatch,
fastvideo_args: FastVideoArgs) -> ForwardBatch:
batch = cast(PreprocessBatch, batch)
assert isinstance(batch.fps, list)
assert isinstance(batch.num_frames, list)
if batch.data_type != "video":
return batch
if len(batch.video_loader) == 0:
raise ValueError("Video loader is not set")
video_pixel_batch = []
for i in range(len(batch.video_loader)):
frame_interval = batch.fps[i] / self.train_fps
start_frame_idx = 0
frame_indices = np.arange(start_frame_idx, batch.num_frames[i],
frame_interval).astype(int)
if len(frame_indices) > self.num_frames:
if self.temporal_sample_fn is not None:
begin_index, end_index = self.temporal_sample_fn(
len(frame_indices))
frame_indices = frame_indices[begin_index:end_index]
else:
frame_indices = frame_indices[:self.num_frames]
video = batch.video_loader[i].get_frames_at(frame_indices).data
video = self.video_transform(video)
video_pixel_batch.append(video)
video_pixel_values = torch.stack(video_pixel_batch)
video_pixel_values = rearrange(video_pixel_values,
"b t c h w -> b c t h w")
video_pixel_values = video_pixel_values.to(torch.uint8)
video_pixel_values = video_pixel_values.float() / 255.0
batch.latents = video_pixel_values
return cast(ForwardBatch, batch)
fastvideo/pipelines/preprocess/v1_preprocessing_new.py
+27
View File
@@ -0,0 +1,27 @@
import os
from fastvideo.distributed import (
maybe_init_distributed_environment_and_model_parallel)
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.logger import init_logger
from fastvideo.utils import FlexibleArgumentParser
from fastvideo.workflow.preprocess.preprocess_workflow_t2v import (
PreprocessWorkflowT2V)
logger = init_logger(__name__)
def main(fastvideo_args: FastVideoArgs) -> None:
maybe_init_distributed_environment_and_model_parallel(1, 1)
num_gpus = int(os.environ["WORLD_SIZE"])
assert num_gpus == 1, "Only support 1 GPU"
preprocess_workflow = PreprocessWorkflowT2V(fastvideo_args)
preprocess_workflow.run()
if __name__ == "__main__":
parser = FlexibleArgumentParser()
parser = FastVideoArgs.add_cli_args(parser)
args = parser.parse_args()
fastvideo_args = FastVideoArgs.from_cli_args(args)
main(fastvideo_args)
fastvideo/pipelines/preprocess/wan/__init__.py
View File
fastvideo/pipelines/preprocess/wan/wan_preprocess_pipelines.py
+53
View File
@@ -0,0 +1,53 @@
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
from fastvideo.pipelines.preprocess.preprocess_stages import VideoTransformStage
from fastvideo.pipelines.stages import (EncodingStage, ImageEncodingStage,
TextEncodingStage)
class I2VPreprocessPipeline(ComposedPipelineBase):
_required_config_modules = [
"image_encoder", "image_processor", "text_encoder", "tokenizer", "vae"
]
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
self.add_stage(stage_name="prompt_encoding_stage",
stage=TextEncodingStage(
text_encoders=[self.get_module("text_encoder")],
tokenizers=[self.get_module("tokenizer")],
))
self.add_stage(stage_name="image_encoding_stage",
stage=ImageEncodingStage(
image_encoder=self.get_module("image_encoder"),
image_processor=self.get_module("image_processor"),
))
self.add_stage(stage_name="video_encoding_stage",
stage=EncodingStage(vae=self.get_module("vae"), ))
class T2VPreprocessPipeline(ComposedPipelineBase):
_required_config_modules = ["text_encoder", "tokenizer", "vae"]
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
assert fastvideo_args.preprocess_config is not None
self.add_stage(stage_name="prompt_encoding_stage",
stage=TextEncodingStage(
text_encoders=[self.get_module("text_encoder")],
tokenizers=[self.get_module("tokenizer")],
))
self.add_stage(
stage_name="video_transform_stage",
stage=VideoTransformStage(
train_fps=fastvideo_args.preprocess_config.train_fps,
num_frames=fastvideo_args.preprocess_config.num_frames,
max_height=fastvideo_args.preprocess_config.max_height,
max_width=fastvideo_args.preprocess_config.max_width,
do_temporal_sample=fastvideo_args.preprocess_config.
do_temporal_sample,
))
self.add_stage(stage_name="video_encoding_stage",
stage=EncodingStage(vae=self.get_module("vae"), ))
EntryClass = [I2VPreprocessPipeline, T2VPreprocessPipeline]
fastvideo/pipelines/stages/__init__.py
+3 -1
View File
@@ -12,7 +12,8 @@ from fastvideo.pipelines.stages.decoding import DecodingStage
from fastvideo.pipelines.stages.denoising import (DenoisingStage,
DmdDenoisingStage)
from fastvideo.pipelines.stages.encoding import EncodingStage
from fastvideo.pipelines.stages.image_encoding import ImageEncodingStage
from fastvideo.pipelines.stages.image_encoding import (ImageEncodingStage,
ImageVAEEncodingStage)
from fastvideo.pipelines.stages.input_validation import InputValidationStage
from fastvideo.pipelines.stages.latent_preparation import LatentPreparationStage
from fastvideo.pipelines.stages.stepvideo_encoding import (
@@ -32,6 +33,7 @@ __all__ = [
"EncodingStage",
"DecodingStage",
"ImageEncodingStage",
"ImageVAEEncodingStage",
"TextEncodingStage",
"StepvideoPromptEncodingStage",
]
fastvideo/pipelines/stages/base.py
+2
View File
@@ -155,6 +155,8 @@ class PipelineStage(ABC):
execution_time = time.perf_counter() - start_time
logger.info("[%s] Execution completed in %s ms", stage_name,
execution_time * 1000)
batch.logging_info.add_stage_execution_time(
stage_name, execution_time)
except Exception as e:
execution_time = time.perf_counter() - start_time
logger.error("[%s] Error during execution after %s ms: %s",
fastvideo/pipelines/stages/decoding.py
-3
View File
@@ -3,7 +3,6 @@
Decoding stage for diffusion pipelines.
"""
import gc
import weakref
import torch
@@ -140,8 +139,6 @@ class DecodingStage(PipelineStage):
del self.vae
if pipeline is not None and "vae" in pipeline.modules:
del pipeline.modules["vae"]
gc.collect()
torch.mps.empty_cache()
fastvideo_args.model_loaded["vae"] = False
return batch
fastvideo/pipelines/stages/denoising.py
+32 -9
View File
@@ -3,7 +3,6 @@
Denoising stage for diffusion pipelines.
"""
import gc
import inspect
import weakref
from collections.abc import Iterable
@@ -57,9 +56,14 @@ class DenoisingStage(PipelineStage):
the initial noise into the final output.
"""
def __init__(self, transformer, scheduler, pipeline=None) -> None:
def __init__(self,
transformer,
scheduler,
pipeline=None,
transformer_2=None) -> None:
super().__init__()
self.transformer = transformer
self.transformer_2 = transformer_2
self.scheduler = scheduler
self.pipeline = weakref.ref(pipeline) if pipeline else None
attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads
@@ -184,6 +188,12 @@ class DenoisingStage(PipelineStage):
assert neg_prompt_embeds is not None
assert torch.isnan(neg_prompt_embeds[0]).sum() == 0
# (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert
if fastvideo_args.boundary_ratio is not None:
boundary_timestep = fastvideo_args.boundary_ratio * self.scheduler.num_train_timesteps
else:
boundary_timestep = None
# Run denoising loop
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
@@ -191,6 +201,23 @@ class DenoisingStage(PipelineStage):
if hasattr(self, 'interrupt') and self.interrupt:
continue
if boundary_timestep is None or t >= boundary_timestep:
if (fastvideo_args.dit_cpu_offload
and self.transformer_2 is not None and next(
self.transformer_2.parameters()).device.type
== 'cuda'):
self.transformer_2.to('cpu')
current_model = self.transformer
current_guidance_scale = batch.guidance_scale
else:
# low-noise stage in wan2.2
if fastvideo_args.dit_cpu_offload and next(
self.transformer.parameters(
)).device.type == 'cuda':
self.transformer.to('cpu')
current_model = self.transformer_2
current_guidance_scale = batch.guidance_scale_2
# Expand latents for I2V
latent_model_input = latents.to(target_dtype)
if batch.image_latent is not None:
@@ -257,7 +284,7 @@ class DenoisingStage(PipelineStage):
# fastvideo_args=fastvideo_args
):
# Run transformer
noise_pred = self.transformer(
noise_pred = current_model(
latent_model_input,
prompt_embeds,
t_expand,
@@ -276,7 +303,7 @@ class DenoisingStage(PipelineStage):
# fastvideo_args=fastvideo_args
):
# Run transformer
noise_pred_uncond = self.transformer(
noise_pred_uncond = current_model(
latent_model_input,
neg_prompt_embeds,
t_expand,
@@ -285,7 +312,7 @@ class DenoisingStage(PipelineStage):
**neg_cond_kwargs,
)
noise_pred_text = noise_pred
noise_pred = noise_pred_uncond + batch.guidance_scale * (
noise_pred = noise_pred_uncond + current_guidance_scale * (
noise_pred_text - noise_pred_uncond)
# Apply guidance rescale if needed
@@ -296,14 +323,12 @@ class DenoisingStage(PipelineStage):
noise_pred_text,
guidance_rescale=batch.guidance_rescale,
)
# Compute the previous noisy sample
latents = self.scheduler.step(noise_pred,
t,
latents,
**extra_step_kwargs,
return_dict=False)[0]
# Update progress bar
if i == len(timesteps) - 1 or (
(i + 1) > num_warmup_steps and
@@ -329,8 +354,6 @@ class DenoisingStage(PipelineStage):
del self.transformer
if pipeline is not None and "transformer" in pipeline.modules:
del pipeline.modules["transformer"]
gc.collect()
torch.mps.empty_cache()
fastvideo_args.model_loaded["transformer"] = False
logger.info("Memory after deallocating transformer: %s",
torch.mps.current_allocated_memory())
fastvideo/pipelines/stages/encoding.py
+49 -125
View File
@@ -3,15 +3,12 @@
Encoding stage for diffusion pipelines.
"""
import PIL.Image
import torch
from fastvideo.distributed import get_local_torch_device
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.logger import init_logger
from fastvideo.models.vaes.common import ParallelTiledVAE
from fastvideo.models.vision_utils import (get_default_height_width, normalize,
numpy_to_pt, pil_to_numpy, resize)
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
from fastvideo.pipelines.stages.base import PipelineStage
from fastvideo.pipelines.stages.validators import V # Import validators
@@ -23,64 +20,67 @@ logger = init_logger(__name__)
class EncodingStage(PipelineStage):
"""
Stage for encoding pixel representations into latent space.
Stage for encoding pixel space representations into latent space.
This stage handles the encoding of pixel representations into the final
input format (e.g., latents).
This stage handles the encoding of pixel-space video/images into latent
representations for further processing in the diffusion pipeline.
"""
def __init__(self, vae: ParallelTiledVAE) -> None:
self.vae: ParallelTiledVAE = vae
@torch.no_grad()
def verify_input(self, batch: ForwardBatch,
fastvideo_args: FastVideoArgs) -> VerificationResult:
"""Verify encoding stage inputs."""
result = VerificationResult()
# Input video/images for VAE encoding: [batch_size, channels, frames, height, width]
result.add_check("latents", batch.latents,
[V.is_tensor, V.with_dims(5)])
return result
def verify_output(self, batch: ForwardBatch,
fastvideo_args: FastVideoArgs) -> VerificationResult:
"""Verify encoding stage outputs."""
result = VerificationResult()
# Encoded latents: [batch_size, channels, frames, height_latents, width_latents]
result.add_check("latents", batch.latents,
[V.is_tensor, V.with_dims(5)])
return result
def forward(
self,
batch: ForwardBatch,
fastvideo_args: FastVideoArgs,
) -> ForwardBatch:
"""
Encode pixel representations into latent space.
Encode pixel space representations into latent space.
Args:
batch: The current batch information.
fastvideo_args: The inference arguments.
Returns:
The batch with encoded outputs.
The batch with encoded latents.
"""
assert batch.latents is not None and isinstance(batch.latents,
torch.Tensor)
self.vae = self.vae.to(get_local_torch_device())
assert batch.height is not None
assert batch.width is not None
latent_height = batch.height // self.vae.spatial_compression_ratio
latent_width = batch.width // self.vae.spatial_compression_ratio
assert batch.pil_image is not None
image = batch.pil_image
image = self.preprocess(
image,
vae_scale_factor=self.vae.spatial_compression_ratio,
height=batch.height,
width=batch.width).to(get_local_torch_device(), dtype=torch.float32)
image = image.unsqueeze(2)
video_condition = torch.cat([
image,
image.new_zeros(image.shape[0], image.shape[1],
batch.num_frames - 1, batch.height, batch.width)
],
dim=2)
video_condition = video_condition.to(device=get_local_torch_device(),
dtype=torch.float32)
# Setup VAE precision
vae_dtype = PRECISION_TO_TYPE[
fastvideo_args.pipeline_config.vae_precision]
vae_autocast_enabled = (
vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
# Encode Image
# Normalize input to [-1, 1] range (reverse of decoding normalization)
latents = (batch.latents * 2.0 - 1.0).clamp(-1, 1)
# Move to appropriate device and dtype
latents = latents.to(get_local_torch_device())
# Encode image to latents
with torch.autocast(device_type="cuda",
dtype=vae_dtype,
enabled=vae_autocast_enabled):
@@ -89,47 +89,29 @@ class EncodingStage(PipelineStage):
# if fastvideo_args.vae_sp:
# self.vae.enable_parallel()
if not vae_autocast_enabled:
video_condition = video_condition.to(vae_dtype)
encoder_output = self.vae.encode(video_condition)
latents = latents.to(vae_dtype)
latents = self.vae.encode(latents).mean
generator = batch.generator
if generator is None:
raise ValueError("Generator must be provided")
latent_condition = self.retrieve_latents(encoder_output, generator)
# Apply shifting if needed
# Apply shifting if needed (reverse of decoding)
if (hasattr(self.vae, "shift_factor")
and self.vae.shift_factor is not None):
if isinstance(self.vae.shift_factor, torch.Tensor):
latent_condition -= self.vae.shift_factor.to(
latent_condition.device, latent_condition.dtype)
latents -= self.vae.shift_factor.to(latents.device,
latents.dtype)
else:
latent_condition -= self.vae.shift_factor
latents -= self.vae.shift_factor
if isinstance(self.vae.scaling_factor, torch.Tensor):
latent_condition = latent_condition * self.vae.scaling_factor.to(
latent_condition.device, latent_condition.dtype)
else:
latent_condition = latent_condition * self.vae.scaling_factor
# Apply scaling factor
if (hasattr(self.vae, "scaling_factor")
and self.vae.scaling_factor is not None):
if isinstance(self.vae.scaling_factor, torch.Tensor):
latents = latents * self.vae.scaling_factor.to(
latents.device, latents.dtype)
else:
latents = latents * self.vae.scaling_factor
mask_lat_size = torch.ones(1, 1, batch.num_frames, latent_height,
latent_width)
mask_lat_size[:, :, list(range(1, batch.num_frames))] = 0
first_frame_mask = mask_lat_size[:, :, 0:1]
first_frame_mask = torch.repeat_interleave(
first_frame_mask,
dim=2,
repeats=self.vae.temporal_compression_ratio)
mask_lat_size = torch.concat(
[first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
mask_lat_size = mask_lat_size.view(1, -1,
self.vae.temporal_compression_ratio,
latent_height, latent_width)
mask_lat_size = mask_lat_size.transpose(1, 2)
mask_lat_size = mask_lat_size.to(latent_condition.device)
batch.image_latent = torch.concat([mask_lat_size, latent_condition],
dim=1)
# Update batch with encoded latents
batch.latents = latents
# Offload models if needed
if hasattr(self, 'maybe_free_model_hooks'):
@@ -139,61 +121,3 @@ class EncodingStage(PipelineStage):
self.vae.to("cpu")
return batch
def retrieve_latents(self,
encoder_output: torch.Tensor,
generator: torch.Generator | None = None,
sample_mode: str = "sample"):
if sample_mode == "sample":
return encoder_output.sample(generator)
elif sample_mode == "argmax":
return encoder_output.mode()
else:
raise AttributeError(
"Could not access latents of provided encoder_output")
def preprocess(
self,
image: PIL.Image.Image,
vae_scale_factor: int,
height: int | None = None,
width: int | None = None,
resize_mode: str = "default", # "default", "fill", "crop"
) -> torch.Tensor:
image = [image]
height, width = get_default_height_width(image[0], vae_scale_factor,
height, width)
image = [
resize(i, height, width, resize_mode=resize_mode) for i in image
]
image = pil_to_numpy(image) # to np
image = numpy_to_pt(image) # to pt
do_normalize = True
if image.min() < 0:
do_normalize = False
if do_normalize:
image = normalize(image)
return image
def verify_input(self, batch: ForwardBatch,
fastvideo_args: FastVideoArgs) -> VerificationResult:
"""Verify encoding stage inputs."""
result = VerificationResult()
# result.add_check("pil_image", batch.pil_image)
result.add_check("height", batch.height, V.positive_int)
result.add_check("width", batch.width, V.positive_int)
result.add_check("generator", batch.generator,
V.generator_or_list_generators)
result.add_check("num_frames", batch.num_frames, V.positive_int)
return result
def verify_output(self, batch: ForwardBatch,
fastvideo_args: FastVideoArgs) -> VerificationResult:
"""Verify encoding stage outputs."""
result = VerificationResult()
result.add_check("image_latent", batch.image_latent,
[V.is_tensor, V.with_dims(5)])
return result
fastvideo/pipelines/stages/image_encoding.py
+184
View File
@@ -5,16 +5,21 @@ Image encoding stages for I2V diffusion pipelines.
This module contains implementations of image encoding stages for diffusion pipelines.
"""
import PIL
import torch
from fastvideo.distributed import get_local_torch_device
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.forward_context import set_forward_context
from fastvideo.logger import init_logger
from fastvideo.models.vaes.common import ParallelTiledVAE
from fastvideo.models.vision_utils import (get_default_height_width, normalize,
numpy_to_pt, pil_to_numpy, resize)
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
from fastvideo.pipelines.stages.base import PipelineStage
from fastvideo.pipelines.stages.validators import StageValidators as V
from fastvideo.pipelines.stages.validators import VerificationResult
from fastvideo.utils import PRECISION_TO_TYPE
logger = init_logger(__name__)
@@ -87,3 +92,182 @@ class ImageEncodingStage(PipelineStage):
result.add_check("image_embeds", batch.image_embeds,
V.list_of_tensors_dims(3))
return result
class ImageVAEEncodingStage(PipelineStage):
"""
Stage for encoding pixel representations into latent space.
This stage handles the encoding of pixel representations into the final
input format (e.g., latents).
"""
def __init__(self, vae: ParallelTiledVAE) -> None:
self.vae: ParallelTiledVAE = vae
def forward(
self,
batch: ForwardBatch,
fastvideo_args: FastVideoArgs,
) -> ForwardBatch:
"""
Encode pixel representations into latent space.
Args:
batch: The current batch information.
fastvideo_args: The inference arguments.
Returns:
The batch with encoded outputs.
"""
assert batch.height is not None and isinstance(batch.height, int)
assert batch.width is not None and isinstance(batch.width, int)
assert batch.num_frames is not None and isinstance(
batch.num_frames, int)
self.vae = self.vae.to(get_local_torch_device())
latent_height = batch.height // self.vae.spatial_compression_ratio
latent_width = batch.width // self.vae.spatial_compression_ratio
assert batch.pil_image is not None
image = batch.pil_image
image = self.preprocess(
image,
vae_scale_factor=self.vae.spatial_compression_ratio,
height=batch.height,
width=batch.width).to(get_local_torch_device(), dtype=torch.float32)
image = image.unsqueeze(2)
video_condition = torch.cat([
image,
image.new_zeros(image.shape[0], image.shape[1],
batch.num_frames - 1, batch.height, batch.width)
],
dim=2)
video_condition = video_condition.to(device=get_local_torch_device(),
dtype=torch.float32)
# Setup VAE precision
vae_dtype = PRECISION_TO_TYPE[
fastvideo_args.pipeline_config.vae_precision]
vae_autocast_enabled = (
vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
# Encode Image
with torch.autocast(device_type="cuda",
dtype=vae_dtype,
enabled=vae_autocast_enabled):
if fastvideo_args.pipeline_config.vae_tiling:
self.vae.enable_tiling()
# if fastvideo_args.vae_sp:
# self.vae.enable_parallel()
if not vae_autocast_enabled:
video_condition = video_condition.to(vae_dtype)
encoder_output = self.vae.encode(video_condition)
generator = batch.generator
if generator is None:
raise ValueError("Generator must be provided")
latent_condition = self.retrieve_latents(encoder_output, generator)
# Apply shifting if needed
if (hasattr(self.vae, "shift_factor")
and self.vae.shift_factor is not None):
if isinstance(self.vae.shift_factor, torch.Tensor):
latent_condition -= self.vae.shift_factor.to(
latent_condition.device, latent_condition.dtype)
else:
latent_condition -= self.vae.shift_factor
if isinstance(self.vae.scaling_factor, torch.Tensor):
latent_condition = latent_condition * self.vae.scaling_factor.to(
latent_condition.device, latent_condition.dtype)
else:
latent_condition = latent_condition * self.vae.scaling_factor
mask_lat_size = torch.ones(1, 1, batch.num_frames, latent_height,
latent_width)
mask_lat_size[:, :, list(range(1, batch.num_frames))] = 0
first_frame_mask = mask_lat_size[:, :, 0:1]
first_frame_mask = torch.repeat_interleave(
first_frame_mask,
dim=2,
repeats=self.vae.temporal_compression_ratio)
mask_lat_size = torch.concat(
[first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
mask_lat_size = mask_lat_size.view(1, -1,
self.vae.temporal_compression_ratio,
latent_height, latent_width)
mask_lat_size = mask_lat_size.transpose(1, 2)
mask_lat_size = mask_lat_size.to(latent_condition.device)
batch.image_latent = torch.concat([mask_lat_size, latent_condition],
dim=1)
# Offload models if needed
if hasattr(self, 'maybe_free_model_hooks'):
self.maybe_free_model_hooks()
self.vae.to("cpu")
return batch
def retrieve_latents(self,
encoder_output: torch.Tensor,
generator: torch.Generator | None = None,
sample_mode: str = "sample"):
if sample_mode == "sample":
return encoder_output.sample(generator)
elif sample_mode == "argmax":
return encoder_output.mode()
else:
raise AttributeError(
"Could not access latents of provided encoder_output")
def preprocess(
self,
image: PIL.Image.Image,
vae_scale_factor: int,
height: int | None = None,
width: int | None = None,
resize_mode: str = "default", # "default", "fill", "crop"
) -> torch.Tensor:
image = [image]
height, width = get_default_height_width(image[0], vae_scale_factor,
height, width)
image = [
resize(i, height, width, resize_mode=resize_mode) for i in image
]
image = pil_to_numpy(image) # to np
image = numpy_to_pt(image) # to pt
do_normalize = True
if image.min() < 0:
do_normalize = False
if do_normalize:
image = normalize(image)
return image
def verify_input(self, batch: ForwardBatch,
fastvideo_args: FastVideoArgs) -> VerificationResult:
"""Verify encoding stage inputs."""
result = VerificationResult()
# result.add_check("pil_image", batch.pil_image)
result.add_check("height", batch.height, V.positive_int)
result.add_check("width", batch.width, V.positive_int)
result.add_check("generator", batch.generator,
V.generator_or_list_generators)
result.add_check("num_frames", batch.num_frames, V.positive_int)
return result
def verify_output(self, batch: ForwardBatch,
fastvideo_args: FastVideoArgs) -> VerificationResult:
"""Verify encoding stage outputs."""
result = VerificationResult()
result.add_check("image_latent", batch.image_latent,
[V.is_tensor, V.with_dims(5)])
return result
fastvideo/pipelines/stages/input_validation.py
+1
View File
@@ -56,6 +56,7 @@ class InputValidationStage(PipelineStage):
Returns:
The validated batch information.
"""
self._generate_seeds(batch, fastvideo_args)
# Ensure prompt is properly formatted
fastvideo/pipelines/stages/latent_preparation.py
+1 -2
View File
@@ -90,7 +90,6 @@ class LatentPreparationStage(PipelineStage):
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
# Generate or use provided latents
if latents is None:
latents = randn_tensor(shape,
@@ -103,7 +102,6 @@ class LatentPreparationStage(PipelineStage):
# Scale the initial noise if needed
if hasattr(self.scheduler, "init_noise_sigma"):
latents = latents * self.scheduler.init_noise_sigma
# Update batch with prepared latents
batch.latents = latents
batch.raw_latent_shape = latents.shape
@@ -122,6 +120,7 @@ class LatentPreparationStage(PipelineStage):
Returns:
The batch with adjusted video length.
"""
video_length = batch.num_frames
use_temporal_scaling_frames = fastvideo_args.pipeline_config.vae_config.use_temporal_scaling_frames
if use_temporal_scaling_frames:
fastvideo/tests/modal/pr_test.py
+2 -2
View File
@@ -11,7 +11,7 @@ print(f"Using image: {image_tag}")
image = (
modal.Image.from_registry(image_tag, add_python="3.12")
.run_commands("rm -rf /FastVideo")
.apt_install("cmake", "pkg-config", "build-essential", "curl", "libssl-dev")
.apt_install("cmake", "pkg-config", "build-essential", "curl", "libssl-dev", "ffmpeg")
.run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable")
.run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc")
.env({
@@ -100,7 +100,7 @@ def run_precision_tests_STA():
@app.function(gpu="H100:1", image=image, timeout=900)
def run_precision_tests_VSA():
run_test("python csrc/attn/tests/test_block_sparse.py")
run_test("python csrc/attn/tests/test_vsa.py")
@app.function(gpu="L40S:1", image=image, timeout=3600)
def run_inference_lora_tests():
fastvideo/training/distillation_pipeline.py
+4 -4
View File
@@ -56,6 +56,10 @@ class DistillationPipeline(TrainingPipeline):
"scheduler", "transformer", "vae", "real_score_transformer",
"fake_score_transformer"
]
_extra_config_module_map = {
"real_score_transformer": "transformer",
"fake_score_transformer": "transformer"
}
validation_pipeline: ComposedPipelineBase
train_dataloader: StatefulDataLoader
train_loader_iter: Iterator[dict[str, Any]]
@@ -790,13 +794,11 @@ class DistillationPipeline(TrainingPipeline):
# Re-enable gradients for training
transformer.train()
gc.collect()
torch.cuda.empty_cache()
def visualize_intermediate_latents(self, training_batch: TrainingBatch,
training_args: TrainingArgs, step: int):
"""Add visualization data to wandb logging and save frames to disk."""
wandb_loss_dict = {}
torch.cuda.empty_cache()
dmd_latents_vis_dict = training_batch.dmd_latent_vis_dict
fake_score_latents_vis_dict = training_batch.fake_score_latent_vis_dict
fake_score_log_keys = ['generator_pred_video']
@@ -830,7 +832,6 @@ class DistillationPipeline(TrainingPipeline):
video, fps=24, format="mp4") # change to 16 for Wan2.1
# Clean up references
del video, latents
torch.cuda.empty_cache()
# Process DMD training data if available - use decode_stage instead of self.vae.decode
if 'generator_pred_video' in dmd_latents_vis_dict:
@@ -862,7 +863,6 @@ class DistillationPipeline(TrainingPipeline):
video, fps=24, format="mp4") # change to 16 for Wan2.1
# Clean up references
del video, latents
torch.cuda.empty_cache()
# Log to wandb
if self.global_rank == 0:
fastvideo/training/training_pipeline.py
+1 -1
View File
@@ -76,6 +76,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
if self.lora_training and fastvideo_args.lora_rank is None:
raise ValueError("lora rank must be set when using lora training")
set_random_seed(fastvideo_args.seed) # for lora param init
super().__init__(model_path, fastvideo_args, required_config_modules,
loaded_modules) # type: ignore
@@ -703,4 +704,3 @@ class TrainingPipeline(LoRAPipeline, ABC):
# Re-enable gradients for training
training_args.inference_mode = False
transformer.train()
torch.cuda.empty_cache()
fastvideo/training/training_utils.py
-1
View File
@@ -159,7 +159,6 @@ def save_checkpoint(transformer,
local_main_process_only=False)
cpu_state = gather_state_dict_on_cpu_rank0(transformer, device=None)
if rank == 0:
# Save model weights (consolidated)
transformer_save_dir = os.path.join(save_dir, "transformer")
fastvideo/training/wan_distillation_pipeline.py
+3 -3
View File
@@ -41,8 +41,6 @@ class WanDistillationPipeline(DistillationPipeline):
args_copy = deepcopy(training_args)
args_copy.inference_mode = True
args_copy.dit_cpu_offload = True
args_copy.pipeline_config.vae_config.load_encoder = False
validation_pipeline = WanDMDPipeline.from_pretrained(
training_args.model_path,
args=args_copy, # type: ignore
@@ -50,7 +48,9 @@ class WanDistillationPipeline(DistillationPipeline):
loaded_modules={"transformer": self.get_module("transformer")},
tp_size=training_args.tp_size,
sp_size=training_args.sp_size,
num_gpus=training_args.num_gpus)
num_gpus=training_args.num_gpus,
pin_cpu_memory=training_args.pin_cpu_memory,
dit_cpu_offload=True)
self.validation_pipeline = validation_pipeline
fastvideo/training/wan_i2v_distillation_pipeline.py
+1 -3
View File
@@ -7,8 +7,7 @@ import torch
from einops import rearrange
from fastvideo.configs.sample import SamplingParam
from fastvideo.dataset.dataloader.schema import (pyarrow_schema_i2v,
pyarrow_schema_i2v_validation)
from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v
from fastvideo.distributed import get_local_torch_device
from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
from fastvideo.logger import init_logger
@@ -48,7 +47,6 @@ class WanI2VDistillationPipeline(DistillationPipeline):
def set_schemas(self):
self.train_dataset_schema = pyarrow_schema_i2v
self.validation_dataset_schema = pyarrow_schema_i2v_validation
def initialize_validation_pipeline(self, training_args: TrainingArgs):
logger.info("Initializing validation pipeline...")
fastvideo/training/wan_training_pipeline.py
+3 -3
View File
@@ -36,8 +36,6 @@ class WanTrainingPipeline(TrainingPipeline):
args_copy = deepcopy(training_args)
args_copy.inference_mode = True
args_copy.dit_cpu_offload = True
args_copy.pipeline_config.vae_config.load_encoder = False
validation_pipeline = WanPipeline.from_pretrained(
training_args.model_path,
args=args_copy, # type: ignore
@@ -47,7 +45,9 @@ class WanTrainingPipeline(TrainingPipeline):
},
tp_size=training_args.tp_size,
sp_size=training_args.sp_size,
num_gpus=training_args.num_gpus)
num_gpus=training_args.num_gpus,
pin_cpu_memory=training_args.pin_cpu_memory,
dit_cpu_offload=True)
self.validation_pipeline = validation_pipeline
fastvideo/version.py
+1 -1
View File
@@ -1 +1 @@
__version__ = "0.1.4"
__version__ = "0.1.5"
fastvideo/worker/gpu_worker.py
+8 -2
View File
@@ -11,6 +11,7 @@ from typing import Any, TextIO, cast
import psutil
import torch
import fastvideo.envs as envs
from fastvideo.distributed import (
cleanup_dist_env_and_memory,
maybe_init_distributed_environment_and_model_parallel)
@@ -71,7 +72,6 @@ class Worker:
# _check_if_gpu_supports_dtype(self.model_config.dtype)
if current_platform.is_cuda_alike():
torch.cuda.empty_cache()
self.init_gpu_memory = torch.cuda.mem_get_info()[0]
else:
# For MPS, we can't get memory info the same way
@@ -140,7 +140,13 @@ class Worker:
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()})
logging_info = None
if envs.FASTVIDEO_STAGE_LOGGING:
logging_info = output_batch.logging_info
self.pipe.send({
"output_batch": output_batch.output.cpu(),
"logging_info": logging_info
})
elif method_name == 'set_lora_adapter':
lora_nickname = recv_rpc['kwargs']['lora_nickname']
lora_path = recv_rpc['kwargs']['lora_path']
fastvideo/worker/multiproc_executor.py
+17 -3
View File
@@ -8,8 +8,9 @@ import socket
import time
from collections.abc import Callable
from multiprocessing.process import BaseProcess
from typing import Any, cast
from typing import Any
import fastvideo.envs as envs
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.logger import init_logger
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
@@ -40,7 +41,8 @@ class MultiprocExecutor(Executor):
logger.info("Using provided master port: %s", self.master_port)
else:
# Auto-find available port
for port in range(29503, 65535):
import random
for port in range(29503 + random.randint(0, 10000), 65535):
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
if s.connect_ex(('localhost', port)) != 0:
self.master_port = port
@@ -80,7 +82,19 @@ class MultiprocExecutor(Executor):
"forward_batch": forward_batch,
"fastvideo_args": fastvideo_args
})
return cast(ForwardBatch, responses[0]["output_batch"])
output = responses[0]["output_batch"]
logging_info = None
if envs.FASTVIDEO_STAGE_LOGGING:
logging_info = responses[0]["logging_info"]
else:
logging_info = None
result_batch = ForwardBatch(data_type=forward_batch.data_type,
output=output,
logging_info=logging_info)
return result_batch
def set_lora_adapter(self,
lora_nickname: str,
fastvideo/workflow/preprocess/components.py
+467
View File
@@ -0,0 +1,467 @@
import gc
import multiprocessing
import os
import random
from collections.abc import Callable
from concurrent.futures import ProcessPoolExecutor
from typing import Any
import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
import torch
from fastvideo.logger import init_logger
from fastvideo.pipelines.pipeline_batch_info import PreprocessBatch
logger = init_logger(__name__)
class PreprocessingDataValidator:
def __init__(self,
max_height: int = 1024,
max_width: int = 1024,
max_h_div_w_ratio: float = 17 / 16,
min_h_div_w_ratio: float = 8 / 16,
num_frames: int = 16,
train_fps: int = 24,
speed_factor: float = 1.0,
video_length_tolerance_range: float = 5.0,
drop_short_ratio: float = 0.0,
hw_aspect_threshold: float = 1.5):
self.max_height = max_height
self.max_width = max_width
self.max_h_div_w_ratio = max_h_div_w_ratio
self.min_h_div_w_ratio = min_h_div_w_ratio
self.num_frames = num_frames
self.train_fps = train_fps
self.speed_factor = speed_factor
self.video_length_tolerance_range = video_length_tolerance_range
self.drop_short_ratio = drop_short_ratio
self.hw_aspect_threshold = hw_aspect_threshold
self.validators: dict[str, Callable[[dict[str, Any]], bool]] = {}
self.filter_counts: dict[str, int] = {}
self.num_items_before_filtering = 0
self.num_items_after_filtering = 0
self.register_validators()
def register_validators(self) -> None:
self.add_validator("data_type_validator", self._validate_data_type)
self.add_validator("resolution_validator", self._validate_resolution)
self.add_validator("frame_sampling_validator",
self._validate_frame_sampling)
def add_validator(self, name: str, validator: Callable[[dict[str, Any]],
bool]) -> None:
self.validators[name] = validator
self.filter_counts[name] = 0
def __call__(self, batch: dict[str, Any]) -> bool:
"""
Validate whether the preprocessing data batch is valid.
"""
self.num_items_before_filtering += 1
for name, validator in self.validators.items():
if not validator(batch):
self.filter_counts[name] += 1
return False
self.num_items_after_filtering += 1
return True
def _validate_data_type(self, batch: dict[str, Any]) -> bool:
"""Validate basic validity of data items"""
return not (batch["caption"] is None or batch["caption"] == ""
or batch["fps"] is None or batch["fps"] <= 0
or batch["num_frames"] is None or batch["num_frames"] <= 0)
def _validate_resolution(self, batch: dict[str, Any]) -> bool:
"""Validate resolution constraints"""
aspect = self.max_height / self.max_width
if batch["resolution"] is not None:
height = batch["resolution"].get("height", None)
width = batch["resolution"].get("width", None)
if height is None or width is None:
return False
return self._filter_resolution(
height,
width,
max_h_div_w_ratio=self.hw_aspect_threshold * aspect,
min_h_div_w_ratio=1 / self.hw_aspect_threshold * aspect,
)
def _filter_resolution(self, h: int, w: int, max_h_div_w_ratio: float,
min_h_div_w_ratio: float) -> bool:
"""Filter based on aspect ratio"""
return (min_h_div_w_ratio <= h / w <= max_h_div_w_ratio) and (
self.min_h_div_w_ratio <= h / w <= self.max_h_div_w_ratio)
def _validate_frame_sampling(self, batch: dict[str, Any]) -> bool:
"""Validate frame sampling constraints"""
if (batch["num_frames"] / batch["fps"]
> self.video_length_tolerance_range *
(self.num_frames / self.train_fps * self.speed_factor)):
return False
frame_interval = batch["fps"] / self.train_fps
start_frame_idx = 0
frame_indices = np.arange(start_frame_idx, batch["num_frames"],
frame_interval).astype(int)
return not (len(frame_indices) < self.num_frames
and random.random() < self.drop_short_ratio)
def log_validation_stats(self):
info = ""
for name, count in self.filter_counts.items():
info += f"failed in {name}: {count}, "
info += f"number of items before filtering: {self.num_items_before_filtering}, "
info += f"number of items after filtering: {self.num_items_after_filtering}"
logger.info(info)
class VideoForwardBatchBuilder:
def __call__(self, batch: list) -> PreprocessBatch:
forward_batch = PreprocessBatch(
video_loader=[item["video"] for item in batch],
video_file_name=[item["name"] for item in batch],
height=[item["resolution"]["height"] for item in batch],
width=[item["resolution"]["width"] for item in batch],
fps=[item["fps"] for item in batch],
num_frames=[item["num_frames"] for item in batch],
prompt=[
item["caption"]
if isinstance(item["caption"], str) else item["caption"][0]
for item in batch
],
prompt_attention_mask=[],
data_type="video",
)
return forward_batch
def basic_record_creator(
video_name: str,
vae_latent: np.ndarray,
text_embedding: np.ndarray,
batch: PreprocessBatch,
idx: int,
extra_features: dict[str, Any] | None = None) -> dict[str, Any]:
"""Create a record for the Parquet dataset from PreprocessBatch."""
# For batch processing, we need to handle the case where some fields might be single values
# or lists depending on the batch size
# Get caption - handle both single string and list cases
caption = ""
if batch.prompt:
if isinstance(batch.prompt, list):
caption = batch.prompt[idx] if idx < len(
batch.prompt) else batch.prompt[0]
else:
caption = batch.prompt
# Get dimensions - these are typically single values in PreprocessBatch
assert isinstance(batch.width, list)
assert isinstance(batch.height, list)
assert isinstance(batch.fps, list)
assert isinstance(batch.num_frames, list)
width = batch.width[idx] if batch.width is not None else 0
height = batch.height[idx] if batch.height is not None else 0
# Get FPS - single value in PreprocessBatch
fps_val = float(batch.fps[idx]) if batch.fps is not None else 0.0
# For duration, we need to calculate it or use a default since it's not in PreprocessBatch
# duration = num_frames / fps if available
duration_val = 0.0
if batch.num_frames[idx] and batch.fps[idx] and batch.fps[idx] > 0:
duration_val = float(batch.num_frames[idx]) / float(batch.fps[idx])
record = {
"id": video_name,
"vae_latent_bytes": vae_latent.tobytes(),
"vae_latent_shape": list(vae_latent.shape),
"vae_latent_dtype": str(vae_latent.dtype),
"text_embedding_bytes": text_embedding.tobytes(),
"text_embedding_shape": list(text_embedding.shape),
"text_embedding_dtype": str(text_embedding.dtype),
"file_name": video_name,
"caption": caption,
"media_type": "video",
"width": int(width),
"height": int(height),
"num_frames": vae_latent.shape[1] if len(vae_latent.shape) > 1 else 0,
"duration_sec": duration_val,
"fps": fps_val,
}
if extra_features:
record.update(extra_features)
return record
class ParquetDatasetSaver:
"""Component for saving and writing Parquet datasets"""
def __init__(self,
flush_frequency: int,
samples_per_file: int,
schema_fields: list[str],
record_creator: Callable[..., dict[str, Any]] | None = None,
file_writer_fn: Callable | None = None):
"""
Initialize ParquetDatasetSaver
Args:
schema_fields_provider: Function that returns schema fields list
record_creator: Function for creating records
file_writer_fn: Function for writing records to files, uses default implementation if None
"""
self.flush_frequency = flush_frequency
self.samples_per_file = samples_per_file
self.schema_fields = schema_fields
self.create_record = record_creator or basic_record_creator
self.file_writer_fn: Callable[
[tuple], int] = file_writer_fn or self._default_file_writer_fn
self.all_tables: list[pa.Table] = []
self.num_processed_samples: int = 0
self.num_saved_files: int = 0
def save_and_write_parquet_batch(
self,
batch: PreprocessBatch,
output_dir: str,
extra_features: dict[str, Any] | None = None) -> None:
"""
Save and write Parquet dataset batch
Args:
batch: PreprocessBatch containing video and metadata information
output_dir: Output directory
extra_features: Extra features
Returns:
Number of processed samples
"""
assert isinstance(batch.latents, torch.Tensor)
assert isinstance(batch.prompt_embeds, list)
assert isinstance(batch.prompt_attention_mask, list)
prompt_embeds: list[torch.Tensor] = []
# Process non-padded embeddings (if needed)
if batch.prompt_attention_mask is not None:
prompt_embeds = self._process_non_padded_embeddings(
batch.prompt_embeds[0], batch.prompt_attention_mask[0])
else:
raise ValueError("prompt_attention_mask is None")
# Prepare batch data for Parquet dataset
batch_data = []
for idx, video_name in enumerate(batch.video_file_name):
# Get the corresponding latent and info using video name
vae_latent = batch.latents[idx].cpu().numpy()
text_embedding = prompt_embeds[idx].cpu().numpy()
# Get extra features for this sample if needed
sample_extra_features = {}
if extra_features:
for key, value in extra_features.items():
if isinstance(value, torch.Tensor):
sample_extra_features[key] = value[idx].cpu().numpy()
else:
sample_extra_features[key] = value[idx]
# Create record for Parquet dataset
record = self.create_record(video_name=video_name,
vae_latent=vae_latent,
text_embedding=text_embedding,
batch=batch,
idx=idx,
extra_features=sample_extra_features)
batch_data.append(record)
if batch_data:
self.num_processed_samples += len(batch_data)
# Convert batch data to PyArrow arrays
table = self._convert_batch_to_pyarrow_table(batch_data)
# Store the table in a list for later processing
self.all_tables.append(table)
logger.debug("Collected batch with %s samples", len(table))
# If flush is needed
if self.num_processed_samples >= self.flush_frequency:
self.flush_tables(output_dir)
def _process_non_padded_embeddings(
self, prompt_embeds: torch.Tensor,
prompt_attention_mask: torch.Tensor) -> list[torch.Tensor]:
"""Process non-padded embeddings"""
assert isinstance(prompt_embeds, torch.Tensor)
assert isinstance(prompt_attention_mask, torch.Tensor)
assert prompt_embeds.shape[0] == prompt_attention_mask.shape[0]
# Get sequence lengths from attention masks (number of 1s)
seq_lens = prompt_attention_mask.sum(dim=1)
non_padded_embeds = []
# Process each item in the batch
for i in range(prompt_embeds.size(0)):
seq_len = seq_lens[i].item()
# Slice the embeddings and masks to keep only non-padding parts
non_padded_embeds.append(prompt_embeds[i, :seq_len])
return non_padded_embeds
def _convert_batch_to_pyarrow_table(self,
batch_data: list[dict]) -> pa.Table:
"""Convert batch data to PyArrow table"""
arrays = []
for field in self.schema_fields:
if field.endswith('_bytes'):
arrays.append(
pa.array([record[field] for record in batch_data],
type=pa.binary()))
elif field.endswith('_shape'):
arrays.append(
pa.array([record[field] for record in batch_data],
type=pa.list_(pa.int32())))
elif field in ['width', 'height', 'num_frames']:
arrays.append(
pa.array([record[field] for record in batch_data],
type=pa.int32()))
elif field in ['duration_sec', 'fps']:
arrays.append(
pa.array([record[field] for record in batch_data],
type=pa.float32()))
else:
arrays.append(pa.array([record[field]
for record in batch_data]))
return pa.Table.from_arrays(arrays, names=self.schema_fields)
def flush_tables(self, output_dir: str):
"""Flush collected tables to disk"""
if not hasattr(self, 'all_tables') or not self.all_tables:
return
logger.debug("Combining %d batches...", len(self.all_tables))
combined_table = pa.concat_tables(self.all_tables)
assert len(combined_table) == self.num_processed_samples
logger.debug("Total samples collected: %d", len(combined_table))
# Calculate total number of chunks needed, putting remainder into self.all_tables
total_files = max(self.num_processed_samples // self.samples_per_file,
1)
logger.debug("Fixed samples per parquet file: %d",
self.samples_per_file)
logger.debug("Total number of parquet files: %d", total_files)
logger.debug(
"Total samples to be processed: %d (putting %d samples into self.all_tables)",
total_files * self.samples_per_file,
self.num_processed_samples % self.samples_per_file)
# Split work among processes
num_workers = int(min(multiprocessing.cpu_count(), total_files))
files_per_worker = (total_files + num_workers - 1) // num_workers
logger.debug("Using %d workers to process %d files", num_workers,
total_files)
logger.debug("Files per worker: %s", files_per_worker)
# Prepare work ranges
work_ranges = []
for i in range(num_workers):
start_idx = i * files_per_worker
end_idx = min((i + 1) * files_per_worker, total_files)
if start_idx < total_files:
work_ranges.append((start_idx, end_idx, combined_table, i,
output_dir, self.samples_per_file))
total_written = 0
failed_ranges = []
with ProcessPoolExecutor(max_workers=num_workers) as executor:
futures = {
executor.submit(self.file_writer_fn, work_range): work_range
for work_range in work_ranges
}
for future in futures:
try:
written = future.result()
total_written += written
logger.info("Processed file with %s samples", written)
except Exception as e:
work_range = futures[future]
failed_ranges.append(work_range)
logger.error("Failed to process range %s-%s: %s",
work_range[0], work_range[1], str(e))
# Retry failed ranges sequentially
if failed_ranges:
logger.warning("Retrying %s failed ranges sequentially",
len(failed_ranges))
for work_range in failed_ranges:
try:
total_written += self.file_writer_fn(work_range)
except Exception as e:
logger.error(
"Failed to process range %s-%s after retry: %s",
work_range[0], work_range[1], str(e))
self.num_saved_files += total_files
# Clear tables list
self.all_tables = []
if self.num_processed_samples > self.samples_per_file:
saved_samples = total_files * self.samples_per_file
self.all_tables.append(combined_table.slice(saved_samples))
self.num_processed_samples -= saved_samples
else:
self.num_processed_samples = 0
del combined_table
gc.collect()
def clean_up(self) -> None:
"""Clean up all tables"""
self.all_tables = []
self.num_processed_samples = 0
self.num_saved_files = 0
gc.collect()
def _default_file_writer_fn(self, args_tuple: tuple) -> int:
"""Default chunk processing implementation"""
start_idx, end_idx, combined_table, worker_id, output_dir, samples_per_file = args_tuple
written_count = 0
for file_idx in range(start_idx, end_idx):
start_row = file_idx * samples_per_file
end_row = min(start_row + samples_per_file, len(combined_table))
if start_row >= len(combined_table):
break
chunk_table = combined_table.slice(start_row, end_row - start_row)
# Write to file
output_file = os.path.join(
output_dir,
f"chunk_{file_idx + self.num_saved_files:06d}.parquet")
pq.write_table(chunk_table, output_file)
written_count += len(chunk_table)
return written_count
fastvideo/workflow/preprocess/preprocess_workflow.py
+98
View File
@@ -0,0 +1,98 @@
import os
from typing import cast
from datasets import load_dataset
from torch.utils.data import DataLoader
from fastvideo.configs.configs import PreprocessConfig
from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType
from fastvideo.pipelines.pipeline_registry import PipelineType
from fastvideo.workflow.preprocess.components import (
PreprocessingDataValidator, VideoForwardBatchBuilder)
from fastvideo.workflow.workflow_base import WorkflowBase
class PreprocessWorkflow(WorkflowBase):
def register_pipelines(self) -> None:
self.add_pipeline_config("preprocess_pipeline",
(PipelineType.PREPROCESS, self.fastvideo_args))
def register_components(self) -> None:
assert self.fastvideo_args.preprocess_config is not None
preprocess_config: PreprocessConfig = self.fastvideo_args.preprocess_config
# raw data validator
raw_data_validator = PreprocessingDataValidator(
max_height=preprocess_config.max_height,
max_width=preprocess_config.max_width,
num_frames=preprocess_config.num_frames,
train_fps=preprocess_config.train_fps,
speed_factor=preprocess_config.speed_factor,
video_length_tolerance_range=preprocess_config.
video_length_tolerance_range,
drop_short_ratio=preprocess_config.drop_short_ratio,
)
self.add_component("raw_data_validator", raw_data_validator)
# training dataset
training_dataset = load_dataset(preprocess_config.dataset_path,
split="train")
training_dataset = training_dataset.filter(raw_data_validator)
# we do not use collate_fn here because we use iterable-style Dataset
# and want to keep the original type of the dataset
training_dataloader = DataLoader(
training_dataset,
batch_size=preprocess_config.preprocess_video_batch_size,
num_workers=preprocess_config.dataloader_num_workers,
collate_fn=lambda x: x,
)
self.add_component("training_dataloader", training_dataloader)
# validation dataset
validation_dataset = load_dataset(preprocess_config.dataset_path,
split="validation")
validation_dataset = validation_dataset.filter(raw_data_validator)
validation_dataloader = DataLoader(
validation_dataset,
batch_size=preprocess_config.preprocess_video_batch_size,
num_workers=preprocess_config.dataloader_num_workers,
collate_fn=lambda x: x,
)
self.add_component("validation_dataloader", validation_dataloader)
# forward batch builder
video_forward_batch_builder = VideoForwardBatchBuilder()
self.add_component("video_forward_batch_builder",
video_forward_batch_builder)
def prepare_system_environment(self) -> None:
assert self.fastvideo_args.preprocess_config is not None
dataset_output_dir = self.fastvideo_args.preprocess_config.dataset_output_dir
os.makedirs(dataset_output_dir, exist_ok=True)
validation_dataset_output_dir = os.path.join(dataset_output_dir,
"validation_dataset")
os.makedirs(validation_dataset_output_dir, exist_ok=True)
self.validation_dataset_output_dir = validation_dataset_output_dir
training_dataset_output_dir = os.path.join(dataset_output_dir,
"training_dataset")
os.makedirs(training_dataset_output_dir, exist_ok=True)
self.training_dataset_output_dir = training_dataset_output_dir
@classmethod
def get_workflow_cls(cls,
fastvideo_args: FastVideoArgs) -> "PreprocessWorkflow":
if fastvideo_args.workload_type == WorkloadType.T2V:
from fastvideo.workflow.preprocess.preprocess_workflow_t2v import (
PreprocessWorkflowT2V)
return cast(PreprocessWorkflow, PreprocessWorkflowT2V)
elif fastvideo_args.workload_type == WorkloadType.I2V:
from fastvideo.workflow.preprocess.preprocess_workflow_i2v import (
PreprocessWorkflowI2V)
return cast(PreprocessWorkflow, PreprocessWorkflowI2V)
else:
raise ValueError(
f"Workload type: {fastvideo_args.workload_type} is not supported in preprocessing workflow."
)
fastvideo/workflow/preprocess/preprocess_workflow_i2v.py
+14
View File
@@ -0,0 +1,14 @@
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.workflow.preprocess.preprocess_workflow import PreprocessWorkflow
class PreprocessWorkflowI2V(PreprocessWorkflow):
def __init__(self, fastvideo_args: FastVideoArgs):
super().__init__(fastvideo_args)
def register_pipelines(self) -> None:
pass
def register_components(self) -> None:
pass
fastvideo/workflow/preprocess/preprocess_workflow_t2v.py
+69
View File
@@ -0,0 +1,69 @@
from typing import TYPE_CHECKING
from tqdm import tqdm
from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
from fastvideo.pipelines.pipeline_batch_info import PreprocessBatch
from fastvideo.workflow.preprocess.components import ParquetDatasetSaver
from fastvideo.workflow.preprocess.preprocess_workflow import PreprocessWorkflow
if TYPE_CHECKING:
from torch.utils.data import DataLoader
from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
from fastvideo.workflow.preprocess.components import (
VideoForwardBatchBuilder)
class PreprocessWorkflowT2V(PreprocessWorkflow):
training_dataloader: "DataLoader"
validation_dataloader: "DataLoader"
preprocess_pipeline: "ComposedPipelineBase"
processed_dataset_saver: "ParquetDatasetSaver"
video_forward_batch_builder: "VideoForwardBatchBuilder"
def register_components(self) -> None:
assert self.fastvideo_args.preprocess_config is not None
super().register_components()
self.add_component(
"processed_dataset_saver",
ParquetDatasetSaver(
flush_frequency=self.fastvideo_args.preprocess_config.
flush_frequency,
samples_per_file=self.fastvideo_args.preprocess_config.
samples_per_file,
schema_fields=[f.name for f in pyarrow_schema_t2v],
))
def run(self) -> None:
# Training dataset preprocessing
for batch in tqdm(self.training_dataloader,
desc="Preprocessing training dataset",
unit="batch"):
forward_batch: PreprocessBatch = self.video_forward_batch_builder(
batch)
forward_batch = self.preprocess_pipeline.forward(
forward_batch, self.fastvideo_args)
self.processed_dataset_saver.save_and_write_parquet_batch(
forward_batch, self.training_dataset_output_dir)
self.processed_dataset_saver.flush_tables(
self.training_dataset_output_dir)
self.processed_dataset_saver.clean_up()
# Validation dataset preprocessing
for batch in tqdm(self.validation_dataloader,
desc="Preprocessing validation dataset",
unit="batch"):
forward_batch = self.video_forward_batch_builder(batch)
forward_batch = self.preprocess_pipeline.forward(
forward_batch, self.fastvideo_args)
self.processed_dataset_saver.save_and_write_parquet_batch(
forward_batch, self.validation_dataset_output_dir)
self.processed_dataset_saver.flush_tables(
self.validation_dataset_output_dir)
self.processed_dataset_saver.clean_up()
fastvideo/workflow/workflow_base.py
+165 -10
View File
@@ -1,28 +1,183 @@
from abc import ABC, abstractmethod
from typing import Any, Optional
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.fastvideo_args import ExecutionMode, FastVideoArgs
from fastvideo.logger import init_logger
from fastvideo.pipelines import ComposedPipelineBase, build_pipeline
from fastvideo.pipelines.pipeline_registry import PipelineType
logger = init_logger(__name__)
class WorkflowBase(ABC):
pipeline_configs: dict[str, FastVideoArgs] = {}
pipelines: dict[str, ComposedPipelineBase] = {}
"""
Abstract base class for defining video processing workflows.
A workflow serves as the top-level orchestrator that coordinates multiple pipelines
and components to accomplish a specific video processing task. The workflow pattern
provides several key benefits:
1. **Separation of Concerns**: Workflows separate high-level orchestration logic
from low-level processing implementations in pipelines.
2. **Modularity**: Different workflows can be created for different execution modes
(preprocess, inference, etc.) while sharing common pipeline components.
3. **Configuration Management**: Workflows manage the configuration and initialization
of multiple related pipelines and components in a centralized manner.
4. **Environment Setup**: Workflows handle system-level setup and resource
allocation before pipeline execution begins.
5. **Lifecycle Management**: Workflows control the complete lifecycle from
initialization through execution to cleanup.
The workflow acts as a factory and coordinator, creating the appropriate pipelines
based on configuration, setting up the execution environment, and orchestrating
the overall processing flow.
"""
def __init__(self, fastvideo_args: FastVideoArgs):
"""
Initialize the workflow with configuration arguments.
Args:
fastvideo_args: Configuration object containing all parameters
needed for workflow and pipeline setup.
"""
self.fastvideo_args = fastvideo_args
def register_pipelines(self, pipeline_configs: dict[str, FastVideoArgs]):
self.pipeline_configs.update(pipeline_configs)
# TODO: pipeline_config should be: dict[str, PipelineConfig]
# pipeline_type should be included in the PipelineConfig
# pipeline_config[pipeline_name] = (pipeline_type, fastvideo_args)
self._pipeline_configs: dict[str, tuple[PipelineType,
FastVideoArgs]] = {}
self._pipelines: dict[str, ComposedPipelineBase] = {}
self._components: dict[str, Any] = {}
self.register_pipelines()
self.register_components()
def load_pipelines(self):
for pipeline_name, pipeline_config in self.pipeline_configs.items():
pipeline = build_pipeline(pipeline_config)
self.pipelines[pipeline_name] = pipeline
self.prepare_system_environment()
self.load_pipelines()
def load_pipelines(self) -> None:
"""
Create and initialize all registered pipelines.
This method instantiates pipeline objects from their configurations
and makes them available as both dictionary entries and instance
attributes for convenient access.
"""
for pipeline_name, pipeline_config in self._pipeline_configs.items():
pipeline_type, fastvideo_args = pipeline_config
pipeline = build_pipeline(fastvideo_args, pipeline_type)
self._pipelines[pipeline_name] = pipeline
setattr(self, pipeline_name, pipeline)
def add_pipeline_config(
self, pipeline_name: str,
pipeline_config: tuple[PipelineType, FastVideoArgs]) -> None:
"""
Register a pipeline configuration for later instantiation.
Args:
pipeline_name: Unique identifier for the pipeline.
pipeline_config: Tuple containing the pipeline type and
configuration arguments.
"""
self._pipeline_configs[pipeline_name] = pipeline_config
def add_component(self, component_name: str, component: Any) -> None:
"""
Register a component instance with the workflow.
Components are auxiliary objects that may be shared across pipelines
or used for workflow-level functionality (e.g., databases, caches,
external services).
Args:
component_name: Unique identifier for the component.
component: The component instance to register.
"""
self._components[component_name] = component
setattr(self, component_name, component)
def get_component(self, component_name: str) -> Any:
"""
Retrieve a registered component by name.
Args:
component_name: The name of the component to retrieve.
Returns:
The component instance.
"""
return self._components[component_name]
@abstractmethod
def get_components(self):
def register_components(self) -> None:
"""
Register workflow-specific components.
Subclasses must implement this method to register any components
needed for their specific workflow (e.g., databases, external APIs,
shared resources).
"""
pass
@abstractmethod
def register_pipelines(self) -> None:
"""
Register workflow-specific pipelines.
Subclasses must implement this method to define which pipelines
are needed for their specific workflow and how they should be
configured.
"""
pass
@abstractmethod
def prepare_system_environment(self) -> None:
"""
Prepare the system environment for workflow execution.
Subclasses must implement this method to handle any system-level
setup required before pipeline execution (e.g., GPU initialization,
temporary directories, resource allocation).
"""
pass
@abstractmethod
def run(self):
"""
Execute the main workflow logic.
Subclasses must implement this method to define the specific
execution flow for their workflow, coordinating the registered
pipelines and components to accomplish the desired task.
"""
pass
@classmethod
def get_workflow_cls(
cls, fastvideo_args: FastVideoArgs) -> Optional["WorkflowBase"]:
"""
Factory method to get the appropriate workflow class based on execution mode.
This method acts as a workflow factory, returning the appropriate
workflow class implementation based on the specified execution mode
in the configuration arguments.
Args:
fastvideo_args: Configuration object containing the execution mode
and other parameters.
Returns:
The appropriate workflow class for the specified execution mode,
or None if no workflow is available for the given mode.
"""
if fastvideo_args.mode == ExecutionMode.PREPROCESS:
from fastvideo.workflow.preprocess.preprocess_workflow import (
PreprocessWorkflow)
return PreprocessWorkflow.get_workflow_cls(fastvideo_args)
return None
pyproject.toml
+29 -6
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "fastvideo"
version = "0.1.4"
version = "0.1.5"
description = "FastVideo"
readme = "README.md"
requires-python = ">=3.10"
@@ -15,7 +15,7 @@ classifiers = [
dependencies = [
# Core Libraries
"scipy==1.14.1", "six==1.16.0", "h5py==3.12.1",
"scipy==1.14.1", "six==1.16.0", "h5py==3.12.1", "requests>=2.32.2",
# Machine Learning & Transformers
"transformers>=4.46.1", "tokenizers>=0.20.1", "sentencepiece==0.2.0",
@@ -30,11 +30,11 @@ dependencies = [
"imageio-ffmpeg==0.5.1", "einops",
# Experiment Tracking & Logging
"wandb>=0.19.11", "loguru", "test-tube==0.7.5",
"wandb>=0.21.0", "loguru", "test-tube==0.7.5",
# Miscellaneous Utilities
"tqdm", "pytest", "PyYAML==6.0.1", "protobuf>=5.28.3",
"gradio>=5.22.0", "moviepy==1.0.3", "flask",
"gradio>=5.22.0", "moviepy>=2.0.0", "flask",
"flask_restful", "aiohttp", "huggingface_hub", "cloudpickle",
# System & Monitoring Tools
"gpustat", "watch", "remote-pdb",
@@ -45,12 +45,35 @@ dependencies = [
# Training Dependencies
"torchdata",
"pyarrow",
"datasets",
"datasets==4.0.0",
"av",
# Preprocessing Dependencies
"torchcodec"
]
[tool.uv]
extra-index-url = ["https://download.pytorch.org/whl/cu128"]
prerelease = "allow"
[tool.uv.sources]
torch = [
{ index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
]
torchvision = [
{ index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
]
[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true
[project.optional-dependencies]
scripts/dataset_preparation/resize_videos.py
+1 -1
View File
@@ -5,7 +5,7 @@ from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
import numpy as np
from moviepy.editor import VideoFileClip
from moviepy.video.io.VideoFileClip import VideoFileClip
from skimage.transform import resize
from tqdm import tqdm
scripts/distill/v1_distill_dmd_wan.sh
+2 -2
View File
@@ -30,7 +30,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
--max_train_steps 30000 \
--learning_rate 1e-5 \
--mixed_precision "bf16" \
--checkpointing_steps 500 \
--checkpointing_steps 400 \
--validation_steps 100 \
--validation_sampling_steps "3" \
--log_validation \
@@ -43,7 +43,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
--num_width 832 \
--num_frames 77 \
--flow_shift 8 \
--validation_guidance_scale "1.0" \
--validation_guidance_scale "6.0" \
--master_weight_type "fp32" \
--dit_precision "fp32" \
--vae_precision "bf16" \
scripts/distill/v1_distill_dmd_wan_VSA.sh
+2 -2
View File
@@ -31,7 +31,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
--max_train_steps 30000 \
--learning_rate 1e-5 \
--mixed_precision "bf16" \
--checkpointing_steps 500 \
--checkpointing_steps 400 \
--validation_steps 100 \
--validation_sampling_steps "3" \
--log_validation \
@@ -44,7 +44,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
--num_width 832 \
--num_frames 61 \
--flow_shift 8 \
--validation_guidance_scale "1.0" \
--validation_guidance_scale "6.0" \
--master_weight_type "fp32" \
--dit_precision "fp32" \
--vae_precision "bf16" \
scripts/finetune/finetune_v1.sh
+2 -2
View File
@@ -31,7 +31,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
--learning_rate=1e-5\
--mixed_precision="bf16"\
--checkpointing_steps=6000 \
--validation_steps 50\
--validation_steps 200\
--validation_sampling_steps "2,4,8" \
--log_validation \
--checkpoints_total_limit 3\
@@ -42,7 +42,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
--num_height 480 \
--num_width 832 \
--num_frames 81 \
--validation_guidance_scale "1.0" \
--validation_guidance_scale "6.0" \
--num_euler_timesteps 50 \
--multi_phased_distill_schedule "4000-1" \
--weight_decay 0.01 \