rockerBOO/flux2-klein-4b-nvfp4-convrot

FLUX.2-klein-4B — quantized

Quantized variants of black-forest-labs/FLUX.2-klein-4B, the distilled 4B-parameter FLUX.2 Klein diffusion transformer. Quantized from the repo's native/ComfyUI-format checkpoint (flux-2-klein-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-4b_convrot_int8.safetensors: broadest compatibility — any modern GPU with usable INT8 tensor-core throughput, no Blackwell requirement.
• —flux2-klein-4b_nvfp4_convrot_int8.safetensors: requires a Blackwell GPU (SM ≥ 10.0/12.0) for NVFP4 inference support.
• —flux2-klein-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-4b.safetensors -o flux2-klein-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): modulation (double_stream_modulation_img/txt, single_stream_modulation), time embedding (time_in), image/text input embedders (img_in, txt_in), and the final layer (final_layer). 80 of 149 tensors were quantized; the remainder are these excluded weights plus non-weight tensors (norms, scales) that aren't quantization targets.

NVFP4 + ConvRot INT8

ctq -i flux-2-klein-4b.safetensors -o flux2-klein-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.

All three use learned-rounding optimization (SVD/AdaRound) — none are --simple/RTN builds.

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

# 2. Quantize the extracted subset directly to INT4 W4A4 ConvRot.
ctq -i flux2-klein-4b_nvfp4_candidates_bf16.safetensors \
    -o flux2-klein-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-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

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

Verification

All three files confirmed to load and generate successfully end-to-end in ComfyUI (Qwen3-4B text encoder, FLUX.2 VAE, res_multistep sampler, 20 steps). 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).

License

This model is a derivative of black-forest-labs/FLUX.2-klein-4B, released under the Apache License 2.0. The full license text is included as LICENSE.md in this repo. These quantized derivatives are 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). These files have been modified from the original via the quantization methods described above.