Home

Awesome

<p align="center" style="border-radius: 10px"> <img src="asset/logo.png" width="35%" alt="logo"/> </p>

⚑️Sana: Efficient High-Resolution Image Synthesis with Linear Diffusion Transformer

<div align="center"> <a href="https://nvlabs.github.io/Sana/"><img src="https://img.shields.io/static/v1?label=Project&message=Github&color=blue&logo=github-pages"></a> &ensp; <a href="https://hanlab.mit.edu/projects/sana/"><img src="https://img.shields.io/static/v1?label=Page&message=MIT&color=darkred&logo=github-pages"></a> &ensp; <a href="https://arxiv.org/abs/2410.10629"><img src="https://img.shields.io/static/v1?label=Arxiv&message=Sana&color=red&logo=arxiv"></a> &ensp; <a href="https://nv-sana.mit.edu/"><img src="https://img.shields.io/static/v1?label=Demo:8x3090&message=MIT&color=yellow"></a> &ensp; <a href="https://replicate.com/chenxwh/sana"><img src="https://img.shields.io/static/v1?label=API:H100&message=Replicate&color=pink"></a> &ensp; <a href="https://discord.gg/rde6eaE5Ta"><img src="https://img.shields.io/static/v1?label=Discuss&message=Discord&color=purple&logo=discord"></a> &ensp; </div> <p align="center" border-raduis="10px"> <img src="asset/Sana.jpg" width="90%" alt="teaser_page1"/> </p>

πŸ’‘ Introduction

We introduce Sana, a text-to-image framework that can efficiently generate images up to 4096 Γ— 4096 resolution. Sana can synthesize high-resolution, high-quality images with strong text-image alignment at a remarkably fast speed, deployable on laptop GPU. Core designs include:

(1) DC-AE: unlike traditional AEs, which compress images only 8Γ—, we trained an AE that can compress images 32Γ—, effectively reducing the number of latent tokens.
(2) Linear DiT: we replace all vanilla attention in DiT with linear attention, which is more efficient at high resolutions without sacrificing quality.
(3) Decoder-only text encoder: we replaced T5 with modern decoder-only small LLM as the text encoder and designed complex human instruction with in-context learning to enhance the image-text alignment.
(4) Efficient training and sampling: we propose Flow-DPM-Solver to reduce sampling steps, with efficient caption labeling and selection to accelerate convergence.

As a result, Sana-0.6B is very competitive with modern giant diffusion model (e.g. Flux-12B), being 20 times smaller and 100+ times faster in measured throughput. Moreover, Sana-0.6B can be deployed on a 16GB laptop GPU, taking less than 1 second to generate a 1024 Γ— 1024 resolution image. Sana enables content creation at low cost.

<p align="center" border-raduis="10px"> <img src="asset/model-incremental.jpg" width="90%" alt="teaser_page2"/> </p>

πŸ”₯πŸ”₯ News

Performance

Methods (1024x1024)Throughput (samples/s)Latency (s)Params (B)SpeedupFID πŸ‘‡CLIP πŸ‘†GenEval πŸ‘†DPG πŸ‘†
FLUX-dev0.0423.012.01.0Γ—10.1527.470.6784.0
Sana-0.6B1.70.90.639.5Γ—5.8128.360.6483.6
Sana-0.6B-MultiLing1.70.90.639.5Γ—5.61<u>28.80</u><u>0.68</u>84.2
Sana-1.6B1.01.21.623.3Γ—<u>5.76</u>28.670.6684.8
Sana-1.6B-MultiLing1.01.21.623.3Γ—5.9228.940.69<u>84.5</u>
<details> <summary><h3>Click to show all</h3></summary>
MethodsThroughput (samples/s)Latency (s)Params (B)SpeedupFID πŸ‘†CLIP πŸ‘†GenEval πŸ‘†DPG πŸ‘†
512 Γ— 512 resolution
PixArt-Ξ±1.51.20.61.0Γ—6.1427.550.4871.6
PixArt-Ξ£1.51.20.61.0Γ—6.3427.62<u>0.52</u>79.5
Sana-0.6B6.70.80.65.0Γ—<u>5.67</u><u>27.92</u>0.64<u>84.3</u>
Sana-1.6B3.80.61.62.5Γ—5.1628.190.6685.5
1024 Γ— 1024 resolution
LUMINA-Next0.129.12.02.8Γ—7.5826.840.4674.6
SDXL0.156.52.63.5Γ—6.6329.030.5574.7
PlayGroundv2.50.215.32.64.9Γ—6.0929.130.5675.5
Hunyuan-DiT0.0518.21.51.2Γ—6.5428.190.6378.9
PixArt-Ξ£0.42.70.69.3Γ—6.1528.260.5480.5
DALLE3------0.6783.5
SD3-medium0.284.42.06.5Γ—11.9227.830.62<u>84.1</u>
FLUX-dev0.0423.012.01.0Γ—10.1527.470.6784.0
FLUX-schnell0.52.112.011.6Γ—7.9428.140.7184.8
Sana-0.6B1.70.90.639.5Γ—<u>5.81</u>28.360.6483.6
Sana-1.6B1.01.21.623.3Γ—5.76<u>28.67</u><u>0.66</u>84.8
</details>

Contents

πŸ”§ 1. Dependencies and Installation

git clone https://github.com/NVlabs/Sana.git
cd Sana

./environment_setup.sh sana
# or you can install each components step by step following environment_setup.sh

πŸ’» 2. How to Play with Sana (Inference)

πŸ’°Hardware requirement

πŸ”› Choose your model: Model card

πŸ”› Quick start with Gradio

# official online demo
DEMO_PORT=15432 \
python app/app_sana.py \
    --share \
    --config=configs/sana_config/1024ms/Sana_1600M_img1024.yaml \
    --model_path=hf://Efficient-Large-Model/Sana_1600M_1024px/checkpoints/Sana_1600M_1024px.pth \
    --image_size=1024

1. How to use SanaPipeline with 🧨diffusers

[!IMPORTANT] Upgrade your diffusers>=0.32.0.dev to make the SanaPipeline and SanaPAGPipeline available!

pip install git+https://github.com/huggingface/diffusers

Make sure to specify pipe.transformer to default torch_dtype and variant according to Model Card.

Set pipe.text_encoder to BF16 and pipe.vae to FP32 or BF16. For more info, docs are here.

# run `pip install git+https://github.com/huggingface/diffusers` before use Sana in diffusers
import torch
from diffusers import SanaPipeline

pipe = SanaPipeline.from_pretrained(
    "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers",
    variant="bf16",
    torch_dtype=torch.bfloat16,
)
pipe.to("cuda")

pipe.vae.to(torch.bfloat16)
pipe.text_encoder.to(torch.bfloat16)

prompt = 'a cyberpunk cat with a neon sign that says "Sana"'
image = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    guidance_scale=4.5,
    num_inference_steps=20,
    generator=torch.Generator(device="cuda").manual_seed(42),
)[0]

image[0].save("sana.png")

2. How to use SanaPAGPipeline with 🧨diffusers

# run `pip install git+https://github.com/huggingface/diffusers` before use Sana in diffusers
import torch
from diffusers import SanaPAGPipeline

pipe = SanaPAGPipeline.from_pretrained(
  "Efficient-Large-Model/Sana_1600M_1024px_diffusers",
  variant="fp16",
  torch_dtype=torch.float16,
  pag_applied_layers="transformer_blocks.8",
)
pipe.to("cuda")

pipe.text_encoder.to(torch.bfloat16)
pipe.vae.to(torch.bfloat16)

prompt = 'a cyberpunk cat with a neon sign that says "Sana"'
image = pipe(
    prompt=prompt,
    guidance_scale=5.0,
    pag_scale=2.0,
    num_inference_steps=20,
    generator=torch.Generator(device="cuda").manual_seed(42),
)[0]
image[0].save('sana.png')
<details> <summary><h3>3. How to use Sana in this repo</h3></summary>
import torch
from app.sana_pipeline import SanaPipeline
from torchvision.utils import save_image

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
generator = torch.Generator(device=device).manual_seed(42)

sana = SanaPipeline("configs/sana_config/1024ms/Sana_1600M_img1024.yaml")
sana.from_pretrained("hf://Efficient-Large-Model/Sana_1600M_1024px_BF16/checkpoints/Sana_1600M_1024px_BF16.pth")
prompt = 'a cyberpunk cat with a neon sign that says "Sana"'

image = sana(
    prompt=prompt,
    height=1024,
    width=1024,
    guidance_scale=5.0,
    pag_guidance_scale=2.0,
    num_inference_steps=18,
    generator=generator,
)
save_image(image, 'output/sana.png', nrow=1, normalize=True, value_range=(-1, 1))
</details> <details> <summary><h3>4. Run Sana (Inference) with Docker</h3></summary>
# Pull related models
huggingface-cli download google/gemma-2b-it
huggingface-cli download google/shieldgemma-2b
huggingface-cli download mit-han-lab/dc-ae-f32c32-sana-1.0
huggingface-cli download Efficient-Large-Model/Sana_1600M_1024px

# Run with docker
docker build . -t sana
docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \
    -v ~/.cache:/root/.cache \
    sana
</details>

πŸ”› Run inference with TXT or JSON files

# Run samples in a txt file
python scripts/inference.py \
      --config=configs/sana_config/1024ms/Sana_1600M_img1024.yaml \
      --model_path=hf://Efficient-Large-Model/Sana_1600M_1024px/checkpoints/Sana_1600M_1024px.pth \
      --txt_file=asset/samples_mini.txt

# Run samples in a json file
python scripts/inference.py \
      --config=configs/sana_config/1024ms/Sana_1600M_img1024.yaml \
      --model_path=hf://Efficient-Large-Model/Sana_1600M_1024px/checkpoints/Sana_1600M_1024px.pth \
      --json_file=asset/samples_mini.json

where each line of asset/samples_mini.txt contains a prompt to generate

πŸ”₯ 3. How to Train Sana

πŸ’°Hardware requirement

We provide a training example here and you can also select your desired config file from config files dir based on your data structure.

To launch Sana training, you will first need to prepare data in the following formats. Here is an example for the data structure for reference.

asset/example_data
β”œβ”€β”€ AAA.txt
β”œβ”€β”€ AAA.png
β”œβ”€β”€ BCC.txt
β”œβ”€β”€ BCC.png
β”œβ”€β”€ ......
β”œβ”€β”€ CCC.txt
└── CCC.png

Then Sana's training can be launched via

# Example of training Sana 0.6B with 512x512 resolution from scratch
bash train_scripts/train.sh \
  configs/sana_config/512ms/Sana_600M_img512.yaml \
  --data.data_dir="[asset/example_data]" \
  --data.type=SanaImgDataset \
  --model.multi_scale=false \
  --train.train_batch_size=32

# Example of fine-tuning Sana 1.6B with 1024x1024 resolution
bash train_scripts/train.sh \
  configs/sana_config/1024ms/Sana_1600M_img1024.yaml \
  --data.data_dir="[asset/example_data]" \
  --data.type=SanaImgDataset \
  --model.load_from=hf://Efficient-Large-Model/Sana_1600M_1024px/checkpoints/Sana_1600M_1024px.pth \
  --model.multi_scale=false \
  --train.train_batch_size=8

πŸ’» 4. Metric toolkit

Refer to Toolkit Manual.

πŸ’ͺTo-Do List

We will try our best to release

πŸ€—Acknowledgements

πŸ“–BibTeX

@misc{xie2024sana,
      title={Sana: Efficient High-Resolution Image Synthesis with Linear Diffusion Transformer},
      author={Enze Xie and Junsong Chen and Junyu Chen and Han Cai and Haotian Tang and Yujun Lin and Zhekai Zhang and Muyang Li and Ligeng Zhu and Yao Lu and Song Han},
      year={2024},
      eprint={2410.10629},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2410.10629},
    }