Update model card for PSA: Pyramid Sparse Attention

README.md

@@ -1,288 +1,88 @@
 ---
 license: apache-2.0
+pipeline_tag: text-to-video
+library_name: diffusers
 ---
-# BLADE: Block-Sparse Attention Meets Step Distillation for Efficient Video Generation
 
-<div align="center">
+# PSA: Pyramid Sparse Attention for Efficient Video Understanding and Generation
 
-[📖 Paper](https://arxiv.org/abs/2508.10774) | [🚀 Homepage](http://ziplab.co/BLADE-Homepage/) | [💾 Models](https://huggingface.co/GYP666/BLADE) | [📖 中文阅读](README_zh.md)  
+[📖 Paper](https://huggingface.co/papers/2512.04025) | [🚀 Project Page](http://ziplab.co/PSA) | [💻 Code](https://github.com/ziplab/Pyramid-Sparse-Attention)
 
-</div>
+Official PyTorch implementation of [PSA: Pyramid Sparse Attention for Efficient Video Understanding and Generation](https://huggingface.co/papers/2512.04025).
 
-BLADE is a data-free framework for efficient video generation. By jointly training an adaptive sparse attention mechanism with a step distillation technique, it achieves a significant acceleration in video generation models. This project combines a block-sparse attention mechanism with step distillation, reducing the number of inference steps from 50 to just 8 while maintaining high-quality generation.
+<p align="center">
+  <img src="https://github.com/ziplab/Pyramid-Sparse-Attention/raw/main/figures/prompt007comparison.jpg" width="100%">
+</p>
 
-## 📢 News
+<p align="center"><em>Visual comparison of sparse attention methods at similar sparsity levels (~90%). PSA maintains visual fidelity close to full attention while other methods show noticeable artifacts.</em></p>
 
-  - **[Aug 2025]** 🎉 The code and pre-trained models for BLADE have been released\!
-  - **[Aug 2025]** 📝 Support for two mainstream video generation models, CogVideoX-5B and WanX-1.3B, is now available.
-  - **[Aug 2025]** ⚡ Achieved high-quality video generation in just 8 steps, a significant speedup compared to the 50-step baseline.
+Pyramid Sparse Attention (PSA) is a versatile attention module designed to overcome the quadratic complexity bottleneck of attention mechanisms in foundation models. It introduces multi-level pooled Key-Value (KV) representations, enabling a finer mask granularity than traditional binary masking approaches. This design allows critical KV blocks to receive full resolution attention while less important blocks utilize progressively pooled representations, creating an informative interpolation between full retention and complete pruning. This approach effectively mitigates information loss and preserves computational efficiency. PSA is applicable to both video understanding and generation tasks, consistently outperforming or achieving comparable performance to existing sparse attention baselines with superior efficiency-quality trade-offs.
 
-## ✨ Key Features
+> **Note:** This release focuses on **inference-only** with **bidirectional attention**. Support for causal attention masks and backward propagation (training) is still under optimization and will be released in a future update.
 
-  - 🚀 **Efficient Inference**: Reduces the number of inference steps from 50 to 8 while preserving generation quality.
-  - 🎯 **Adaptive Sparse Attention**: Employs a block-sparse attention mechanism to significantly reduce computational complexity.
-  - 📈 **Step Distillation**: Utilizes the Trajectory Distillation Method (TDM), enabling training without the need for video data.
-  - 🎮 **Plug-and-Play**: Supports CogVideoX-5B and WanX-1.3B models without requiring modifications to their original architectures.
+## Installation
 
-## 🛠️ Environment Setup
-
-### System Requirements
-
-  - Python \>= 3.11 (Recommended)
-  - CUDA \>= 11.6 (Recommended)
-  - GPU Memory \>= 24GB (for Inference)
-  - GPU Memory \>= 80GB (for Training)
-
-### Installation Steps
-
-1.  **Clone the repository**
-
-    ```bash
-    git clone https://github.com/Tacossp/BLADE
-    cd BLADE
-    ```
-
-2.  **Install dependencies**
-
-    ```bash
-    # Install using uv (Recommended)
-    uv pip install -r requirements.txt
-
-    # Or use pip
-    pip install -r requirements.txt
-    ```
-
-3.  **Compile the Block-Sparse-Attention library**
-
-    ```bash
-    git clone https://github.com/mit-han-lab/Block-Sparse-Attention.git
-    cd Block-Sparse-Attention
-    pip install packaging
-    pip install ninja
-    python setup.py install
-    cd ..
-    ```
-
-## 📥 Model Weights Download
-
-### Base Model Weights
-
-Please download the following base model weights and place them in the specified directories:
-
-1.  **CogVideoX-5B Model**
-
-    ```bash
-    # Download from Hugging Face
-    git lfs install
-    git clone https://huggingface.co/zai-org/CogVideoX-5b cogvideox/CogVideoX-5b
-    ```
-
-2.  **WanX-1.3B Model**
-
-    ```bash
-    # Download from Hugging Face
-    git clone https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers wanx/wan1.3b
-    ```
-
-### Pre-trained BLADE Weights
-
-We provide pre-trained weights for BLADE:
+### Using uv (Recommended)
 
 ```bash
-# Download pre-trained weights
-git clone https://huggingface.co/GYP666/BLADE pretrained_weights
+uv venv --python 3.11
+source .venv/bin/activate
+uv pip install -e .
 ```
 
-### Weight Directory Structure
-
-Ensure your directory structure for weights is as follows:
-
-```
-BLADE/
-├── cogvideox/
-│   └── CogVideoX-5b/           # Base model weights for CogVideoX
-├── wanx/
-│   └── wan1.3b/               # Base model weights for WanX
-└── pretrained_weights/         # Pre-trained weights for BLADE
-    ├── BLADE_cogvideox_weight/
-    └── BLADE_wanx_weight/
-```
-
-## 🚀 Quick Start - Inference
-
-### CogVideoX Inference
+### Using pip
 
 ```bash
-cd cogvideox
-python train/inference.py \
-    --lora_path ../pretrained_weights/cogvideox_checkpoints/your_checkpoint \
-    --gpu 0
+python -m venv .venv
+source .venv/bin/activate
+pip install -e .
 ```
 
-**Argument Descriptions**:
-
-  - `--lora_path`: Path to the LoRA weights file.
-  - `--gpu`: The ID of the GPU device to use (Default: 0).
+> For best performance, we recommend using PyTorch nightly version.
 
-**Output**: The generated videos will be saved in the `cogvideox/outputs/inference/` directory.
+## Download Weights
 
-### WanX Inference
+### CogVideoX-5B LoRA (4-step)
 
 ```bash
-cd wanx
-python train/inference.py \
-    --lora_path ../pretrained_weights/wanx_checkpoints/your_checkpoint \
-    --gpu 0
+huggingface-cli download GYP666/BLADE cogvideox-5b-psa-lora/pytorch_lora_weights.safetensors --local-dir ./weights
 ```
 
-**Output**: The generated videos will be saved in the `wanx/outputs/` directory.
-
-## 🔧 Training Process
+**Note:** After downloading, update the `lora_path` in `examples/configs/model_configs.py` to point to your weights directory.
 
-### Step 1: Prompt Preprocessing
+## Quick Start (Inference)
 
-Before training, you need to preprocess the text prompts to generate embeddings.
-
-#### CogVideoX Preprocessing
-
-```bash
-cd utils
-python process_prompts_cogvideox.py \
-    --input_file your_prompts.txt \
-    --output_dir ../cogvideox/prompts \
-    --model_path ../cogvideox/CogVideoX-5b \
-    --batch_size 32 \
-    --save_separate
-```
-
-**Argument Descriptions**:
-
-  - `--input_file`: A `.txt` file containing prompts, with one prompt per line.
-  - `--output_dir`: The directory to save the output embeddings.
-  - `--model_path`: Path to the CogVideoX model.
-  - `--batch_size`: The batch size for processing.
-  - `--save_separate`: Whether to save each embedding as a separate file.
-
-#### WanX Preprocessing
-
-```bash
-cd utils
-python process_prompts_wanx.py
-```
-
-This script will automatically process the prompts in `utils/all_dimension_aug_wanx.txt` and generate the corresponding embeddings.
-
-### Step 2: Start Training
-
-#### CogVideoX Training
+### CogVideoX1.5-5B
 
 ```bash
-cd cogvideox
-bash train_tdm_1.sh
+python examples/inference/cogvideo/cogvideo_5b.py \
+    --model cogvideo1.5_5b \
+    --prompt "your prompt here" \
+    --use_psa
 ```
 
-**Core Training Parameters**:
+### Wan2.1-1.3B
 
 ```bash
-# If not training with 8 GPUs, you must modify CUDA_VISIBLE_DEVICES and the num_processes in config.yaml
-CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch \
-    --config_file train/config.yaml \
-    train/train_cogvideo_tdm.py \
-    --pretrained_model_name_or_path CogVideoX-5b \        # Path to the base model
-    --mixed_precision bf16 \                              # Use mixed-precision for reduced memory usage
-    --train_batch_size 5 \                                # Training batch size
-    --gradient_accumulation_steps 4 \                     # Number of gradient accumulation steps
-    --learning_rate 1e-4 \                                # Learning rate for the student model
-    --learning_rate_g 1e-4 \                              
-    --learning_rate_fake 5e-4 \                           # Learning rate for the fake model
-    --lambda_reg 0.5 \                                    # Regularization weight
-    --k_step 8 \                                          # Target number of steps for distillation
-    --cfg 3.5 \                                           # Classifier-Free Guidance scale
-    --eta 0.9 \                                           # ETA parameter for DDIM
-    --use_sparsity true \                                 # Enable sparse attention
-    --rank 64 \
-    --lora_alpha 64 \                                     # LoRA configuration
-    --max_train_steps 300 \                               # Maximum number of training steps
-    --checkpointing_steps 15 \                            # Interval for saving checkpoints
-    --gradient_checkpointing \                            # Use gradient checkpointing to save memory
-    --enable_slicing \
-    --enable_tiling                                       # VAE memory optimization
+python examples/inference/wan21/wan21_1.3b.py \
+    --prompt "your prompt here" \
+    --use_psa --no_warmup
 ```
 
-#### WanX Training
+For more inference examples, see [examples/README.md](https://github.com/ziplab/Pyramid-Sparse-Attention/blob/main/examples/README.md).
 
-```bash
-cd wanx
-bash train_wanx_tdm.sh
-```
-
-## 📊 Project Structure
-
-```
-BLADE/
-├── README.md                   # Project documentation
-├── requirements.txt           # List of Python dependencies
-│
-├── cogvideox/                 # Code related to CogVideoX
-│   ├── CogVideoX-5b/         # Directory for base model weights
-│   ├── train/                # Training scripts
-│   │   ├── inference.py      # Inference script
-│   │   ├── train_cogvideo_tdm.py  # Training script
-│   │   ├── train_tdm_1.sh    # Script to launch training
-│   │   ├── modify_cogvideo.py # Model modification script
-│   │   └── config.yaml       # Training configuration file
-│   ├── prompts/              # Preprocessed prompts and embeddings
-│   └── outputs/              # Output from training and inference
-│
-├── wanx/                     # Code related to WanX
-│   ├── wan1.3b/             # Directory for base model weights
-│   ├── train/               # Training scripts
-│   │   ├── inference.py     # Inference script
-│   │   ├── train_wanx_tdm.py # Training script
-│   │   ├── train_wanx_tdm.sh # Script to launch training
-│   │   └── modify_wan.py    # Model modification script
-│   ├── prompts/             # Preprocessed prompts and embeddings
-│   └── outputs/             # Output from training and inference
-│
-├── utils/                   # Utility scripts
-│   ├── process_prompts_cogvideox.py  # Data preprocessing for CogVideoX
-│   ├── process_prompts_wanx.py       # Data preprocessing for WanX
-│   └── all_dimension_aug_wanx.txt    # Training prompts for WanX
-│
-├── Block-Sparse-Attention/  # Sparse attention library
-│   ├── setup.py            # Compilation and installation script
-│   ├── block_sparse_attn/  # Core library code
-│   └── README.md           # Library usage instructions
-│
-└── ds_config.json          # DeepSpeed configuration file
-```
-
-## 🤝 Acknowledgements
+## Citation
 
-  - [FlashAttention](https://github.com/Dao-AILab/flash-attention), [Block-Sparse-Attention](https://github.com/mit-han-lab/Block-Sparse-Attention): For the foundational work on sparse attention.
-  - [CogVideoX](https://github.com/THUDM/CogVideo), [Wan2.1](https://github.com/Wan-Video/Wan2.1): For the supported models.
-  - [TDM](https://www.google.com/search?q=https://github.com/Luo-Yihong/TDM): For the foundational work on distillation implementation.
-  - [Diffusers](https://github.com/huggingface/diffusers): For the invaluable diffusion models library.
-
-## 📄 Citation
-
-If you use BLADE in your research, please cite our work:
+If you find this work useful, please cite our paper:
 
 ```bibtex
-@misc{gu2025videobladeblocksparseattentionmeets,
-    title={BLADE: Block-Sparse Attention Meets Step Distillation for Efficient Video Generation},
-    author={Youping Gu and Xiaolong Li and Yuhao Hu and Bohan Zhuang},
-    year={2025},
-    eprint={2508.10774},
-    archivePrefix={arXiv},
-    primaryClass={cs.CV},
-    url={https://arxiv.org/abs/2508.10774},
+@misc{li2025psapyramidsparseattention,
+      title={PSA: Pyramid Sparse Attention for Efficient Video Understanding and Generation}, 
+      author={Xiaolong Li and Youping Gu and Xi Lin and Weijie Wang and Bohan Zhuang},
+      year={2025},
+      eprint={2512.04025},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2512.04025}, 
 }
-```
-
-## 📧 Contact
-
-For any questions or suggestions, feel free to:
-
-  - Contact Youping Gu at youpgu71@gmail.com.
-  - Submit an issue on our [Github page](https://github.com/ziplab/BLADE/issues).
+```