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

File	Size	What it is	Derived from
flux2-klein-4b_convrot_int8.safetensors	~4.1 GB	Every quantization-eligible weight quantized to INT8 with ConvRot (group-wise Hadamard rotation, dynamic group size, 256 floor); modulation, embedding, and final-layer weights kept at source BF16 precision	flux-2-klein-4b.safetensors (native/BFL format, BF16)
flux2-klein-4b_nvfp4_convrot_int8.safetensors	~2.9 GB	Bulk weights (the middle double_blocks 1-3 and single_blocks 2-17) quantized to NVFP4; the first/last block of each stack (double_blocks 0,4 and single_blocks 0,1,18,19 — structurally the most precision-sensitive per common transformer-quantization heuristics) quantized to INT8 ConvRot instead of plain FP8; modulation, embedding, and final-layer weights kept at source BF16 precision	flux-2-klein-4b.safetensors (native/BFL format, BF16)
flux2-klein-4b_mixed_int4_int8_convrot.safetensors	~2.7 GB	Same layer split as the NVFP4 variant, but the 56 bulk layers (formerly NVFP4) are INT4 W4A4 ConvRot instead; the 24 edge-block layers stay INT8 ConvRot	flux-2-klein-4b.safetensors (native/BFL format, BF16) — the 56 bulk layers were quantized standalone from this BF16 source (not from the NVFP4 file) and spliced into a copy of flux2-klein-4b_nvfp4_convrot_int8.safetensors, replacing its NVFP4 tensors, to avoid compounding quantization error

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.