flux2-klein-base-4b-nvfp4-convrot

FLUX.2-klein-base-4B — quantized

Quantized variants of black-forest-labs/FLUX.2-klein-base-4B, the base (non-distilled) 4B-parameter FLUX.2 Klein diffusion transformer. Quantized from the repo's native/ComfyUI-format checkpoint (flux-2-klein-base-4b.safetensors, double_blocks.* key naming) — not the diffusers-format transformer/diffusion_pytorch_model.safetensors in the same repo, which uses different tensor names (x_embedder, context_embedder, etc.) and is not what ComfyUI's loader expects.

Files

Hardware requirements

• flux2-klein-base-4b_convrot_int8.safetensors: broadest compatibility — any modern GPU with usable INT8 tensor-core throughput, no Blackwell requirement.
• flux2-klein-base-4b_nvfp4_convrot_int8.safetensors: requires a Blackwell GPU (SM ≥ 10.0/12.0) for NVFP4 inference support.
• flux2-klein-base-4b_mixed_int4_int8_convrot.safetensors: no Blackwell dependency, but INT4 tensor-core throughput varies significantly by GPU generation — verify actual runtime behavior on target hardware rather than assuming uniform benefit. Quality has not been evaluated (see Verification below); treat this variant as experimental.

Quantization method

Quantized with convert_to_quant (ctq).

ConvRot INT8

ctq -i flux-2-klein-base-4b.safetensors -o flux2-klein-base-4b_convrot_int8.safetensors \
    --int8 --scaling-mode row --dynamic-convrot --convrot-group-size 256 \
    --flux2 --comfy_quant --save-quant-metadata

The --flux2 exclusion preset kept the following layers at source BF16 precision (not quantized): double_stream_modulation_img.lin.weight, double_stream_modulation_txt.lin.weight, final_layer.adaLN_modulation.1.weight, final_layer.linear.weight, img_in.weight, single_stream_modulation.lin.weight, time_in.in_layer.weight, time_in.out_layer.weight, txt_in.weight — modulation, the final output layer, the image input embedder, time embedding, and the text input embedder. 80 of 149 tensors were quantized; the remainder are these 9 excluded weights plus non-weight tensors (norms, scales) that aren't quantization targets. Uses learned-rounding optimization (SVD/AdaRound) — not a --simple/RTN build.

NVFP4 + ConvRot INT8

ctq -i flux-2-klein-base-4b.safetensors -o flux2-klein-base-4b_nvfp4_convrot_int8.safetensors \
    --nvfp4 \
    --custom-layers 'double_blocks\.(0|4)\..*\.weight$|single_blocks\.(0|1|18|19)\.linear[12]\.weight$' \
    --custom-type int8 --custom-scaling-mode row --custom-convrot --custom-convrot-group-size 256 \
    --flux2 --comfy_quant --save-quant-metadata

double_blocks and single_blocks in this architecture have unequal role-separation: double_blocks keep attention (*_attn.qkv/*_attn.proj) and MLP (*_mlp.0/*_mlp.2) as separate tensors, but single_blocks.*.linear1/linear2 fuse attention and MLP into single tensors (standard FLUX single-stream design), so a clean attention-vs-MLP split isn't possible there. This build instead routes by block position: the first/last block of each stack (structurally the most precision-sensitive per common transformer heuristics) stays INT8 ConvRot; the middle blocks go NVFP4. 56 of 80 quantization-eligible tensors went to NVFP4, 24 to INT8 ConvRot; the same 9 layers as the ConvRot INT8 build stayed BF16. Uses learned-rounding optimization (SVD/AdaRound) — not a --simple/RTN build.

A known bug in the convert_to_quant checkout used for this build (--layer-config silently maps "format": "nvfp4" to FP8 instead) meant this file was built with --custom-layers/--custom-type rather than --layer-config; a fix is proposed in silveroxides/convert_to_quant#56.

Mixed INT4/INT8 ConvRot

Built via extract → quantize → splice, reusing the same 56-vs-24 layer split as the NVFP4 variant above but replacing the 56 bulk (formerly NVFP4) layers with INT4 W4A4 ConvRot instead:

# 1. Extract just the 56 target tensors from the original BF16 source (not the NVFP4 file --
#    quantizing from BF16 avoids compounding error).
python extract_nvfp4_candidates.py \
    --src flux-2-klein-base-4b.safetensors \
    --out flux2-klein-base-4b_nvfp4_candidates_bf16.safetensors

# 2. Quantize the extracted subset directly to INT4 W4A4 ConvRot.
ctq -i flux2-klein-base-4b_nvfp4_candidates_bf16.safetensors \
    -o flux2-klein-base-4b_nvfp4_candidates_int4.safetensors \
    --int4 --dynamic-convrot --convrot-group-size 256 \
    --comfy_quant --save-quant-metadata

# 3. Splice the INT4 tensors into a copy of flux2-klein-base-4b_nvfp4_convrot_int8.safetensors,
#    replacing the NVFP4 tensors at those keys. Updates both the per-tensor .comfy_quant blob
#    and the global __metadata__["_quantization_metadata"]["layers"] entry for every replaced
#    key -- ComfyUI's loader treats the global blob as authoritative for format detection.
python splice_int4_into_nvfp4_convrot_int8.py \
    --base flux2-klein-base-4b_nvfp4_convrot_int8.safetensors \
    --int4 flux2-klein-base-4b_nvfp4_candidates_int4.safetensors \
    --out flux2-klein-base-4b_mixed_int4_int8_convrot.safetensors

Verification

All three files confirmed to load and generate successfully end-to-end in ComfyUI (Qwen3-4B text encoder, FLUX.2 VAE), with the intended per-layer quantization format actually dispatched per each output file's own _quantization_metadata (not just a successful command exit). Output quality has not been separately assessed against the BF16 source for any variant — only functional correctness (loads, runs, produces a coherent image matching the prompt). Visually, the mixed INT4/INT8 output showed a noticeable loss of fine detail (e.g. mountain texture) compared to the INT8 and NVFP4 variants in a spot-check generation — consistent with the standing caveat that INT4 ConvRot quality depends heavily on layer-selection and rounding method and hasn't been formally evaluated.

License

This model is a derivative of black-forest-labs/FLUX.2-klein-base-4B, released under the Apache License 2.0. The full license text is included as LICENSE.md in this repo. This quantized derivative is also distributed under Apache 2.0, per the terms of the source license. No files from the original NOTICE were required (the source repo does not include a NOTICE file). This file has been modified from the original via the quantization method described above.