LocalMuseAI/mlx-flux2-klein-4b-aio-4bit

FLUX.2 Klein 4B AIO — MLX INT4 for LocalMuse

Source AIO repack by SeeSeeLP / SeeSee21: https://civitai.com/models/2327389?modelVersionId=2618128 This is the author's repack of the official Black Forest Labs Klein 4B step-distilled model, not a new fine-tune. Model, text encoder and VAE are Apache-2.0. Preserve LICENSE, NOTICE and the source attribution.

Source BF16 checkpoint: 16,132,318,096 bytes; SHA-256 f9f60d53bb21970e19d31d85eefb294f9b0b5e36669ec81ba2bf52f102a7679e. Exact source metadata and permissions are in source.json.

Runtime contract

flux2Klein4B4Bit (ordinary affine INT4, group 64), separate from Bonsai's ternary INT2 layout. DiT attention/FFN and Qwen3 are INT4; sensitive DiT layers remain BF16. Qwen raw hidden-state taps 9/18/27 are concatenated; unconsumed layers 27–35 and the language head are omitted.

Text to Image and Image to Image (reference editing), batch one, 1024 × 1024, Euler / Simple, 4–6 steps (default 4), distilled CFG 1, no negative prompt. For editing, the source is encoded by the AIO's FP32 VAE encoder, 2×2-patchified, BN-normalized and appended to the image stream as 4096 reference tokens at RoPE time 10, as in Diffusers' Flux2KleinPipeline; generation still starts from noise, so there is no strength. Text encoder, image encoder, transformer and FP32 VAE decoder load sequentially. The full original decoder uses 384-pixel overlapping windows for mobile memory. Live previews project the clean latent on CPU without another neural network.

VAE projections

The AIO stores the learned 1×1 VAE projections in Black Forest Labs' nested layout (vae.encoder.quant_conv, vae.decoder.post_quant_conv) rather than Diffusers' top-level names. The first conversion looked only for the top-level decoder key and restored post_quant_conv from the matching official black-forest-labs/FLUX.2-klein-4B VAE at revision e7b7dc27f91deacad38e78976d1f2b499d76a294; SHA-256 ca70d2202afe6415bdbcb8793ba8cd99fd159cfe6192381504d6c4d3036e0f04. Every checked FP32 AIO VAE value, including both nested projections, rounds exactly to that official BF16 file, so the restored projection is the same learned tensor. An identity projection is invalid and produces mosaic. The encoder export uses the AIO's own FP32 quant_conv. Encoder, decoder and batch-normalization statistics keep FP32 precision; this is not a merge or training.

Files and verification

The seven application runtime files total 4,560,019,137 bytes. manifest.json authenticates every byte size and SHA-256. Tokenizer files are pinned to the same official BFL revision. Install by immutable HF commit, never by main.

A standalone .safetensors file in this repository is a prequantized tensor component, not a ComfyUI AIO checkpoint. Requires LocalMuse's Klein INT4 runtime or a compatible MLX implementation; do not load it with Bonsai INT2.

Native validation

Reference edit of the sample below with the same library: 1024 × 1024, four steps, seed 1729, prompt “Change the cream sweater to a bright red wool sweater, keep the face, pose, window and lighting unchanged.” On the M1 Pro Mac with Metal API validation it took 173.81 seconds with a peak MLX allocation of 4,349,946,880 bytes and reported three preparation stages. Against independent Diffusers code (validation/edit-reference.json), the FP32 encoder and reference tokens match to relative L2 0.0000026; with the reference appended, dequantized INT4 DiT in Diffusers versus native Swift has cosine 0.999943 (relative L2 0.010691), and source BF16 DiT cosine 0.995487.

The sample above was generated by the same Swift library used in LocalMuse: 1024 × 1024, Euler / Simple, CFG 1, four steps, seed 1729. Prompt: “Candid photograph of an adult woman in a cream sweater by a studio window, soft evening light, natural skin texture, analog grain.”

On an Apple M1 Pro Mac with 16 GiB RAM and Metal API validation enabled, this run completed in 99.47 seconds, with a peak MLX allocation of 3,606,375,424 bytes (3.36 GiB). It reported both preparation stages and four CPU-projected preview frames. These are Mac measurements; physical iPhone validation is pending and will be performed by the app user. The current LocalMuse catalog requires the nominal 12 GB memory class.

Independent Transformers/Diffusers comparison: source BF16 encoder cosine 0.995114; source BF16 DiT cosine 0.993426; Diffusers with independently unpacked INT4 matrices versus native Swift cosine 0.999953 and relative L2 0.009666. Restored FP32 VAE relative L2: 0.00000255. Token IDs match exactly and source BN statistics are preserved exactly. See validation/reference.json.

These are component comparisons, not proof of identical complete BF16 and INT4 images. MLX and ComfyUI use different noise RNG implementations. INT4 and windowed VAE decoding can also change the final image. The additional BF16 DiT reference uses the same INT4 Qwen3 conditioning and initial noise as the native run; see validation/vae-repair.json.

Application download manifests pin the seven runtime files to commit 646b51828f39b2541fbfcc643d6984d199c6f6a5. The first six are byte-identical to the original c2abb5051572bcda02841b3e24af3e3fe9911d57 release; later documentation-only commits do not change runtime artifacts. verification.json records every remote size/hash and anonymous pinned-URL check.