fix jinzhe comment

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

README.md

@@ -1,29 +1,23 @@
-<div align="center">
-<img src=assets/logos/logo.svg width="30%"/>
-</div>
 
-<p align="center">
+| **[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) |
-     | <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>
 
 **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>
+- `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)
-<summary>More</summary>
+- `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)
+### More News
-- ```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/).
 
-</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                                                  |
+| 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-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) |
-| [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

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:

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,

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:
+    # Only show process-awareness warnings when actually running distributed
-        logger.warning(
+    if is_distributed:
-            '%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',
+        if not _warned_local_main_process and local_main_process_only:
-            GREEN,
+            logger.warning(
-            RESET,
+                '%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,
-        _warned_local_main_process = True
+                RESET,
-    if not _warned_main_process and main_process_only:
+            )
-        logger.warning(
+            _warned_local_main_process = True
-            '%s is_main_process_only is set to True, logging only from the main (RANK==0) process.%s',
+        if not _warned_main_process and main_process_only:
-            GREEN,
+            logger.warning(
-            RESET,
+                '%s is_main_process_only is set to True, logging only from the main (RANK==0) process.%s',
-        )
+                GREEN,
-        _warned_main_process = True
+                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)

fastvideo/pipelines/stages/denoising.py

@@ -509,7 +509,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
@@ -602,8 +602,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
+        STA_mode = fastvideo_args.pipeline_config.STA_mode
-        skip_time_steps = fastvideo_args.skip_time_steps
+        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]

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