diff --git a/README.md b/README.md index f2e34ff2c..716b64b37 100644 --- a/README.md +++ b/README.md @@ -1,29 +1,23 @@ -
- -
-

- | Documentation | Quick Start | Weekly Dev Meeting | 🟣💬 Slack | 🟣💬 WeChat | -

+| **[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/). -
-More +- `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 -
+- `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/) - ## 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 diff --git a/fastvideo/configs/models/dits/matrixgame.py b/fastvideo/configs/models/dits/matrixgame.py index 4b39f1601..9996d2909 100644 --- a/fastvideo/configs/models/dits/matrixgame.py +++ b/fastvideo/configs/models/dits/matrixgame.py @@ -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]) diff --git a/fastvideo/configs/pipelines/base.py b/fastvideo/configs/pipelines/base.py index a057628a3..7a6e623ed 100644 --- a/fastvideo/configs/pipelines/base.py +++ b/fastvideo/configs/pipelines/base.py @@ -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: diff --git a/fastvideo/dataset/preprocessing_datasets.py b/fastvideo/dataset/preprocessing_datasets.py index 15ed7723c..cba721ce4 100644 --- a/fastvideo/dataset/preprocessing_datasets.py +++ b/fastvideo/dataset/preprocessing_datasets.py @@ -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) diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index d586d9215..9b82ed374 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -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 diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 96fac3913..b17df7047 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -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, diff --git a/fastvideo/logger.py b/fastvideo/logger.py index 51d424e8f..b2f0e68d5 100644 --- a/fastvideo/logger.py +++ b/fastvideo/logger.py @@ -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) diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index d9a7f6419..ca92f1343 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -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: diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/pipelines/stages/decoding.py index 7450ff226..267b3bf27 100644 --- a/fastvideo/pipelines/stages/decoding.py +++ b/fastvideo/pipelines/stages/decoding.py @@ -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( diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 09ee67567..f89d1075a 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -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] diff --git a/fastvideo/pipelines/stages/sr_denoising.py b/fastvideo/pipelines/stages/sr_denoising.py index 6311930bd..123d8b0b0 100644 --- a/fastvideo/pipelines/stages/sr_denoising.py +++ b/fastvideo/pipelines/stages/sr_denoising.py @@ -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 diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index c9ba23588..2fd735dc9 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -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 diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 07c09b005..7d89c3b96 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -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 diff --git a/fastvideo/worker/multiproc_executor.py b/fastvideo/worker/multiproc_executor.py index cc09c5898..22047f9ba 100644 --- a/fastvideo/worker/multiproc_executor.py +++ b/fastvideo/worker/multiproc_executor.py @@ -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, }) diff --git a/fastvideo/workflow/preprocess/components.py b/fastvideo/workflow/preprocess/components.py index 22e5a8526..ddcfda9a9 100644 --- a/fastvideo/workflow/preprocess/components.py +++ b/fastvideo/workflow/preprocess/components.py @@ -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"""