Nunchaku chega ao Diffusers: inferência de difusão 4-bit com aceleração real em GPUs Blackwell

Quantização que acelera de verdade — não só economiza VRAM

Modelos de difusão modernos — como FLUX, Stable Diffusion e ERNIE-Image — produzem imagens impressionantes, mas carregar um modelo texto-para-imagem em precisão BF16 exige 20 a 30 GB de VRAM, colocando-os fora do alcance da maioria das GPUs de consumidor. A quantização é uma solução poderosa, e o Diffusers já integrava backends como bitsandbytes, GGUF, torchao e Quanto. O problema é que a maioria desses backends é weight-only: armazenam os pesos em baixa precisão, mas os dequantizam de volta para alta precisão no momento do cálculo. Isso reduz o uso de memória, mas geralmente não acelera a inferência — e pode até adicionar latência.

O SVDQuant, método de quantização por trás do popular motor de inferência Nunchaku, adota uma abordagem diferente. Ele roda as camadas principais do transformer com pesos e ativações em 4 bits (W4A4), reduzindo memória e acelerando o loop de denoising. Até agora, porém, usar esses checkpoints exigia uma biblioteca de inferência separada. Isso mudou.

Nunchaku nativo no Diffusers: from_pretrained() e pronto

A partir da versão mais recente do Diffusers, carregar um checkpoint Nunchaku é tão simples quanto chamar from_pretrained() — sem compilação CUDA local, graças ao pacote kernels. O ecossistema inclui ainda o diffuse-compressor, um toolkit que permite quantizar novas arquiteturas você mesmo e publicá-las como repositórios Diffusers comuns.

Para começar, instale os requisitos:

pip install -U diffusers transformers accelerate kernels bitsandbytes

Depois carregue um pipeline pré-quantizado como qualquer modelo Diffusers:

import torch
from diffusers import ErnieImagePipeline

pipe = ErnieImagePipeline.from_pretrained(
    "lite-infer/ERNIE-Image-Turbo-nunchaku-lite-nvfp4_r32-bnb4-text-encoder",
    torch_dtype=torch.bfloat16,
).to("cuda")

image = pipe(
    prompt="A cinematic portrait of a red fox in a misty forest at sunrise, "
           "detailed fur, volumetric light",
    height=1024,
    width=1024,
    num_inference_steps=8,
    guidance_scale=1.0,
    generator=torch.Generator("cuda").manual_seed(42),
).images[0]
image.save("output.png")

Nada de classe de pipeline customizada, nada de motor de inferência separado, nada para compilar localmente. Os kernels NVFP4 são baixados automaticamente do Hub na primeira execução. Este checkpoint combina um transformer Nunchaku NVFP4 com um codificador de texto bitsandbytes NF4 e gera uma imagem 1024×1024 em ~1,7 segundos em uma RTX 5090, com pico de memória de ~12 GB (contra ~24 GB do pipeline BF16).

Suporte de hardware: Blackwell, Ampere e Ada Lovelace

Para GPUs de gerações anteriores (RTX 30 e 40), use as variantes INT4. Para Blackwell (RTX 50), o NVFP4 oferece o melhor desempenho.

Ganhos reais de desempenho

Os benchmarks mostram reduções significativas tanto em latência quanto em memória. Em uma RTX 5090, o pipeline Nunchaku reduz a latência de inferência de ponta a ponta em comparação com BF16, enquanto corta o uso de VRAM pela metade. A qualidade de imagem se mantém comparável ao modelo original em precisão total, graças à decomposição SVD que preserva os componentes mais importantes dos pesos.

Quantizando seu próprio modelo

O diffuse-compressor permite quantizar arquiteturas personalizadas. O fluxo é direto:

# 1. Inspecionar o que será quantizado
python examples/text_to_image/quantize_hf.py black-forest-labs/FLUX.2-klein-4B \
  --precision int4 --rank 32 --inspect-config

# 2. Executar a quantização
python examples/text_to_image/quantize_hf.py black-forest-labs/FLUX.2-klein-4B \
  --precision int4 --rank 32

# 3. Empacotar como pipeline Diffusers
python examples/text_to_image/quantize_hf.py black-forest-labs/FLUX.2-klein-4B \
  --precision int4 --rank 32 --push-to-hub my-org/my-quantized-model

Checkpoints prontos para uso já estão disponíveis no Hub para ERNIE-Image Turbo, Krea 2 Turbo, FLUX.2 Klein e outros modelos populares.

Por que isso importa agora

A integração nativa do Nunchaku ao Diffusers elimina a fricção que mantinha a quantização W4A4 longe do ecossistema HuggingFace. Com metade da VRAM e inferência mais rápida, GPUs de consumidor como a RTX 5090 conseguem rodar modelos de difusão de última geração com qualidade comparável ao BF16. Para quem não tem uma Blackwell, os checkpoints INT4 garantem acesso com GPUs Ampere e Ada Lovelace. O resultado prático: modelos de difusão de ponta acessíveis a muito mais desenvolvedores, sem precisar de uma H100 para cada experimento.