FLUX.2 Klein-4B — JAX-native unified checkpoint
Unofficial, JAX-native repackaging of black-forest-labs/FLUX.2-klein-4B for text-to-image generation. The diffusion transformer, the autoencoder decoder, and the Qwen3-4B text encoder are converted from their original PyTorch checkpoints into a single bundle of Orbax parameter pytrees. There is no PyTorch dependency anywhere in this bundle or in the conversion that produced it.
This repository contains weights only. It does not contain a model implementation, a sampling loop, or an inference pipeline.
No safety filtering of any kind
This bundle does not include any content moderation, NSFW classifier, or prompt/output filtering. The reference FLUX.2 pipeline's safety layer (a separate large language model used to screen prompts and generated images) was intentionally not ported. Nothing in these weights or in any accompanying code restricts what can be generated. Anyone using this checkpoint is solely responsible for how they use it.
Not compatible with PyTorch, Transformers, or Diffusers
The parameter files in this repository are stored in Orbax's checkpoint
format, not safetensors, and cannot be loaded with from_pretrained,
safetensors.torch.load_file, DiffusionPipeline, or any other
PyTorch-based loader. Loading this bundle requires jax and
orbax-checkpoint only. transformers is only needed if you want to
use its tokenizer loader for convenience; it is not required and does
not need a torch backend for that purpose.
Repository layout
.
├── manifest.json # conversion metadata, see below
├── params/
│ ├── text_encoder/ # Orbax checkpoint, ~5.8 GiB
│ ├── transformer/ # Orbax checkpoint, ~7.2 GiB
│ └── vae/ # Orbax checkpoint, ~0.19 GiB
└── tokenizer/
├── vocab.json
├── merges.txt
├── tokenizer_config.json
├── special_tokens_map.json
├── added_tokens.json
└── chat_template.jinja
Each of the three params/ subdirectories is an independent Orbax
checkpoint and can be restored on its own; you do not need to load the
other two components to use just one.
Size and parameter counts
All figures below are computed directly from each source checkpoint's own tensor shapes, not estimated.
| Component | Parameters | Precision | Size |
|---|---|---|---|
| Diffusion transformer | 3,875,544,576 | bf16 (unchanged) | 7.22 GiB |
| Text encoder (truncated, 27 of 36 layers) | 3,114,088,192 | bf16 (unchanged) | 5.80 GiB |
| Autoencoder (decoder only) | 49,620,259 | fp32 (unchanged) | 0.185 GiB |
| Bundle total | 7,039,253,027 | — | ≈ 13.2 GiB |
For comparison, the three original, unconverted checkpoints together total approximately 15.0 GiB (full 36-layer text encoder plus encoder and decoder for the autoencoder). The difference is entirely accounted for by dropping tensors that are provably never read for text-to-image generation, not by reduced precision anywhere. Actual on-disk size may differ slightly from the raw parameter byte counts above due to Orbax's own checkpoint metadata overhead.
What was changed, and what was not
Every transformation applied during conversion is lossless: no value was approximated, rounded, or quantized. Precision matches the reference implementation exactly, tensor by tensor.
| Component | Source | What changed |
|---|---|---|
| Diffusion transformer | black-forest-labs/FLUX.2-klein-4B, flux-2-klein-4b.safetensors |
Linear weight matrices transposed from PyTorch's (out, in) to JAX-idiomatic (in, out). The 5 double-stream blocks and 20 single-stream blocks are each stacked along a new leading axis, for use with jax.lax.scan. No values changed. |
| Text encoder (Qwen3-4B) | Comfy-Org/vae-text-encorder-for-flux-klein-4b, qwen_3_4b.safetensors |
Truncated from 36 to 27 transformer layers. FLUX.2 Klein conditioning is built only from hidden states after layers 9, 18 and 27; producing the deepest of those requires exactly 27 executed layers, so layers 27–35, the final norm, and the (tied, and therefore already-absent) lm_head are all provably dead weight for this use case and were dropped. Linear weights transposed to (in, out) as above. The embedding table keeps its original (vocab_size, hidden_size) row layout, since it is read by lookup, not by matrix multiplication. |
| Autoencoder | Comfy-Org/vae-text-encorder-for-flux-klein-4b, flux2-vae.safetensors |
Decoder only; the encoder is never used for text-to-image generation and was dropped entirely. Convolution kernels reordered from PyTorch's OIHW to JAX's HWIO layout. The BatchNorm layer used to de-normalise the latent (frozen, evaluation mode, no learnable affine) was collapsed into two constant vectors, scale and shift. The 1x1 post_quant_conv was reduced from a four-dimensional kernel to a plain (32, 32) matrix, an exact reformulation rather than an approximation, since a 1x1 convolution has no spatial receptive field. |
| Tokenizer | black-forest-labs/FLUX.2-klein-4B, tokenizer/ |
Copied verbatim. |
Exact tensor schema
Shapes below are as they exist after conversion, inside this bundle. Every shape was read directly from each source checkpoint's own header; none are approximate.
Diffusion transformer (params/transformer)
No tensor in the diffusion transformer has a bias; every projection is
bias-free, and the RMSNorm-style query/key normalisation layers use a
single scale vector each.
Stacked over a leading axis of 5 (double-stream blocks):
Key (after double_blocks_ prefix) |
Shape |
|---|---|
img_attn_qkv_weight |
(5, 3072, 9216) |
img_attn_proj_weight |
(5, 3072, 3072) |
img_attn_norm_query_norm_scale |
(5, 128) |
img_attn_norm_key_norm_scale |
(5, 128) |
img_mlp_0_weight |
(5, 3072, 18432) |
img_mlp_2_weight |
(5, 9216, 3072) |
txt_attn_qkv_weight |
(5, 3072, 9216) |
txt_attn_proj_weight |
(5, 3072, 3072) |
txt_attn_norm_query_norm_scale |
(5, 128) |
txt_attn_norm_key_norm_scale |
(5, 128) |
txt_mlp_0_weight |
(5, 3072, 18432) |
txt_mlp_2_weight |
(5, 9216, 3072) |
Stacked over a leading axis of 20 (single-stream, parallel attention+MLP
blocks; linear1 fuses QKV and the MLP's up-projection, linear2 fuses
the attention output and MLP down-projection):
Key (after single_blocks_ prefix) |
Shape |
|---|---|
linear1_weight |
(20, 3072, 27648) |
linear2_weight |
(20, 12288, 3072) |
norm_query_norm_scale |
(20, 128) |
norm_key_norm_scale |
(20, 128) |
Global tensors, used once per forward pass, no leading axis:
| Key | Shape |
|---|---|
img_in_weight |
(128, 3072) |
txt_in_weight |
(7680, 3072) |
time_in_in_layer_weight |
(256, 3072) |
time_in_out_layer_weight |
(3072, 3072) |
double_stream_modulation_img_lin_weight |
(3072, 18432) |
double_stream_modulation_txt_lin_weight |
(3072, 18432) |
single_stream_modulation_lin_weight |
(3072, 9216) |
final_layer_adaLN_modulation_1_weight |
(3072, 6144) |
final_layer_linear_weight |
(3072, 128) |
Text encoder (params/text_encoder)
embed_tokens.weight (151936, 2560) # unchanged layout
layers.input_layernorm_weight (27, 2560)
layers.self_attn_q_proj_weight (27, 2560, 4096)
layers.self_attn_k_proj_weight (27, 2560, 1024)
layers.self_attn_v_proj_weight (27, 2560, 1024)
layers.self_attn_o_proj_weight (27, 4096, 2560)
layers.self_attn_q_norm_weight (27, 128)
layers.self_attn_k_norm_weight (27, 128)
layers.post_attention_layernorm_weight (27, 2560)
layers.mlp_gate_proj_weight (27, 2560, 9728)
layers.mlp_up_proj_weight (27, 2560, 9728)
layers.mlp_down_proj_weight (27, 9728, 2560)
Autoencoder decoder (params/vae)
post_quant_conv and latent_denormalize are top-level entries, not
part of decoder:
post_quant_conv.weight (32, 32) # squeezed and transposed, see above
post_quant_conv.bias (32,)
latent_denormalize.scale (128,) # folded BatchNorm
latent_denormalize.shift (128,)
Everything else lives under decoder, with the original decoder.
prefix stripped and remaining dots replaced with underscores (for
example mid.block_1.conv1.weight becomes the key
mid_block_1_conv1_weight). Convolution kernels are HWIO; every 1x1
convolution that is not post_quant_conv (the four attention
projections inside the middle block) keeps its four-dimensional shape
rather than being squeezed to a matrix — only post_quant_conv receives
that special-cased treatment.
Fixed stem and head:
| Key | Shape (HWIO) |
|---|---|
conv_in_weight / conv_in_bias |
(3, 3, 32, 512) / (512,) |
mid_block_1_conv1_weight / ..._conv2_weight |
(3, 3, 512, 512) each |
mid_attn_1_q_weight / _k_ / _v_ / _proj_out_weight |
(1, 1, 512, 512) each |
norm_out_weight / norm_out_bias |
(128,) |
conv_out_weight / conv_out_bias |
(3, 3, 128, 3) / (3,) |
Four up levels (indexed 0–3, level 0 nearest the output), each with
three residual blocks (block_0, block_1, block_2) and an upsample
convolution except at level 0:
| Level | In → out channels | block_0 has nin_shortcut |
Has upsample |
|---|---|---|---|
| 0 | 256 → 128 | yes (channels change) | no |
| 1 | 512 → 256 | yes | yes |
| 2 | 512 → 512 | no | yes |
| 3 | 512 → 512 | no | yes |
block_1 and block_2 at every level keep that level's output channel
count on both sides (no shortcut needed). Naming follows the pattern
up_{level}_block_{index}_{conv1|conv2|norm1|norm2}_{weight|bias}, plus
up_{level}_block_0_nin_shortcut_{weight|bias} where applicable and
up_{level}_upsample_conv_{weight|bias} where applicable.
manifest.json
A small JSON file recording what was converted and from where, written alongside the parameters:
{
"conversion_timestamp_utc": "2026-08-17T15:03:38+00:00",
"source": {
"dit_repo_id": "black-forest-labs/FLUX.2-klein-4B",
"dit_filename": "flux-2-klein-4b.safetensors",
"vae_repo_id": "Comfy-Org/vae-text-encorder-for-flux-klein-4b",
"vae_filename": "split_files/vae/flux2-vae.safetensors",
"text_encoder_repo_id": "Comfy-Org/vae-text-encorder-for-flux-klein-4b",
"text_encoder_filename": "split_files/text_encoders/qwen_3_4b.safetensors",
"tokenizer_repo_id": "black-forest-labs/FLUX.2-klein-4B"
},
"text_encoder": {
"hidden_states_output_layers": [9, 18, 27],
"num_layers_kept": 27
},
"vae": {
"latent_channels": 32,
"patch_size": [2, 2],
"batch_norm_eps": 0.0001
},
"dit": {
"num_double_blocks": 5,
"num_single_blocks": 20,
"hidden_size": 3072
}
}
Architecture summary
- Diffusion transformer: hidden size 3072, 5 double-stream blocks
followed by 20 single-stream (parallel attention+MLP) blocks, 4-axis
rotary position embeddings, joint attention over text and image
tokens. No convolutions, no biases. Guidance is not embedded
(
guidance_embed=False); this is the 4-step, distillation-guided variant of FLUX.2 Klein, not the 50-step base model. The number of sampling steps and the guidance value are fixed by distillation in the reference implementation and are not meant to be changed. - Text encoder: Qwen3-4B, hidden size 2560, 27 of the original 36 transformer layers. Conditioning is built by concatenating the hidden states after layers 9, 18 and 27 into a 7680-dimensional context vector.
- Autoencoder: decoder-only, 32 latent channels, 2x2 patchify (128 channels after patchify), f8 spatial downscale. Reference decoding runs in float32, which this bundle preserves.
Recommended resolutions
The reference sampling schedule has a discontinuity at 4300 image tokens: above that point the step-count-aware schedule formula is replaced by a different one, which this checkpoint's 4-step distillation was not tuned against. Resolutions at or below that threshold are recommended:
| Resolution | Aspect ratio | Image tokens |
|---|---|---|
| 1024 x 1024 | 1:1 | 4096 |
| 1360 x 768 | 16:9 | 4080 |
| 768 x 1360 | 9:16 | 4080 |
Loading
pip install jax orbax-checkpoint huggingface_hub transformers
from huggingface_hub import snapshot_download
import orbax.checkpoint as ocp
local_path = snapshot_download(repo_id="johaness14/flux2-klein-4b-jax")
checkpointer = ocp.StandardCheckpointer()
text_encoder_params = checkpointer.restore(f"{local_path}/params/text_encoder")
transformer_params = checkpointer.restore(f"{local_path}/params/transformer")
vae_params = checkpointer.restore(f"{local_path}/params/vae")
# Tokenizer loading does not require PyTorch or a torch backend.
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(f"{local_path}/tokenizer")
To inspect the exact tensor names and shapes inside any component yourself, rather than relying on the tables above:
import jax
jax.tree_util.tree_map(lambda x: x.shape, transformer_params)
What this repository does not include
There is no model forward pass, no sampling loop, and no image pre/post-processing code here. Restoring these pytrees gives you parameters; running inference with them requires a separate JAX implementation of the diffusion transformer, the text encoder, and the decoder, matching the tensor layouts described above.
License and attribution
Released under Apache-2.0, matching every source component. This is an independent, unofficial conversion and is not affiliated with or endorsed by Black Forest Labs. All original weights are credited to:
Status
Personal, experimental project. Provided as-is, with no warranty and no support commitment.
Model tree for johaness14/flux2-klein-4b-jax
Base model
black-forest-labs/FLUX.2-klein-4B