flux2-klein-4b-jax

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.

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.

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):

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):

Global tensors, used once per forward pass, no leading axis:

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:

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:

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:

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:

• black-forest-labs/FLUX.2-klein-4B
• Comfy-Org/vae-text-encorder-for-flux-klein-4b

Status

Personal, experimental project. Provided as-is, with no warranty and no support commitment.