@@ -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",
diff --git a/csrc/attn/README.md b/csrc/attn/README.md
index b7d226546..30af4bf34 100644
--- a/csrc/attn/README.md
+++ b/csrc/attn/README.md
@@ -84,7 +84,7 @@ out = sliding_tile_attention(q, k, v, window_size, 0, False)
### Test
```bash
python tests/test_sta.py # test STA
-python tests/test_block_sparse.py # test VSA
+python tests/test_vsa.py # test VSA
```
### Benchmark
```bash
diff --git a/csrc/attn/benchmarks/bench_sta.py b/csrc/attn/benchmarks/bench_sta.py
index 534fa2627..8e0c8afc1 100644
--- a/csrc/attn/benchmarks/bench_sta.py
+++ b/csrc/attn/benchmarks/bench_sta.py
@@ -86,8 +86,6 @@ def benchmark_attention(configurations):
# print(f"Average TFLOPS: {tflops_bwd}")
# print("=" * 60)
- torch.cuda.empty_cache()
-
return results
diff --git a/csrc/attn/setup_vsa.py b/csrc/attn/setup_vsa.py
index 4ffb71b8f..7e9026269 100644
--- a/csrc/attn/setup_vsa.py
+++ b/csrc/attn/setup_vsa.py
@@ -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'])
+]
diff --git a/csrc/attn/vsa/__init__.py b/csrc/attn/vsa/__init__.py
index 2faa9e18d..930e981c7 100644
--- a/csrc/attn/vsa/__init__.py
+++ b/csrc/attn/vsa/__init__.py
@@ -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
diff --git a/csrc/attn/vsa/block_sparse_wrapper.py b/csrc/attn/vsa/block_sparse_wrapper.py
index 5a7418339..bb588c44c 100644
--- a/csrc/attn/vsa/block_sparse_wrapper.py
+++ b/csrc/attn/vsa/block_sparse_wrapper.py
@@ -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)
\ No newline at end of file
+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)
\ No newline at end of file
diff --git a/docs/source/distillation/dmd.md b/docs/source/distillation/dmd.md
index ea5ab5572..74a3cbc4f 100644
--- a/docs/source/distillation/dmd.md
+++ b/docs/source/distillation/dmd.md
@@ -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
diff --git a/docs/source/index.md b/docs/source/index.md
index ac55e5e13..4f803b66e 100644
--- a/docs/source/index.md
+++ b/docs/source/index.md
@@ -9,7 +9,7 @@
:::{raw} html
-FastVideo is An unified inference and post-training framework for accelerated video generation.
+FastVideo is a unified inference and post-training framework for accelerated video generation.
diff --git a/docs/source/sliding_tile_attention/installation.md b/docs/source/sliding_tile_attention/installation.md
index 9069bcdd0..46bde6b6f 100644
--- a/docs/source/sliding_tile_attention/installation.md
+++ b/docs/source/sliding_tile_attention/installation.md
@@ -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
+```
diff --git a/docs/source/video_sparse_attention/installation.md b/docs/source/video_sparse_attention/installation.md
index 78a9a907c..00192802b 100644
--- a/docs/source/video_sparse_attention/installation.md
+++ b/docs/source/video_sparse_attention/installation.md
@@ -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
+```
diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.slurm b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.slurm
index 9e057f03d..a3bcdbd2f 100644
--- a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.slurm
+++ b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.slurm
@@ -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
diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slurm b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slurm
index 01ac5ccb1..de3bb87d5 100644
--- a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slurm
+++ b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slurm
@@ -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
diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.slurm b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.slurm
index 906389f2d..2912e3dba 100644
--- a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.slurm
+++ b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.slurm
@@ -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
diff --git a/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/README.md b/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/README.md
new file mode 100644
index 000000000..f709a92da
--- /dev/null
+++ b/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/README.md
@@ -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.
\ No newline at end of file
diff --git a/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B.sh b/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B.sh
index b262ca07c..a6f3d033a 100644
--- a/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B.sh
+++ b/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B.sh
@@ -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
diff --git a/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B_VSA.sh b/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B_VSA.sh
index 40e74e7e6..06841e55d 100644
--- a/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B_VSA.sh
+++ b/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B_VSA.sh
@@ -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
diff --git a/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/README.md b/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/README.md
index bb05d2a48..0c71a0005 100644
--- a/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/README.md
+++ b/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/README.md
@@ -9,4 +9,14 @@ git submodule update --init --recursive
python setup_vsa.py install
```
-### TODO
\ No newline at end of file
+### 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/
+```
diff --git a/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh b/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh
index 7321bbd58..ee942605b 100644
--- a/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh
+++ b/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh
@@ -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
diff --git a/examples/inference/basic/basic.py b/examples/inference/basic/basic.py
index 337b89550..f12a4969e 100644
--- a/examples/inference/basic/basic.py
+++ b/examples/inference/basic/basic.py
@@ -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,
)
diff --git a/examples/inference/basic/basic_dmd.py b/examples/inference/basic/basic_dmd.py
index ed30e6e0a..48b2965ce 100644
--- a/examples/inference/basic/basic_dmd.py
+++ b/examples/inference/basic/basic_dmd.py
@@ -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,
diff --git a/examples/inference/basic/basic_wan2_2.py b/examples/inference/basic/basic_wan2_2.py
new file mode 100644
index 000000000..b1420b7da
--- /dev/null
+++ b/examples/inference/basic/basic_wan2_2.py
@@ -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()
\ No newline at end of file
diff --git a/examples/inference/basic/basic_wan2_2_i2v.py b/examples/inference/basic/basic_wan2_2_i2v.py
new file mode 100644
index 000000000..9c55694e3
--- /dev/null
+++ b/examples/inference/basic/basic_wan2_2_i2v.py
@@ -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()
\ No newline at end of file
diff --git a/examples/inference/gradio/README.md b/examples/inference/gradio/README.md
deleted file mode 100644
index c7ff9f5fd..000000000
--- a/examples/inference/gradio/README.md
+++ /dev/null
@@ -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
diff --git a/examples/inference/gradio/gradio_demo.py b/examples/inference/gradio/gradio_demo.py
deleted file mode 100644
index 06340bae6..000000000
--- a/examples/inference/gradio/gradio_demo.py
+++ /dev/null
@@ -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)
diff --git a/examples/inference/gradio/gradio_frontend.py b/examples/inference/gradio/gradio_frontend.py
new file mode 100644
index 000000000..fab308800
--- /dev/null
+++ b/examples/inference/gradio/gradio_frontend.py
@@ -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"""
+
+ """)
+
+ with gr.Accordion("🎥 What Is FastVideo?", open=False):
+ gr.HTML("""
+
+
+ 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.
+
The compute for this demo is generously provided by GMI Cloud. 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.
+
+ """)
+
+ 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"""
+
+
+
+
+
+
+ FastWan
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ """
+
+ 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()
\ No newline at end of file
diff --git a/examples/inference/gradio/ray_serve_backend.py b/examples/inference/gradio/ray_serve_backend.py
new file mode 100644
index 000000000..67192fd80
--- /dev/null
+++ b/examples/inference/gradio/ray_serve_backend.py
@@ -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)
\ No newline at end of file
diff --git a/examples/inference/gradio/start.sh b/examples/inference/gradio/start.sh
new file mode 100644
index 000000000..6eabdd9ce
--- /dev/null
+++ b/examples/inference/gradio/start.sh
@@ -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"
\ No newline at end of file
diff --git a/examples/inference/gradio/start_ray_serve_app.py b/examples/inference/gradio/start_ray_serve_app.py
new file mode 100644
index 000000000..437121132
--- /dev/null
+++ b/examples/inference/gradio/start_ray_serve_app.py
@@ -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()
\ No newline at end of file
diff --git a/examples/inference/lora/wan_lora_inference.py b/examples/inference/lora/wan_lora_inference.py
index 5471a8980..bc3d264d9 100644
--- a/examples/inference/lora/wan_lora_inference.py
+++ b/examples/inference/lora/wan_lora_inference.py
@@ -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,
}
diff --git a/examples/inference/lora/wan_lora_inference_from_ckpt.py b/examples/inference/lora/wan_lora_inference_from_ckpt.py
index 444fda6e2..519abf0f8 100644
--- a/examples/inference/lora/wan_lora_inference_from_ckpt.py
+++ b/examples/inference/lora/wan_lora_inference_from_ckpt.py
@@ -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()
\ No newline at end of file
diff --git a/examples/inference/optimizations/attention_example.py b/examples/inference/optimizations/attention_example.py
index ce727212f..558a001d2 100644
--- a/examples/inference/optimizations/attention_example.py
+++ b/examples/inference/optimizations/attention_example.py
@@ -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")
diff --git a/examples/inference/optimizations/teacache_example.py b/examples/inference/optimizations/teacache_example.py
index ba2e353ea..77e0582a6 100644
--- a/examples/inference/optimizations/teacache_example.py
+++ b/examples/inference/optimizations/teacache_example.py
@@ -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")
diff --git a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh
index 50a51743b..11c389a76 100644
--- a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh
+++ b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh
@@ -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
diff --git a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm
index a3953f250..cfd5226b5 100644
--- a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm
+++ b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm
@@ -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
diff --git a/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.sh b/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.sh
index dc8fdbdfd..ec52b8493 100644
--- a/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.sh
+++ b/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.sh
@@ -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
diff --git a/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.slurm b/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.slurm
index b3b5d7af5..f24071c5c 100644
--- a/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.slurm
+++ b/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.slurm
@@ -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
diff --git a/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v_lora.sh b/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v_lora.sh
index 4ba9449bf..781870e3f 100644
--- a/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v_lora.sh
+++ b/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v_lora.sh
@@ -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
diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh
index cd1e640ec..c7e2c251e 100644
--- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh
+++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh
@@ -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 \
diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.slurm b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.slurm
index 77918427c..414e28447 100644
--- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.slurm
+++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.slurm
@@ -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
)
diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_lora.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_lora.sh
index 28729d1f6..0f396a687 100644
--- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_lora.sh
+++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_lora.sh
@@ -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
)
diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh
new file mode 100644
index 000000000..ced763f51
--- /dev/null
+++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh
@@ -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
diff --git a/fastvideo/configs/models/dits/wanvideo.py b/fastvideo/configs/models/dits/wanvideo.py
index c952b7641..8cd25cf6c 100644
--- a/fastvideo/configs/models/dits/wanvideo.py
+++ b/fastvideo/configs/models/dits/wanvideo.py
@@ -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
diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py
index bf72e9cdf..cd8fb9938 100644
--- a/fastvideo/configs/pipelines/registry.py
+++ b/fastvideo/configs/pipelines/registry.py
@@ -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
}
diff --git a/fastvideo/configs/sample/base.py b/fastvideo/configs/sample/base.py
index 2da5f05d7..32b096561 100644
--- a/fastvideo/configs/sample/base.py
+++ b/fastvideo/configs/sample/base.py
@@ -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
diff --git a/fastvideo/configs/sample/registry.py b/fastvideo/configs/sample/registry.py
index 44786d42a..740cbbde9 100644
--- a/fastvideo/configs/sample/registry.py
+++ b/fastvideo/configs/sample/registry.py
@@ -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
}
diff --git a/fastvideo/configs/sample/wan.py b/fastvideo/configs/sample/wan.py
index 84a0ac362..23cb0cdbf 100644
--- a/fastvideo/configs/sample/wan.py
+++ b/fastvideo/configs/sample/wan.py
@@ -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
diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py
index 2e6f63696..a24ea4554 100644
--- a/fastvideo/entrypoints/video_generator.py
+++ b/fastvideo/entrypoints/video_generator.py
@@ -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,
diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py
index 12006cd96..cce97ce2b 100644
--- a/fastvideo/fastvideo_args.py
+++ b/fastvideo/fastvideo_args.py
@@ -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
diff --git a/fastvideo/layers/lora/linear.py b/fastvideo/layers/lora/linear.py
index 2f63de200..a0a67f807 100644
--- a/fastvideo/layers/lora/linear.py
+++ b/fastvideo/layers/lora/linear.py
@@ -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
diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py
index 2ca840622..1cae2da01 100644
--- a/fastvideo/models/dits/wanvideo.py
+++ b/fastvideo/models/dits/wanvideo.py
@@ -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)
diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py
index 5d95a1859..1eebd9f6f 100644
--- a/fastvideo/models/loader/component_loader.py
+++ b/fastvideo/models/loader/component_loader.py
@@ -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,
diff --git a/fastvideo/models/schedulers/base.py b/fastvideo/models/schedulers/base.py
index 8a516ec0b..4c647efc4 100644
--- a/fastvideo/models/schedulers/base.py
+++ b/fastvideo/models/schedulers/base.py
@@ -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):
diff --git a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py
index a1c4d3e27..73ce8c954 100644
--- a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py
+++ b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py
@@ -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:
"""
diff --git a/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py b/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py
index 945ffd132..ce111ddc5 100644
--- a/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py
+++ b/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py
@@ -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):
"""
diff --git a/fastvideo/models/schedulers/scheduling_unipc_multistep.py b/fastvideo/models/schedulers/scheduling_unipc_multistep.py
index 16efe2c1a..b44160eee 100644
--- a/fastvideo/models/schedulers/scheduling_unipc_multistep.py
+++ b/fastvideo/models/schedulers/scheduling_unipc_multistep.py
@@ -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
diff --git a/fastvideo/pipelines/basic/wan/wan_i2v_pipeline.py b/fastvideo/pipelines/basic/wan/wan_i2v_pipeline.py
index 29c6824aa..2050ad71c 100644
--- a/fastvideo/pipelines/basic/wan/wan_i2v_pipeline.py
+++ b/fastvideo/pipelines/basic/wan/wan_i2v_pipeline.py
@@ -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",
diff --git a/fastvideo/pipelines/basic/wan/wan_pipeline.py b/fastvideo/pipelines/basic/wan/wan_pipeline.py
index 8574c2d29..caabd7073 100644
--- a/fastvideo/pipelines/basic/wan/wan_pipeline.py
+++ b/fastvideo/pipelines/basic/wan/wan_pipeline.py
@@ -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))
diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py
index b083482f4..5db168c77 100644
--- a/fastvideo/pipelines/composed_pipeline_base.py
+++ b/fastvideo/pipelines/composed_pipeline_base.py
@@ -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,
diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py
index 75d28d0dc..eb89f62f5 100644
--- a/fastvideo/pipelines/pipeline_batch_info.py
+++ b/fastvideo/pipelines/pipeline_batch_info.py
@@ -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)
\ No newline at end of file
+ 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)
diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py
index dd8220506..fcaa7edcb 100644
--- a/fastvideo/pipelines/pipeline_registry.py
+++ b/fastvideo/pipelines/pipeline_registry.py
@@ -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:
diff --git a/fastvideo/pipelines/preprocess/__init__.py b/fastvideo/pipelines/preprocess/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/fastvideo/pipelines/preprocess/preprocess_stages.py b/fastvideo/pipelines/preprocess/preprocess_stages.py
new file mode 100644
index 000000000..fafc811de
--- /dev/null
+++ b/fastvideo/pipelines/preprocess/preprocess_stages.py
@@ -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)
diff --git a/fastvideo/pipelines/preprocess/v1_preprocessing_new.py b/fastvideo/pipelines/preprocess/v1_preprocessing_new.py
new file mode 100644
index 000000000..c1614ad4f
--- /dev/null
+++ b/fastvideo/pipelines/preprocess/v1_preprocessing_new.py
@@ -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)
diff --git a/fastvideo/pipelines/preprocess/wan/__init__.py b/fastvideo/pipelines/preprocess/wan/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/fastvideo/pipelines/preprocess/wan/wan_preprocess_pipelines.py b/fastvideo/pipelines/preprocess/wan/wan_preprocess_pipelines.py
new file mode 100644
index 000000000..c0349c45f
--- /dev/null
+++ b/fastvideo/pipelines/preprocess/wan/wan_preprocess_pipelines.py
@@ -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]
diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py
index 0a7845c01..1852edc49 100644
--- a/fastvideo/pipelines/stages/__init__.py
+++ b/fastvideo/pipelines/stages/__init__.py
@@ -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",
]
diff --git a/fastvideo/pipelines/stages/base.py b/fastvideo/pipelines/stages/base.py
index 10b4f690d..97d818044 100644
--- a/fastvideo/pipelines/stages/base.py
+++ b/fastvideo/pipelines/stages/base.py
@@ -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",
diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/pipelines/stages/decoding.py
index 8e692c06a..8f226b5c4 100644
--- a/fastvideo/pipelines/stages/decoding.py
+++ b/fastvideo/pipelines/stages/decoding.py
@@ -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
diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py
index 176fbdb9e..8a2d9f0b6 100644
--- a/fastvideo/pipelines/stages/denoising.py
+++ b/fastvideo/pipelines/stages/denoising.py
@@ -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())
diff --git a/fastvideo/pipelines/stages/encoding.py b/fastvideo/pipelines/stages/encoding.py
index c8cdfab36..1d3beb1f6 100644
--- a/fastvideo/pipelines/stages/encoding.py
+++ b/fastvideo/pipelines/stages/encoding.py
@@ -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
diff --git a/fastvideo/pipelines/stages/image_encoding.py b/fastvideo/pipelines/stages/image_encoding.py
index 889c01198..d4074a5d1 100644
--- a/fastvideo/pipelines/stages/image_encoding.py
+++ b/fastvideo/pipelines/stages/image_encoding.py
@@ -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
diff --git a/fastvideo/pipelines/stages/input_validation.py b/fastvideo/pipelines/stages/input_validation.py
index 466ece75c..2f8cb4811 100644
--- a/fastvideo/pipelines/stages/input_validation.py
+++ b/fastvideo/pipelines/stages/input_validation.py
@@ -56,6 +56,7 @@ class InputValidationStage(PipelineStage):
Returns:
The validated batch information.
"""
+
self._generate_seeds(batch, fastvideo_args)
# Ensure prompt is properly formatted
diff --git a/fastvideo/pipelines/stages/latent_preparation.py b/fastvideo/pipelines/stages/latent_preparation.py
index e4f6c0448..ea23a7daa 100644
--- a/fastvideo/pipelines/stages/latent_preparation.py
+++ b/fastvideo/pipelines/stages/latent_preparation.py
@@ -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:
diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py
index 3b4e8585b..6dea8bf1d 100644
--- a/fastvideo/tests/modal/pr_test.py
+++ b/fastvideo/tests/modal/pr_test.py
@@ -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():
diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py
index 6223ade0d..970e3225d 100644
--- a/fastvideo/training/distillation_pipeline.py
+++ b/fastvideo/training/distillation_pipeline.py
@@ -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:
diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py
index 32e45c315..0615ed24f 100644
--- a/fastvideo/training/training_pipeline.py
+++ b/fastvideo/training/training_pipeline.py
@@ -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()
diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py
index d0bf0961a..dd794a197 100644
--- a/fastvideo/training/training_utils.py
+++ b/fastvideo/training/training_utils.py
@@ -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")
diff --git a/fastvideo/training/wan_distillation_pipeline.py b/fastvideo/training/wan_distillation_pipeline.py
index ebbdd9321..246d3f0c5 100644
--- a/fastvideo/training/wan_distillation_pipeline.py
+++ b/fastvideo/training/wan_distillation_pipeline.py
@@ -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
diff --git a/fastvideo/training/wan_i2v_distillation_pipeline.py b/fastvideo/training/wan_i2v_distillation_pipeline.py
index 75fbe86c1..47b5d73de 100644
--- a/fastvideo/training/wan_i2v_distillation_pipeline.py
+++ b/fastvideo/training/wan_i2v_distillation_pipeline.py
@@ -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...")
diff --git a/fastvideo/training/wan_training_pipeline.py b/fastvideo/training/wan_training_pipeline.py
index 5af79519a..1b3785197 100644
--- a/fastvideo/training/wan_training_pipeline.py
+++ b/fastvideo/training/wan_training_pipeline.py
@@ -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
diff --git a/fastvideo/version.py b/fastvideo/version.py
index bbab0242f..1276d0254 100644
--- a/fastvideo/version.py
+++ b/fastvideo/version.py
@@ -1 +1 @@
-__version__ = "0.1.4"
+__version__ = "0.1.5"
diff --git a/fastvideo/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py
index 3d639c5ad..6f89d1772 100644
--- a/fastvideo/worker/gpu_worker.py
+++ b/fastvideo/worker/gpu_worker.py
@@ -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']
diff --git a/fastvideo/worker/multiproc_executor.py b/fastvideo/worker/multiproc_executor.py
index 1d770d322..b3b1398df 100644
--- a/fastvideo/worker/multiproc_executor.py
+++ b/fastvideo/worker/multiproc_executor.py
@@ -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,
diff --git a/fastvideo/workflow/preprocess/components.py b/fastvideo/workflow/preprocess/components.py
new file mode 100644
index 000000000..fcd34201b
--- /dev/null
+++ b/fastvideo/workflow/preprocess/components.py
@@ -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
diff --git a/fastvideo/workflow/preprocess/preprocess_workflow.py b/fastvideo/workflow/preprocess/preprocess_workflow.py
new file mode 100644
index 000000000..e93321cb6
--- /dev/null
+++ b/fastvideo/workflow/preprocess/preprocess_workflow.py
@@ -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."
+ )
diff --git a/fastvideo/workflow/preprocess/preprocess_workflow_i2v.py b/fastvideo/workflow/preprocess/preprocess_workflow_i2v.py
new file mode 100644
index 000000000..62d19374b
--- /dev/null
+++ b/fastvideo/workflow/preprocess/preprocess_workflow_i2v.py
@@ -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
diff --git a/fastvideo/workflow/preprocess/preprocess_workflow_t2v.py b/fastvideo/workflow/preprocess/preprocess_workflow_t2v.py
new file mode 100644
index 000000000..54dee8fe2
--- /dev/null
+++ b/fastvideo/workflow/preprocess/preprocess_workflow_t2v.py
@@ -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()
diff --git a/fastvideo/workflow/workflow_base.py b/fastvideo/workflow/workflow_base.py
index e00c4dbcd..ef74f846e 100644
--- a/fastvideo/workflow/workflow_base.py
+++ b/fastvideo/workflow/workflow_base.py
@@ -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
diff --git a/pyproject.toml b/pyproject.toml
index 09ec2b4df..8c3051941 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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]
diff --git a/scripts/dataset_preparation/resize_videos.py b/scripts/dataset_preparation/resize_videos.py
index 5e37d1aeb..dd6a6b4ef 100644
--- a/scripts/dataset_preparation/resize_videos.py
+++ b/scripts/dataset_preparation/resize_videos.py
@@ -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
diff --git a/scripts/distill/v1_distill_dmd_wan.sh b/scripts/distill/v1_distill_dmd_wan.sh
index b645f43f6..5520dabbd 100644
--- a/scripts/distill/v1_distill_dmd_wan.sh
+++ b/scripts/distill/v1_distill_dmd_wan.sh
@@ -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" \
diff --git a/scripts/distill/v1_distill_dmd_wan_VSA.sh b/scripts/distill/v1_distill_dmd_wan_VSA.sh
index f06f158d3..442177556 100644
--- a/scripts/distill/v1_distill_dmd_wan_VSA.sh
+++ b/scripts/distill/v1_distill_dmd_wan_VSA.sh
@@ -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" \
diff --git a/scripts/finetune/finetune_v1.sh b/scripts/finetune/finetune_v1.sh
index 0778d6dba..d25ed27ab 100644
--- a/scripts/finetune/finetune_v1.sh
+++ b/scripts/finetune/finetune_v1.sh
@@ -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 \