edit to attention

README.md

@@ -2,62 +2,109 @@
 <img src=assets/logo.jpg width="30%"/>
 </div>
 
-FastVideo is a lightweight framework for accelerating large video diffusion models.
+**FastVideo is a unified framework for accelerated video generation.**
+
+It features a clean, consistent API that works across popular video models, making it easier for developers to author new models and incorporate system- or kernel-level optimizations.
+With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.
 
 <p align="center">
-    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> <b>Slack</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/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> <b>Slack</b> </a> |
 </p>
 
-https://github.com/user-attachments/assets/79af5fb8-707c-4263-b153-9ab2a01d3ac1
+<div align="center">
+<img src=assets/perf.png width="90%"/>
+</div>
 
-FastVideo currently offers: (with more to come)
+## Key Features
 
-- [NEW!] V1 inference API available. Full announcement coming soon!
-- [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
-- FastHunyuan and FastMochi: consistency distilled video diffusion models for 8x inference speedup.
-- First open distillation recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
-- Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
+FastVideo has the following features:
+- State-of-the-art performance optimizations for inference
+  - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
+  - [TeaCache](https://arxiv.org/pdf/2411.19108)
+  - [Sage Attention](https://arxiv.org/abs/2410.02367)
+- Cutting edge models
+  - Wan2.1 T2V, I2V
+  - HunyuanVideo
+  - FastHunyuan: consistency distilled video diffusion models for 8x inference speedup.
+  - StepVideo T2V
+- Distillation support
+  - Recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
+  - Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
 - Scalable training with FSDP, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs.
 - Memory efficient finetuning with LoRA, precomputed latent, and precomputed text embeddings.
 
-Dev in progress and highly experimental.
-
-## Change Log
-- ```2025/02/20```: FastVideo now supports STA on [StepVideo](https://github.com/stepfun-ai/Step-Video-T2V) with 3.4X speedup!
-- ```2025/02/18```: Release the inference code and kernel for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
-- ```2025/01/13```: Support Lora finetuning for HunyuanVideo.
-- ```2024/12/25```: Enable single 4090 inference for `FastHunyuan`, please rerun the installation steps to update the environment.
-- ```2024/12/17```: `FastVideo` v0.0.1 is released.
-
 ## Getting Started
+We recommend using an environment manager such as `Conda` to create a clean environment:
+
+```bash
+# Create and activate a new conda environment
+conda create -n fastvideo python=3.12
+conda activate fastvideo
+
+# Install FastVideo
+pip install fastvideo
+```
+
+Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html) for more detailed installation instructions.
+
+## Inference
+### Generating Your First Video
+Here's a minimal example to generate a video using the default settings. Create a file called `example.py` with the following code:
+
+```python
+from fastvideo import VideoGenerator
+
+def main():
+    # Create a video generator with a pre-trained model
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,  # Adjust based on your hardware
+    )
+
+    # Define a prompt for your video
+    prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest."
+
+    # Generate the video
+    video = generator.generate_video(
+        prompt,
+        return_frames=True,  # Also return frames from this call (defaults to False)
+        output_path="my_videos/",  # Controls where videos are saved
+        save_video=True
+    )
+
+if __name__ == '__main__':
+    main()
+```
+
+Run the script with:
+
+```bash
+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.html).
+
+### Other docs:
 
-- [Install FastVideo](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html)
 - [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)
 
-### Inference
-- [Quick Start](https://hao-ai-lab.github.io/FastVideo/inference/examples/basic.html)
-- V1 Inference API Guide (Coming soon!)
-
-### Distillation and Finetuning
+## Distillation and Finetuning
 - [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/training/distillation.html)
 - [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetuning.html)
 
-### Deprecated APIs
-- [V0 Inference (Deprecated)](https://hao-ai-lab.github.io/FastVideo/inference/v0_inference.html)
-
 ## 📑 Development Plan
 
 <!-- - More distillation methods -->
   <!-- - [ ] Add Distribution Matching Distillation -->
 - More models support
   <!-- - [ ] Add CogvideoX model -->
-  - [ ] Add StepVideo to V1
+  - [x] Add StepVideo to V1
 - Optimization features
-  - [ ] Teacache in V1
-  - [ ] SageAttention in V1
+  - [x] Teacache in V1
+  - [x] SageAttention in V1
 - Code updates
-  - [ ] V1 Configuration API
+  - [x] V1 Configuration API
   - [ ] Support Training in V1
   <!-- - [ ] fp8 support -->
   <!-- - [ ] faster load model and save model support -->

docs/source/index.md

@@ -9,7 +9,7 @@
 
 :::{raw} html
 <p style="text-align:center">
-<strong>FastVideo is a lightweight framework for accelerating large video diffusion models.
+<strong>FastVideo is a unified framework for accelerated video generation.
 </strong>
 </p>
 
@@ -21,27 +21,31 @@
 </p>
 :::
 
-FastVideo is a lightweight framework for accelerating large video diffusion models developed by the [Hao AI Lab](https://hao-ai-lab.github.io/).
+It features a clean, consistent API that works across popular video models, making it easier for developers to author new models and incorporate system- or kernel-level optimizations.
+With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.
 
 <div style="text-align: center;">
-  <video controls width="800">
-    <source src="https://github.com/user-attachments/assets/79af5fb8-707c-4263-b153-9ab2a01d3ac1" type="video/mp4">
-    Your browser does not support the video tag.
-  </video>
+  <img src=_static/images/perf.png width="100%"/>
 </div>
 
-FastVideo currently offers: (with more to come)
+## Key Features
 
-- [NEW!] V1 inference API available. Full announcement coming soon!
-- [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
-- FastHunyuan and FastMochi: consistency distilled video diffusion models for 8x inference speedup.
-- First open distillation recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
-- Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
+FastVideo has the following features:
+- State-of-the-art performance optimizations for inference
+  - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
+  - [TeaCache](https://arxiv.org/pdf/2411.19108)
+  - [Sage Attention](https://arxiv.org/abs/2410.02367)
+- Cutting edge models
+  - Wan2.1 T2V, I2V
+  - HunyuanVideo
+  - FastHunyuan: consistency distilled video diffusion models for 8x inference speedup.
+  - StepVideo T2V
+- Distillation support
+  - Recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
+  - Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
 - Scalable training with FSDP, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs.
 - Memory efficient finetuning with LoRA, precomputed latent, and precomputed text embeddings.
 
-Dev in progress and highly experimental.
-
 ## Documentation
 
 % How to start using FastVideo?
@@ -63,6 +67,7 @@ inference/configuration
 inference/optimizations
 inference/support_matrix
 inference/examples/examples_inference_index
+inference/cli
 inference/add_pipeline
 inference/v0_inference
 :::

docs/source/inference/configuration.md

@@ -0,0 +1,77 @@
+(inference-configuration)=
+# Configuration
+
+## Multi-GPU Setup
+
+FastVideo automatically distributes the generation process when multiple GPUs are specified:
+
+```python
+# Will use 4 GPUs in parallel for faster generation
+generator = VideoGenerator.from_pretrained(
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    num_gpus=4,
+)
+```
+
+## Customizing Generation
+
+- `PipelineConfig`: Initialization time parameters
+- `SamplingParam`: Generation time parameters
+
+You can customize various parameters when generating videos using the `PipelineConfig` and `SamplingParam` class:
+
+```python
+from fastvideo import VideoGenerator, SamplingParam, PipelineConfig
+
+def main():
+    model_name = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+    config = PipelineConfig.from_pretrained(model_name)
+    config.vae_precision = "fp16"
+    config.use_cpu_offload = True
+
+    # Create the generator
+    generator = VideoGenerator.from_pretrained(
+        model_name,
+        num_gpus=1,
+        pipeline_config=config
+    )
+
+    # Create and customize sampling parameters
+    sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    
+    # How many frames to generate
+    sampling_param.num_frames = 45
+    
+    # Video resolution (width, height)
+    sampling_param.width = 1024
+    sampling_param.height = 576
+    
+    # How many steps we denoise the video (higher = better quality, slower generation)
+    sampling_param.num_inference_steps = 30
+    
+    # How strongly the video conforms to the prompt (higher = more faithful to prompt)
+    sampling_param.guidance_scale = 7.5
+    
+    # Random seed for reproducibility
+    sampling_param.seed = 42  # Optional, leave unset for random results
+
+    # Generate video with custom parameters
+    prompt = "A beautiful sunset over a calm ocean, with gentle waves."
+    video = generator.generate_video(
+        prompt, 
+        sampling_param=sampling_param, 
+        output_path="my_videos/",  # Controls where videos are saved
+        return_frames=True,  # Also return frames from this call (defaults to False)
+        save_video=True
+    )
+
+    # If return_frames=True, video contains the generated frames as a NumPy array
+    print(f"Generated {len(video)} frames")
+
+if __name__ == '__main__':
+    main()
+```
+
+## Performance Optimization
+
+For configuring optimizations, please see our [optimizations guide](#inference-optimizations)

docs/source/inference/inference_quick_start.md

@@ -2,19 +2,11 @@
 
 This page contains step-by-step instructions to get you quickly started with video generation using FastVideo.
 
-## Table of Contents
-- [Generating Your First Video](#generating-your-first-video)
-- [Customizing Generation](#customizing-generation)
-- [Available Models](#available-models)
-- [Image-to-Video Generation](#image-to-video-generation)
-- [Troubleshooting](#troubleshooting)
-- [Advanced Configuration](#advanced-configuration)
-- [Next Steps](#next-steps)
-
-## Software Requirements
+## Requirements
 - **OS**: Linux (Tested on Ubuntu 22.04+)
 - **Python**: 3.10-3.12
 - **CUDA**: 12.4
+- **GPU**: At least one NVIDIA GPU
 
 ## Installation
 
@@ -78,18 +70,25 @@ You can generate a video starting from an initial image:
 ```python
 from fastvideo import VideoGenerator, SamplingParam
 
-# Create the generator
-generator = VideoGenerator.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+def main():
+    # Create the generator
+    model_name = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+    generator = VideoGenerator.from_pretrained(model_name, num_gpus=1)
 
-# Set up parameters with an initial image
-sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
-sampling_param.image_path = "path/to/your/image.jpg"
-sampling_param.num_frames = 24
-sampling_param.image_strength = 0.8  # How much to preserve the original image (0-1)
+    # Set up parameters with an initial image
+    sampling_param = SamplingParam.from_pretrained(model_name)
+    sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    sampling_param.num_frames = 107
+    sampling_param.image_strength = 0.8  # How much to preserve the original image (0-1)
 
-# Generate video based on the image
-prompt = "A photograph coming to life with gentle movement"
-video = generator.generate_video(prompt, sampling_param=sampling_param)
+    # Generate video based on the image
+    prompt = "A photograph coming to life with gentle movement"
+    generator.generate_video(prompt, sampling_param=sampling_param,
+                             output_path="my_videos/",
+                             save_video=True)
+
+if __name__ == '__main__':
+    main()
 ```
 
 ## Troubleshooting
@@ -100,7 +99,7 @@ Common issues and their solutions:
 If you encounter CUDA out of memory errors:
 - Reduce `num_frames` or video resolution
 - Enable memory optimization with `enable_model_cpu_offload`
-- Try a smaller model or use quantized versions
+- Try a smaller model or use distilled versions
 - Use `num_gpus` > 1 if multiple GPUs are available
 
 ### Slow Generation
@@ -116,14 +115,10 @@ If the generated video doesn't match your prompt:
 - Experiment with different random seeds
 - Try a different model
 
-## Advanced Configuration
-
-## Optimizations
-
 ## Next Steps
 
-- Explore the [API Reference](../api/index.md) for detailed documentation
-- Learn about [Advanced Inference Options](../inference/overview_back.md)
-- See [Examples](../examples/index.md) for more usage scenarios
-- Check out the [Model Training](../training/overview.md) guide to fine-tune models
-- Join our [Community Discord](https://discord.gg/fastvideo) for support and sharing
+- Learn about [Advanced Inference Configurations](#inference-configuration)
+- Learn about using [Optimizations](#inference-optimizations)
+- See [Examples](../examples/examples_inference_index.md) for more usage scenarios
+- Join our [Community Discord](https://discord.gg/JA7cksDz86).
+- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg).

docs/source/inference/optimizations.md

@@ -1,3 +1,4 @@
+(inference-optimizations)=
 # Optimizations
 
 This page describes the various options for speeding up generation times in FastVideo.

docs/source/inference/support_matrix.md

@@ -49,14 +49,14 @@ The `HuggingFace Model ID` can be directly pass to `from_pretrained()` methods a
   * ❌
   * ✅
   * ✅
-- * Wan T2V 1.4B
+- * Wan T2V 1.3B
   * `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`
   * 480P
   * ✅
   * ✅*
   * ✅
 - * Wan T2V 14B
-  * `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`
+  * `Wan-AI/Wan2.1-T2V-14B-Diffusers`
   * 480P, 720P
   * ✅
   * ✅*
@@ -67,18 +67,18 @@ The `HuggingFace Model ID` can be directly pass to `from_pretrained()` methods a
   * ✅
   * ✅*
   * ✅
-- * Wan T2V 720P
-  * `Wan-AI/Wan2.1-T2V-14B-Diffusers`
+- * Wan I2V 720P
+  * `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers`
   * 720P
   * ✅
   * ✅*
   * ✅
 - * StepVideo T2V
-  * Coming soon!
+  * `FastVideo/stepvideo-t2v-diffusers`
   * 768px768px204f<br>544px992px204f<br>544px992px136f
-  *
-  *
-  *
+  * ❌
+  * ❌
+  * ✅
 :::
 
 **Note**: there are some known quality issues with Wan2.1 + Sliding Tile Attn. We are working on fixing this issue.

fastvideo/models/stepvideo/text_encoder/clip.py

@@ -16,7 +16,7 @@ class HunyuanClip(nn.Module):
 
         self.max_length = max_length
         self.tokenizer = BertTokenizer.from_pretrained(os.path.join(model_dir, 'tokenizer'))
-        self.text_encoder = BertModel.from_pretrained(os.path.join(model_dir, 'clip_text_encoder'))
+        self.text_encoder = BertModel.from_pretrained(os.path.join(model_dir, 'text_encoder'))
 
     @torch.no_grad
     def forward(self, prompts, with_mask=True):
@@ -31,6 +31,6 @@ class HunyuanClip(nn.Module):
         )
         prompt_embeds = self.text_encoder(
             text_inputs.input_ids.to(self.device),
-            attention_mask=text_inputs.attention_mask.to(self.device) if with_mask else None,
+            # attention_mask=text_inputs.attention_mask.to(self.device) if with_mask else None,
         )
-        return prompt_embeds.last_hidden_state, prompt_embeds.pooler_output
+        return prompt_embeds.last_hidden_state, text_inputs

fastvideo/v1/configs/models/dits/base.py

@@ -2,7 +2,7 @@ from dataclasses import dataclass, field
 from typing import Any, Optional, Tuple
 
 from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
-from fastvideo.v1.configs.quantization import QuantizationConfig
+from fastvideo.v1.layers.quantization import QuantizationConfig
 from fastvideo.v1.platforms import _Backend
 
 

fastvideo/v1/configs/models/encoders/__init__.py

@@ -6,9 +6,10 @@ from fastvideo.v1.configs.models.encoders.clip import (CLIPTextConfig,
                                                        CLIPVisionConfig)
 from fastvideo.v1.configs.models.encoders.llama import LlamaConfig
 from fastvideo.v1.configs.models.encoders.t5 import T5Config
+from fastvideo.v1.configs.models.encoders.bert import BertConfig
 
 __all__ = [
     "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig",
     "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig", "LlamaConfig",
-    "T5Config"
+    "T5Config", "BertConfig"
 ]

fastvideo/v1/configs/models/encoders/base.py

@@ -4,7 +4,7 @@ from typing import Any, Dict, List, Optional, Tuple
 import torch
 
 from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
-from fastvideo.v1.configs.quantization import QuantizationConfig
+from fastvideo.v1.layers.quantization import QuantizationConfig
 from fastvideo.v1.platforms import _Backend
 
 

fastvideo/v1/configs/models/encoders/bert.py

@@ -0,0 +1,48 @@
+# fastvideo/v1/configs/models/encoders/bert.py
+from dataclasses import dataclass, field
+from typing import Optional
+
+from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
+                                                       TextEncoderConfig)
+
+# ---------- Architecture-level hyper-parameters -----------------
+@dataclass
+class BertArchConfig(TextEncoderArchConfig):
+    vocab_size: int = 47020
+    hidden_size: int = 1024
+    intermediate_size: int = 4096
+    num_hidden_layers: int = 24
+    num_attention_heads: int = 16
+    
+    hidden_act: str = "gelu"
+    max_position_embeddings: int = 512
+    type_vocab_size: int = 2
+    position_embedding_type: str = "absolute"
+    layer_norm_eps: float = 1e-12
+    initializer_range: float = 0.02
+
+    attention_probs_dropout_prob: float = 0.1
+    hidden_dropout_prob: float = 0.1
+    use_cache: bool = True                       # KV cache toggle
+    pad_token_id: int = 0
+    bos_token_id: int = 0
+    eos_token_id: int = 2
+
+    text_len: int = 77
+
+    # === Pooler / NSP head ===================================================
+    pooler_fc_size: int = 768
+    pooler_num_attention_heads: int = 12
+    pooler_num_fc_layers: int = 3
+    pooler_size_per_head: int = 128
+    pooler_type: str = "first_token_transform"
+    
+    classifier_dropout: Optional[float] = None
+    directionality: str = "bidi"
+    output_past: bool = True 
+
+# ---------- Top-level config wrapper (YAML-friendly) ------------
+@dataclass
+class BertConfig(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(default_factory=BertArchConfig)
+    prefix: str = "bert"

fastvideo/v1/configs/pipelines/hunyuan.py

@@ -68,6 +68,9 @@ class HunyuanConfig(PipelineConfig):
     embedded_cfg_scale: int = 6
     flow_shift: int = 7
 
+    # Video parameters
+    use_cpu_offload: bool = True
+
     # Text encoding stage
     text_encoder_configs: Tuple[EncoderConfig, ...] = field(
         default_factory=lambda: (LlamaConfig(), CLIPTextConfig()))

fastvideo/v1/configs/pipelines/stepvideo.py

@@ -18,6 +18,9 @@ class StepVideoT2VConfig(PipelineConfig):
     vae_tiling: bool = False
     vae_sp: bool = False
 
+    # Video parameters
+    use_cpu_offload: bool = True
+
     # Denoising stage
     flow_shift: int = 13
     timesteps_scale: bool = False

fastvideo/v1/configs/quantization/__init__.py

@@ -1,3 +0,0 @@
-from fastvideo.v1.configs.quantization.base import QuantizationConfig
-
-__all__ = ["QuantizationConfig"]

fastvideo/v1/configs/quantization/base.py

@@ -1,6 +0,0 @@
-from dataclasses import dataclass
-
-
-@dataclass
-class QuantizationConfig:
-    pass

fastvideo/v1/entrypoints/cli/generate.py

@@ -11,8 +11,11 @@ from fastvideo.v1.configs.sample.base import SamplingParam
 from fastvideo.v1.entrypoints.cli.cli_types import CLISubcommand
 from fastvideo.v1.entrypoints.cli.utils import RaiseNotImplementedAction
 from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import FlexibleArgumentParser
 
+logger = init_logger(__name__)
+
 
 class GenerateSubcommand(CLISubcommand):
     """The `generate` subcommand for the FastVideo CLI"""
@@ -35,12 +38,22 @@ class GenerateSubcommand(CLISubcommand):
         excluded_args = ['subparser', 'config', 'dispatch_function']
 
         FastVideoArgs.from_cli_args(args)
-        filtered_args = {}
-        for k, v in vars(args).items():
-            if k not in excluded_args and v is not None:
-                filtered_args[k] = v
 
-        merged_args = {**filtered_args}
+        provided_args = {}
+        for k, v in vars(args).items():
+            if (k not in excluded_args and v is not None
+                    and hasattr(args, '_provided') and k in args._provided):
+                provided_args[k] = v
+
+        if 'model_path' in vars(args) and args.model_path is not None:
+            provided_args['model_path'] = args.model_path
+
+        if 'prompt' in vars(args) and args.prompt is not None:
+            provided_args['prompt'] = args.prompt
+
+        merged_args = {**provided_args}
+
+        logger.info('CLI Args: %s', merged_args)
 
         if 'model_path' not in merged_args or not merged_args['model_path']:
             raise ValueError(

fastvideo/v1/entrypoints/video_generator.py

@@ -7,6 +7,7 @@ diffusion models.
 """
 
 import gc
+import math
 import os
 import time
 from typing import Any, Dict, List, Optional, Union
@@ -180,12 +181,46 @@ class VideoGenerator:
                 f"height={sampling_param.height}, width={sampling_param.width}, "
                 f"num_frames={sampling_param.num_frames}")
 
-        if (
-                sampling_param.num_frames - 1
-        ) % fastvideo_args.vae_config.arch_config.temporal_compression_ratio != 0:
-            raise ValueError(
-                f"num_frames-1 must be a multiple of {fastvideo_args.vae_config.arch_config.temporal_compression_ratio}, got {sampling_param.num_frames}"
-            )
+        temporal_scale_factor = fastvideo_args.vae_config.arch_config.temporal_compression_ratio
+        num_frames = sampling_param.num_frames
+        num_gpus = fastvideo_args.num_gpus
+        use_temporal_scaling_frames = fastvideo_args.vae_config.use_temporal_scaling_frames
+
+        # Adjust number of frames based on number of GPUs
+        if use_temporal_scaling_frames:
+            orig_latent_num_frames = (num_frames -
+                                      1) // temporal_scale_factor + 1
+        else:  # stepvideo only
+            orig_latent_num_frames = sampling_param.num_frames // 17 * 3
+
+        if orig_latent_num_frames % fastvideo_args.num_gpus != 0:
+            # Adjust latent frames to be divisible by number of GPUs
+            if sampling_param.num_frames_round_down:
+                # Ensure we have at least 1 batch per GPU
+                new_latent_num_frames = max(
+                    1, (orig_latent_num_frames // num_gpus)) * num_gpus
+            else:
+                new_latent_num_frames = math.ceil(
+                    orig_latent_num_frames / num_gpus) * num_gpus
+
+            if use_temporal_scaling_frames:
+                # Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor
+                new_num_frames = (new_latent_num_frames -
+                                  1) * temporal_scale_factor + 1
+            else:  # stepvideo only
+                # Find the least common multiple of 3 and num_gpus
+                divisor = math.lcm(3, num_gpus)
+                # Round up to the nearest multiple of this LCM
+                new_latent_num_frames = (
+                    (new_latent_num_frames + divisor - 1) // divisor) * divisor
+                # Convert back to actual frames using the StepVideo formula
+                new_num_frames = new_latent_num_frames // 3 * 17
+
+            logger.info(
+                "Adjusting number of frames from %s to %s based on number of GPUs (%s)",
+                sampling_param.num_frames, new_num_frames,
+                fastvideo_args.num_gpus)
+            sampling_param.num_frames = new_num_frames
 
         # Calculate sizes
         target_height = align_to(sampling_param.height, 16)

fastvideo/v1/layers/activation.py

@@ -9,7 +9,6 @@ import torch.nn.functional as F
 
 # TODO (will): remove this dependency
 from fastvideo.v1.layers.custom_op import CustomOp
-from fastvideo.v1.platforms import current_platform
 
 
 @CustomOp.register("silu_and_mul")
@@ -25,21 +24,12 @@ class SiluAndMul(CustomOp):
 
     def __init__(self) -> None:
         super().__init__()
-        if current_platform.is_cuda_alike() or current_platform.is_cpu():
-            self.op = torch.ops._C.silu_and_mul
 
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
         d = x.shape[-1] // 2
         return F.silu(x[..., :d]) * x[..., d:]
 
-    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        d = x.shape[-1] // 2
-        output_shape = (x.shape[:-1] + (d, ))
-        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
-        self.op(out, x)
-        return out
-
 
 @CustomOp.register("gelu_and_mul")
 class GeluAndMul(CustomOp):
@@ -57,24 +47,12 @@ class GeluAndMul(CustomOp):
         self.approximate = approximate
         if approximate not in ("none", "tanh"):
             raise ValueError(f"Unknown approximate mode: {approximate}")
-        if current_platform.is_cuda_alike() or current_platform.is_cpu():
-            if approximate == "none":
-                self.op = torch.ops._C.gelu_and_mul
-            elif approximate == "tanh":
-                self.op = torch.ops._C.gelu_tanh_and_mul
 
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
         d = x.shape[-1] // 2
         return F.gelu(x[..., :d], approximate=self.approximate) * x[..., d:]
 
-    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        d = x.shape[-1] // 2
-        output_shape = (x.shape[:-1] + (d, ))
-        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
-        self.op(out, x)
-        return out
-
     def extra_repr(self) -> str:
         return f'approximate={repr(self.approximate)}'
 
@@ -84,8 +62,6 @@ class NewGELU(CustomOp):
 
     def __init__(self):
         super().__init__()
-        if current_platform.is_cuda_alike() or current_platform.is_cpu():
-            self.op = torch.ops._C.gelu_new
 
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
@@ -93,32 +69,17 @@ class NewGELU(CustomOp):
         return 0.5 * x * (1.0 + torch.tanh(c *
                                            (x + 0.044715 * torch.pow(x, 3.0))))
 
-    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        out = torch.empty_like(x)
-        self.op(out, x)
-        return out
-
-    def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
-        return self.op(x)
-
 
 @CustomOp.register("quick_gelu")
 class QuickGELU(CustomOp):
     # https://github.com/huggingface/transformers/blob/main/src/transformers/activations.py#L90
     def __init__(self):
         super().__init__()
-        if current_platform.is_cuda_alike() or current_platform.is_cpu():
-            self.op = torch.ops._C.gelu_quick
 
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
         return x * torch.sigmoid(1.702 * x)
 
-    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        out = torch.empty_like(x)
-        self.op(out, x)
-        return out
-
 
 _ACTIVATION_REGISTRY = {
     "gelu": nn.GELU,

fastvideo/v1/layers/custom_op.py

@@ -50,6 +50,10 @@ class CustomOp(nn.Module):
         return self.forward_native(*args, **kwargs)
 
     def dispatch_forward(self) -> Callable:
+        # FIXME(will): for now, we always use the native implementation, since
+        # forward_cuda is using vllm's custom ops and it doesn't support
+        # backwards. We should add our own custom ops that support backwards.
+        return self.forward_native
         # NOTE(woosuk): Here we assume that vLLM was built for only one
         # specific backend. Currently, we do not support dynamic dispatching.
         enabled = self.enabled()

fastvideo/v1/layers/layernorm.py

@@ -75,33 +75,6 @@ class RMSNorm(CustomOp):
         else:
             return x, residual
 
-    def forward_cuda(
-        self,
-        x: torch.Tensor,
-        residual: Optional[torch.Tensor] = None,
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-        if self.variance_size_override is not None:
-            return self.forward_native(x, residual)
-
-        from vllm import _custom_ops as ops
-
-        if residual is not None:
-            ops.fused_add_rms_norm(
-                x,
-                residual,
-                self.weight.data,
-                self.variance_epsilon,
-            )
-            return x, residual
-        out = torch.empty_like(x)
-        ops.rms_norm(
-            out,
-            x,
-            self.weight.data,
-            self.variance_epsilon,
-        )
-        return out
-
     def extra_repr(self) -> str:
         s = f"hidden_size={self.weight.data.size(0)}"
         s += f", eps={self.variance_epsilon}"

fastvideo/v1/layers/linear.py

@@ -7,16 +7,14 @@ from typing import Optional, Union
 import torch
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter
-# TODO(will): remove this import by copying the definition from vLLM then
-# manually import each quantization method we want to use. Refer to SGLang
-from vllm.model_executor.layers.quantization.base_config import (
-    QuantizationConfig, QuantizeMethodBase)
 
 from fastvideo.v1.distributed import (divide, get_tensor_model_parallel_rank,
                                       get_tensor_model_parallel_world_size,
                                       split_tensor_along_last_dim,
                                       tensor_model_parallel_all_gather,
                                       tensor_model_parallel_all_reduce)
+from fastvideo.v1.layers.quantization.base_config import (QuantizationConfig,
+                                                          QuantizeMethodBase)
 from fastvideo.v1.logger import init_logger
 # yapf: disable
 from fastvideo.v1.models.parameter import (BasevLLMParameter,
@@ -545,19 +543,21 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
         tp_size = get_tensor_model_parallel_world_size()
 
         if isinstance(param, BlockQuantScaleParameter):
-            from vllm.model_executor.layers.quantization.fp8 import (
-                Fp8LinearMethod, Fp8MoEMethod)
-            assert self.quant_method is not None
-            assert isinstance(self.quant_method,
-                              (Fp8LinearMethod, Fp8MoEMethod))
-            weight_block_size = self.quant_method.quant_config.weight_block_size
-            assert weight_block_size is not None
-            block_n, _ = weight_block_size[0], weight_block_size[1]
-            shard_offset = (
-                (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) //
-                block_n) // tp_size
-            shard_size = ((self.output_sizes[loaded_shard_id] + block_n - 1) //
-                          block_n // tp_size)
+            raise NotImplementedError("FP8 is not implemented yet")
+            # FIXME(will): add fp8 support
+            # from vllm.model_executor.layers.quantization.fp8 import (
+            #     Fp8LinearMethod, Fp8MoEMethod)
+            # assert self.quant_method is not None
+            # assert isinstance(self.quant_method,
+            #                   (Fp8LinearMethod, Fp8MoEMethod))
+            # weight_block_size = self.quant_method.quant_config.weight_block_size
+            # assert weight_block_size is not None
+            # block_n, _ = weight_block_size[0], weight_block_size[1]
+            # shard_offset = (
+            #     (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) //
+            #     block_n) // tp_size
+            # shard_size = ((self.output_sizes[loaded_shard_id] + block_n - 1) //
+            #               block_n // tp_size)
         else:
             shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size
             shard_size = self.output_sizes[loaded_shard_id] // tp_size

fastvideo/v1/layers/quantization/__init__.py

@@ -0,0 +1,65 @@
+from typing import Literal, get_args
+
+from fastvideo.v1.layers.quantization.base_config import QuantizationConfig
+
+QuantizationMethods = Literal[None]
+
+QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
+
+# The customized quantization methods which will be added to this dict.
+_CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {}
+
+
+def register_quantization_config(quantization: str):
+    """Register a customized vllm quantization config.
+
+    When a quantization method is not supported by vllm, you can register a customized
+    quantization config to support it.
+
+    Args:
+        quantization (str): The quantization method name.
+
+    Examples:
+        >>> from fastvideo.v1.layers.quantization import register_quantization_config
+        >>> from fastvideo.v1.layers.quantization import get_quantization_config
+        >>> from fastvideo.v1.layers.quantization.base_config import QuantizationConfig
+        >>>
+        >>> @register_quantization_config("my_quant")
+        ... class MyQuantConfig(QuantizationConfig):
+        ...     pass
+        >>>
+        >>> get_quantization_config("my_quant")
+        <class 'MyQuantConfig'>
+    """  # noqa: E501
+
+    def _wrapper(quant_config_cls):
+        if quantization in QUANTIZATION_METHODS:
+            raise ValueError(
+                f"The quantization method `{quantization}` is already exists.")
+        if not issubclass(quant_config_cls, QuantizationConfig):
+            raise ValueError("The quantization config must be a subclass of "
+                             "`QuantizationConfig`.")
+        _CUSTOMIZED_METHOD_TO_QUANT_CONFIG[quantization] = quant_config_cls
+        QUANTIZATION_METHODS.append(quantization)
+        return quant_config_cls
+
+    return _wrapper
+
+
+def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
+    if quantization not in QUANTIZATION_METHODS:
+        raise ValueError(f"Invalid quantization method: {quantization}")
+
+    method_to_config: dict[str, type[QuantizationConfig]] = {}
+    # Update the `method_to_config` with customized quantization methods.
+    method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
+
+    return method_to_config[quantization]
+
+
+all = [
+    "QuantizationMethods",
+    "QuantizationConfig",
+    "get_quantization_config",
+    "QUANTIZATION_METHODS",
+]

fastvideo/v1/layers/quantization/base_config.py

@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/quantization/base_config.py
+
+import inspect
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Optional
+
+import torch
+from torch import nn
+
+if TYPE_CHECKING:
+    from fastvideo.v1.layers.quantization import QuantizationMethods
+else:
+    QuantizationMethods = str
+
+
+class QuantizeMethodBase(ABC):
+    """Base class for different quantized methods."""
+
+    @abstractmethod
+    def create_weights(self, layer: torch.nn.Module, *weight_args,
+                       **extra_weight_attrs):
+        """Create weights for a layer.
+
+        The weights will be set as attributes of the layer."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor:
+        """Apply the weights in layer to the input tensor.
+
+        Expects create_weights to have been called before on the layer."""
+        raise NotImplementedError
+
+    # Not required functions
+    def embedding(self, layer: torch.nn.Module, *args,
+                  **kwargs) -> torch.Tensor:
+        """Gather embeddings in the layer based on indices in the input tensor.
+
+        Expects create_weights to have been called before on the layer."""
+        raise NotImplementedError
+
+    def process_weights_after_loading(self, layer: nn.Module) -> None:
+        """Process the weight after loading.
+
+        This can be used for example, to transpose weights for computation.
+        """
+        return
+
+
+def method_has_implemented_embedding(
+        method_class: type[QuantizeMethodBase]) -> bool:
+    """
+    Not all quant methods have embedding implemented, so we need to check that
+    it exists for our given method. We check this by making sure the function
+    has been changed from the base implementation.
+    """
+    base_embedding = inspect.getattr_static(QuantizeMethodBase, "embedding",
+                                            None)
+    class_embedding = inspect.getattr_static(method_class, "embedding", None)
+
+    return (class_embedding is not None
+            and class_embedding is not base_embedding)
+
+
+class QuantizationConfig(ABC):
+    """Base class for quantization configs."""
+
+    def __init__(self):
+        super().__init__()
+        # mapping is updated by models as they initialize
+        self.packed_modules_mapping: dict[str, list[str]] = dict()
+
+    @abstractmethod
+    def get_name(self) -> QuantizationMethods:
+        """Name of the quantization method."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_supported_act_dtypes(self) -> list[torch.dtype]:
+        """List of supported activation dtypes."""
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def get_min_capability(cls) -> int:
+        """Minimum GPU capability to support the quantization method.
+
+        E.g., 70 for Volta, 75 for Turing, 80 for Ampere.
+        This requirement is due to the custom CUDA kernels used by the
+        quantization method.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_config_filenames() -> list[str]:
+        """List of filenames to search for in the model directory."""
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def from_config(cls, config: dict[str, Any]) -> "QuantizationConfig":
+        """Create a config class from the model's quantization config."""
+        raise NotImplementedError
+
+    @classmethod
+    def override_quantization_method(
+            cls, hf_quant_cfg, user_quant) -> Optional[QuantizationMethods]:
+        """
+           Detects if this quantization method can support a given checkpoint
+           format by overriding the user specified quantization method -- 
+           this method should only be overwritten by subclasses in exceptional 
+           circumstances
+        """
+        return None
+
+    @staticmethod
+    def get_from_keys(config: dict[str, Any], keys: list[str]) -> Any:
+        """Get a value from the model's quantization config."""
+        for key in keys:
+            if key in config:
+                return config[key]
+        raise ValueError(f"Cannot find any of {keys} in the model's "
+                         "quantization config.")
+
+    @staticmethod
+    def get_from_keys_or(config: dict[str, Any], keys: list[str],
+                         default: Any) -> Any:
+        """Get a optional value from the model's quantization config."""
+        try:
+            return QuantizationConfig.get_from_keys(config, keys)
+        except ValueError:
+            return default
+
+    @abstractmethod
+    def get_quant_method(self, layer: torch.nn.Module,
+                         prefix: str) -> Optional[QuantizeMethodBase]:
+        """Get the quantize method to use for the quantized layer.
+        
+        Args:
+            layer: The layer for the quant method.
+            prefix: The full name of the layer in the state dict
+        Returns:
+            The quantize method. None if the given layer doesn't support quant
+            method.
+        """
+        raise NotImplementedError
+
+    def get_cache_scale(self, name: str) -> Optional[str]:
+        return None

fastvideo/v1/layers/rotary_embedding.py

@@ -152,28 +152,6 @@ class RotaryEmbedding(CustomOp):
         key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
         return query, key
 
-    def forward_cuda(
-        self,
-        positions: torch.Tensor,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        offsets: Optional[torch.Tensor] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        from vllm import _custom_ops as ops
-
-        self.cos_sin_cache = self.cos_sin_cache.to(query.device,
-                                                   dtype=query.dtype)
-        # ops.rotary_embedding()/batched_rotary_embedding()
-        # are in-place operations that update the query and key tensors.
-        if offsets is not None:
-            ops.batched_rotary_embedding(positions, query, key, self.head_size,
-                                         self.cos_sin_cache, self.is_neox_style,
-                                         self.rotary_dim, offsets)
-        else:
-            ops.rotary_embedding(positions, query, key, self.head_size,
-                                 self.cos_sin_cache, self.is_neox_style)
-        return query, key
-
     def extra_repr(self) -> str:
         s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
         s += f", max_position_embeddings={self.max_position_embeddings}"

fastvideo/v1/layers/vocab_parallel_embedding.py

@@ -6,12 +6,12 @@ from typing import List, Optional, Sequence, Tuple
 import torch
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter, UninitializedParameter
-from vllm.model_executor.layers.quantization.base_config import (
-    QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding)
 
 from fastvideo.v1.distributed import (divide, get_tensor_model_parallel_rank,
                                       get_tensor_model_parallel_world_size,
                                       tensor_model_parallel_all_reduce)
+from fastvideo.v1.layers.quantization.base_config import (
+    QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding)
 from fastvideo.v1.models.parameter import BasevLLMParameter
 from fastvideo.v1.models.utils import set_weight_attrs
 from fastvideo.v1.platforms import current_platform

fastvideo/v1/models/dits/base.py

@@ -95,8 +95,6 @@ class CachableDiT(BaseDiT):
         if self.config.prefix == "wan":
             self.use_ret_steps = self.config.cache_config.use_ret_steps
             self.is_even = False
-            self.previous_e0_even: torch.Tensor | None = None
-            self.previous_e0_odd: torch.Tensor | None = None
             self.previous_residual_even: torch.Tensor | None = None
             self.previous_residual_odd: torch.Tensor | None = None
             self.accumulated_rel_l1_distance_even = 0
@@ -106,7 +104,9 @@ class CachableDiT(BaseDiT):
         else:
             self.accumulated_rel_l1_distance = 0
             self.previous_modulated_input = None
-            self.previous_residual = None
+            self.previous_resiual = None
+        self.previous_e0_even: torch.Tensor | None = None
+        self.previous_e0_odd: torch.Tensor | None = None
 
     def maybe_cache_states(self, hidden_states: torch.Tensor,
                            original_hidden_states: torch.Tensor) -> None:

fastvideo/v1/models/dits/hunyuanvideo.py

@@ -562,8 +562,7 @@ class HunyuanVideoTransformer3DModel(CachableDiT):
         """
         forward_context = get_forward_context()
         forward_batch = forward_context.forward_batch
-        assert forward_batch is not None
-        enable_teacache = forward_batch.enable_teacache
+        enable_teacache = forward_batch is not None and forward_batch.enable_teacache
 
         if guidance is None:
             guidance = torch.tensor([6016.0],
@@ -661,7 +660,8 @@ class HunyuanVideoTransformer3DModel(CachableDiT):
 
         forward_context = get_forward_context()
         forward_batch = forward_context.forward_batch
-        assert forward_batch is not None
+        if forward_batch is None:
+            return False
         current_timestep = forward_context.current_timestep
         enable_teacache = forward_batch.enable_teacache
 

fastvideo/v1/models/dits/wanvideo.py

@@ -424,8 +424,7 @@ class WanTransformer3DModel(CachableDiT):
                 guidance=None,
                 **kwargs) -> torch.Tensor:
         forward_batch = get_forward_context().forward_batch
-        assert forward_batch is not None
-        enable_teacache = forward_batch.enable_teacache
+        enable_teacache = forward_batch is not None and forward_batch.enable_teacache
 
         orig_dtype = hidden_states.dtype
         if not isinstance(encoder_hidden_states, torch.Tensor):
@@ -525,8 +524,7 @@ class WanTransformer3DModel(CachableDiT):
 
         forward_context = get_forward_context()
         forward_batch = forward_context.forward_batch
-        assert forward_batch is not None
-        if not forward_batch.enable_teacache:
+        if forward_batch is None or not forward_batch.enable_teacache:
             return False
         teacache_params = forward_batch.teacache_params
         assert teacache_params is not None, "teacache_params is not initialized"

fastvideo/v1/models/encoders/bert.py

@@ -1,40 +1,411 @@
-# type: ignore
-import os
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Iterable, Optional, Set, Tuple
 
 import torch
-import torch.nn as nn
-from transformers import BertModel, BertTokenizer
+from torch import nn
+from transformers import BertConfig
+
+# from vllm.attention import Attention, AttentionType
+from fastvideo.v1.attention import LocalAttention
+# from vllm.compilation.decorators import support_torch_compile
+# from vllm.config import CacheConfig, PoolerConfig, VllmConfig
+# CacheConfig is used for attention, could be replaced with LocalAttention
+# VllmConfig is replaced with BertConfig (needs to implement)
+from fastvideo.v1.configs.models.encoders.bert import BertConfig
+# from vllm.distributed import get_tensor_model_parallel_world_size
+from fastvideo.v1.distributed import (divide,
+                                      get_tensor_model_parallel_world_size)
+# from vllm.forward_context import get_forward_context
+from fastvideo.v1.forward_context import get_forward_context
+# from vllm.model_executor.layers.activation import get_act_fn
+from fastvideo.v1.layers.activation import get_act_fn
+# from vllm.model_executor.layers.linear import (ColumnParallelLinear,
+#                                                QKVParallelLinear,
+#                                                RowParallelLinear)
+from fastvideo.v1.layers.linear import (ColumnParallelLinear, QKVParallelLinear,
+                                        RowParallelLinear)
+# from vllm.model_executor.layers.pooler import (CrossEncodingPooler, Pooler,
+#                                                PoolingType)
+# from vllm.model_executor.layers.quantization import QuantizationConfig
+from fastvideo.v1.layers.quantization import QuantizationConfig
+# from vllm.model_executor.layers.vocab_parallel_embedding import (
+#     VocabParallelEmbedding)
+from fastvideo.v1.layers.vocab_parallel_embedding import VocabParallelEmbedding
+# from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from fastvideo.v1.models.loader.weight_utils import default_weight_loader
+# from vllm.model_executor.pooling_metadata import PoolingMetadata
+# from vllm.sequence import PoolerOutput
+# from vllm.transformers_utils.config import (
+#     get_cross_encoder_activation_function)
+from fastvideo.v1.models.encoders.base import TextEncoder
+# from .interfaces import SupportsCrossEncoding, SupportsQuant, SupportsV0Only
+# from .utils import WeightsMapper, maybe_prefix
+
+from fastvideo.v1.platforms import _Backend
+class BertEmbedding(nn.Module):
+
+    def __init__(self, config: BertConfig):
+
+        super().__init__()
+        self.size = config.hidden_size
+        self.word_embeddings = VocabParallelEmbedding(config.vocab_size,
+                                                      config.hidden_size)
+        self.position_embeddings = VocabParallelEmbedding(
+            config.max_position_embeddings, config.hidden_size)
+        self.token_type_embeddings = VocabParallelEmbedding(
+            config.type_vocab_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size,
+                                      eps=config.layer_norm_eps)
+        # self.position_ids = nn.Parameter(
+        #     torch.empty((1, config.max_position_embeddings)), )
+
+        self.position_embedding_type = config.position_embedding_type
+        if self.position_embedding_type != "absolute":
+            raise ValueError("Only 'absolute' position_embedding_type" +
+                             " is supported")
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor,
+        token_type_ids: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        input_shape = input_ids.size()
+
+        # Input embeddings.
+        inputs_embeds = self.word_embeddings(input_ids)
+
+        # Position embeddings.
+        position_embeddings = self.position_embeddings(position_ids)
+
+        if token_type_ids is None:
+            token_type_ids = torch.zeros(input_shape,
+                                         dtype=torch.long,
+                                         device=inputs_embeds.device)
+
+        token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+        embeddings = inputs_embeds + token_type_embeddings + position_embeddings
+        embeddings = self.LayerNorm(embeddings)
+        return embeddings
 
 
-class HunyuanClip(nn.Module):
-    """
-        Hunyuan clip code copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py
-        hunyuan's clip used BertModel and BertTokenizer, so we copy it.
-    """
+class BertPooler(nn.Module):
 
-    def __init__(self, model_dir, max_length=77):
+    def __init__(self, config: BertConfig):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        first_token_tensor = hidden_states[0, :]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+# @support_torch_compile
+class BertEncoder(nn.Module):
+
+    def __init__(self, config: BertConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = ""):
+        super().__init__()
+        # config = vllm_config.model_config.hf_config
+        # cache_config = vllm_config.cache_config
+        # quant_config = vllm_config.quant_config
+        self.layer = nn.ModuleList([
+            BertLayer(config=config,
+                    #   cache_config=cache_config,
+                      quant_config=quant_config,
+                      prefix=f"{prefix}.layer.{layer_idx}")
+            for layer_idx in range(config.num_hidden_layers)
+        ])
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        for layer in self.layer:
+            hidden_states = layer(hidden_states)
+        print("hidden_states shape", hidden_states.shape)
+        return hidden_states
+
+
+class BertLayer(nn.Module):
+
+    def __init__(self,
+                 config: BertConfig,
+                #  cache_config: Optional[CacheConfig] = None,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 prefix: str = ""):
         super().__init__()
 
-        self.max_length = max_length
-        self.tokenizer = BertTokenizer.from_pretrained(
-            os.path.join(model_dir, 'tokenizer'))
-        self.text_encoder = BertModel.from_pretrained(
-            os.path.join(model_dir, 'clip_text_encoder'))
+        self.attention = BertAttention(
+            hidden_size=config.hidden_size,
+            num_attention_heads=config.num_attention_heads,
+            layer_norm_eps=config.layer_norm_eps,
+            # cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attention")
 
-    @torch.no_grad
-    def forward(self, prompts, with_mask=True):
-        self.device = next(self.text_encoder.parameters()).device
-        text_inputs = self.tokenizer(
-            prompts,
-            padding="max_length",
-            max_length=self.max_length,
-            truncation=True,
-            return_attention_mask=True,
-            return_tensors="pt",
-        )
-        prompt_embeds = self.text_encoder(
-            text_inputs.input_ids.to(self.device),
-            attention_mask=text_inputs.attention_mask.to(self.device)
-            if with_mask else None,
-        )
-        return prompt_embeds.last_hidden_state, prompt_embeds.pooler_output
+        self.intermediate = BertIntermediate(
+            hidden_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            prefix=f"{prefix}.intermediate")
+
+        self.output = BertOutput(hidden_size=config.hidden_size,
+                                 intermediate_size=config.intermediate_size,
+                                 layer_norm_eps=config.layer_norm_eps,
+                                 quant_config=quant_config,
+                                 prefix=f"{prefix}.output")
+
+    def forward(self, hidden_states: torch.Tensor):
+        attn_output = self.attention(hidden_states)
+        intermediate_output = self.intermediate(attn_output)
+        output = self.output(intermediate_output, attn_output)
+        return output
+
+
+class BertAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_attention_heads: int,
+        layer_norm_eps: float,
+        # cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        self.self = BertSelfAttention(hidden_size=hidden_size,
+                                      num_attention_heads=num_attention_heads,
+                                    #   cache_config=cache_config,
+                                      quant_config=quant_config,
+                                      prefix=f"{prefix}.output")
+
+        self.output = BertSelfOutput(hidden_size=hidden_size,
+                                     layer_norm_eps=layer_norm_eps,
+                                     quant_config=quant_config,
+                                     prefix=f"{prefix}.output")
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        self_output = self.self(hidden_states)
+        return self.output(self_output, hidden_states)
+
+
+class BertSelfAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_attention_heads: int,
+        # cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+
+        self.total_num_heads = num_attention_heads
+        assert self.total_num_heads % tp_size == 0
+
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = self.total_num_heads
+        self.head_dim = self.hidden_size // self.total_num_heads
+        assert self.head_dim * self.total_num_heads == self.hidden_size
+
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=self.hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj")
+
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.num_heads_per_partition = divide(self.num_heads, self.tp_size)
+        
+        self.attn = LocalAttention(num_heads=self.num_heads,
+                              head_size=self.head_dim,
+                              num_kv_heads=self.num_kv_heads,
+                            #   softmax_scale=self.scaling,
+                            #   cache_config=cache_config,
+                            #   quant_config=quant_config,
+                            #   prefix=f"{prefix}.attn",
+                            #   attn_type=AttentionType.ENCODER_ONLY
+                            causal=False,
+                            supported_attention_backends=(_Backend.FLASH_ATTN,
+                                                           _Backend.TORCH_SDPA)
+                            ) # River TODO, fix hardcoded backend
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        B, L, _ = q.shape
+        q = q.reshape(B, L, self.num_heads, self.head_dim)
+        k = k.reshape(B, L, self.num_heads, self.head_dim)
+        v = v.reshape(B, L, self.num_heads, self.head_dim)
+        output = self.attn(q, k, v)
+        output = output.reshape(B, L, self.hidden_size)
+        return output
+    # def forward(
+    #     self,
+    #     hidden_states: torch.Tensor,
+    # ) -> torch.Tensor:
+    #     qkv, _ = self.qkv_proj(hidden_states)
+    #     q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+    #     output = self.attn(q, k, v)
+    #     return output
+
+class BertSelfOutput(nn.Module):
+
+    def __init__(self,
+                 hidden_size: int,
+                 layer_norm_eps: float,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 prefix: str = ""):
+        super().__init__()
+        self.dense = RowParallelLinear(input_size=hidden_size,
+                                       output_size=hidden_size,
+                                       bias=True,
+                                       quant_config=quant_config,
+                                       prefix=f"{prefix}.dense")
+        self.LayerNorm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
+
+    def forward(self, hidden_states: torch.Tensor,
+                input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = self.dense(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BertIntermediate(nn.Module):
+
+    def __init__(self,
+                 hidden_size: int,
+                 intermediate_size: int,
+                 hidden_act: str,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 prefix: str = ""):
+        super().__init__()
+        self.dense = ColumnParallelLinear(input_size=hidden_size,
+                                          output_size=intermediate_size,
+                                          bias=True,
+                                          quant_config=quant_config,
+                                          prefix=f"{prefix}.dense")
+        self.intermediate_act_fn = get_act_fn(hidden_act)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = self.dense(hidden_states)
+        hidden_states = self.intermediate_act_fn(hidden_states)
+        return hidden_states
+
+
+class BertOutput(nn.Module):
+
+    def __init__(self,
+                 hidden_size: int,
+                 intermediate_size: int,
+                 layer_norm_eps: float,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 prefix: str = ""):
+        super().__init__()
+
+        self.dense = RowParallelLinear(input_size=intermediate_size,
+                                       output_size=hidden_size,
+                                       bias=True,
+                                       quant_config=quant_config,
+                                       prefix=f"{prefix}.dense")
+
+        self.LayerNorm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
+
+    def forward(self, hidden_states: torch.Tensor,
+                input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = self.dense(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BertModel(TextEncoder):
+    packed_modules_mapping = {"qkv_proj": ["query", "key", "value"]}
+
+    def __init__(
+        self,
+        config: BertConfig,
+    ) -> None:
+        super().__init__(config)
+        self.embeddings = BertEmbedding(config=config)
+        print("prefix here ",config.prefix)
+        self.encoder = BertEncoder(config=config,
+                                prefix=f"{config.prefix}.encoder")
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        if inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        else:
+            hidden_states = self.embeddings(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                token_type_ids=token_type_ids)
+        return self.encoder(hidden_states)
+
+    def load_weights(self, weights: Iterable[Tuple[str,
+                                                   torch.Tensor]]) -> Set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "query", "q"),
+            ("qkv_proj", "key", "k"),
+            ("qkv_proj", "value", "v"),
+        ]
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: Set[str] = set()
+        for name, loaded_weight in weights:
+            name = name[len("bert."):] if name.startswith("bert.") else name
+            for (param_name, weight_name, shard_id) in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                # Skip loading extra bias for GPTQ models.
+                if name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader",
+                                        default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params

fastvideo/v1/models/encoders/clip.py

@@ -13,12 +13,12 @@ from fastvideo.v1.attention import LocalAttention
 from fastvideo.v1.configs.models.encoders import (BaseEncoderOutput,
                                                   CLIPTextConfig,
                                                   CLIPVisionConfig)
-from fastvideo.v1.configs.quantization import QuantizationConfig
 from fastvideo.v1.distributed import (divide,
                                       get_tensor_model_parallel_world_size)
 from fastvideo.v1.layers.activation import get_act_fn
 from fastvideo.v1.layers.linear import (ColumnParallelLinear, QKVParallelLinear,
                                         RowParallelLinear)
+from fastvideo.v1.layers.quantization import QuantizationConfig
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.encoders.base import ImageEncoder, TextEncoder
 from fastvideo.v1.models.encoders.vision import resolve_visual_encoder_outputs

fastvideo/v1/models/encoders/llama.py

@@ -32,12 +32,12 @@ from torch import nn
 from fastvideo.v1.attention import LocalAttention
 # from ..utils import (extract_layer_index)
 from fastvideo.v1.configs.models.encoders import BaseEncoderOutput, LlamaConfig
-from fastvideo.v1.configs.quantization import QuantizationConfig
 from fastvideo.v1.distributed import get_tensor_model_parallel_world_size
 from fastvideo.v1.layers.activation import SiluAndMul
 from fastvideo.v1.layers.layernorm import RMSNorm
 from fastvideo.v1.layers.linear import (MergedColumnParallelLinear,
                                         QKVParallelLinear, RowParallelLinear)
+from fastvideo.v1.layers.quantization import QuantizationConfig
 from fastvideo.v1.layers.rotary_embedding import get_rope
 from fastvideo.v1.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from fastvideo.v1.models.encoders.base import TextEncoder

fastvideo/v1/models/encoders/t5.py

@@ -28,13 +28,13 @@ import torch.nn.functional as F
 from torch import nn
 
 from fastvideo.v1.configs.models.encoders import BaseEncoderOutput, T5Config
-from fastvideo.v1.configs.quantization import QuantizationConfig
 from fastvideo.v1.distributed import (get_tensor_model_parallel_rank,
                                       get_tensor_model_parallel_world_size)
 from fastvideo.v1.layers.activation import get_act_fn
 from fastvideo.v1.layers.layernorm import RMSNorm
 from fastvideo.v1.layers.linear import (MergedColumnParallelLinear,
                                         QKVParallelLinear, RowParallelLinear)
+from fastvideo.v1.layers.quantization import QuantizationConfig
 from fastvideo.v1.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from fastvideo.v1.models.encoders.base import TextEncoder
 from fastvideo.v1.models.loader.weight_utils import default_weight_loader

fastvideo/v1/models/registry.py

@@ -36,7 +36,7 @@ _TEXT_ENCODER_MODELS = {
     "LlamaModel": ("encoders", "llama", "LlamaModel"),
     "UMT5EncoderModel": ("encoders", "t5", "UMT5EncoderModel"),
     "STEP1TextEncoder": ("encoders", "stepllm", "STEP1TextEncoder"),
-    "BertModel": ("encoders", "clip", "CLIPTextModel"),
+    "BertModel": ("encoders", "bert", "BertModel"),
 }
 
 _IMAGE_ENCODER_MODELS: dict[str, tuple] = {
@@ -295,7 +295,7 @@ class _ModelRegistry:
             model_cls = self._try_load_model_cls(arch)
             if model_cls is not None:
                 return (model_cls, arch)
-
+        print("unsupported ",architectures)
         return self._raise_for_unsupported(architectures)
 
 

fastvideo/v1/pipelines/stages/input_validation.py

@@ -3,8 +3,6 @@
 Input validation stage for diffusion pipelines.
 """
 
-import math
-
 import torch
 
 from fastvideo.v1.fastvideo_args import FastVideoArgs
@@ -88,24 +86,4 @@ class InputValidationStage(PipelineStage):
                 f"Guidance scale must be positive, but got {batch.guidance_scale}"
             )
 
-        # Adjust number of frames based on number of GPUs
-        temporal_scale_factor = fastvideo_args.vae_config.arch_config.temporal_compression_ratio
-        orig_latent_num_frames = (batch.num_frames -
-                                  1) // temporal_scale_factor + 1
-        if orig_latent_num_frames % fastvideo_args.num_gpus != 0:
-            if batch.num_frames_round_down:
-                new_latent_num_frames = math.floor(
-                    orig_latent_num_frames /
-                    fastvideo_args.num_gpus) * fastvideo_args.num_gpus
-            else:
-                new_latent_num_frames = math.ceil(
-                    orig_latent_num_frames /
-                    fastvideo_args.num_gpus) * fastvideo_args.num_gpus
-            new_num_frames = (new_latent_num_frames -
-                              1) * temporal_scale_factor + 1
-            logger.info(
-                "Adjusting number of frames from %s to %s based on number of GPUs (%s)",
-                batch.num_frames, new_num_frames, fastvideo_args.num_gpus)
-            batch.num_frames = new_num_frames
-
         return batch

fastvideo/v1/pipelines/stepvideo/stepvideo_pipeline.py

@@ -16,7 +16,7 @@ from huggingface_hub import hf_hub_download
 
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.models.encoders.bert import HunyuanClip  # type: ignore
+# from fastvideo.v1.models.encoders.bert_old import HunyuanClip  # type: ignore
 from fastvideo.v1.models.encoders.stepllm import STEP1TextEncoder
 from fastvideo.v1.models.loader.component_loader import PipelineComponentLoader
 from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
@@ -31,7 +31,7 @@ logger = init_logger(__name__)
 
 class StepVideoPipeline(ComposedPipelineBase):
 
-    _required_config_modules = ["transformer", "scheduler", "vae"]
+    _required_config_modules = ["transformer", "scheduler", "vae", "tokenizer", "text_encoder"]
 
     def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
         """Set up pipeline stages with proper dependency injection."""
@@ -41,8 +41,8 @@ class StepVideoPipeline(ComposedPipelineBase):
 
         self.add_stage(stage_name="prompt_encoding_stage",
                        stage=StepvideoPromptEncodingStage(
-                           stepllm=self.get_module("text_encoder"),
-                           clip=self.get_module("text_encoder_2"),
+                           stepllm=self.get_module("text_encoder_2"),
+                           clip=self.get_module("text_encoder"),
                        ))
 
         self.add_stage(stage_name="timestep_preparation_stage",
@@ -80,9 +80,9 @@ class StepVideoPipeline(ComposedPipelineBase):
         llm_dir = os.path.join(self.model_path, "step_llm")
         clip_dir = os.path.join(self.model_path, "hunyuan_clip")
         text_enc = self.build_llm(llm_dir, target_device)
-        clip_enc = self.build_clip(clip_dir, target_device)
-        self.add_module("text_encoder", text_enc)
-        self.add_module("text_encoder_2", clip_enc)
+        # clip_enc = self.build_clip(clip_dir, target_device)
+        self.add_module("text_encoder_2", text_enc)
+        # self.add_module("text_encoder", clip_enc)
         lib_path = (
             os.path.join(
                 fastvideo_args.model_path,
@@ -111,7 +111,7 @@ class StepVideoPipeline(ComposedPipelineBase):
             modules_config
         ) > 1, "model_index.json must contain at least one pipeline module"
 
-        required_modules = ["transformer", "scheduler", "vae"]
+        required_modules = ["transformer", "scheduler", "tokenizer", "text_encoder", "vae"]
         for module_name in required_modules:
             if module_name not in modules_config:
                 raise ValueError(

fastvideo/v1/platforms/__init__.py

@@ -21,10 +21,10 @@ def cuda_platform_plugin() -> Optional[str]:
         pynvml = import_pynvml()  # type: ignore[no-untyped-call]
         pynvml.nvmlInit()
         try:
-            # NOTE: Edge case: vllm cpu build on a GPU machine.
+            # NOTE: Edge case: fastvideo cpu build on a GPU machine.
             # Third-party pynvml can be imported in cpu build,
-            # we need to check if vllm is built with cpu too.
-            # Otherwise, vllm will always activate cuda plugin
+            # we need to check if fastvideo is built with cpu too.
+            # Otherwise, fastvideo will always activate cuda plugin
             # on a GPU machine, even if in a cpu build.
             is_cuda = (pynvml.nvmlDeviceGetCount() > 0)
         finally:
@@ -72,12 +72,12 @@ if TYPE_CHECKING:
 def __getattr__(name: str):
     if name == 'current_platform':
         # lazy init current_platform.
-        # 1. out-of-tree platform plugins need `from vllm.platforms import
+        # 1. out-of-tree platform plugins need `from fastvideo.platforms import
         #    Platform` so that they can inherit `Platform` class. Therefore,
         #    we cannot resolve `current_platform` during the import of
-        #    `vllm.platforms`.
+        #    `fastvideo.platforms`.
         # 2. when users use out-of-tree platform plugins, they might run
-        #    `import vllm`, some vllm internal code might access
+        #    `import fastvideo`, some fastvideo internal code might access
         #    `current_platform` during the import, and we need to make sure
         #    `current_platform` is only resolved after the plugins are loaded
         #    (we have tests for this, if any developer violate this, they will

fastvideo/v1/platforms/cuda.py

@@ -9,10 +9,6 @@ from functools import lru_cache, wraps
 from typing import Callable, List, Optional, Tuple, TypeVar, Union
 
 import torch
-# NOTE(will): this import is necessary to trigger the registration of the custom
-# ops from vllm, which we use
-# import custom ops, trigger op registration
-import vllm._C  # noqa
 from typing_extensions import ParamSpec
 
 import fastvideo.v1.envs as envs

fastvideo/v1/tests/encoders/test_bert_encoder.py

@@ -0,0 +1,158 @@
+import sys
+sys.argv = [sys.argv[0]] 
+import os
+from fastvideo.v1.models.loader.component_loader import VAELoader, TokenizerLoader, TextEncoderLoader
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+import torch
+from fastvideo.v1.configs.models.vaes import StepVideoVAEConfig
+from fastvideo.models.stepvideo.text_encoder.clip import HunyuanClip
+from fastvideo.v1.configs.models.encoders.bert import BertConfig
+import torch.distributed as dist
+from fastvideo.v1.distributed.parallel_state import initialize_model_parallel
+import pytest
+from fastvideo.v1.forward_context import set_forward_context
+from types import SimpleNamespace
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29500")
+# if not dist.is_initialized():
+#     dist.init_process_group("gloo", rank=0, world_size=1)
+# initialize_model_parallel(tensor_model_parallel_size=1)
+
+model_path = 'cache/hub/models--FastVideo--stepvideo-t2v-diffusers/snapshots/572e0ce299de9fe2f8b843afe5afce5facb23c13'
+VAE_PATH = model_path + '/vae'
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+precision = torch.bfloat16
+precision_str = "bf16"
+args = FastVideoArgs(model_path=model_path,
+                        text_encoder_precisions=("fp16",),
+                        text_encoder_configs=(BertConfig(),))
+args.device_str = "cuda:0"
+# args = FastVideoArgs(model_path=model_path, vae_precision=precision_str)
+args.vae_config = StepVideoVAEConfig()
+PROMPTS = [
+    "A beautiful sunset over the mountains",
+    "A futuristic city skyline at night",
+    "A serene beach with palm trees and clear water",
+    "A bustling market street in a vibrant city",
+    "A cozy cabin in the woods during winter"
+]
+
+@pytest.mark.usefixtures("distributed_setup")
+def build_fastvideo_encoder(old_model_path):
+    """Instantiate the FV tokenizer + text encoder the same way your loader will
+    inside the pipeline, but *outside* the full pipeline for easier unit test."""
+
+    tokenizer = TokenizerLoader().load(
+        model_path=os.path.join(old_model_path, "tokenizer"),
+        architecture="bert",
+        fastvideo_args=args,
+    )
+
+    encoder = TextEncoderLoader().load(
+        model_path=os.path.join(old_model_path, "text_encoder"),
+        architecture="bert",
+        fastvideo_args=args,
+    ).to(dtype=precision, device=device)
+
+    return tokenizer, encoder  
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_hidden_and_pooled_equivalence():
+    # ----- original HF path
+    old_model_path = os.path.join(model_path, "hunyuan_clip")
+    legacy = HunyuanClip(old_model_path).to(dtype=precision, device=device).eval()
+    with torch.no_grad():
+        h_legacy, text_inputs = legacy(PROMPTS)
+
+    # ----- fastvideo path
+    tok, fv_enc = build_fastvideo_encoder(old_model_path)
+    tok_out = tok(
+        PROMPTS,
+        padding="max_length",
+        max_length=legacy.max_length,
+        truncation=True,
+        return_attention_mask=True,
+        return_tensors="pt",
+    )
+    # 2) Get the “fastvideo” tokenizer outputs
+    fv_tok = tok_out  # already a BatchEncoding from your TokenizerLoader
+
+    # 3) Make sure they have the same keys
+    assert set(text_inputs.keys()) == set(fv_tok.keys()), \
+        f"Mismatch in keys: {text_inputs.keys()} vs {fv_tok.keys()}"
+
+    # 4) For each tensor, compare shape and contents
+    for key in text_inputs.keys():
+        a = text_inputs[key]
+        b = fv_tok[key]
+        # Move both to CPU so the comparison is unambiguous
+        a_cpu = a.cpu()
+        b_cpu = b.cpu()
+        # Check shapes
+        assert a_cpu.shape == b_cpu.shape, f"Shape mismatch for `{key}`: {a_cpu.shape} vs {b_cpu.shape}"
+        # Check values exactly
+        if a_cpu.dtype.is_floating_point:
+            # floats: allow exact equality since token IDs/masks are integers
+            assert torch.equal(a_cpu, b_cpu), f"Floating values differ for `{key}`"
+            print(key, torch.equal(a_cpu, b_cpu))
+        else:
+            # ints (input_ids, attention_mask): must be identical
+            assert torch.equal(a_cpu, b_cpu), f"Integer tensors differ for `{key}`"
+            print(key, torch.equal(a_cpu, b_cpu))
+        
+        
+    seq_len      = tok_out.input_ids.size(1)               # 77
+    position_ids = torch.arange(seq_len, device=device).unsqueeze(0)
+    position_ids = position_ids.expand(tok_out.input_ids.size(0), -1)
+
+
+    with set_forward_context(current_timestep=0, attn_metadata=None):
+        fv_out = fv_enc(
+            input_ids=tok_out.input_ids.to(device),
+            position_ids=position_ids.to(device),
+            attention_mask=tok_out.attention_mask.to(device),
+        )
+    h_fv= fv_out
+    
+    # assume h_legacy, h_fv are Float/BFloat tensors on the same device
+    diff = (h_legacy - h_fv).abs()
+
+    # overall stats
+    max_diff  = diff.max().item()
+    mean_diff = diff.mean().item()
+    print(f"[DEBUG] overall max abs diff   = {max_diff:.6f}")
+    print(f"[DEBUG] overall mean abs diff  = {mean_diff:.6f}")
+
+    # per‐position mean diff (averaged over hidden dim)
+    per_token = diff.mean(dim=-1)      # shape [batch, seq_len]
+    print("[DEBUG] per-token mean diff:", per_token[0].tolist())
+
+    # if you only want to compare *valid* tokens (mask==1)
+    mask = tok_out.attention_mask.to(device).unsqueeze(-1)   # [B, L, 1]
+    valid_diff = (h_legacy - h_fv).abs()[mask.bool().expand_as(h_legacy)]
+    print(f"[DEBUG] valid-token max diff = {valid_diff.max().item():.6f}")
+    print(f"[DEBUG] valid-token mean diff= {valid_diff.mean().item():.6f}")
+
+    # after you have `legacy` and `fv_enc` on the same device
+    legacy_state = {n: p.detach().cpu() for n, p in legacy.text_encoder.named_parameters()}
+    fv_state     = {n: p.detach().cpu() for n, p in fv_enc.named_parameters()}
+    print(len(legacy_state.keys()), len(fv_state.keys()))
+    print([k for k in fv_state.keys() if not k.startswith('layer.')])
+    for name in ["embeddings.word_embeddings.weight",
+                "embeddings.position_embeddings.weight"]:
+        a = legacy_state[name]
+        b = fv_state[name][:47020,:]
+        print(f"for {name}: {a.shape}, {b.shape}")
+        print(f"{name}  max abs diff = {(a - b).abs().max():.6f}")
+
+    for k in legacy_state.keys():
+        if k not in fv_state:
+            print(f"missing in fv_state: {k}")
+    # ----- 5.  assertions
+    assert h_legacy.shape == h_fv.shape
+    # assert p_legacy.shape == p_fv.shape
+    assert torch.allclose(h_legacy, h_fv, atol=1e-4, rtol=1e-3)
+    # assert torch.allclose(p_legacy, p_fv)
+
+if __name__ == "__main__":
+    test_hidden_and_pooled_equivalence()

fastvideo/v1/utils.py

@@ -180,8 +180,33 @@ class FlexibleArgumentParser(argparse.ArgumentParser):
             else:
                 processed_args.append(arg)
 
-        return super().parse_args(  # type: ignore[no-any-return]
-            processed_args, namespace)
+        namespace = super().parse_args(processed_args, namespace)
+
+        # Track which arguments were explicitly provided
+        namespace._provided = set()
+
+        i = 0
+        while i < len(args):
+            arg = args[i]
+            if arg.startswith('--'):
+                # Handle --key=value format
+                if '=' in arg:
+                    key = arg.split('=')[0][2:].replace('-', '_')
+                    namespace._provided.add(key)
+                    i += 1
+                # Handle --key value format
+                else:
+                    key = arg[2:].replace('-', '_')
+                    namespace._provided.add(key)
+                    # Skip the value if there is one
+                    if i + 1 < len(args) and not args[i + 1].startswith('-'):
+                        i += 2
+                    else:
+                        i += 1
+            else:
+                i += 1
+
+        return namespace  # type: ignore[no-any-return]
 
     def _pull_args_from_config(self, args: List[str]) -> List[str]:
         """Method to pull arguments specified in the config file

fastvideo/v1/worker/multiproc_executor.py

@@ -42,7 +42,6 @@ class MultiprocExecutor(Executor):
                                             pipe=worker_pipe))
             worker.start()
             self.workers.append(worker)
-        logger.info("Workers: %s", self.workers)
 
         # Wait for all workers to be ready
         for idx, pipe in enumerate(self.worker_pipes):

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fastvideo"
-version = "0.0.5"
+version = "0.1.0"
 description = "FastVideo"
 readme = "README.md"
 requires-python = ">=3.8"
@@ -22,9 +22,6 @@ dependencies = [
     "timm==1.0.11", "peft==0.13.2", "diffusers>=0.33.0", "bitsandbytes",
     "torch==2.6.0", "torchvision",
 
-    # vLLM
-    "vllm>=0.7.3",
-    
     # Acceleration & Optimization
     "accelerate==1.0.1",