Nunchaku llega a Diffusers y cambia la forma en que ejecutas transformadores de difusión: menos memoria, inferencia más rápida y sin compilar CUDA en tu máquina. ¿Suena bien? Aquí te explico qué hay detrás, cómo usarlo y cuándo te conviene cuantizar tus modelos.
Qué es Nunchaku y por qué importa
SVDQuant, la técnica detrás del motor Nunchaku, no es la típica cuantización "solo pesos". En lugar de guardar pesos en baja precisión y desquantizarlos durante cálculo, SVDQuant representa las matrices más difíciles con una rama de bajo rango a 16 bits y cuantiza el resto a 4 bits. El resultado: W4A4 (pesos y activaciones en 4 bits) que reduce memoria y además acelera el bucle de denoising.
Hasta ahora usar checkpoints Nunchaku requería un motor de inferencia separado. Con Nunchaku Lite, Diffusers puede cargar estos checkpoints directamente con from_pretrained() y sin compilación local: las kernels se descargan desde el Hub a través del paquete kernels.
Nunchaku Lite: integración y modelo de ejecución
¿Cómo funciona en Diffusers? Nunchaku Lite parchea en tiempo de carga los nn.Linear relevantes de un modelo con capas SVDQ/AWQ que implementan SVDQuant y AWQ. Las kernels optimizadas vienen desde el Hub y hay dos familias principales:
svdq_w4a4: W4A4 con corrección low-rank. Disponible en variantesint4ynvfp4.awq_w4a16: pesos en 4 bits con activaciones a 16 bits, útil para proyecciones sensibles a precisión como normalizaciones adaptativas.
La ventaja: un repositorio Nunchaku Lite es un repo Diffusers normal. Schedulers, LoRA, offloading y torch.compile siguen funcionando porque la estructura de módulos se mantiene.
Cómputo, compatibilidad y variantes de GPU
Nunchaku Lite escoge kernels según la GPU y la variante del checkpoint:
nvfp4(NVFP4): requiere GPUs Blackwell (RTX 50, RTX PRO 6000, B200).int4: compatible con Turing / Ampere / Ada (RTX 30/40, A100, L40S).
Volta y Hopper no están soportadas por las kernels 4-bit actualmente. La carga valida la capacidad CUDA y falla con un error claro en GPUs no soportadas.
Rendimiento y memoria: números reales
En la medición de Hugging Face (RTX PRO 6000, 1024x1024 con rootonchair/ERNIE-Image-Turbo) se observó:
- BF16 baseline: 3.00 s, 31.1 GB VRAM.
- Nunchaku Lite NVFP4: 2.27 s (1.35x), 20.6 GB VRAM.
- Nunchaku Lite NVFP4 +
torch.compile: 1.68 s (1.8x), 20.6 GB VRAM. - Nunchaku Lite NVFP4 + NF4 text encoder: 2.29 s, 16.0 GB VRAM.
En resumen: hasta 50% menos VRAM y ~30% de mejora en latencia; torch.compile ayuda a cerrar la brecha con el motor original.
¿Por qué el motor original sigue siendo más rápido para algunos modelos?
Nunchaku (el motor original) usa kernels fusionados específicos por arquitectura: proyecciones QKV fusionadas, GELU+MLP fusionado, etc. Esas optimizaciones requieren reescrituras estructurales del modelo (fusionar to_q/to_k/to_v en un to_qkv), lo que da más velocidad.
Nunchaku Lite es genérico: no infiere por sí solo estas reescrituras. Aun así entrega un acelerón importante sin trabajo de integración por modelo.
Cómo empezar (comandos rápidos)
Instala dependencias:
pip install -U diffusers transformers accelerate kernels bitsandbytes
Carga un pipeline precuantizado como cualquier 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",
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")
No necesitas clases de pipeline custom ni compilar código CUDA localmente: las kernels NVFP4 se descargan la primera vez desde el Hub.
Optimizaciones combinables
torch.compilepuede llevar el speedup de 1.35x a ~1.8x.- Cuantizar el encoder de texto (por ejemplo a NF4 con
bitsandbytes) reduce aún más el peak VRAM (benchmark mostró ~22% adicional). - Offloading (como
enable_model_cpu_offload) funciona con los modelos nunchaku-lite.
Cuantizando tus propios modelos con diffuse-compressor
El flujo completo está disponible en diffuse-compressor: inspección, calibración, cuantización SVDQuant, empaquetado y publicación.
Ejemplo rápido de inspección:
python examples/text_to_image/quantize_hf.py black-forest-labs/FLUX.2-klein-4B \
--precision int4 --rank 32 --inspect-config
Cuantizar el transformer:
python examples/text_to_image/quantize_hf.py black-forest-labs/FLUX.2-klein-4B \
--precision int4 \
--output outputs/checkpoints/svdq-int4_r32-flux-2-klein-4b.safetensors
Convertir a un repo Diffusers listo para publicar:
python examples/convert_nunchaku_lite_diffusers.py \
--checkpoint outputs/checkpoints/svdq-int4_r32-flux-2-klein-4b.safetensors \
--model-id black-forest-labs/FLUX.2-klein-4B \
--bnb4-text-encoder text_encoder \
--compute-dtype bfloat16 \
--output-dir outputs/diffusers/FLUX.2-klein-4B-nunchaku-lite-int4-bnb4-text-encoder
Luego cargas y pruebas:
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained(
"outputs/diffusers/FLUX.2-klein-4B-nunchaku-lite-int4-bnb4-text-encoder",
device_map="cuda",
)
image = pipe("A glass robot in a greenhouse, cinematic lighting", num_inference_steps=4).images[0]
Si todo está bien, pipe.push_to_hub("tu-usuario/tu-model-nunchaku-lite-int4") y otros podrán usarlo con from_pretrained().
Reescrituras estructurales y cuándo necesitas un adapter
La ruta genérica asume que el modelo no requiere reescrituras. Pero modelos como FLUX.1-dev necesitan agrupar Q, K y V en una proyección to_qkv para poder usar el operador fusionado del motor Nunchaku. Eso se hace con una configuración de destino al cuantizar y con pequeños adapters en tiempo de carga.
Diffusers incluye ejemplos y guías para esas reescrituras y adapters cuando son necesarias.
Checkpoints listos y recursos
Algunos repositorios listos para usar:
rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoderrootonchair/ERNIE-Image-Turbo-nunchaku-lite-nvfp4-bnb4-text-encoderOzzyGT/Krea_2_Turbo_nunchaku_lite_nvfp4lite-infer: colecciones y más checkpoints
Además, diffuse-compressor facilita el pipeline para nuevos modelos.
Reflexión final
Nunchaku Lite trae la potencia de SVDQuant al ecosistema Diffusers de forma práctica: menos memoria y latencias mejores sin sacrificar la interoperabilidad del pipeline. ¿Quieres velocidad máxima y tienes recursos para integrar cambios por arquitectura? Entonces el motor Nunchaku con reescrituras fusionadas sigue siendo la mejor opción. Si prefieres una ruta más simple y reproducible, Nunchaku Lite te permite cuantizar y publicar modelos hoy mismo.