⚡ComfyUI-Lightning

Introduction

This repository integrates all the tricks I know to speed up Flux inference:

1. Use TeaCache or FBCache or MBCache or ToCa;
2. Skip some unnessasery blocks;
3. Compile and quantize model;
4. Use fast CuDNN attention kernels;
5. Use SageAttention or SpargeAttn;
6. Fix AttributeError: 'SymInt' object has no attribute 'size' to speed up recompilation after resolution changing.

MBCache extends FBCache and is applied to cache multiple blocks. The codes are modified from SageAttention, ComfyUI-TeaCache, comfyui-flux-accelerator and Comfy-WaveSpeed. More details see above given repositories.

Updates

• [2025/3/10] Add SpargeAttn. For more details, see Usage.

• [2025/2/27] Add ToCa.

• [2025/1/24] Now support Sana. Get your 1024*1024 images within 2s. All the codes are modified from Sana.

Usage

For Flux

You can use XXCache, SageAttention, and torch.compile with the following examples:

More specific:

1. Download Flux diffusion model and VAE image decoder from FLUX.1-dev or FLUX.1-schnell. Put the flux1-dev.safetensors or flux1-schnell.safetensors file into models/diffusion_models and the ae.safetensors file into models/vae;

2. Download Flux text encoder from flux_text_encoders and put all the .safetensors files into models/clip;

3. Run the example workflow.

For Sana

1. Download Sana diffusion model from Model Zoo and put the .pth file into models/diffusion_models;

2. Download Gemma text encoder from google/gemma-2-2b-it, unsloth/gemma-2b-it-bnb-4bit or Efficient-Large-Model/gemma-2-2b-it and put the whole folder into models/text_encoders;

3. Download DCAE image decoder from mit-han-lab/dc-ae-f32c32-sana-1.0 and put the .safetensors file into models/vae;

4. Run the example workflow.

For SpargeAttn

1. First you should follow the steps below to install SpargeAttn. If you have problems installing it, see the original repository;

   git clone https://github.com/thu-ml/SpargeAttn.git
   cd ./SpargeAttn
   pip install -e .
   

2. If you do not have a hyperparameter file, you should perform a few rounds of quality fine-tuning to get one first. You just need to open the enable_tuning_mode of the node Apply SpargeAttn and perform the generation. For example, generate 50 steps of 512*512 images at 10 different prompts (very time-consuming);

   

   • The skip_DoubleStreamBlocks and skip_SingleStreamBlocks arguments are used to skip certain blocks that do not require the use of SpargeAttn, mainly to work with TeaCache and FBCache.
   • If you have a well-tuned hyperparameter file, feel free to share it.

3. Turn off enable_tuning_mode and use the Save Finetuned SpargeAttn Hyperparams node to save your hyperparameter file;

   

4. Remove or disable the Save Finetuned SpargeAttn Hyperparams node and place the saved hyperparameter file in the models/checkpoints folder. Load this hyperparameter file in the Apply SpargeAttn node;

   

5. Enjoy yourself.