From 8d483fe4aaae2c0083ca0a69a8bbbf642960454e Mon Sep 17 00:00:00 2001 From: Yongqi Chen <144848849+BrianChen1129@users.noreply.github.com> Date: Mon, 4 Aug 2025 18:54:06 -0400 Subject: [PATCH 01/16] [Bugfix] Fix neg_prompt bug when training from local cp (#708) --- README.md | 2 +- fastvideo/configs/sample/base.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index ea7fa445a..7476df996 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,7 @@ FastVideo features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.

- | Documentation | Quick Start | πŸ€— FastWan2.1 | πŸ€— FastWan2.2 | πŸŸ£πŸ’¬ Slack | πŸŸ£πŸ’¬ WeChat | + | Documentation | Quick Start | πŸ€— FastWan2.1 | πŸ€— FastWan2.2 | πŸŸ£πŸ’¬ Slack | πŸŸ£πŸ’¬ WeChat |

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 From 7d8cf38dbe849b099c75a4e46daac49065dd030f Mon Sep 17 00:00:00 2001 From: Yongqi Chen <144848849+BrianChen1129@users.noreply.github.com> Date: Mon, 4 Aug 2025 23:21:21 -0400 Subject: [PATCH 02/16] Fix typo (#709) --- docs/source/index.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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.

From 1d018acb3e8a5d7ccf5ce3bb49b1fb22f2802143 Mon Sep 17 00:00:00 2001 From: Yongqi Chen <144848849+BrianChen1129@users.noreply.github.com> Date: Tue, 5 Aug 2025 14:27:39 -0400 Subject: [PATCH 03/16] [Feature]Add Data-free distillation readme (#710) --- docs/source/distillation/dmd.md | 2 +- docs/source/sliding_tile_attention/installation.md | 6 ++++++ docs/source/video_sparse_attention/installation.md | 6 ++++++ .../Wan2.2-TI2V-5B-Diffusers/Data-free/README.md | 13 +++++++++++++ .../Wan2.2-TI2V-5B-Diffusers/crush_smol/README.md | 12 +++++++++++- 5 files changed, 37 insertions(+), 2 deletions(-) create mode 100644 examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/README.md 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/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.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/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/ +``` From 0ceff110da072ab6be645499e9761ad0e6f9ff28 Mon Sep 17 00:00:00 2001 From: William Lin Date: Wed, 6 Aug 2025 13:07:52 -0700 Subject: [PATCH 04/16] [chore] Release 0.1.5 (#717) --- fastvideo/version.py | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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/pyproject.toml b/pyproject.toml index 09ec2b4df..a7bc92749 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" From 98ce9034f0a3ce3b0228cc07eb893c3c56700814 Mon Sep 17 00:00:00 2001 From: Zhang Peiyuan Date: Wed, 6 Aug 2025 19:29:40 -0700 Subject: [PATCH 05/16] [Chore] Include our demo in the readme. (#720) --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 7476df996..67c9fe051 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,7 @@ FastVideo features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.

- | Documentation | Quick Start | πŸ€— FastWan2.1 | πŸ€— FastWan2.2 | πŸŸ£πŸ’¬ Slack | πŸŸ£πŸ’¬ WeChat | + | πŸ•ΉοΈ Online Demo | Documentation | Quick Start | πŸ€— FastWan | πŸŸ£πŸ’¬ Slack | πŸŸ£πŸ’¬ WeChat |

From c41305ad1871a03866197bf0386ec26a09c1694a Mon Sep 17 00:00:00 2001 From: Wei Zhou <69577934+JerryZhou54@users.noreply.github.com> Date: Wed, 6 Aug 2025 20:31:03 -0700 Subject: [PATCH 06/16] [Feat] Add Wan2.2 14B MoE (#688) Co-authored-by: SolitaryThinker --- README.md | 5 +- examples/inference/basic/basic_wan2_2.py | 48 +++++++++++++++++++ examples/inference/basic/basic_wan2_2_i2v.py | 30 ++++++++++++ fastvideo/configs/models/dits/wanvideo.py | 2 +- fastvideo/configs/pipelines/registry.py | 4 +- fastvideo/configs/sample/registry.py | 25 ++++++---- fastvideo/configs/sample/wan.py | 10 +++- fastvideo/fastvideo_args.py | 3 ++ fastvideo/models/loader/component_loader.py | 3 +- fastvideo/models/schedulers/base.py | 3 +- .../scheduling_flow_match_euler_discrete.py | 3 ++ .../scheduling_flow_unipc_multistep.py | 3 ++ .../schedulers/scheduling_unipc_multistep.py | 1 + .../pipelines/basic/wan/wan_i2v_pipeline.py | 14 ++++-- fastvideo/pipelines/basic/wan/wan_pipeline.py | 1 + fastvideo/pipelines/composed_pipeline_base.py | 45 +++++++++++++---- fastvideo/pipelines/pipeline_batch_info.py | 3 ++ fastvideo/pipelines/stages/denoising.py | 36 ++++++++++++-- fastvideo/training/distillation_pipeline.py | 4 ++ 19 files changed, 206 insertions(+), 37 deletions(-) create mode 100644 examples/inference/basic/basic_wan2_2.py create mode 100644 examples/inference/basic/basic_wan2_2_i2v.py diff --git a/README.md b/README.md index 67c9fe051..b8b1adaf1 100644 --- a/README.md +++ b/README.md @@ -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/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/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/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/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/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 5d95a1859..fe2c8ca54 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"), 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..27458fc81 100644 --- a/fastvideo/pipelines/basic/wan/wan_i2v_pipeline.py +++ b/fastvideo/pipelines/basic/wan/wan_i2v_pipeline.py @@ -48,11 +48,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()) @@ -72,6 +75,7 @@ class WanImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase): 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..a695d6b03 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] = {} @@ -239,6 +240,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 +261,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 +284,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 +293,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..4c5b2335d 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -93,6 +93,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 @@ -136,6 +137,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) diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 176fbdb9e..c8fb76ce4 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -57,9 +57,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 +189,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 +202,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 +285,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 +304,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 +313,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 diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 6223ade0d..086c75288 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]] From 47bd1e01782078d23852288b76d8d1e86f0a356b Mon Sep 17 00:00:00 2001 From: Zhang Peiyuan Date: Wed, 6 Aug 2025 21:54:09 -0700 Subject: [PATCH 07/16] [Misc] change installation logic of vsa (#721) --- csrc/attn/setup_vsa.py | 23 ++-- csrc/attn/vsa/__init__.py | 7 +- csrc/attn/vsa/block_sparse_wrapper.py | 183 +++++++++++++------------- 3 files changed, 107 insertions(+), 106 deletions(-) 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 From 86604a684b510dc8a4b832271212a0fbc30416d9 Mon Sep 17 00:00:00 2001 From: Jinzhe Pan <48981407+Eigensystem@users.noreply.github.com> Date: Thu, 7 Aug 2025 16:49:07 +0800 Subject: [PATCH 08/16] [3/3][Preprocess] add preprocessing workflows (#645) --- .../crush_smol/preprocess_wan_data_t2v_new.sh | 24 + .../pipelines/basic/wan/wan_i2v_pipeline.py | 12 +- fastvideo/pipelines/pipeline_batch_info.py | 21 +- fastvideo/pipelines/pipeline_registry.py | 6 +- fastvideo/pipelines/preprocess/__init__.py | 0 .../pipelines/preprocess/preprocess_stages.py | 73 +++ .../preprocess/v1_preprocessing_new.py | 27 + .../pipelines/preprocess/wan/__init__.py | 0 .../wan/wan_preprocess_pipelines.py | 53 ++ fastvideo/pipelines/stages/__init__.py | 4 +- fastvideo/pipelines/stages/encoding.py | 174 ++----- fastvideo/pipelines/stages/image_encoding.py | 184 +++++++ .../pipelines/stages/input_validation.py | 1 + .../pipelines/stages/latent_preparation.py | 1 + fastvideo/tests/modal/pr_test.py | 4 +- fastvideo/workflow/preprocess/components.py | 467 ++++++++++++++++++ .../preprocess/preprocess_workflow.py | 98 ++++ .../preprocess/preprocess_workflow_i2v.py | 14 + .../preprocess/preprocess_workflow_t2v.py | 69 +++ fastvideo/workflow/workflow_base.py | 175 ++++++- pyproject.toml | 11 +- scripts/dataset_preparation/resize_videos.py | 2 +- 22 files changed, 1260 insertions(+), 160 deletions(-) create mode 100644 examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh create mode 100644 fastvideo/pipelines/preprocess/__init__.py create mode 100644 fastvideo/pipelines/preprocess/preprocess_stages.py create mode 100644 fastvideo/pipelines/preprocess/v1_preprocessing_new.py create mode 100644 fastvideo/pipelines/preprocess/wan/__init__.py create mode 100644 fastvideo/pipelines/preprocess/wan/wan_preprocess_pipelines.py create mode 100644 fastvideo/workflow/preprocess/components.py create mode 100644 fastvideo/workflow/preprocess/preprocess_workflow.py create mode 100644 fastvideo/workflow/preprocess/preprocess_workflow_i2v.py create mode 100644 fastvideo/workflow/preprocess/preprocess_workflow_t2v.py 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/pipelines/basic/wan/wan_i2v_pipeline.py b/fastvideo/pipelines/basic/wan/wan_i2v_pipeline.py index 27458fc81..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) @@ -70,7 +68,7 @@ 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( diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 4c5b2335d..b5acc76f6 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -13,6 +13,7 @@ from typing import Any import PIL.Image import torch +from torchcodec.decoders import VideoDecoder from fastvideo.attention import AttentionMetadata from fastvideo.configs.sample.teacache import TeaCacheParams, WanTeaCacheParams @@ -75,15 +76,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 @@ -192,4 +193,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/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..5c7e54f1a 100644 --- a/fastvideo/pipelines/stages/latent_preparation.py +++ b/fastvideo/pipelines/stages/latent_preparation.py @@ -122,6 +122,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..05e81e971 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({ @@ -52,7 +52,7 @@ def run_test(pytest_command: str): git clone {git_repo} /FastVideo && cd /FastVideo && {checkout_command} && - uv pip install -e .[test] && + uv pip install -e .[test] --index-strategy unsafe-best-match && {pytest_command} """ 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 a7bc92749..9e7f58a81 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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,8 +45,11 @@ dependencies = [ # Training Dependencies "torchdata", "pyarrow", - "datasets", + "datasets==4.0.0", "av", + + # Preprocessing Dependencies + "torchcodec" ] [tool.uv] 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 From f2bc037fff57ab8f5ffa55e0cfbd72654ba4581a Mon Sep 17 00:00:00 2001 From: Jinzhe Pan <48981407+Eigensystem@users.noreply.github.com> Date: Thu, 7 Aug 2025 17:31:20 +0800 Subject: [PATCH 09/16] [Fix] training pipeline pin_cpu_memory issue (#692) --- fastvideo/models/loader/component_loader.py | 1 + fastvideo/pipelines/composed_pipeline_base.py | 2 -- fastvideo/training/wan_distillation_pipeline.py | 6 +++--- fastvideo/training/wan_training_pipeline.py | 6 +++--- 4 files changed, 7 insertions(+), 8 deletions(-) diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index fe2c8ca54..1eebd9f6f 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -451,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/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index a695d6b03..5db168c77 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -149,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) 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_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 From f79b08de81e4dee240b1b2502e9cb0a58a166792 Mon Sep 17 00:00:00 2001 From: Jiali Chen <90408393+Gary-ChenJL@users.noreply.github.com> Date: Thu, 7 Aug 2025 18:53:05 -0700 Subject: [PATCH 10/16] add cicd workflow for publishing VSA kernel (#723) --- .github/workflows/vsa-publish.yml | 249 ++++++++++++++++++++++++++++++ 1 file changed, 249 insertions(+) create mode 100644 .github/workflows/vsa-publish.yml diff --git a/.github/workflows/vsa-publish.yml b/.github/workflows/vsa-publish.yml new file mode 100644 index 000000000..901785f7c --- /dev/null +++ b/.github/workflows/vsa-publish.yml @@ -0,0 +1,249 @@ +name: Publish Video Sparse Attention Kernel to PyPI on Version Change + +on: + push: + branches: + - main + paths: + - "csrc/attn/setup_vsa.py" + workflow_dispatch: + +jobs: + check-version-change: + runs-on: ubuntu-latest + outputs: + version-changed: ${{ steps.check-version.outputs.changed }} + new-version: ${{ steps.check-version.outputs.new-version }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + fetch-depth: 2 + + - name: Check if version changed + id: check-version + run: | + cd csrc/attn + # Get current commit's version + NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup_vsa.py) + echo "New version: $NEW_VERSION" + + # Get previous version from git history + OLD_VERSION=$(git show HEAD~1:./setup_vsa.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0") + echo "Old version: $OLD_VERSION" + + if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then + echo "Version changed from $OLD_VERSION to $NEW_VERSION" + echo "changed=true" >> $GITHUB_OUTPUT + echo "new-version=$NEW_VERSION" >> $GITHUB_OUTPUT + else + echo "Version did not change" + echo "changed=false" >> $GITHUB_OUTPUT + fi + + build_wheels: + name: Build Wheel + needs: check-version-change + if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }} + runs-on: ${{ matrix.os }} + + strategy: + fail-fast: false + matrix: + # Using ubuntu-20.04 instead of 22.04 for more compatibility (glibc). Ideally we'd use the + # manylinux docker image, but I haven't figured out how to install CUDA on manylinux. + os: [ubuntu-22.04] + python-version: ['3.10', '3.11', '3.12', '3.13'] + torch-version: ['2.5.1', '2.6.0', '2.7.1'] + cuda-version: ['12.4.1', '12.5.1', '12.6.3', '12.8.0'] + + steps: + - name: Free up disk space + run: | + echo "Initial disk space:" + df -h + + # Remove large directories + sudo rm -rf /usr/share/dotnet + sudo rm -rf /usr/local/lib/android + sudo rm -rf /opt/ghc + sudo rm -rf /usr/local/share/boost + sudo rm -rf /usr/share/swift + sudo rm -rf /usr/local/lib/node_modules + sudo rm -rf /usr/local/share/powershell + sudo rm -rf /usr/share/rust + sudo rm -rf /usr/local/.ghcup + + # Remove cached files + sudo rm -rf /var/lib/apt/lists/* + sudo rm -rf /var/cache/apt/archives/* + + echo "Disk space after cleanup:" + df -h + + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install CUDA ${{ matrix.cuda-version }} + uses: Jimver/cuda-toolkit@v0.2.21 + id: cuda-toolkit + with: + cuda: ${{ matrix.cuda-version }} + linux-local-args: '["--toolkit"]' + method: 'network' + + - name: Install dependencies (GCC, Clang, CUDA Paths, Git) + run: | + sudo apt update + sudo apt install -y git patchelf gcc-11 g++-11 clang-11 + sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11 + + # Allow Git to Access Safe Directory + git config --global --add safe.directory /__w/FastVideo/FastVideo + + # Set CUDA environment variables + export CUDA_HOME=/usr/local/cuda-${{ matrix.cuda-version }} + export PATH=${CUDA_HOME}/bin:${PATH} + export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH + + # Verify installation + gcc --version + g++ --version + clang-11 --version + nvcc --version + + - name: Install PyTorch ${{ matrix.torch-version }}+cu${{ matrix.cuda-version }} + run: | + pip install --upgrade pip + # With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error + # AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable + pip install typing-extensions==4.12.2 + # We want to figure out the CUDA version to download pytorch + # e.g. we can have system CUDA version being 11.7 but if torch==1.12 then we need to download the wheel from cu116 + # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix + export TORCH_CUDA_VERSION=124 + pip install --no-cache-dir torch==${{ matrix.torch-version }} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION} + nvcc --version + python --version + python -c "import torch; print('PyTorch:', torch.__version__)" + python -c "import torch; print('CUDA:', torch.version.cuda)" + python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)" + + - name: Build wheel + run: | + export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH + + # We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA version is 11.7 and pytorch cuda version is 11.6 + # https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/utils/cpp_extension.py#L810 + # However this still fails so I'm using a newer version of setuptools + pip install setuptools + pip install ninja packaging wheel + + cd csrc/attn # Move into the correct folder + git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized + python setup_vsa.py bdist_wheel --dist-dir=dist + + - name: Rename wheel file + run: | + cd csrc/attn + + CUDA_SHORT_VERSION=$(echo ${{ matrix.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g') + TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-version }} | cut -d. -f1,2) + # Get the correct version format + tmpname=cu${CUDA_SHORT_VERSION}torch${TORCH_SHORT_VERSION} + wheel_name=$(ls dist/*whl | xargs -n 1 basename | sed "s/-/+$tmpname-/2") + # Rename with version information + ls dist/*whl |xargs -I {} mv {} dist/${wheel_name} + echo "wheel_name=${wheel_name}" >> $GITHUB_ENV + + - name: Upload wheel artifact + uses: actions/upload-artifact@v4 + with: + name: ${{ env.wheel_name }} + path: csrc/attn/dist/*.whl + retention-days: 90 + + publish_package: + name: Publish package + needs: [build_wheels, check-version-change] + if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }} + runs-on: ubuntu-22.04 + permissions: + id-token: write # Needed for OIDC Trusted Publishing + + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Install CUDA 12.4.1 + uses: Jimver/cuda-toolkit@v0.2.21 + id: cuda-toolkit + with: + cuda: 12.4.1 + linux-local-args: '["--toolkit"]' + method: 'network' + sub-packages: '["nvcc"]' + + - name: Install dependencies (GCC, Clang, CUDA Paths, Git) + run: | + sudo apt update + sudo apt install -y git patchelf gcc-11 g++-11 clang-11 + sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11 + + # Allow Git to Access Safe Directory + git config --global --add safe.directory /__w/FastVideo/FastVideo + + # Set CUDA environment variables + export CUDA_HOME=/usr/local/cuda-12.4.1 + export PATH=${CUDA_HOME}/bin:${PATH} + export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH + + # Verify installation + gcc --version + g++ --version + clang-11 --version + nvcc --version + + - name: Install PyTorch 2.5.1+cu12.4.1 + run: | + pip install --upgrade pip + # With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error + # AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable + pip install typing-extensions==4.12.2 + # We want to figure out the CUDA version to download pytorch + # e.g. we can have system CUDA version being 11.7 but if torch==1.12 then we need to download the wheel from cu116 + # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix + export TORCH_CUDA_VERSION=124 + pip install --no-cache-dir torch==2.5.1 --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION} + nvcc --version + python --version + python -c "import torch; print('PyTorch:', torch.__version__)" + python -c "import torch; print('CUDA:', torch.version.cuda)" + python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)" + + - name: Build source distribution + run: | + export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH + + # We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA version is 11.7 and pytorch cuda version is 11.6 + # https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/utils/cpp_extension.py#L810 + # However this still fails so I'm using a newer version of setuptools + pip install setuptools + pip install ninja packaging wheel + + cd csrc/attn # Move into the correct folder + git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized + python setup_vsa.py sdist --dist-dir=dist + + - name: Publish release distributions to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + packages-dir: csrc/attn/dist/ From 6c6bcd914df863abf2cfe31498fb269c96a0a2eb Mon Sep 17 00:00:00 2001 From: Wenxuan Tan Date: Fri, 8 Aug 2025 00:50:38 -0500 Subject: [PATCH 11/16] Remove all empty_cache (#713) --- csrc/attn/benchmarks/bench_sta.py | 2 -- .../Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.slurm | 2 +- .../Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slurm | 2 +- .../Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.slurm | 2 +- .../Data-free/distill_dmd_t2v_5B.sh | 2 +- .../Data-free/distill_dmd_t2v_5B_VSA.sh | 2 +- .../crush_smol/distill_dmd_VSA_t2v_5B.sh | 2 +- .../inference/lora/wan_lora_inference_from_ckpt.py | 12 +++++++++--- .../Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh | 2 +- .../crush_smol/finetune_i2v.slurm | 2 +- .../wan_i2v_14B_480p/crush_smol/finetune_i2v.sh | 2 +- .../wan_i2v_14B_480p/crush_smol/finetune_i2v.slurm | 2 +- .../wan_i2v_14B_480p/crush_smol/finetune_i2v_lora.sh | 2 +- .../finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh | 4 ++-- .../wan_t2v_1.3B/crush_smol/finetune_t2v.slurm | 4 ++-- .../wan_t2v_1.3B/crush_smol/finetune_t2v_lora.sh | 6 +++--- fastvideo/pipelines/stages/decoding.py | 3 --- fastvideo/pipelines/stages/denoising.py | 3 --- fastvideo/training/distillation_pipeline.py | 4 ---- fastvideo/training/training_pipeline.py | 1 - fastvideo/worker/gpu_worker.py | 1 - scripts/distill/v1_distill_dmd_wan.sh | 4 ++-- scripts/distill/v1_distill_dmd_wan_VSA.sh | 4 ++-- scripts/finetune/finetune_v1.sh | 4 ++-- 24 files changed, 33 insertions(+), 41 deletions(-) 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/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/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/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/lora/wan_lora_inference_from_ckpt.py b/examples/inference/lora/wan_lora_inference_from_ckpt.py index 444fda6e2..c2ad8598d 100644 --- a/examples/inference/lora/wan_lora_inference_from_ckpt.py +++ b/examples/inference/lora/wan_lora_inference_from_ckpt.py @@ -14,14 +14,14 @@ def main(): vae_cpu_offload=False, text_encoder_cpu_offload=True, pin_cpu_memory=False, - lora_path="checkpoints/wan_t2v_finetune_lora/checkpoint-1250/transformer", + 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, } @@ -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/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..3ac715007 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 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/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 c8fb76ce4..592ee352d 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 @@ -357,8 +356,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/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 086c75288..970e3225d 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -794,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'] @@ -834,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: @@ -866,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..8293f10ca 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -703,4 +703,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/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py index 3d639c5ad..f3eebcbb7 100644 --- a/fastvideo/worker/gpu_worker.py +++ b/fastvideo/worker/gpu_worker.py @@ -71,7 +71,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 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 \ From 3d75ba8251f2b5bc28cb985be4ca34ecf90492cb Mon Sep 17 00:00:00 2001 From: Jiali Chen <90408393+Gary-ChenJL@users.noreply.github.com> Date: Fri, 8 Aug 2025 13:05:16 -0700 Subject: [PATCH 12/16] update version selection for VSA workflow (#725) --- .github/workflows/vsa-publish.yml | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/.github/workflows/vsa-publish.yml b/.github/workflows/vsa-publish.yml index 901785f7c..e3f08fdc4 100644 --- a/.github/workflows/vsa-publish.yml +++ b/.github/workflows/vsa-publish.yml @@ -54,8 +54,17 @@ jobs: # manylinux docker image, but I haven't figured out how to install CUDA on manylinux. os: [ubuntu-22.04] python-version: ['3.10', '3.11', '3.12', '3.13'] - torch-version: ['2.5.1', '2.6.0', '2.7.1'] - cuda-version: ['12.4.1', '12.5.1', '12.6.3', '12.8.0'] + # For version reference https://pytorch.org/get-started/previous-versions/ + torch-cuda: + - torch-version: '2.5.1' + cuda-version: '12.4.1' + torch-cuda-short: 'cu124' + - torch-version: '2.6.0' + cuda-version: '12.6.3' + torch-cuda-short: 'cu126' + - torch-version: '2.7.1' + cuda-version: '12.8.0' + torch-cuda-short: 'cu128' steps: - name: Free up disk space @@ -89,11 +98,11 @@ jobs: with: python-version: ${{ matrix.python-version }} - - name: Install CUDA ${{ matrix.cuda-version }} + - name: Install CUDA ${{ matrix.torch-cuda.cuda-version }} uses: Jimver/cuda-toolkit@v0.2.21 id: cuda-toolkit with: - cuda: ${{ matrix.cuda-version }} + cuda: ${{ matrix.torch-cuda.cuda-version }} linux-local-args: '["--toolkit"]' method: 'network' @@ -107,7 +116,7 @@ jobs: git config --global --add safe.directory /__w/FastVideo/FastVideo # Set CUDA environment variables - export CUDA_HOME=/usr/local/cuda-${{ matrix.cuda-version }} + export CUDA_HOME=/usr/local/cuda-${{ matrix.torch-cuda.cuda-version }} export PATH=${CUDA_HOME}/bin:${PATH} export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH @@ -117,7 +126,7 @@ jobs: clang-11 --version nvcc --version - - name: Install PyTorch ${{ matrix.torch-version }}+cu${{ matrix.cuda-version }} + - name: Install PyTorch ${{ matrix.torch-cuda.torch-version }}+cu${{ matrix.torch-cuda.cuda-version }} run: | pip install --upgrade pip # With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error @@ -126,8 +135,7 @@ jobs: # We want to figure out the CUDA version to download pytorch # e.g. we can have system CUDA version being 11.7 but if torch==1.12 then we need to download the wheel from cu116 # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix - export TORCH_CUDA_VERSION=124 - pip install --no-cache-dir torch==${{ matrix.torch-version }} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION} + pip install --no-cache-dir torch==${{ matrix.torch-cuda.torch-version }} --index-url https://download.pytorch.org/whl/${{matrix.torch-cuda.torch-cuda-short}} nvcc --version python --version python -c "import torch; print('PyTorch:', torch.__version__)" @@ -152,8 +160,8 @@ jobs: run: | cd csrc/attn - CUDA_SHORT_VERSION=$(echo ${{ matrix.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g') - TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-version }} | cut -d. -f1,2) + CUDA_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g') + TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.torch-version }} | cut -d. -f1,2) # Get the correct version format tmpname=cu${CUDA_SHORT_VERSION}torch${TORCH_SHORT_VERSION} wheel_name=$(ls dist/*whl | xargs -n 1 basename | sed "s/-/+$tmpname-/2") From 60f61c810190f4e5bb8c29716fc2d25caa62ea39 Mon Sep 17 00:00:00 2001 From: William Lin Date: Fri, 8 Aug 2025 18:45:03 -0700 Subject: [PATCH 13/16] [bugfix] fix pyproject install and VSA precision test (#726) --- .buildkite/pipeline.yml | 1 + csrc/attn/README.md | 2 +- fastvideo/pipelines/pipeline_batch_info.py | 8 +++++--- fastvideo/tests/modal/pr_test.py | 4 ++-- pyproject.toml | 22 +++++++++++++++++++++- 5 files changed, 30 insertions(+), 7 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 01f4ee505..9f3850b84 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -163,6 +163,7 @@ steps: - path: - "csrc/attn/vsa/**" - "csrc/attn/tk/**" + - "csrc/attn/tests/test_vsa.py" - "csrc/attn/setup_vsa.py" - "csrc/attn/config_vsa.py" - "csrc/attn/vsa.cpp" diff --git a/csrc/attn/README.md b/csrc/attn/README.md index b7d226546..30af4bf34 100644 --- a/csrc/attn/README.md +++ b/csrc/attn/README.md @@ -84,7 +84,7 @@ out = sliding_tile_attention(q, k, v, window_size, 0, False) ### Test ```bash python tests/test_sta.py # test STA -python tests/test_block_sparse.py # test VSA +python tests/test_vsa.py # test VSA ``` ### Benchmark ```bash diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index b5acc76f6..ba79b95d0 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -9,11 +9,13 @@ in a functional manner, reducing the need for explicit parameter passing. import pprint from dataclasses import asdict, dataclass, field -from typing import Any +from typing import TYPE_CHECKING, Any import PIL.Image import torch -from torchcodec.decoders import VideoDecoder + +if TYPE_CHECKING: + from torchcodec.decoders import VideoDecoder from fastvideo.attention import AttentionMetadata from fastvideo.configs.sample.teacache import TeaCacheParams, WanTeaCacheParams @@ -198,5 +200,5 @@ class TrainingBatch: @dataclass class PreprocessBatch(ForwardBatch): - video_loader: list[VideoDecoder] = field(default_factory=list) + video_loader: list["VideoDecoder"] = field(default_factory=list) video_file_name: list[str] = field(default_factory=list) diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index 05e81e971..6dea8bf1d 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -52,7 +52,7 @@ def run_test(pytest_command: str): git clone {git_repo} /FastVideo && cd /FastVideo && {checkout_command} && - uv pip install -e .[test] --index-strategy unsafe-best-match && + uv pip install -e .[test] && {pytest_command} """ @@ -100,7 +100,7 @@ def run_precision_tests_STA(): @app.function(gpu="H100:1", image=image, timeout=900) def run_precision_tests_VSA(): - run_test("python csrc/attn/tests/test_block_sparse.py") + run_test("python csrc/attn/tests/test_vsa.py") @app.function(gpu="L40S:1", image=image, timeout=3600) def run_inference_lora_tests(): diff --git a/pyproject.toml b/pyproject.toml index 9e7f58a81..8c3051941 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,7 +53,27 @@ dependencies = [ ] [tool.uv] -extra-index-url = ["https://download.pytorch.org/whl/cu128"] +prerelease = "allow" + +[tool.uv.sources] +torch = [ + { index = "pytorch-cpu", marker = "sys_platform != 'linux'" }, + { index = "pytorch-cu128", marker = "sys_platform == 'linux'" }, +] +torchvision = [ + { index = "pytorch-cpu", marker = "sys_platform != 'linux'" }, + { index = "pytorch-cu128", marker = "sys_platform == 'linux'" }, +] + +[[tool.uv.index]] +name = "pytorch-cpu" +url = "https://download.pytorch.org/whl/cpu" +explicit = true + +[[tool.uv.index]] +name = "pytorch-cu128" +url = "https://download.pytorch.org/whl/cu128" +explicit = true [project.optional-dependencies] From 5357f63327791ab6ac0aaa1021dd7129a75f1a4c Mon Sep 17 00:00:00 2001 From: Wenxuan Tan Date: Sat, 9 Aug 2025 20:46:00 -0500 Subject: [PATCH 14/16] Fix LoRA load from training checkpoint (#719) --- examples/inference/basic/basic.py | 3 +-- examples/inference/basic/basic_dmd.py | 2 +- examples/inference/lora/wan_lora_inference.py | 4 ++-- examples/inference/lora/wan_lora_inference_from_ckpt.py | 6 +++--- examples/inference/optimizations/attention_example.py | 2 +- examples/inference/optimizations/teacache_example.py | 2 +- .../finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh | 1 + fastvideo/layers/lora/linear.py | 6 +++--- fastvideo/models/dits/wanvideo.py | 2 -- fastvideo/pipelines/stages/denoising.py | 2 -- fastvideo/pipelines/stages/latent_preparation.py | 2 -- fastvideo/training/training_pipeline.py | 1 + fastvideo/training/training_utils.py | 1 - 13 files changed, 14 insertions(+), 20 deletions(-) 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/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 c2ad8598d..519abf0f8 100644 --- a/examples/inference/lora/wan_lora_inference_from_ckpt.py +++ b/examples/inference/lora/wan_lora_inference_from_ckpt.py @@ -11,9 +11,9 @@ 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, + 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" ) @@ -26,7 +26,7 @@ def main(): "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, 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/wan_t2v_1.3B/crush_smol/finetune_t2v.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh index 3ac715007..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 @@ -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/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/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 592ee352d..8a2d9f0b6 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -323,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 diff --git a/fastvideo/pipelines/stages/latent_preparation.py b/fastvideo/pipelines/stages/latent_preparation.py index 5c7e54f1a..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 diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 8293f10ca..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 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") From 2a773fa34e73b907164589f6fc7f9b7eb04f8f9e Mon Sep 17 00:00:00 2001 From: William Lin Date: Sat, 9 Aug 2025 20:47:42 -0700 Subject: [PATCH 15/16] [bugfix] [distill] remove i2v validation schema import in distill (#728) --- fastvideo/training/wan_i2v_distillation_pipeline.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) 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...") From 8cae5e470818c486da06862185e05ab60edc9a01 Mon Sep 17 00:00:00 2001 From: William Lin Date: Sun, 10 Aug 2025 15:34:03 -0700 Subject: [PATCH 16/16] [feature] add Gradio live serving demo code (#727) Co-authored-by: RandNMR73 --- examples/inference/gradio/README.md | 59 -- examples/inference/gradio/gradio_demo.py | 169 ---- examples/inference/gradio/gradio_frontend.py | 730 ++++++++++++++++++ .../inference/gradio/ray_serve_backend.py | 379 +++++++++ examples/inference/gradio/start.sh | 3 + .../inference/gradio/start_ray_serve_app.py | 257 ++++++ fastvideo/entrypoints/video_generator.py | 11 +- fastvideo/pipelines/pipeline_batch_info.py | 41 + fastvideo/pipelines/stages/base.py | 2 + fastvideo/worker/gpu_worker.py | 9 +- fastvideo/worker/multiproc_executor.py | 20 +- 11 files changed, 1446 insertions(+), 234 deletions(-) delete mode 100644 examples/inference/gradio/README.md delete mode 100644 examples/inference/gradio/gradio_demo.py create mode 100644 examples/inference/gradio/gradio_frontend.py create mode 100644 examples/inference/gradio/ray_serve_backend.py create mode 100644 examples/inference/gradio/start.sh create mode 100644 examples/inference/gradio/start_ray_serve_app.py 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""" +
+

⏱️ Timing Breakdown

+
+
+
πŸš€
+
DiT Denoising
+
{dit_denoising_time}
+
+
+
🧠
+
E2E (w. vae/text encoder)
+
{inference_time:.2f}s
+
+
+
🎬
+
Video Encoding
+
{encoding_time:.2f}s
+
+
+
🌐
+
Network Transfer
+
{network_time:.2f}s
+
+
+
πŸ“Š
+
Total Processing
+
{total_time:.2f}s
+
+
""" + + if inference_time > 0: + fps = num_frames / inference_time + timing_html += f""" +
+ Generation Speed: + {fps:.1f} frames/second +
""" + + return timing_html + "
" + + +def load_example_prompts(): + def contains_chinese(text): + return any('\u4e00' <= char <= '\u9fff' for char in text) + + def load_from_file(filepath): + prompts, labels = [], [] + try: + with open(filepath, "r", encoding='utf-8') as f: + for line in f: + line = line.strip() + if line and not contains_chinese(line): + label = line[:100] + "..." if len(line) > 100 else line + labels.append(label) + prompts.append(line) + except Exception as e: + print(f"Warning: Could not read {filepath}: {e}") + return prompts, labels + + examples, example_labels = load_from_file("prompts/prompts_final.txt") + + if not examples: + examples = ["A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background."] + example_labels = ["Crowded rooftop bar at night"] + + return examples, example_labels + + +def create_gradio_interface(backend_url: str, default_params: dict[str, SamplingParam]): + + client = RayServeClient(backend_url) + + def generate_video( + prompt, negative_prompt, use_negative_prompt, seed, guidance_scale, + num_frames, height, width, randomize_seed, model_selection, progress + ): + if not client.check_health(): + return None, f"Backend is not available. Please check if Ray Serve is running at {backend_url}", "" + + # Validate dimensions + max_pixels = 720 * 1280 + if height * width > max_pixels: + return None, f"Video dimensions too large. Maximum: 720x1280 pixels", "" + + if progress: + progress(0.1, desc="Checking backend health...") + + request_data = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "use_negative_prompt": use_negative_prompt, + "seed": seed, + "guidance_scale": guidance_scale, + "num_frames": num_frames, + "height": height, + "width": width, + "randomize_seed": randomize_seed, + "return_frames": False, + "image_path": None, + "model_path": MODEL_PATH_MAPPING.get(model_selection, "FastVideo/FastWan2.1-T2V-1.3B-Diffusers") + } + + if progress: + progress(0.4, desc="Generating video...") + + response = client.generate_video(request_data) + + if progress: + progress(0.8, desc="Processing response...") + + if response.get("success", False): + video_data = response.get("video_data", "") + used_seed = response.get("seed", seed) + inference_time = response.get("inference_time", 0.0) + encoding_time = response.get("encoding_time", 0.0) + total_time = response.get("total_time", 0.0) + network_time = response.get("network_time", 0.0) + stage_execution_times = response.get("stage_execution_times", []) + + timing_details = create_timing_display( + inference_time, encoding_time, network_time, total_time, + stage_execution_times, num_frames + ) + + if video_data: + if progress: + progress(0.9, desc="Saving video...") + + video_path = save_video_from_base64(video_data, "outputs", prompt) + + if progress: + progress(1.0, desc="Generation complete!") + + if video_path and os.path.exists(video_path): + return video_path, used_seed, timing_details + else: + return None, "Failed to save video", "" + else: + return None, "No video data received from backend", "" + else: + error_msg = response.get("error_message", "Unknown error occurred") + return None, f"Generation failed: {error_msg}", "" + + examples, example_labels = load_example_prompts() + + theme = gr.themes.Base().set( + button_primary_background_fill="#2563eb", + button_primary_background_fill_hover="#1d4ed8", + button_primary_text_color="white", + slider_color="#2563eb", + checkbox_background_color_selected="#2563eb", + ) + + def get_default_values(model_name): + model_path = MODEL_PATH_MAPPING.get(model_name) + if model_path and model_path in default_params: + params = default_params[model_path] + return { + 'height': params.height, + 'width': params.width, + 'num_frames': params.num_frames, + 'guidance_scale': params.guidance_scale, + 'seed': params.seed, + } + + return { + 'height': 448, + 'width': 832, + 'num_frames': 61, + 'guidance_scale': 3.0, + 'seed': 1024, + } + + initial_values = get_default_values("FastWan2.1-T2V-1.3B") + + with gr.Blocks(title="FastWan", theme=theme) as demo: + gr.Image("fastvideo-logos/main/svg/full.svg", show_label=False, container=False, height=80) + gr.HTML(""" +
+

Make Video Generation Go Blurrrrrrr

+

Code | Blog | Docs

+
+ """) + + 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. +

+
+ """) + + with gr.Row(): + model_selection = gr.Dropdown( + choices=list(MODEL_PATH_MAPPING.keys()), + value="FastWan2.1-T2V-1.3B", + label="Select Model", + interactive=True + ) + + with gr.Row(): + example_dropdown = gr.Dropdown( + choices=example_labels, + label="Example Prompts", + value=None, + interactive=True, + allow_custom_value=False + ) + + with gr.Row(): + with gr.Column(scale=6): + prompt = gr.Text( + label="Prompt", + show_label=False, + max_lines=3, + placeholder="Describe your scene...", + container=False, + lines=3, + autofocus=True, + ) + with gr.Column(scale=1, min_width=120, elem_classes="center-button"): + run_button = gr.Button("Run", variant="primary", size="lg") + + with gr.Row(): + with gr.Column(): + error_output = gr.Text(label="Error", visible=False) + timing_display = gr.Markdown(label="Timing Breakdown", visible=False) + + with gr.Row(equal_height=True, elem_classes="main-content-row"): + with gr.Column(scale=1, elem_classes="advanced-options-column"): + with gr.Group(): + gr.HTML("
Advanced Options
") + with gr.Row(): + height = gr.Number( + label="Height", + value=initial_values['height'], + interactive=False, + container=True + ) + width = gr.Number( + label="Width", + value=initial_values['width'], + interactive=False, + container=True + ) + + with gr.Row(): + num_frames = gr.Number( + label="Number of Frames", + value=initial_values['num_frames'], + interactive=False, + container=True + ) + guidance_scale = gr.Slider( + label="Guidance Scale", + minimum=1, + maximum=12, + value=initial_values['guidance_scale'], + ) + + with gr.Row(): + use_negative_prompt = gr.Checkbox( + label="Use negative prompt", value=False) + negative_prompt = gr.Text( + label="Negative prompt", + max_lines=3, + lines=3, + placeholder="Enter a negative prompt", + visible=False, + ) + + seed = gr.Slider( + label="Seed", + minimum=0, + maximum=1000000, + step=1, + value=initial_values['seed'], + ) + randomize_seed = gr.Checkbox(label="Randomize seed", value=False) + seed_output = gr.Number(label="Used Seed") + + with gr.Column(scale=1, elem_classes="video-column"): + result = gr.Video( + label="Generated Video", + show_label=True, + height=466, + width=600, + container=True, + elem_classes="video-component" + ) + + gr.HTML(""" + + """) + + def on_example_select(example_label): + if example_label and example_label in example_labels: + index = example_labels.index(example_label) + return examples[index] + return "" + + example_dropdown.change( + fn=on_example_select, + inputs=example_dropdown, + outputs=prompt, + ) + + gr.HTML(""" +
+

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/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/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index ba79b95d0..eb89f62f5 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -17,10 +17,47 @@ 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: """ @@ -132,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.""" 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/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py index f3eebcbb7..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) @@ -139,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,