Kuantala

Kuantala quantizes diffusers-format generative models to NVFP4 (or FP8) using NVIDIA Model Optimizer. It can also convert the quantized output to ComfyUI-compatible format, evaluate quantization quality, and inspect model components.

NVFP4 is W4A4 (4-bit weights, 4-bit activations). Weights are stored as FP4 E2M1 packed in uint8, with FP8 per-block scales. Activation scales are calibrated during quantization, enabling native FP4 tensor core math on Blackwell+ GPUs for both memory savings and compute speedups.

FP8 is W8A8 (8-bit weights, 8-bit activations). Weights are stored as float8 E4M3 with per-tensor scales. Activation scales are calibrated during quantization, enabling native FP8 tensor core math on Hopper+ GPUs.

Installation

Requires Python ≤ 3.12 due to nvidia-modelopt constraints. Tested on Linux.

python -m venv .venv
source .venv/bin/activate

# First install PyTorch with CUDA from https://pytorch.org
pip install torch --index-url https://download.pytorch.org/whl/cu130

git clone https://github.com/metebalci/kuantala.git
cd kuantala
pip install -e .

Kuantala requires models in diffusers format (with model_index.json). If a model has both a raw and a diffusers variant on HuggingFace, use the diffusers one (typically suffixed with -Diffusers).

Quick Start

# Show supported formats and default keep presets
kuantala info

# Quantize to FP8 (~50% size, ~2x inference speed on Hopper+/Blackwell)
kuantala quantize Wan-AI/Wan2.2-I2V-A14B-Diffusers --dtype FP8

# Quantize to NVFP4 (~75% size, fastest on Blackwell)
kuantala quantize ./local-model --dtype NVFP4

# Inspect model components
kuantala components Wan-AI/Wan2.2-I2V-A14B-Diffusers

# Show model architecture from config (no weights downloaded)
kuantala config Wan-AI/Wan2.2-I2V-A14B-Diffusers

# Convert NVFP4 output to ComfyUI format
kuantala convert ./output-wan/transformer-NVFP4.safetensors --remap-keys wan

# Evaluate quantization quality (compare original vs quantized)
kuantala eval Wan-AI/Wan2.2-I2V-A14B-Diffusers -q ./output-Wan-AI-Wan2.2-I2V-A14B-Diffusers

# Inspect tensors in a quantized file
kuantala tensors ./output/transformer-FP8.safetensors

How It Works

Kuantala uses NVIDIA Model Optimizer to quantize model components:

1. Load the full diffusers pipeline (transformer, text encoder, VAE, etc.) in BF16 on CUDA
2. Quantize with Model Optimizer — inserts quantizer nodes and runs calibration by executing the pipeline with prompts from HuggingFace datasets (see Prompt Sources), producing realistic activations for optimal scale factor estimation
3. Compress — converts fake-quantized weights to real low-precision (FP8 or packed FP4)
4. Save as safetensors with actual quantized weights

The output files are genuinely smaller and load into VRAM at low precision. FP8 uses float8_e4m3fn tensors, NVFP4 uses packed uint8 with block-wise FP8 scales.

ComfyUI Conversion

Kuantala's NVFP4 output uses Model Optimizer's internal format. To use NVFP4 models in ComfyUI, convert them first:

kuantala convert ./output/transformer-NVFP4.safetensors
# Creates ./output/transformer-NVFP4-comfyui.safetensors

# Custom output path
kuantala convert ./output/transformer-NVFP4.safetensors -o ./comfy-model.safetensors

For models where diffusers uses different key names than the original (e.g. Wan), use --remap-keys to translate back to original names that ComfyUI expects:

kuantala convert ./wan-nvfp4/transformer-NVFP4.safetensors --remap-keys wan

The conversion performs:

• Nibble swap — reorders packed FP4 byte layout to match ComfyUI expectations
• Scale tiling — converts block scales from plain to cuBLAS tiled layout
• Quantizer key renaming — translates Model Optimizer quantizer names to ComfyUI conventions
• Layer key remapping — translates diffusers layer names to original names (with --remap-keys)
• Metadata injection — adds .comfy_quant entries required by ComfyUI

Full workflow (Wan example):

kuantala quantize Wan-AI/Wan2.2-I2V-A14B-Diffusers --dtype NVFP4
kuantala convert ./output-Wan-AI-Wan2.2-I2V-A14B-Diffusers/transformer-NVFP4.safetensors --remap-keys wan
# Load transformer-NVFP4-comfyui.safetensors in ComfyUI

CLI Reference

kuantala components

kuantala components [OPTIONS] MODEL

kuantala config

kuantala config MODEL

Shows model architecture from config: full module hierarchy with layer types, shapes, and parameter counts. Only downloads config.json files — no model weights are downloaded.

kuantala convert

kuantala convert [OPTIONS] INPUT

kuantala eval

kuantala eval [OPTIONS] MODEL
kuantala evaluate [OPTIONS] MODEL

Compares original vs quantized pipeline outputs using PSNR and SSIM metrics. Runs the pipeline with fixed seeds on both original and quantized models to measure quality loss. Eval prompts come from HuggingFace datasets selected by --psrc (see Prompt Sources).

# Basic eval
kuantala eval Wan-AI/Wan2.2-I2V-A14B-Diffusers -q ./output-wan

# With pixel-space comparison
kuantala eval Wan-AI/Wan2.2-I2V-A14B-Diffusers -q ./output-wan --decode

# Custom prompts
kuantala eval ./local-model -q ./output --prompts eval_prompts.txt --nprompts 8

kuantala info

kuantala info

Shows supported quantization formats, default keep presets, and known models.

kuantala quantize

kuantala quantize [OPTIONS] MODEL

Quantization config presets (--cfg):

• default — Standard quantization with max calibration. Fastest, good baseline. Works with FP8 and NVFP4.
• awq_lite — Activation-aware weight quantization with lightweight search. Better quality than default, moderate overhead. NVFP4 only.
• awq_clip — AWQ with clipping-based optimization. Good quality/speed tradeoff. NVFP4 only.
• awq_full — AWQ with full search (lite + clip). Slowest, best quality. NVFP4 only. Use --alpha-step to control search granularity (default: 0.1).

Approximate calibration speed relative to default (measured with Z-Image 6B on RTX Pro 6000): awq_lite and awq_clip are ~8x slower, awq_full is ~16x slower.

kuantala tensors

kuantala tensors FILE_PATH

Shows per-tensor detail: name, dtype, shape, and parameter count. Also shows a dtype summary.

Global Options

Per-Component Control

Kuantala detects and quantizes the following component types from diffusers models:

Schedulers, tokenizers, and other non-neural components are always skipped. VAE is skipped by default because quantizing it causes visible artifacts.

kuantala quantize black-forest-labs/FLUX.1-dev \
    --dtype FP8 \
    --vae-dtype skip \
    --te-dtype FP8

Layer-Level Control

Use --keep to disable quantization on specific layers by glob pattern. Matched layers stay at their original precision. Time embeddings, conditioning projections, and input/output layers are small but sensitive — keeping them unquantized has negligible size impact but helps quality.

Default Keeps

For known HuggingFace model IDs, kuantala automatically applies preset keep patterns that disable quantization on sensitive layers. These are based on NVIDIA Model Optimizer's example settings.

Use --use-default-keeps <preset> to explicitly select a preset (e.g. for local paths). Use --no-default-keeps to disable auto-detection.

# Wan 2.2 I2V / T2V 14B (keeps auto-detected)
kuantala quantize Wan-AI/Wan2.2-I2V-A14B-Diffusers --dtype NVFP4

# FLUX.1 dev
kuantala quantize black-forest-labs/FLUX.1-dev --dtype NVFP4

# FLUX.2 dev
kuantala quantize black-forest-labs/FLUX.2-dev --dtype NVFP4

# LTX-2
kuantala quantize Lightricks/LTX-2 --dtype NVFP4

# Z-Image
kuantala quantize Tongyi-MAI/Z-Image --dtype NVFP4

# Qwen-Image-2512
kuantala quantize Qwen/Qwen-Image-2512 --dtype NVFP4

# CogVideoX-5b
kuantala quantize zai-org/CogVideoX-5b --dtype FP8

# SDXL 1.0
kuantala quantize stabilityai/stable-diffusion-xl-base-1.0 --dtype FP8

# With AWQ config preset for better quality
kuantala quantize Wan-AI/Wan2.2-I2V-A14B-Diffusers --dtype NVFP4 --cfg awq_lite

Norms (FP32LayerNorm, RMSNorm) and embeddings (WanRotaryPosEmbed) are already kept at original precision by Model Optimizer.

Prompt Sources

Calibration and evaluation prompts are loaded from HuggingFace datasets. Use --psrc to select the source, or let kuantala auto-detect it from known model IDs.

Calibration uses the first N entries from each dataset (default: 32), evaluation uses entries starting at offset 1024 to avoid overlap. For I2V and TI2I sources, the dataset provides both text prompts and conditioning images.

Datasets are downloaded on demand from HuggingFace and cached locally by the datasets library. Custom prompts can be provided via --prompts FILE (one prompt per line) to override the default dataset. For image-input models, append image:/path/to/image.png to the line:

a cat sitting on a table
a dog running in a park image:/data/images/dog.jpg
sunset over the ocean image:/data/images/sunset.png

Python API

from pathlib import Path
from kuantala import QuantConfig, quantize

config = QuantConfig(
    model_source="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
    dtype="NVFP4",
    output_dir=Path("./output-wan"),
    keep=[
        "*patch_embedding*", "*condition_embedder*", "*proj_out*",
        "*blocks.0.*", "*blocks.1.*", "*blocks.2.*",
        "*blocks.37.*", "*blocks.38.*", "*blocks.39.*",
    ],
)
output_files = quantize(config)

Development

python -m venv .venv
source .venv/bin/activate
pip install torch --index-url https://download.pytorch.org/whl/cu130
pip install -e ".[dev]"
pytest tests/

License

This project is licensed under the GNU General Public License v3.0.