|
|
--- |
|
|
license: other |
|
|
library_name: diffusers |
|
|
pipeline_tag: text-to-image |
|
|
tags: |
|
|
- diffusers |
|
|
- image-generation |
|
|
- quantization |
|
|
- int8 |
|
|
- torchao |
|
|
- amd |
|
|
- rocm |
|
|
base_model: black-forest-labs/FLUX.2-dev |
|
|
--- |
|
|
|
|
|
# FLUX.2-dev — Attention-only INT8 Weight-Only Transformer (ROCm) |
|
|
|
|
|
This repository provides an **INT8 weight-only quantized transformer** for |
|
|
[`black-forest-labs/FLUX.2-dev`](https://huggingface.co/black-forest-labs/FLUX.2-dev). |
|
|
|
|
|
It is designed to be: |
|
|
|
|
|
- ✅ **ROCm-compatible** |
|
|
- ✅ **Stable on AMD Instinct MI210** |
|
|
- ✅ **Image-quality preserving** |
|
|
|
|
|
Only **attention Linear layers (Q/K/V + projections)** are quantized. |
|
|
All other components remain in **BF16**. |
|
|
|
|
|
--- |
|
|
|
|
|
## 🔍 What is included |
|
|
|
|
|
- ✅ Transformer with **attention-only INT8 weight-only quantization** |
|
|
- ✅ TorchAO-based quantization (no bitsandbytes) |
|
|
- ✅ Compatible with **Diffusers standard pipelines** |
|
|
|
|
|
--- |
|
|
|
|
|
## ❌ What is NOT included |
|
|
|
|
|
- ❌ VAE |
|
|
- ❌ Text encoders |
|
|
- ❌ Scheduler |
|
|
|
|
|
These components are automatically loaded from the base FLUX.2 model. |
|
|
|
|
|
--- |
|
|
|
|
|
## 💡 Why attention-only INT8? |
|
|
|
|
|
Full INT8 quantization of FLUX.2 introduces visible artifacts on ROCm. |
|
|
Quantizing **only attention layers** provides: |
|
|
|
|
|
- Significant VRAM reduction |
|
|
- Stable generation |
|
|
- No "confetti noise" artifacts |
|
|
- Safe inference on MI210 (64 GB) |
|
|
|
|
|
--- |
|
|
|
|
|
## 🚀 Usage (Diffusers) |
|
|
|
|
|
```python |
|
|
import torch |
|
|
from diffusers import Flux2Pipeline, AutoModel |
|
|
|
|
|
BASE_MODEL = "black-forest-labs/FLUX.2-dev" |
|
|
ATTN_INT8 = "AmdGoose/FLUX.2-dev-transformer-attn-int8wo" |
|
|
|
|
|
dtype = torch.bfloat16 |
|
|
device = "cuda" # ROCm uses "cuda" in PyTorch |
|
|
|
|
|
transformer = AutoModel.from_pretrained( |
|
|
ATTN_INT8, |
|
|
subfolder="transformer_attn_int8wo", |
|
|
torch_dtype=dtype, |
|
|
use_safetensors=False, |
|
|
).to(device) |
|
|
|
|
|
pipe = Flux2Pipeline.from_pretrained( |
|
|
BASE_MODEL, |
|
|
transformer=transformer, |
|
|
torch_dtype=dtype, |
|
|
) |
|
|
|
|
|
pipe.enable_attention_slicing() |
|
|
pipe.vae.enable_tiling() |
|
|
pipe.enable_model_cpu_offload() |
|
|
|
|
|
image = pipe( |
|
|
prompt="A realistic starter pack figurine in a blister box, studio lighting", |
|
|
num_inference_steps=28, |
|
|
guidance_scale=4, |
|
|
height=1024, |
|
|
width=1024, |
|
|
).images[0] |
|
|
|
|
|
image.save("out.png") |
|
|
|
