Merge branch 'hao-ai-lab:main' into hunyuan-gamecraft-v2-port
This commit is contained in:
@@ -1,29 +1,23 @@
|
||||
<div align="center">
|
||||
<img src=assets/logos/logo.svg width="30%"/>
|
||||
</div>
|
||||
|
||||
<p align="center">
|
||||
| <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/"><b> Quick Start</b></a> | <a href="https://github.com/hao-ai-lab/FastVideo/discussions/982" target="_blank"><b>Weekly Dev Meeting</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ" target="_blank"> <b>Slack</b> </a> | 🟣💬 <a href="https://ibb.co/sv3MMKyv" target="_blank"> <b> WeChat </b> </a> |
|
||||
</p>
|
||||
| **[Documentation](https://hao-ai-lab.github.io/FastVideo)** | **[Quick Start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/)** | **[Weekly Dev Meeting](https://github.com/hao-ai-lab/FastVideo/discussions/982)** | 🟣💬 **[Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ)** |
|
||||
|
||||
**FastVideo is a unified post-training and inference framework for accelerated video generation.**
|
||||
|
||||
## NEWS
|
||||
- ```2025/11/19```: Release [CausalWan2.2 I2V A14B Preview](https://huggingface.co/FastVideo/CausalWan2.2-I2V-A14B-Preview-Diffusers) models, [Blog](https://hao-ai-lab.github.io/blogs/fastvideo_causalwan_preview/) and [Inference Code!](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py)
|
||||
- ```2025/08/04```: Release [FastWan](https://hao-ai-lab.github.io/FastVideo/distillation/dmd) models and [Sparse-Distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).
|
||||
|
||||
<details>
|
||||
<summary>More</summary>
|
||||
- `2025/11/19`: Release [CausalWan2.2 I2V A14B Preview](https://huggingface.co/FastVideo/CausalWan2.2-I2V-A14B-Preview-Diffusers) models, [Blog](https://hao-ai-lab.github.io/blogs/fastvideo_causalwan_preview/) and [Inference Code!](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py)
|
||||
- `2025/08/04`: Release [FastWan](https://hao-ai-lab.github.io/FastVideo/distillation/dmd) models and [Sparse-Distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).
|
||||
|
||||
- ```2025/06/14```: Release finetuning and inference code for [VSA](https://arxiv.org/pdf/2505.13389)
|
||||
- ```2025/04/24```: [FastVideo V1](https://hao-ai-lab.github.io/blogs/fastvideo/) is released!
|
||||
- ```2025/02/18```: Release the inference code for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
|
||||
### More News
|
||||
|
||||
</details>
|
||||
- `2025/06/14`: Release finetuning and inference code for [VSA](https://arxiv.org/pdf/2505.13389)
|
||||
- `2025/04/24`: [FastVideo V1](https://hao-ai-lab.github.io/blogs/fastvideo/) is released!
|
||||
- `2025/02/18`: Release the inference code for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
|
||||
|
||||
## Key Features
|
||||
|
||||
FastVideo has the following features:
|
||||
|
||||
- End-to-end post-training support for bidirectional and autoregressive models:
|
||||
- Support full finetuning and LoRA finetuning for state-of-the-art open video DiTs
|
||||
- Data preprocessing pipeline for video, image, and text data
|
||||
@@ -44,6 +38,7 @@ FastVideo has the following features:
|
||||
- See this [page](https://hao-ai-lab.github.io/FastVideo/inference/hardware_support/) for full list of supported hardware and OS.
|
||||
|
||||
## Getting Started
|
||||
|
||||
We recommend using an environment manager such as `Conda` to create a clean environment:
|
||||
|
||||
```bash
|
||||
@@ -58,18 +53,20 @@ pip install fastvideo
|
||||
Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/) for more detailed installation instructions.
|
||||
|
||||
## Sparse Distillation
|
||||
|
||||
For our sparse distillation techniques, please see our [distillation docs](https://hao-ai-lab.github.io/FastVideo/distillation/dmd/) and check out our [blog](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).
|
||||
|
||||
See below for recipes and datasets:
|
||||
|
||||
| Model | Sparse Distillation | Dataset |
|
||||
|:-------------------------------------------------------------------------------------------: |:---------------------------------------------------------------------------------------------------------------: |:--------------------------------------------------------------------------------------------------------: |
|
||||
| [FastWan2.1-T2V-1.3B](https://huggingface.co/FastVideo/FastWan2.1-T2V-1.3B-Diffusers) | [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P) | [FastVideo Synthetic Wan2.1 480P](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x448x832_600k) |
|
||||
| [FastWan2.1-T2V-14B-Preview](https://huggingface.co/FastVideo/FastWan2.1-T2V-14B-Diffusers) | Coming soon! | [FastVideo Synthetic Wan2.1 720P](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x768x1280_250k) |
|
||||
| [FastWan2.2-TI2V-5B](https://huggingface.co/FastVideo/FastWan2.2-TI2V-5B-Diffusers) | [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free) | [FastVideo Synthetic Wan2.2 720P](https://huggingface.co/datasets/FastVideo/Wan2.2-Syn-121x704x1280_32k) |
|
||||
| Model | Sparse Distillation | Dataset |
|
||||
| ------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------- |
|
||||
| [FastWan2.1-T2V-1.3B](https://huggingface.co/FastVideo/FastWan2.1-T2V-1.3B-Diffusers) | [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P) | [FastVideo Synthetic Wan2.1 480P](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x448x832_600k) |
|
||||
| [FastWan2.2-TI2V-5B](https://huggingface.co/FastVideo/FastWan2.2-TI2V-5B-Diffusers) | [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free) | [FastVideo Synthetic Wan2.2 720P](https://huggingface.co/datasets/FastVideo/Wan2.2-Syn-121x704x1280_32k) |
|
||||
|
||||
## Inference
|
||||
|
||||
### Generating Your First Video
|
||||
|
||||
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/). Create a file called `example.py` with the following code:
|
||||
|
||||
```python
|
||||
@@ -113,7 +110,6 @@ For a more detailed guide, please see our [inference quick start](https://hao-ai
|
||||
- [Design Overview](https://hao-ai-lab.github.io/FastVideo/design/overview/)
|
||||
- [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/distillation/dmd/)
|
||||
- [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/contributing/overview/)
|
||||
<!-- - [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetune.html) -->
|
||||
|
||||
## Awesome work using FastVideo or our research projects
|
||||
|
||||
@@ -131,9 +127,11 @@ We welcome all contributions. Please check out our guide [here](https://hao-ai-l
|
||||
See details in [development roadmap](https://github.com/hao-ai-lab/FastVideo/issues/899).
|
||||
|
||||
## Acknowledgement
|
||||
|
||||
We learned the design and reused code from the following projects: [Wan-Video](https://github.com/Wan-Video), [ThunderKittens](https://github.com/HazyResearch/ThunderKittens), [DMD2](https://github.com/tianweiy/DMD2), [diffusers](https://github.com/huggingface/diffusers), [xDiT](https://github.com/xdit-project/xDiT), [vLLM](https://github.com/vllm-project/vllm), [SGLang](https://github.com/sgl-project/sglang). We thank [MBZUAI](https://ifm.mbzuai.ac.ae/), [Anyscale](https://www.anyscale.com/), and [GMI Cloud](https://www.gmicloud.ai/) for their support throughout this project.
|
||||
|
||||
## Citation
|
||||
|
||||
If you find FastVideo useful, please consider citing our research work:
|
||||
|
||||
```bibtex
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import torch
|
||||
from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig
|
||||
|
||||
|
||||
@@ -76,8 +78,14 @@ class MatrixGameWanVideoArchConfig(WanVideoArchConfig):
|
||||
image_dim: int = 1280
|
||||
|
||||
|
||||
def _is_transformer_block(param_name: str, module: torch.nn.Module) -> bool:
|
||||
return bool("blocks" in param_name and param_name.split(".")[-1].isdigit())
|
||||
|
||||
|
||||
@dataclass
|
||||
class MatrixGameWanVideoConfig(WanVideoConfig):
|
||||
arch_config: MatrixGameWanVideoArchConfig = field(
|
||||
default_factory=MatrixGameWanVideoArchConfig)
|
||||
prefix: str = "Wan"
|
||||
_compile_conditions: list = field(
|
||||
default_factory=lambda: [_is_transformer_block])
|
||||
|
||||
@@ -219,6 +219,24 @@ class PipelineConfig:
|
||||
"Comma-separated list of denoising steps (e.g., '1000,757,522')",
|
||||
)
|
||||
|
||||
# STA (Sliding Tile Attention) parameters
|
||||
parser.add_argument(
|
||||
f"--{prefix_with_dot}STA-mode",
|
||||
type=str,
|
||||
dest=f"{prefix_with_dot.replace('-', '_')}STA_mode",
|
||||
default=PipelineConfig.STA_mode.value,
|
||||
choices=[mode.value for mode in STA_Mode],
|
||||
help=
|
||||
"STA mode: STA_inference, STA_searching, STA_tuning, STA_tuning_cfg, None",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix_with_dot}skip-time-steps",
|
||||
type=int,
|
||||
dest=f"{prefix_with_dot.replace('-', '_')}skip_time_steps",
|
||||
default=PipelineConfig.skip_time_steps,
|
||||
help="Number of time steps to warmup (full attention) for STA",
|
||||
)
|
||||
|
||||
# Add VAE configuration arguments
|
||||
from fastvideo.configs.models.vaes.base import VAEConfig
|
||||
VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config")
|
||||
@@ -299,6 +317,12 @@ class PipelineConfig:
|
||||
# 4. Update PipelineConfig from CLI arguments if provided
|
||||
kwargs[prefix_with_dot + 'model_path'] = model_path
|
||||
pipeline_config.update_config_from_dict(kwargs, config_cli_prefix)
|
||||
|
||||
# Convert STA_mode string to enum if necessary
|
||||
if isinstance(pipeline_config.STA_mode, str) and not isinstance(
|
||||
pipeline_config.STA_mode, STA_Mode):
|
||||
pipeline_config.STA_mode = STA_Mode(pipeline_config.STA_mode)
|
||||
|
||||
return pipeline_config
|
||||
|
||||
def check_pipeline_config(self) -> None:
|
||||
|
||||
@@ -148,62 +148,6 @@ class DataValidationStage(DatasetFilterStage):
|
||||
return batch
|
||||
|
||||
|
||||
class ResolutionFilterStage(DatasetFilterStage):
|
||||
"""Stage for filtering data items based on resolution constraints."""
|
||||
|
||||
def __init__(self,
|
||||
max_h_div_w_ratio: float = 17 / 16,
|
||||
min_h_div_w_ratio: float = 8 / 16,
|
||||
max_height: int = 1024,
|
||||
max_width: int = 1024):
|
||||
self.max_h_div_w_ratio = max_h_div_w_ratio
|
||||
self.min_h_div_w_ratio = min_h_div_w_ratio
|
||||
self.max_height = max_height
|
||||
self.max_width = max_width
|
||||
|
||||
def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
|
||||
"""
|
||||
Check if data item passes resolution filtering.
|
||||
|
||||
Args:
|
||||
batch: Dataset batch with resolution information
|
||||
|
||||
Returns:
|
||||
True if passes filter, False otherwise
|
||||
"""
|
||||
# Only apply to videos
|
||||
if not batch.is_video:
|
||||
return True
|
||||
|
||||
if batch.resolution is None:
|
||||
return False
|
||||
|
||||
height = batch.resolution.get("height", None)
|
||||
width = batch.resolution.get("width", None)
|
||||
if height is None or width is None:
|
||||
return False
|
||||
|
||||
# Check aspect ratio
|
||||
aspect = self.max_height / self.max_width
|
||||
hw_aspect_thr = 1.5
|
||||
|
||||
return self.filter_resolution(
|
||||
height,
|
||||
width,
|
||||
max_h_div_w_ratio=hw_aspect_thr * aspect,
|
||||
min_h_div_w_ratio=1 / hw_aspect_thr * aspect,
|
||||
)
|
||||
|
||||
def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
|
||||
"""Process does nothing for resolution filtering - filtering is handled by should_keep."""
|
||||
return batch
|
||||
|
||||
def filter_resolution(self, h: int, w: int, max_h_div_w_ratio: float,
|
||||
min_h_div_w_ratio: float) -> bool:
|
||||
"""Filter based on height/width ratio."""
|
||||
return h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio
|
||||
|
||||
|
||||
class FrameSamplingStage(DatasetFilterStage):
|
||||
"""Stage for temporal frame sampling and indexing."""
|
||||
|
||||
@@ -328,11 +272,6 @@ class VideoTransformStage(DatasetStage):
|
||||
video = rearrange(video, "t c h w -> c t h w")
|
||||
video = video.to(torch.uint8)
|
||||
|
||||
h, w = video.shape[-2:]
|
||||
assert (
|
||||
h / w <= 17 / 16 and h / w >= 8 / 16
|
||||
), f"Only videos with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But video ({batch.path}) found ratio is {round(h / w, 2)} with the shape of {video.shape}"
|
||||
|
||||
video = video.float() / 127.5 - 1.0
|
||||
batch.pixel_values = video
|
||||
return batch
|
||||
@@ -489,8 +428,6 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
|
||||
tokenizer) -> None:
|
||||
"""Initialize all processing stages."""
|
||||
self.validation_stage = DataValidationStage()
|
||||
self.resolution_filter_stage = ResolutionFilterStage(
|
||||
max_height=args.max_height, max_width=args.max_width)
|
||||
self.frame_sampling_stage = FrameSamplingStage(
|
||||
num_frames=args.num_frames,
|
||||
train_fps=args.train_fps,
|
||||
@@ -543,7 +480,6 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
|
||||
# Initialize counters
|
||||
filter_counts = {
|
||||
"validation_failed": 0,
|
||||
"resolution_failed": 0,
|
||||
"frame_sampling_failed": 0
|
||||
}
|
||||
sample_num_frames: list[int] = []
|
||||
@@ -579,10 +515,6 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
|
||||
filter_counts["validation_failed"] += 1
|
||||
return False
|
||||
|
||||
if not self.resolution_filter_stage.should_keep(batch):
|
||||
filter_counts["resolution_failed"] += 1
|
||||
return False
|
||||
|
||||
if not self.frame_sampling_stage.should_keep(batch):
|
||||
filter_counts["frame_sampling_failed"] += 1
|
||||
return False
|
||||
@@ -594,10 +526,9 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
|
||||
after_count: int):
|
||||
"""Log filtering statistics."""
|
||||
logger.info(
|
||||
"validation_failed: %d, resolution_failed: %d, frame_sampling_failed: %d, "
|
||||
"validation_failed: %d, frame_sampling_failed: %d, "
|
||||
"Counter(sample_num_frames): %s, before filter: %d, after filter: %d",
|
||||
filter_counts['validation_failed'],
|
||||
filter_counts['resolution_failed'],
|
||||
filter_counts['frame_sampling_failed'], Counter(sample_num_frames),
|
||||
before_count, after_count)
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@ diffusion models.
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
from copy import deepcopy
|
||||
from typing import Any
|
||||
@@ -31,6 +32,21 @@ from fastvideo.worker.executor import Executor
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _infer_latent_batch_size(batch: ForwardBatch) -> int:
|
||||
if isinstance(batch.prompt, list):
|
||||
latent_batch_size = len(batch.prompt)
|
||||
elif batch.prompt is not None:
|
||||
latent_batch_size = 1
|
||||
elif batch.prompt_embeds is not None and len(batch.prompt_embeds) > 0:
|
||||
latent_batch_size = batch.prompt_embeds[0].shape[0]
|
||||
else:
|
||||
raise ValueError(
|
||||
"Cannot infer batch size from batch; no prompt or prompt_embeds found"
|
||||
)
|
||||
latent_batch_size *= batch.num_videos_per_prompt
|
||||
return latent_batch_size
|
||||
|
||||
|
||||
class VideoGenerator:
|
||||
"""
|
||||
A unified class for generating videos using diffusion models.
|
||||
@@ -372,8 +388,31 @@ class VideoGenerator:
|
||||
|
||||
# Run inference
|
||||
start_time = time.perf_counter()
|
||||
output_batch = self.executor.execute_forward(batch, fastvideo_args)
|
||||
samples = output_batch.output
|
||||
|
||||
# Execute forward pass in a new thread for non-blocking tensor allocation
|
||||
result_container = {}
|
||||
|
||||
def execute_forward_thread():
|
||||
result_container['output_batch'] = self.executor.execute_forward(
|
||||
batch, fastvideo_args)
|
||||
|
||||
thread = threading.Thread(target=execute_forward_thread)
|
||||
thread.start()
|
||||
latent_batch_size = _infer_latent_batch_size(batch)
|
||||
samples = torch.empty((latent_batch_size, 3, sampling_param.num_frames,
|
||||
sampling_param.height, sampling_param.width),
|
||||
device='cpu',
|
||||
pin_memory=fastvideo_args.pin_cpu_memory)
|
||||
thread.join()
|
||||
|
||||
output_batch = result_container['output_batch']
|
||||
if output_batch.output.shape == samples.shape:
|
||||
samples.copy_(output_batch.output)
|
||||
else:
|
||||
logger.warning(
|
||||
"Output shape %s does not match expected shape %s; use slow path",
|
||||
output_batch.output.shape, samples.shape)
|
||||
samples = output_batch.output.cpu()
|
||||
logging_info = output_batch.logging_info
|
||||
|
||||
gen_time = time.perf_counter() - start_time
|
||||
|
||||
@@ -10,7 +10,7 @@ from enum import Enum
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
from fastvideo.configs.configs import PreprocessConfig
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig, STA_Mode
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig
|
||||
from fastvideo.configs.utils import clean_cli_args
|
||||
from fastvideo.layers.quantization import QUANTIZATION_METHODS, QuantizationMethods
|
||||
from fastvideo.logger import init_logger
|
||||
@@ -141,8 +141,6 @@ class FastVideoArgs:
|
||||
|
||||
# STA (Sliding Tile Attention) parameters
|
||||
mask_strategy_file_path: str | None = None
|
||||
STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
|
||||
skip_time_steps: int = 15
|
||||
|
||||
# Compilation
|
||||
enable_torch_compile: bool = False
|
||||
@@ -465,20 +463,6 @@ class FastVideoArgs:
|
||||
)
|
||||
|
||||
# STA (Sliding Tile Attention) parameters
|
||||
parser.add_argument(
|
||||
"--STA-mode",
|
||||
type=str,
|
||||
default=FastVideoArgs.STA_mode.value,
|
||||
choices=[mode.value for mode in STA_Mode],
|
||||
help=
|
||||
"STA mode contains STA_inference, STA_searching, STA_tuning, STA_tuning_cfg, None",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip-time-steps",
|
||||
type=int,
|
||||
default=FastVideoArgs.skip_time_steps,
|
||||
help="Number of time steps to warmup (full attention) for STA",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mask-strategy-file-path",
|
||||
type=str,
|
||||
|
||||
+17
-14
@@ -102,6 +102,7 @@ def _info(logger: Logger,
|
||||
- When both are False, the message will be logged from all processes
|
||||
- By default, only logs from processes with LOCAL_RANK=0
|
||||
"""
|
||||
is_distributed = int(os.environ.get("WORLD_SIZE", 1)) > 1
|
||||
try:
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
rank = int(os.environ["RANK"])
|
||||
@@ -118,20 +119,22 @@ def _info(logger: Logger,
|
||||
|
||||
global _warned_local_main_process, _warned_main_process
|
||||
|
||||
if not _warned_local_main_process and local_main_process_only:
|
||||
logger.warning(
|
||||
'%s By default, logger.info(..) will only log from the local main process. Set logger.info(..., is_local_main_process=False) to log from all processes.%s',
|
||||
GREEN,
|
||||
RESET,
|
||||
)
|
||||
_warned_local_main_process = True
|
||||
if not _warned_main_process and main_process_only:
|
||||
logger.warning(
|
||||
'%s is_main_process_only is set to True, logging only from the main (RANK==0) process.%s',
|
||||
GREEN,
|
||||
RESET,
|
||||
)
|
||||
_warned_main_process = True
|
||||
# Only show process-awareness warnings when actually running distributed
|
||||
if is_distributed:
|
||||
if not _warned_local_main_process and local_main_process_only:
|
||||
logger.warning(
|
||||
'%s By default, logger.info(..) will only log from the local main process. Set logger.info(..., is_local_main_process=False) to log from all processes.%s',
|
||||
GREEN,
|
||||
RESET,
|
||||
)
|
||||
_warned_local_main_process = True
|
||||
if not _warned_main_process and main_process_only:
|
||||
logger.warning(
|
||||
'%s is_main_process_only is set to True, logging only from the main (RANK==0) process.%s',
|
||||
GREEN,
|
||||
RESET,
|
||||
)
|
||||
_warned_main_process = True
|
||||
|
||||
if not main_process_only and not local_main_process_only:
|
||||
logger.log(logging.INFO, msg, *args, stacklevel=2, **kwargs)
|
||||
|
||||
@@ -101,6 +101,58 @@ class ComposedPipelineBase(ABC):
|
||||
module.requires_grad_(True)
|
||||
module.train()
|
||||
|
||||
@staticmethod
|
||||
def _compile_with_conditions(
|
||||
module: torch.nn.Module,
|
||||
compile_kwargs: dict[str, Any],
|
||||
) -> int:
|
||||
"""Compile submodules that match module._compile_conditions."""
|
||||
compile_conditions = getattr(module, "_compile_conditions", None)
|
||||
if not compile_conditions:
|
||||
return 0
|
||||
|
||||
compiled_count = 0
|
||||
for name, submodule in module.named_modules():
|
||||
if not name:
|
||||
continue
|
||||
if any(cond(name, submodule) for cond in compile_conditions):
|
||||
submodule.forward = torch.compile(submodule.forward,
|
||||
**compile_kwargs)
|
||||
compiled_count += 1
|
||||
return compiled_count
|
||||
|
||||
def _maybe_compile_pipeline_module(
|
||||
self,
|
||||
module_name: str,
|
||||
fsdp_module_cls: type | None,
|
||||
compile_kwargs: dict[str, Any],
|
||||
) -> None:
|
||||
if module_name not in self.modules:
|
||||
return
|
||||
|
||||
module = self.modules[module_name]
|
||||
if fsdp_module_cls is not None and isinstance(module, fsdp_module_cls):
|
||||
logger.info(
|
||||
"%s is already FSDP-wrapped; skipping torch.compile in pipeline",
|
||||
module_name.capitalize(),
|
||||
)
|
||||
return
|
||||
|
||||
compiled_count = self._compile_with_conditions(module, compile_kwargs)
|
||||
if compiled_count > 0:
|
||||
logger.info(
|
||||
"Enabled torch.compile for %d submodules in %s via _compile_conditions with kwargs=%s",
|
||||
compiled_count,
|
||||
module_name,
|
||||
compile_kwargs,
|
||||
)
|
||||
return
|
||||
|
||||
# Backward-compatible fallback: compile full module if no condition matched.
|
||||
logger.info("Enabling torch.compile for %s with kwargs=%s", module_name,
|
||||
compile_kwargs)
|
||||
self.modules[module_name] = torch.compile(module, **compile_kwargs)
|
||||
|
||||
def post_init(self) -> None:
|
||||
assert self.fastvideo_args is not None, "fastvideo_args must be set"
|
||||
if self.post_init_called:
|
||||
@@ -116,7 +168,6 @@ class ComposedPipelineBase(ABC):
|
||||
|
||||
self.initialize_pipeline(self.fastvideo_args)
|
||||
if self.fastvideo_args.enable_torch_compile:
|
||||
transformer_module = self.modules["transformer"]
|
||||
if self.fastvideo_args.training_mode:
|
||||
logger.info(
|
||||
"Torch Compile enabled via FSDP loader for training; skipping additional pipeline compile"
|
||||
@@ -130,29 +181,16 @@ class ComposedPipelineBase(ABC):
|
||||
fsdp_module_cls = None
|
||||
|
||||
compile_kwargs = self.fastvideo_args.torch_compile_kwargs or {}
|
||||
if fsdp_module_cls is not None and isinstance(
|
||||
transformer_module, fsdp_module_cls):
|
||||
logger.info(
|
||||
"Transformer is already FSDP-wrapped; skipping torch.compile in pipeline"
|
||||
)
|
||||
else:
|
||||
logger.info("Enabling torch.compile for DiT with kwargs=%s",
|
||||
compile_kwargs)
|
||||
self.modules["transformer"] = torch.compile(
|
||||
transformer_module, **compile_kwargs)
|
||||
if "transformer_2" in self.modules:
|
||||
transformer_module_2 = self.modules["transformer_2"]
|
||||
if fsdp_module_cls is not None and isinstance(
|
||||
transformer_module_2, fsdp_module_cls):
|
||||
logger.info(
|
||||
"Transformer_2 is already FSDP-wrapped; skipping torch.compile in pipeline"
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"Enabling torch.compile for Transformer_2 with kwargs=%s",
|
||||
compile_kwargs)
|
||||
self.modules["transformer_2"] = torch.compile(
|
||||
transformer_module_2, **compile_kwargs)
|
||||
self._maybe_compile_pipeline_module(
|
||||
module_name="transformer",
|
||||
fsdp_module_cls=fsdp_module_cls,
|
||||
compile_kwargs=compile_kwargs,
|
||||
)
|
||||
self._maybe_compile_pipeline_module(
|
||||
module_name="transformer_2",
|
||||
fsdp_module_cls=fsdp_module_cls,
|
||||
compile_kwargs=compile_kwargs,
|
||||
)
|
||||
logger.info("Torch Compile enabled for DiT")
|
||||
|
||||
if not self.fastvideo_args.training_mode:
|
||||
|
||||
@@ -242,8 +242,8 @@ class DecodingStage(PipelineStage):
|
||||
decoded_frames = self.decode(cur_latent, fastvideo_args)
|
||||
batch.trajectory_decoded.append(decoded_frames.cpu().float())
|
||||
|
||||
# Convert to CPU float32 for compatibility
|
||||
frames = frames.cpu().float()
|
||||
# Convert to float32 for compatibility
|
||||
frames = frames.to(torch.float32)
|
||||
|
||||
# Crop padding if this is a LongCat refinement
|
||||
if hasattr(batch, 'num_cond_frames_added') and hasattr(
|
||||
|
||||
@@ -518,7 +518,7 @@ class DenoisingStage(PipelineStage):
|
||||
mgr2.release_all()
|
||||
|
||||
# Save STA mask search results if needed
|
||||
if st_attn_available and self.attn_backend == SlidingTileAttentionBackend and fastvideo_args.STA_mode == STA_Mode.STA_SEARCHING:
|
||||
if st_attn_available and self.attn_backend == SlidingTileAttentionBackend and fastvideo_args.pipeline_config.STA_mode == STA_Mode.STA_SEARCHING:
|
||||
self.save_sta_search_results(batch)
|
||||
|
||||
# deallocate transformer if on mps
|
||||
@@ -611,8 +611,8 @@ class DenoisingStage(PipelineStage):
|
||||
"""
|
||||
# TODO(kevin): STA mask search, currently only support Wan2.1 with 69x768x1280
|
||||
from fastvideo.attention.backends.STA_configuration import configure_sta
|
||||
STA_mode = fastvideo_args.STA_mode
|
||||
skip_time_steps = fastvideo_args.skip_time_steps
|
||||
STA_mode = fastvideo_args.pipeline_config.STA_mode
|
||||
skip_time_steps = fastvideo_args.pipeline_config.skip_time_steps
|
||||
if batch.timesteps is None:
|
||||
raise ValueError("Timesteps must be provided")
|
||||
timesteps_num = batch.timesteps.shape[0]
|
||||
|
||||
@@ -343,7 +343,7 @@ class SRDenoisingStage(PipelineStage):
|
||||
batch.latents = latents
|
||||
|
||||
# Save STA mask search results if needed
|
||||
if st_attn_available and self.attn_backend == SlidingTileAttentionBackend and fastvideo_args.STA_mode == STA_Mode.STA_SEARCHING:
|
||||
if st_attn_available and self.attn_backend == SlidingTileAttentionBackend and fastvideo_args.pipeline_config.STA_mode == STA_Mode.STA_SEARCHING:
|
||||
self.save_sta_search_results(batch)
|
||||
|
||||
# deallocate transformer if on mps
|
||||
|
||||
@@ -1267,7 +1267,7 @@ class DistillationPipeline(TrainingPipeline):
|
||||
with torch.no_grad():
|
||||
output_batch = self.validation_pipeline.forward(
|
||||
batch, training_args)
|
||||
samples = output_batch.output
|
||||
samples = output_batch.output.cpu()
|
||||
if self.rank_in_sp_group != 0:
|
||||
continue
|
||||
|
||||
|
||||
@@ -857,7 +857,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
|
||||
# Run validation inference
|
||||
output_batch = self.validation_pipeline.forward(
|
||||
batch, training_args)
|
||||
samples = output_batch.output
|
||||
samples = output_batch.output.cpu()
|
||||
|
||||
if self.rank_in_sp_group != 0:
|
||||
continue
|
||||
|
||||
@@ -650,8 +650,10 @@ class WorkerMultiprocProc:
|
||||
logging_info = None
|
||||
if envs.FASTVIDEO_STAGE_LOGGING:
|
||||
logging_info = output_batch.logging_info
|
||||
# result tensor shared by CUDA IPC to avoid serialization overhead
|
||||
result = output_batch.output
|
||||
self.pipe.send({
|
||||
"output_batch": output_batch.output.cpu(),
|
||||
"output_batch": result,
|
||||
"logging_info": logging_info,
|
||||
"extra": output_batch.extra,
|
||||
})
|
||||
|
||||
@@ -26,24 +26,18 @@ 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):
|
||||
drop_short_ratio: float = 0.0):
|
||||
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] = {}
|
||||
|
||||
@@ -86,26 +80,11 @@ class PreprocessingDataValidator:
|
||||
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)
|
||||
return not (height is None or width is None)
|
||||
|
||||
def _validate_frame_sampling(self, batch: dict[str, Any]) -> bool:
|
||||
"""Validate frame sampling constraints"""
|
||||
|
||||
Reference in New Issue
Block a user