update

README.md

@@ -7,7 +7,7 @@
 FastVideo features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.
 
 <p align="center">
-    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <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://huggingface.co/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</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/XcY0Cpv" target="_blank"> <b> WeChat </b> </a> |
+    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <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.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3csdw1isz-Euq8_Q8~baewG8hxjXs2gQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/tMwknPLY" target="_blank"> <b> WeChat </b> </a> |
 </p>
 
 <div align="center">
@@ -49,10 +49,10 @@ conda activate fastvideo
 pip install fastvideo
 ```
 
-Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/) for more detailed installation instructions.
+Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html) 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/).
+For our sparse distillation techniques, please see our [distillation docs](https://hao-ai-lab.github.io/FastVideo/distillation/dmd.html) and check out our [blog](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).
 
 See below for recipes and datasets:
 
@@ -64,7 +64,7 @@ See below for recipes and datasets:
 
 ## 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:
+Here's a minimal example to generate a video using the default settings. Make sure VSA kernels are [installed](https://hao-ai-lab.github.io/FastVideo/video_sparse_attention/installation.html). Create a file called `example.py` with the following code:
 
 ```python
 import os
@@ -100,15 +100,15 @@ Run the script with:
 python example.py
 ```
 
-For a more detailed guide, please see our [inference quick start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/).
+For a more detailed guide, please see our [inference quick start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html).
 
 ### Other docs:
 
-- [Design Overview](https://hao-ai-lab.github.io/FastVideo/design/overview/)
-- [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/)
+- [Design Overview](https://hao-ai-lab.github.io/FastVideo/design/overview.html)
+- [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html)
 
 ## Distillation and Finetuning
-- [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/distillation/dmd/)
+- [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/distillation/dmd.html)
 <!-- - [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetune.html) -->
 
 ## 📑 Development Plan
@@ -127,7 +127,7 @@ See details in [development roadmap](https://github.com/hao-ai-lab/FastVideo/iss
 
 ## 🤝 Contributing
 
-We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/contributing/overview/)
+We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/contributing/overview.html)
 
 ## Acknowledgement
 We learned and reused code from the following projects:

csrc/attn/video_sparse_attn/vsa/block_sparse_attn_triton.py

@@ -247,11 +247,11 @@ def _attn_bwd_dq(dq, q, K, V,  #
 
     kv_blocks = tl.load(q2k_num  + meta_base)                 # int32
     kv_ptr    = q2k_index + meta_base * max_kv_blks           # ptr to list
-    block_size = tl.load(variable_block_sizes + q_blk)
     
     
     for blk_idx in range(kv_blocks*2):
         block_sparse_offset = (tl.load(kv_ptr + blk_idx//2).to(tl.int32)*2 + blk_idx%2) *step_n * stride_tok
+        block_size = tl.load(variable_block_sizes + blk_idx//2) - (blk_idx%2) * step_n
         kT = tl.load(kT_ptrs + block_sparse_offset)
         vT = tl.load(vT_ptrs + block_sparse_offset)
         qk = tl.dot(q, kT)

docker/Dockerfile.yaml

@@ -0,0 +1,80 @@
+FROM nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+SHELL ["/bin/bash", "-c"]
+
+WORKDIR /FastVideo
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    wget \
+    git \
+    ca-certificates \
+    openssh-server \
+    zsh \
+    vim \
+    curl \
+    gcc-11 \
+    g++-11 \
+    clang-11 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set up C++20 compilers for ThunderKittens
+RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+# Set CUDA environment variables
+ENV CUDA_HOME=/usr/local/cuda-12.8
+ENV PATH=${CUDA_HOME}/bin:${PATH}
+ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+
+# Install Miniconda and create conda environment
+COPY Miniconda3-latest-Linux-x86_64.sh /tmp/miniconda.sh
+RUN bash /tmp/miniconda.sh -b -p /opt/conda && rm /tmp/miniconda.sh
+ENV PATH=/opt/conda/bin:${PATH}
+RUN /opt/conda/bin/conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main
+RUN /opt/conda/bin/conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r
+RUN /opt/conda/bin/conda update -y -n base conda && \
+    /opt/conda/bin/conda create -y -n fastvideo python=3.12 && \
+    /opt/conda/bin/conda clean -afy
+ENV PATH=/opt/conda/envs/fastvideo/bin:/opt/conda/bin:${PATH}
+
+# Copy just the pyproject.toml first to leverage Docker cache
+COPY pyproject.toml ./
+
+# Create a dummy README to satisfy the installation
+RUN echo "# Placeholder" > README.md
+
+# Install project dependencies inside the conda environment
+RUN source /opt/conda/etc/profile.d/conda.sh && \
+    conda activate fastvideo && \
+    pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir .[dev] && \
+    pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation
+
+COPY . .
+
+# Install dependencies using pip inside conda env and set up shell configuration
+RUN source /opt/conda/etc/profile.d/conda.sh && \
+    conda activate fastvideo && \
+    pip install --no-cache-dir -e .[dev] && \
+    git config --unset-all http.https://github.com/.extraheader || true && \
+    echo 'source /opt/conda/etc/profile.d/conda.sh && conda activate fastvideo' >> /root/.bashrc && \
+    echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
+
+# Install STA (Sliding Tile Attention)
+RUN source /opt/conda/etc/profile.d/conda.sh && \
+    conda activate fastvideo && \
+    cd csrc/attn/sliding_tile_attn && \
+    git submodule update --init --recursive && \
+    python setup.py install
+
+# Install VSA 
+RUN source /opt/conda/etc/profile.d/conda.sh && \
+    conda activate fastvideo && \
+    cd csrc/attn/video_sparse_attn && \
+    git submodule update --init --recursive && \
+    python setup.py install
+
+EXPOSE 22
+
+ENTRYPOINT ["/bin/bash", "-lc", "source /opt/conda/etc/profile.d/conda.sh && conda activate fastvideo && exec /FastVideo/examples/inference/gradio/start.sh"]

docs/design/overview.md

@@ -106,8 +106,6 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward
     return batch
 ```
 
-![Pipeline execution and data flow](../assets/images/pipeline.png)
-
 ### ForwardBatch
 
 Defined in `fastvideo/pipelines/pipeline_batch_info.py`, `ForwardBatch` encapsulates the data payload passed between pipeline stages. It typically holds:
@@ -210,10 +208,6 @@ def step(
     return prev_sample
 ```
 
-This diagram shows how models are discovered, validated, and loaded across entrypoints, executors, pipelines, and model loaders.
-
-![Model loading flow](../assets/images/load_models.png)
-
 ## Optimized Attention
 
 The `fastvideo/attention/` directory contains optimized attention implementations crucial for efficient video diffusion:
@@ -237,8 +231,6 @@ self.attn = LocalAttention(
 # export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
 ```
 
-![Attention backend selector design](../assets/images/attention_backend.png)
-
 ### Attention Patterns
 Supports various patterns with memory optimization techniques:
 - **Cross/Self/Temporal/Global-Local Attention**

examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py

@@ -1,46 +0,0 @@
-# NOTE: This is still a work in progress, and the checkpoints are not released yet.
-
-from fastvideo import VideoGenerator, SamplingParam
-import json
-# from fastvideo.configs.sample import SamplingParam
-
-OUTPUT_PATH = "video_samples_self_forcing_causal_wan2_2_14B_i2v"
-def main():
-    # FastVideo will automatically use the optimal default arguments for the
-    # model.
-    # If a local path is provided, FastVideo will make a best effort
-    # attempt to identify the optimal arguments.
-    generator = VideoGenerator.from_pretrained(
-        "rand0nmr/SFWan2.2-T2V-A14B-Diffusers",
-        # FastVideo will automatically handle distributed setup
-        num_gpus=1,
-        use_fsdp_inference=True,
-        dit_cpu_offload=True, # DiT need to be offloaded for MoE
-        dit_precision="fp32",
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=True,
-        dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125],
-        # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
-        pin_cpu_memory=True,
-        init_weights_from_safetensors="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_x0_CacheAppend_1333/checkpoint-442_weight_only/ema/generator_ema.safetensors",
-        init_weights_from_safetensors_2="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_x0_CacheAppend_1333/checkpoint-442_weight_only/ema/generator_ema_2.safetensors",
-        # image_encoder_cpu_offload=False,
-    )
-
-    sampling_param = SamplingParam.from_pretrained("rand0nmr/SFWan2.2-T2V-A14B-Diffusers")
-    sampling_param.num_frames = 73
-    sampling_param.width = 832
-    sampling_param.height = 480
-    sampling_param.seed = 1000
-
-    with open("prompts/mixkit_i2v.jsonl", "r") as f:
-        prompt_image_pairs = json.load(f)
-
-    for prompt_image_pair in prompt_image_pairs:
-        prompt = prompt_image_pair["prompt"]
-        image_path = prompt_image_pair["image_path"]
-        _ = generator.generate_video(prompt, image_path=image_path, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
-
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py

@@ -1,43 +0,0 @@
-# NOTE: This is still a work in progress, and the checkpoints are not released yet.
-
-from fastvideo import VideoGenerator
-
-# from fastvideo.configs.sample import SamplingParam
-
-OUTPUT_PATH = "video_samples_self_forcing_causal_wan2_2_14B_t2v"
-def main():
-    # FastVideo will automatically use the optimal default arguments for the
-    # model.
-    # If a local path is provided, FastVideo will make a best effort
-    # attempt to identify the optimal arguments.
-    generator = VideoGenerator.from_pretrained(
-        "rand0nmr/SFWan2.2-T2V-A14B-Diffusers",
-        # FastVideo will automatically handle distributed setup
-        num_gpus=1,
-        use_fsdp_inference=True,
-        dit_cpu_offload=True, # DiT need to be offloaded for MoE
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=True,
-        dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125],
-        # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
-        pin_cpu_memory=True,
-        init_weights_from_safetensors="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_inference_transformer/",
-        init_weights_from_safetensors_2="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_2_inference_transformer/",
-        num_frame_per_block=7,
-        # image_encoder_cpu_offload=False,
-    )
-
-    # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
-    # sampling_param.num_frames = 45
-    # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
-    # Generate videos with the same simple API, regardless of GPU count
-    prompt = (
-        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
-        "wide with interest. The playful yet serene atmosphere is complemented by soft "
-        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
-    )
-    _ = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, num_frames=81)
-
-
-if __name__ == "__main__":
-    main()

examples/inference/gradio/serving/ray_serve_backend.py

@@ -200,7 +200,7 @@ class BaseModelDeployment:
 
 
 @serve.deployment(
-    ray_actor_options={"num_cpus": 2, "num_gpus": 1, "runtime_env": {"conda": "fv"}},
+    ray_actor_options={"num_cpus": 15, "num_gpus": 1, "runtime_env": {"conda": "fastvideo"}},
 )
 class T2VModelDeployment(BaseModelDeployment):
     def __init__(self, t2v_model_path: str, output_path: str = "outputs"):
@@ -210,7 +210,7 @@ class T2VModelDeployment(BaseModelDeployment):
 
 
 @serve.deployment(
-    ray_actor_options={"num_cpus": 16, "num_gpus": 1, "runtime_env": {"conda": "fv"}},
+    ray_actor_options={"num_cpus": 16, "num_gpus": 1, "runtime_env": {"conda": "fastvideo"}},
 )
 class T2V14BModelDeployment(BaseModelDeployment):
     def __init__(self, t2v_14b_model_path: str, output_path: str = "outputs"):
@@ -227,7 +227,7 @@ app.state.limiter = limiter
 app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
 
 
-@serve.deployment(num_replicas=50, ray_actor_options={"num_cpus": 2})
+@serve.deployment(num_replicas=1, ray_actor_options={"num_cpus": 2})
 @serve.ingress(app)
 class FastVideoAPI:
 

examples/inference/gradio/serving/start.sh

@@ -1,3 +1,3 @@
-python examples/inference/gradio/start_ray_serve_app.py \
---t2v_model_paths "FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers" \
---t2v_model_replicas "4,4"
+python examples/inference/gradio/serving/start_ray_serve_app.py \
+--t2v_model_paths "FastVideo/FastWan2.1-T2V-1.3B-Diffusers" \
+--t2v_model_replicas "1"

fastvideo/configs/pipelines/registry.py

@@ -14,7 +14,7 @@ from fastvideo.configs.pipelines.wan import (
     FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config,
     Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config, Wan2_2_TI2V_5B_Config,
     WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig,
-    SelfForcingWanT2V480PConfig, WANV2VConfig, SelfForcingWan2_2_T2V480PConfig)
+    SelfForcingWanT2V480PConfig, WANV2VConfig)
 # isort: on
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
@@ -38,7 +38,6 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig,
     "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers": WanT2V720PConfig,
     "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig,
-    "rand0nmr/SFWan2.2-T2V-A14B-Diffusers": SelfForcingWan2_2_T2V480PConfig,
     "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_Config,
     "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_Config,
     "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_Config,

fastvideo/configs/pipelines/wan.py

@@ -176,17 +176,3 @@ class SelfForcingWanT2V480PConfig(WanT2V480PConfig):
     dmd_denoising_steps: list[int] | None = field(
         default_factory=lambda: [1000, 750, 500, 250])
     warp_denoising_step: bool = True
-
-
-@dataclass
-class SelfForcingWan2_2_T2V480PConfig(Wan2_2_T2V_A14B_Config):
-    is_causal: bool = True
-    flow_shift: float | None = 12.0
-    boundary_ratio: float | None = 0.875
-    dmd_denoising_steps: list[int] | None = field(
-        default_factory=lambda: [1000, 850, 700, 550, 350, 275, 200, 125])
-    warp_denoising_step: bool = True
-
-    def __post_init__(self) -> None:
-        self.vae_config.load_encoder = True
-        self.vae_config.load_decoder = True

fastvideo/configs/sample/registry.py

@@ -11,7 +11,7 @@ from fastvideo.configs.sample.cosmos import Cosmos_Predict2_2B_Video2World_Sampl
 
 # isort: off
 from fastvideo.configs.sample.wan import (
-    FastWanT2V480P_SamplingParam,
+    FastWanT2V480PConfig,
     Wan2_1_Fun_1_3B_InP_SamplingParam,
     Wan2_2_I2V_A14B_SamplingParam,
     Wan2_2_T2V_A14B_SamplingParam,
@@ -21,8 +21,7 @@ from fastvideo.configs.sample.wan import (
     WanT2V_1_3B_SamplingParam,
     WanT2V_14B_SamplingParam,
     Wan2_1_Fun_1_3B_Control_SamplingParam,
-    SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam,
-    SelfForcingWan2_2_T2V_A14B_480P_SamplingParam,
+    SelfForcingWanT2V480PConfig,
 )
 # isort: on
 from fastvideo.logger import init_logger
@@ -65,7 +64,7 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
 
     # FastWan2.1
     "FastVideo/FastWan2.1-T2V-1.3B-Diffusers":
-    FastWanT2V480P_SamplingParam,
+    FastWanT2V480PConfig,
 
     # FastWan2.2
     "FastVideo/FastWan2.2-TI2V-5B-Diffusers":
@@ -73,16 +72,11 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
 
     # Causal Self-Forcing Wan2.1
     "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers":
-    SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam,
-
-    # Causal Self-Forcing Wan2.2
-    "rand0nmr/SFWan2.2-T2V-A14B-Diffusers":
-    SelfForcingWan2_2_T2V_A14B_480P_SamplingParam,
+    SelfForcingWanT2V480PConfig,
 
     # Cosmos2
     "nvidia/Cosmos-Predict2-2B-Video2World":
     Cosmos_Predict2_2B_Video2World_SamplingParam,
-
     # Add other specific weight variants
 }
 
@@ -92,8 +86,6 @@ SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = {
     "wanpipeline": lambda id: "wanpipeline" in id.lower(),
     "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
     "stepvideo": lambda id: "stepvideo" in id.lower(),
-    "wandmdpipeline": lambda id: "wandmdpipeline" in id.lower(),
-    "wancausaldmdpipeline": lambda id: "wancausaldmdpipeline" in id.lower(),
     # Add other pipeline architecture detectors
 }
 
@@ -104,9 +96,7 @@ SAMPLING_FALLBACK_PARAM: dict[str, Any] = {
     "wanpipeline":
     WanT2V_1_3B_SamplingParam,  # Base Wan config as fallback for any Wan variant
     "wanimagetovideo": WanI2V_14B_480P_SamplingParam,
-    "wandmdpipeline": FastWanT2V480P_SamplingParam,
-    "wancausaldmdpipeline": SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam,
-    "stepvideo": StepVideoT2VSamplingParam,
+    "stepvideo": StepVideoT2VSamplingParam
     # Other fallbacks by architecture
 }
 

fastvideo/configs/sample/wan.py

@@ -97,7 +97,7 @@ class WanI2V_14B_720P_SamplingParam(WanT2V_14B_SamplingParam):
 
 
 @dataclass
-class FastWanT2V480P_SamplingParam(WanT2V_1_3B_SamplingParam):
+class FastWanT2V480PConfig(WanT2V_1_3B_SamplingParam):
     # DMD parameters
     # dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 757, 522])
     num_inference_steps: int = 3
@@ -183,16 +183,5 @@ class Wan2_2_Fun_A14B_Control_SamplingParam(
 # ============= Causal Self-Forcing =============
 # =============================================
 @dataclass
-class SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam(
-        Wan2_1_Fun_1_3B_InP_SamplingParam):
+class SelfForcingWanT2V480PConfig(WanT2V_1_3B_SamplingParam):
     pass
-
-
-@dataclass
-class SelfForcingWan2_2_T2V_A14B_480P_SamplingParam(
-        Wan2_2_T2V_A14B_SamplingParam):
-    num_inference_steps: int = 8
-    num_frames: int = 81
-    height: int = 448
-    width: int = 832
-    fps: int = 16

fastvideo/models/loader/fsdp_load.py

@@ -4,8 +4,6 @@
 # Copyright 2024 The TorchTune Authors.
 # Copyright 2025 The FastVideo Authors.
 
-from __future__ import annotations
-import os
 import contextlib
 from collections.abc import Callable, Generator
 from itertools import chain
@@ -199,11 +197,10 @@ def shard_model(
     Raises:
         ValueError: If no layer modules were sharded, indicating that no shard_condition was triggered.
     """
-    # Check if we should use size-based filtering
-    use_size_filtering = os.environ.get("FASTVIDEO_FSDP2_AUTOWRAP", "0") == "1"
-    
-    if not fsdp_shard_conditions:
-        logger.warning("No FSDP shard conditions provided; nothing will be sharded.")
+    if fsdp_shard_conditions is None or len(fsdp_shard_conditions) == 0:
+        logger.warning(
+            "The FSDP shard condition list is empty or None. No modules will be sharded in %s",
+            type(model).__name__)
         return
 
     fsdp_kwargs = {
@@ -218,38 +215,20 @@ def shard_model(
     # iterating in reverse to start with
     # lowest-level modules first
     num_layers_sharded = 0
-    
-    if use_size_filtering:
-        # Size-based filtering mode
-        min_params = int(os.environ.get("FASTVIDEO_FSDP2_MIN_PARAMS", "10000000"))
-        logger.info("Using size-based filtering with threshold: %.2fM", min_params / 1e6)
-        
-        for n, m in reversed(list(model.named_modules())):
-            if any([shard_condition(n, m) for shard_condition in fsdp_shard_conditions]):
-                # Count all parameters
-                param_count = sum(p.numel() for p in m.parameters(recurse=True))
-                
-                # Skip small modules
-                if param_count < min_params:
-                    logger.info("Skipping module %s (%.2fM params < %.2fM threshold)", 
-                               n, param_count / 1e6, min_params / 1e6)
-                    continue
-                
-                # Shard this module
-                logger.info("Sharding module %s (%.2fM params)", n, param_count / 1e6)
-                fully_shard(m, **fsdp_kwargs)
-                num_layers_sharded += 1
-    else:
-        # Shard all modules matching conditions        
-        for n, m in reversed(list(model.named_modules())):
-            if any([shard_condition(n, m) for shard_condition in fsdp_shard_conditions]):
-                fully_shard(m, **fsdp_kwargs)
-                num_layers_sharded += 1
-        
-        if num_layers_sharded == 0:
-            raise ValueError(
-                "No layer modules were sharded. Please check if shard conditions are working as expected."
-            )
+    # TODO(will): don't reshard after forward for the last layer to save on the
+    # all-gather that will immediately happen Shard the model with FSDP,
+    for n, m in reversed(list(model.named_modules())):
+        if any([
+                shard_condition(n, m)
+                for shard_condition in fsdp_shard_conditions
+        ]):
+            fully_shard(m, **fsdp_kwargs)
+            num_layers_sharded += 1
+
+    if num_layers_sharded == 0:
+        raise ValueError(
+            "No layer modules were sharded. Please check if shard conditions are working as expected."
+        )
 
     # Finally shard the entire model to account for any stragglers
     fully_shard(model, **fsdp_kwargs)

fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py

@@ -86,12 +86,6 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):
             return (prev_sample, )
         return SelfForcingFlowMatchSchedulerOutput(prev_sample=prev_sample)
 
-    @staticmethod
-    def calculate_alpha_beta_high(sigma, sigma_bound):
-        alpha = (1 - sigma) / (1 - sigma_bound)
-        beta = torch.sqrt(sigma ** 2 - (alpha * sigma_bound) ** 2)
-        return alpha, beta
-        
     def add_noise(self, original_samples, noise, timestep):
         """
         Diffusion forward corruption process.
@@ -111,32 +105,6 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):
         sample = (1 - sigma) * original_samples + sigma * noise
         return sample.type_as(noise)
 
-    def add_noise_high(self, original_samples, noise, timestep, boundary_timestep):
-        """
-        Diffusion forward corruption process.
-        Input:
-            - clean_latent: the clean latent with shape [B*T, C, H, W]
-            - noise: the noise with shape [B*T, C, H, W]
-            - timestep: the timestep with shape [B*T]
-        Output: the corrupted latent with shape [B*T, C, H, W]
-        """
-        if timestep.ndim == 2:
-            timestep = timestep.flatten(0, 1)
-        if boundary_timestep.ndim == 2:
-            boundary_timestep = boundary_timestep.flatten(0, 1)
-        self.sigmas = self.sigmas.to(noise.device)
-        self.timesteps = self.timesteps.to(noise.device)
-        timestep_id = torch.argmin(
-            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
-        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
-
-        boundary_timestep_id = torch.argmin(
-            (self.timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1)
-        sigma_boundary = self.sigmas[boundary_timestep_id].reshape(-1, 1, 1, 1)
-        alpha, beta = self.calculate_alpha_beta_high(sigma, sigma_boundary)
-        sample = alpha * original_samples + beta * noise
-        return sample.type_as(noise)
-
     def training_target(self, sample, noise, timestep):
         target = noise - sample
         return target

fastvideo/models/utils.py

@@ -180,51 +180,3 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor,
     sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
     pred_video = noise_input_latent - sigma_t * pred_noise
     return pred_video.to(dtype)
-
-def pred_noise_to_x_bound(pred_noise: torch.Tensor,
-                             noise_input_latent: torch.Tensor,
-                             timestep: torch.Tensor,
-                             boundary_timestep: torch.Tensor,
-                             scheduler: Any) -> torch.Tensor:
-    """
-    Convert predicted noise to clean latent.
-
-    Args:
-    pred_noise: the predicted noise with shape [B, C, H, W]
-        where B is batch_size or batch_size * num_frames
-    noise_input_latent: the noisy latent with shape [B, C, H, W],
-    timestep: the timestep with shape [1] or [bs * num_frames] or [bs, num_frames]
-    boundary_timestep: the boundary timestep with shape [1] or [bs * num_frames] or [bs, num_frames]
-    scheduler: the scheduler
-
-    Returns:
-        the predicted video with shape [B, C, H, W]
-    """
-    # If timestep is [bs, num_frames]
-    if timestep.ndim == 2:
-        timestep = timestep.flatten(0, 1)
-        assert timestep.numel() == noise_input_latent.shape[0]
-    elif timestep.ndim == 1:
-        # If timestep is [1]
-        if timestep.shape[0] == 1:
-            timestep = timestep.expand(noise_input_latent.shape[0])
-        else:
-            assert timestep.numel() == noise_input_latent.shape[0]
-    else:
-        raise ValueError(f"[pred_noise_to_pred_video] Invalid timestep shape: {timestep.shape}")
-    # timestep shape should be [B]
-    dtype = pred_noise.dtype
-    device = pred_noise.device
-    pred_noise = pred_noise.double().to(device)
-    noise_input_latent = noise_input_latent.double().to(device)
-    sigmas = scheduler.sigmas.double().to(device)
-    timesteps = scheduler.timesteps.double().to(device)
-    timestep_id = torch.argmin(
-        (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
-    sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
-
-    boundary_timestep_id = torch.argmin(
-        (timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1)
-    sigma_t_boundary = sigmas[boundary_timestep_id].reshape(-1, 1, 1, 1)
-    pred_video = noise_input_latent - (sigma_t - sigma_t_boundary) * pred_noise
-    return pred_video.to(dtype)

fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py

@@ -50,8 +50,7 @@ class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
                        stage=CausalDMDDenosingStage(
                            transformer=self.get_module("transformer"),
                            transformer_2=self.get_module("transformer_2", None),
-                           scheduler=self.get_module("scheduler"),
-                           vae=self.get_module("vae")))
+                           scheduler=self.get_module("scheduler")))
 
         self.add_stage(stage_name="decoding_stage",
                        stage=DecodingStage(vae=self.get_module("vae")))

fastvideo/pipelines/stages/causal_denoising.py

@@ -4,7 +4,7 @@ from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
-from fastvideo.models.utils import pred_noise_to_pred_video, pred_noise_to_x_bound
+from fastvideo.models.utils import pred_noise_to_pred_video
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages.denoising import DenoisingStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
@@ -34,16 +34,13 @@ class CausalDMDDenosingStage(DenoisingStage):
     Denoising stage for causal diffusion.
     """
 
-    def __init__(self,
-                 transformer,
-                 scheduler,
-                 transformer_2=None,
-                 vae=None) -> None:
+    def __init__(self, transformer, scheduler, transformer_2=None) -> None:
         super().__init__(transformer, scheduler, transformer_2)
         # KV and cross-attention cache state (initialized on first forward)
         self.transformer = transformer
         self.transformer_2 = transformer_2
-        self.vae = vae
+        self.kv_cache1: list | None = None
+        self.crossattn_cache: list | None = None
         # Model-dependent constants (aligned with causal_inference.py assumptions)
         self.num_transformer_blocks = len(self.transformer.blocks)
         self.num_frames_per_block = self.transformer.config.arch_config.num_frames_per_block
@@ -83,13 +80,6 @@ class CausalDMDDenosingStage(DenoisingStage):
             timesteps = scheduler_timesteps[1000 - timesteps]
         timesteps = timesteps.to(get_local_torch_device())
 
-        if fastvideo_args.pipeline_config.dit_config.boundary_ratio is not None:
-            boundary_timestep = fastvideo_args.pipeline_config.dit_config.boundary_ratio * self.scheduler.num_train_timesteps
-        else:
-            boundary_timestep = None
-
-        high_noise_timesteps = timesteps[timesteps >= boundary_timestep]
-
         # Image kwargs (kept empty unless caller provides compatible args)
         image_kwargs: dict = {}
 
@@ -113,101 +103,113 @@ class CausalDMDDenosingStage(DenoisingStage):
         assert torch.isnan(prompt_embeds[0]).sum() == 0
 
         # Initialize or reset caches
-        kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0],
-                                              dtype=target_dtype,
-                                              device=latents.device)
-        kv_cache2 = None
-        if boundary_timestep is not None:
-            # Initialize the low noise kv cache
-            kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0],
-                                                  dtype=target_dtype,
-                                                  device=latents.device)
+        if self.kv_cache1 is None:
+            self._initialize_kv_cache(batch_size=latents.shape[0],
+                                      dtype=target_dtype,
+                                      device=latents.device)
+            self._initialize_crossattn_cache(
+                batch_size=latents.shape[0],
+                max_text_len=fastvideo_args.pipeline_config.
+                text_encoder_configs[0].arch_config.text_len,
+                dtype=target_dtype,
+                device=latents.device)
+        else:
+            assert self.crossattn_cache is not None
+            # reset cross-attention cache
+            for block_index in range(self.num_transformer_blocks):
+                self.crossattn_cache[block_index][
+                    "is_init"] = False  # type: ignore
+            # reset kv cache pointers
+            for block_index in range(len(self.kv_cache1)):
+                self.kv_cache1[block_index][
+                    "global_end_index"] = torch.tensor(  # type: ignore
+                        [0],
+                        dtype=torch.long,
+                        device=latents.device)
+                self.kv_cache1[block_index][
+                    "local_end_index"] = torch.tensor(  # type: ignore
+                        [0],
+                        dtype=torch.long,
+                        device=latents.device)
 
-        def _get_kv_cache(timestep: float) -> dict:
-            if boundary_timestep is not None:
-                if timestep >= boundary_timestep:
-                    return kv_cache1
-                else:
-                    return kv_cache2
-            return kv_cache1
-
-        crossattn_cache = self._initialize_crossattn_cache(
-            batch_size=latents.shape[0],
-            max_text_len=fastvideo_args.pipeline_config.text_encoder_configs[0].
-            arch_config.text_len,
-            dtype=target_dtype,
-            device=latents.device)
-
-        pos_start_base = 0
-
-        # Determine block sizes
-        block_sizes = [self.num_frames_per_block] * 7
-        block_sizes[0] = 1
-        start_index = 0
-
-        first_frame_latent = None
-        if batch.pil_image is not None:
-            # Causal video gen directly replaces the first frame of the latent with
-            # the image latent instead of appending along the channel dim
-            assert self.vae is not None, "VAE is not provided for causal video gen task"
-            self.vae = self.vae.to(get_local_torch_device())
-            first_frame_latent = self.vae.encode(batch.pil_image).mean.float()
-            if (hasattr(self.vae, "shift_factor")
-                    and self.vae.shift_factor is not None):
-                if isinstance(self.vae.shift_factor, torch.Tensor):
-                    first_frame_latent -= self.vae.shift_factor.to(
-                        first_frame_latent.device, first_frame_latent.dtype)
-                else:
-                    first_frame_latent -= self.vae.shift_factor
-
-            if isinstance(self.vae.scaling_factor, torch.Tensor):
-                first_frame_latent = first_frame_latent * self.vae.scaling_factor.to(
-                    first_frame_latent.device, first_frame_latent.dtype)
-            else:
-                first_frame_latent = first_frame_latent * self.vae.scaling_factor
-
-            if fastvideo_args.vae_cpu_offload:
-                self.vae = self.vae.to("cpu")
-
-            # Fill the low noise and high noise kv cache with first_frame_latent and timestep 0
-            t_zero = torch.zeros([latents.shape[0], 1],
+        # Optional: cache context features from provided image latents prior to generation
+        current_start_frame = 0
+        if getattr(batch, "image_latent", None) is not None:
+            image_latent = batch.image_latent
+            assert image_latent is not None
+            input_frames = image_latent.shape[2]
+            # timestep zero (or configured context noise) for cache warm-up
+            t_zero = torch.zeros([latents.shape[0]],
                                  device=latents.device,
                                  dtype=torch.long)
-            with torch.autocast(device_type="cuda",
-                                dtype=target_dtype,
-                                enabled=autocast_enabled), \
-                set_forward_context(current_timestep=0,
-                                    attn_metadata=None,
-                                    forward_batch=batch):
-                self.transformer(
-                    first_frame_latent.to(target_dtype),
-                    prompt_embeds,
-                    t_zero,
-                    kv_cache=kv_cache1,
-                    crossattn_cache=crossattn_cache,
-                    current_start=(pos_start_base + start_index) *
-                    self.frame_seq_length,
-                    start_frame=start_index,
-                    **image_kwargs,
-                    **pos_cond_kwargs,
-                )
-                if boundary_timestep is not None:
-                    self.transformer_2(
-                        first_frame_latent.to(target_dtype),
+            if independent_first_frame and input_frames >= 1:
+                # warm-up with the very first frame independently
+                image_first_btchw = image_latent[:, :, :1, :, :].to(
+                    target_dtype).permute(0, 2, 1, 3, 4)
+                with torch.autocast(device_type="cuda",
+                                    dtype=target_dtype,
+                                    enabled=autocast_enabled):
+                    _ = self.transformer(
+                        image_first_btchw,
                         prompt_embeds,
                         t_zero,
-                        kv_cache=kv_cache2,
-                        crossattn_cache=crossattn_cache,
-                        current_start=(pos_start_base + start_index) *
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=current_start_frame *
                         self.frame_seq_length,
-                        start_frame=start_index,
                         **image_kwargs,
                         **pos_cond_kwargs,
                     )
+                current_start_frame += 1
+                remaining_frames = input_frames - 1
+            else:
+                remaining_frames = input_frames
 
-            start_index += 1
-            block_sizes.pop(0)
-            latents[:, :, :1, :, :] = first_frame_latent
+            # process remaining input frames in blocks of num_frame_per_block
+            while remaining_frames > 0:
+                block = min(self.num_frames_per_block, remaining_frames)
+                ref_btchw = image_latent[:, :, current_start_frame:
+                                         current_start_frame +
+                                         block, :, :].to(target_dtype).permute(
+                                             0, 2, 1, 3, 4)
+                with torch.autocast(device_type="cuda",
+                                    dtype=target_dtype,
+                                    enabled=autocast_enabled):
+                    _ = self.transformer(
+                        ref_btchw,
+                        prompt_embeds,
+                        t_zero,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=current_start_frame *
+                        self.frame_seq_length,
+                        **image_kwargs,
+                        **pos_cond_kwargs,
+                    )
+                current_start_frame += block
+                remaining_frames -= block
+
+        # Base position offset from any cache warm-up
+        pos_start_base = current_start_frame
+
+        # Determine block sizes
+        if not independent_first_frame or (independent_first_frame
+                                           and batch.image_latent is not None):
+            if t % self.num_frames_per_block != 0:
+                raise ValueError(
+                    "num_frames must be divisible by num_frames_per_block for causal DMD denoising"
+                )
+            num_blocks = t // self.num_frames_per_block
+            block_sizes = [self.num_frames_per_block] * num_blocks
+            start_index = 0
+        else:
+            if (t - 1) % self.num_frames_per_block != 0:
+                raise ValueError(
+                    "(num_frames - 1) must be divisible by num_frame_per_block when independent_first_frame=True"
+                )
+            num_blocks = (t - 1) // self.num_frames_per_block
+            block_sizes = [1] + [self.num_frames_per_block] * num_blocks
+            start_index = 0
 
         # DMD loop in causal blocks
         with self.progress_bar(total=len(block_sizes) *
@@ -220,7 +222,7 @@ class CausalDMDDenosingStage(DenoisingStage):
                 video_raw_latent_shape = noise_latents_btchw.shape
 
                 for i, t_cur in enumerate(timesteps):
-                    if boundary_timestep is not None and t_cur < boundary_timestep:
+                    if self.transformer_2 is not None and fastvideo_args.pipeline_config.dit_config.boundary_ratio is not None and t_cur < fastvideo_args.pipeline_config.dit_config.boundary_ratio * self.scheduler.num_train_timesteps:
                         current_model = self.transformer_2
                     else:
                         current_model = self.transformer
@@ -278,8 +280,8 @@ class CausalDMDDenosingStage(DenoisingStage):
                             latent_model_input,
                             prompt_embeds,
                             t_expanded_noise,
-                            kv_cache=_get_kv_cache(t_cur),
-                            crossattn_cache=crossattn_cache,
+                            kv_cache=self.kv_cache1,
+                            crossattn_cache=self.crossattn_cache,
                             current_start=(pos_start_base + start_index) *
                             self.frame_seq_length,
                             start_frame=start_index,
@@ -288,22 +290,12 @@ class CausalDMDDenosingStage(DenoisingStage):
                         ).permute(0, 2, 1, 3, 4)
 
                     # Convert pred noise to pred video with FM Euler scheduler utilities
-                    if boundary_timestep is not None and t_cur >= boundary_timestep:
-                        pred_video_btchw = pred_noise_to_x_bound(
-                            pred_noise=pred_noise_btchw.flatten(0, 1),
-                            noise_input_latent=noise_latents.flatten(0, 1),
-                            timestep=t_expand,
-                            boundary_timestep=torch.ones_like(t_expand) *
-                            boundary_timestep,
-                            scheduler=self.scheduler).unflatten(
-                                0, pred_noise_btchw.shape[:2])
-                    else:
-                        pred_video_btchw = pred_noise_to_pred_video(
-                            pred_noise=pred_noise_btchw.flatten(0, 1),
-                            noise_input_latent=noise_latents.flatten(0, 1),
-                            timestep=t_expand,
-                            scheduler=self.scheduler).unflatten(
-                                0, pred_noise_btchw.shape[:2])
+                    pred_video_btchw = pred_noise_to_pred_video(
+                        pred_noise=pred_noise_btchw.flatten(0, 1),
+                        noise_input_latent=noise_latents.flatten(0, 1),
+                        timestep=t_expand,
+                        scheduler=self.scheduler).unflatten(
+                            0, pred_noise_btchw.shape[:2])
 
                     if i < len(timesteps) - 1:
                         next_timestep = timesteps[i + 1] * torch.ones(
@@ -317,23 +309,11 @@ class CausalDMDDenosingStage(DenoisingStage):
                                 batch.generator, list) else
                                        batch.generator)).to(self.device)
                         noise_btchw = noise
-                        if boundary_timestep is not None and i < len(
-                                high_noise_timesteps) - 1:
-                            noise_latents_btchw = self.scheduler.add_noise_high(
-                                pred_video_btchw.flatten(0, 1),
-                                noise_btchw.flatten(0, 1), next_timestep,
-                                torch.ones_like(next_timestep) *
-                                boundary_timestep).unflatten(
-                                    0, pred_video_btchw.shape[:2])
-                        elif boundary_timestep is not None and i == len(
-                                high_noise_timesteps) - 1:
-                            noise_latents_btchw = pred_video_btchw
-                        else:
-                            noise_latents_btchw = self.scheduler.add_noise(
-                                pred_video_btchw.flatten(0, 1),
-                                noise_btchw.flatten(0, 1),
-                                next_timestep).unflatten(
-                                    0, pred_video_btchw.shape[:2])
+                        noise_latents_btchw = self.scheduler.add_noise(
+                            pred_video_btchw.flatten(0, 1),
+                            noise_btchw.flatten(0, 1),
+                            next_timestep).unflatten(0,
+                                                     pred_video_btchw.shape[:2])
                         current_latents = noise_latents_btchw.permute(
                             0, 2, 1, 3, 4)
                     else:
@@ -361,34 +341,18 @@ class CausalDMDDenosingStage(DenoisingStage):
                                         attn_metadata=attn_metadata,
                                         forward_batch=batch):
                     t_expanded_context = t_context.unsqueeze(1)
-
-                    if boundary_timestep is not None:
-                        self.transformer_2(
-                            context_bcthw,
-                            prompt_embeds,
-                            t_expanded_context,
-                            kv_cache=kv_cache2,
-                            crossattn_cache=crossattn_cache,
-                            current_start=(pos_start_base + start_index) *
-                            self.frame_seq_length,
-                            start_frame=start_index,
-                            **image_kwargs,
-                            **pos_cond_kwargs,
-                        )
-
-                    self.transformer(
+                    _ = current_model(
                         context_bcthw,
                         prompt_embeds,
                         t_expanded_context,
-                        kv_cache=kv_cache1,
-                        crossattn_cache=crossattn_cache,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
                         current_start=(pos_start_base + start_index) *
                         self.frame_seq_length,
                         start_frame=start_index,
                         **image_kwargs,
                         **pos_cond_kwargs,
                     )
-
                 start_index += current_num_frames
 
         batch.latents = latents
@@ -428,7 +392,7 @@ class CausalDMDDenosingStage(DenoisingStage):
                 torch.tensor([0], dtype=torch.long, device=device),
             })
 
-        return kv_cache1
+        self.kv_cache1 = kv_cache1
 
     def _initialize_crossattn_cache(self, batch_size, max_text_len, dtype,
                                     device) -> None:
@@ -457,7 +421,7 @@ class CausalDMDDenosingStage(DenoisingStage):
                 "is_init":
                 False,
             })
-        return crossattn_cache
+        self.crossattn_cache = crossattn_cache
 
     def verify_input(self, batch: ForwardBatch,
                      fastvideo_args: FastVideoArgs) -> VerificationResult:

fastvideo/pipelines/stages/input_validation.py

@@ -5,6 +5,7 @@ Input validation stage for diffusion pipelines.
 
 import torch
 import torchvision.transforms.functional as TF
+from PIL import Image
 
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
@@ -13,6 +14,7 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.validators import (StageValidators,
                                                    VerificationResult)
+from fastvideo.utils import best_output_size
 
 logger = init_logger(__name__)
 
@@ -104,31 +106,27 @@ class InputValidationStage(PipelineStage):
             batch.pil_image = image
 
         # further processing for ti2v task
-        if (fastvideo_args.pipeline_config.ti2v_task
-                or fastvideo_args.pipeline_config.is_causal
-            ) and batch.pil_image is not None:
+        if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
             img = batch.pil_image
-            # ih, iw = img.height, img.width
-            # patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size
-            # vae_stride = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial
-            # dh, dw = patch_size[1] * vae_stride, patch_size[2] * vae_stride
-            # max_area = 720 * 1280
-            # ow, oh = best_output_size(iw, ih, dw, dh, max_area)
+            ih, iw = img.height, img.width
+            patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size
+            vae_stride = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial
+            dh, dw = patch_size[1] * vae_stride, patch_size[2] * vae_stride
+            max_area = 704 * 1280
+            ow, oh = best_output_size(iw, ih, dw, dh, max_area)
 
-            # scale = max(ow / iw, oh / ih)
-            # img = img.resize((round(iw * scale), round(ih * scale)),
-            #                  Image.LANCZOS)
-            # logger.info("resized img height: %s, img width: %s", img.height,
-            #             img.width)
+            scale = max(ow / iw, oh / ih)
+            img = img.resize((round(iw * scale), round(ih * scale)),
+                             Image.LANCZOS)
+            logger.info("resized img height: %s, img width: %s", img.height,
+                        img.width)
+
+            # center-crop
+            x1 = (img.width - ow) // 2
+            y1 = (img.height - oh) // 2
+            img = img.crop((x1, y1, x1 + ow, y1 + oh))
+            assert img.width == ow and img.height == oh
 
-            # # center-crop
-            # x1 = (img.width - ow) // 2
-            # y1 = (img.height - oh) // 2
-            # img = img.crop((x1, y1, x1 + ow, y1 + oh))
-            # assert img.width == ow and img.height == oh
-            logger.info("img height: %s, img width: %s", img.height, img.width)
-            oh = img.height
-            ow = img.width
             # to tensor
             img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(
                 self.device).unsqueeze(1)

fastvideo/training/distillation_pipeline.py

@@ -729,6 +729,9 @@ class DistillationPipeline(TrainingPipeline):
                                      self.num_train_timestep, [1],
                                      device=self.device,
                                      dtype=torch.long)
+            world_group = get_world_group()
+            if world_group.world_size > 1:
+                world_group.broadcast(timestep, src=0)
 
             timestep = shift_timestep(
                 timestep,
@@ -841,6 +844,9 @@ class DistillationPipeline(TrainingPipeline):
                                             self.num_train_timestep, [1],
                                             device=self.device,
                                             dtype=torch.long)
+        world_group = get_world_group()
+        if world_group.world_size > 1:
+            world_group.broadcast(fake_score_timestep, src=0)
 
         fake_score_timestep = shift_timestep(
             fake_score_timestep,

fastvideo/training/self_forcing_distillation_pipeline.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 import copy
-import os
 import time
 from collections import deque
 from typing import Any
@@ -46,13 +45,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
 
     def initialize_training_pipeline(self, training_args: TrainingArgs):
         """Initialize the self-forcing training pipeline."""
-        # Check if FSDP2 auto wrap is enabled - not supported for self-forcing distillation
-        if os.environ.get("FASTVIDEO_FSDP2_AUTOWRAP", "0") == "1":
-            raise NotImplementedError(
-                "FASTVIDEO_FSDP2_AUTOWRAP is not implemented for self-forcing distillation. "
-                "Please set FASTVIDEO_FSDP2_AUTOWRAP=0 or unset the environment variable."
-            )
-
         logger.info("Initializing self-forcing distillation pipeline...")
 
         self.generator_ema: EMA_FSDP | None = None

fastvideo/worker/multiproc_executor.py

@@ -467,10 +467,6 @@ class WorkerMultiprocProc:
                             "output_batch": output_batch.output.cpu(),
                             "logging_info": logging_info
                         })
-                    else:
-                        result = self.worker.execute_method(
-                            method, *args, **kwargs)
-                        self.pipe.send(result)
                 else:
                     result = self.worker.execute_method(method, *args, **kwargs)
                     self.pipe.send(result)

mkdocs.yml

@@ -14,7 +14,6 @@ edit_uri: edit/main/docs/
 # Configuration
 theme:
   name: material
-  favicon: assets/logos/icon_simple.svg
   palette:
     - scheme: default
       toggle:

prompts/mixkit_i2v.jsonl

@@ -1,66 +0,0 @@
-[
-    {
-        "prompt": "Young man skating with a skateboard on the ramps with graffiti of a park with trees, on a sunny day.", 
-        "image_path": "images/mixkit-boy-skating-with-a-skateboard-in-a-park-with-ramps-34389.png"
-    },
-    {
-        "prompt": "In the midst of the joyous New Year's Eve celebration, the cheerful group of friends, their spirits lifted by the festivities, decides to immortalize the moment with a vibrant snapshot", 
-        "image_path": "images/mixkit-a-cheerful-group-of-friends-celebrate-new-years-eve-and-51525.png"
-    },
-    {
-        "prompt": "A man and a woman playing in a field with grass, during a bright afternoon, while cars pass by in the distance.", 
-        "image_path": "images/mixkit-a-cute-couple-playing-on-the-grass-4688.png"
-    },
-    {
-        "prompt": "Aerial view of a rocky mountain in the forest at a sunny day drone flight footage", 
-        "image_path": "images/mixkit-aerial-view-of-a-rocky-mountain-in-the-forest-50589.png"
-    },
-    {
-        "prompt": "A little girl wearing a pink security helmet and denim overall discovers the art of cycling amidst the serene park, as the camera captures her graceful progress.", 
-        "image_path": "images/mixkit-a-little-girl-cruises-through-the-forest-path-on-her-50088.png"
-    },
-    {
-        "prompt": "Aerial shot of a beach shore with sea waves. Big rocks on the sand at an alone beach.", 
-        "image_path": "images/mixkit-aerial-shot-of-a-beach-with-sea-waves-1087.png"
-    },
-    {
-        "prompt": "Young woman cleaning her house decorated with plants and decorations, while dancing happily to music in her headphones.", 
-        "image_path": "images/mixkit-woman-cleaning-her-house-dancing-happy-43379.png"
-    },
-    {
-        "prompt": "Aerial tour in a meadow surrounded by hills on the horizon, while some birds fly low over a lake.", 
-        "image_path": "images/mixkit-birds-flying-low-over-a-lake-in-a-meadow-41417.png"
-    },
-    {
-        "prompt": "In the video, a lone rider guides a majestic horse across an expansive, open field as the sun sets in the background. The rider, dressed in a classic blue shirt and wide-brimmed hat, sits confidently in the saddle, silhouetted against the warm glow of the evening sky. The horse moves gracefully, its mane and tail flowing with each step, creating a sense of harmony between horse and rider. Surrounding the pair, towering trees form a natural border, their leaves gently rustling in the breeze. The shadows lengthen on the ground, accentuating the serene and timeless feel of the scene. The distant hills and wooden fences frame the horizon, adding depth to the tranquil landscape. A few horses graze peacefully in the background, blending into the pastoral setting. The overall ambiance evokes a sense of calmness and quietude, capturing a perfect moment in the golden light of dusk.", 
-        "image_path": "images/mixkit-a-rancher-riding-a-horse-at-sunset-1143.png"
-    },
-    {
-        "prompt": "In the video, a martial artist dressed in a traditional white uniform with a black belt demonstrates a series of precise movements against a stark black background. The individual gracefully transitions between stances, embodying a sense of focused discipline and control. Each motion is executed with a deliberate pace, showcasing the fluidity of martial arts techniques. The soft lighting creates subtle highlights on the uniform, adding depth to the figure as it moves. The practitioner begins with an open-hand pose, feet firmly grounded, gradually shifting to a powerful forward punch. The fluidity of the sequence displays a mastery of balance and poise. Every trajectory of the limbs is precise and deliberate, capturing the elegance and strength of martial arts. The serene, isolated setting enhances the intensity and concentration of the practitioner. This visual presentation is an elegant interplay of motion and stillness, displaying the art form's discipline and grace.", 
-        "image_path": "images/mixkit-a-young-man-practicing-his-karate-moves-49635.png"
-    },
-    {
-        "prompt": "In a serene and softly lit yoga studio, three individuals engage in a yoga session, each performing an upward-facing stretch. The central figure is a woman with shoulder-length brown hair, dressed in a light cropped top and green leggings, her posture reflecting grace and concentration. To her right, another participant, a woman in a purple outfit, mirrors the pose with equal poise. On her left, a person with a bun focuses intently, supported slightly by yoga blocks beneath their hands. The warm-colored wooden floor contrasts soothingly with the soft pastel mural on the back wall, featuring an abstract design and partial visage of a serene face. Natural light floods the space from a large window on the right, where lush greens peek through, adding an element of tranquility. In the corner of the room, a collection of meditation instruments, including a gong and a Buddha statue, subtly frame the peaceful setting. The mood is calm yet focused, as all three participants are deeply engaged in their practice. The scene combines elements of balance, harmony, and a shared journey towards mindfulness. This depiction captures the essence of a yoga session that blends personal growth with collective experience.", 
-        "image_path": "images/mixkit-small-group-of-people-doing-yoga-together-43730.png"
-    },
-    {
-        "prompt": "In the deep blue expanse of the ocean, two dolphins glide effortlessly, their sleek bodies reflecting the sunlight filtering through the water. The prominent shadows and caustics create a shimmering effect on their skin, capturing the beauty of their natural habitat. Each dolphin moves with a fluid grace, occasionally interacting with gentle nudges, showcasing their playful and social nature. The scene is vibrant and dynamic, with the clear blue background accentuating the dolphins' movements, making it an ideal subject for AI recreation.", 
-        "image_path": "images/mixkit-dolphins-underwater-4133.png"
-    },
-    {
-        "prompt": "A bustling ski slope comes alive with skiers descending a pristine, snow-covered hill, surrounded by towering, snow-draped evergreens. Several figures stand atop the slope, silhouetted against a clear blue sky, preparing to embark on their ski run. The chair lift on the right continuously drops off eager adventurers, adding to the excitement at the hilltop. Each skier, clad in colorful winter gear, carves distinct paths into the textured snow as they weave their way down. The interplay of sunlight and shadows accentuates the myriad tracks etched into the slope, creating a dynamic visual rhythm. The scene captures a vibrant winter wonderland, full of action and the thrill of a perfect ski day.", 
-        "image_path": "images/mixkit-skiers-on-a-snowy-slope-3327.png"
-    },
-    {
-        "prompt": "A determined climber is scaling a massive rock face, showcasing exceptional strength and skill. The person, clad in a teal shirt and dark pants, climbs with precision, their movements measured and deliberate. They are secured by climbing gear, which includes ropes and a harness, emphasizing their commitment to safety. The rugged texture of the sandy-colored rock provides an imposing backdrop, adding drama and scale to the climb. In the distance, other large rock formations and sparse vegetation can be seen under a bright, overcast sky, contributing to the natural and adventurous atmosphere. The scene captures a moment of focus and challenge, highlighting the climber's tenacity and the breathtaking environment.",
-        "image_path": "images/mixkit-alpinist-climbing-a-huge-rock-in-a-desert-43306.png"
-    },
-    {
-        "prompt": "A silver SUV drives along a winding, snow-covered mountain road, with dense pine trees blanketed in snow lining both sides. The scene is serene, with the vehicle moving smoothly, possibly on a winter journey or vacation. As the SUV disappears around the bend, another, darker SUV follows, creating a sense of motion and perspective on the snow-dusted asphalt. The towering, snow-laden rock formation to the right contrasts with the dark green of the pines, highlighting the peacefulness of the wintry landscape.",
-        "image_path": "images/mixkit-curve-on-a-snowy-forest-road-3317.png"
-    },
-    {
-        "prompt": "A solitary boat glides across the expansive, tranquil expanse of a serene lake. The vessel leaves a gentle wake behind, creating delicate ripples across the mirror-like surface. The water appears a rich shade of teal, seamlessly blending with the sky at the horizon. Silhouettes of distant trees are faintly visible, creating a picturesque backdrop that enhances the solitary journey of the boat. The sky is a calm gradient, shifting from soft oranges near the shore to the pale blues above. In the distance, a few slender poles emerge from the water, remnants of an old structure or natural formation. The mood of the scene is one of peace and solitude, with the boat journeying steadily through the quiet landscape. There is a sense of endless possibilities as the boat moves toward the unseen beyond the frame. The simplicity and stillness of the scene invite contemplation and reflection, encapsulating a perfect moment of quietude on the water.",
-        "image_path": "images/mixkit-motorboat-on-a-large-lake-with-turquoise-blue-waters-4996.png"
-    }
-]