Merge branch 'hao-ai-lab:main' into hunyuan-gamecraft-v2-port

This commit is contained in:
Mihir Jagtap
2026-02-06 19:17:21 -08:00
committed by GitHub
15 changed files with 185 additions and 179 deletions
+18 -20
View File
@@ -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])
+24
View File
@@ -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:
+1 -70
View File
@@ -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)
+41 -2
View File
@@ -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
+1 -17
View File
@@ -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
View File
@@ -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)
+62 -24
View File
@@ -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:
+2 -2
View File
@@ -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(
+3 -3
View File
@@ -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]
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+3 -1
View File
@@ -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,
})
+2 -23
View File
@@ -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"""