flux2-klein-p100a

FLUX.2 [klein] 4B, ported to one Tenstorrent Blackhole p100a. Weights: black-forest-labs/FLUX.2-klein-4B (e7b7dc2) · Port: SunandBean/tt-flux2-klein

Runs on p100a (mesh 1×1). Configuration: Tensix (worker) dispatch, 1 command queue, 11×10 compute grid. The Qwen3-4B text encoder, the DiT and the VAE all run on the card; only the scheduler and the image post-processing stay on the host.

What this port adds — the 4-step klein port on one chip, GroupNorm-grid and non-square fixes in the tt_dit VAE, and the host helpers shared with the Z-Image port. What it builds on — the changh95/qwen-image-2.1-p150 single-chip port and tt-metal's tt_dit VAE, both Apache-2.0; diffusers 0.37.1 Flux2KleinPipeline is the reference implementation.

The fastest of the models on this card: a 1024×1024 image in about 3.2 s.

Quickstart (Python)

Prerequisite: a tt-metal / ttnn environment. ttnn is not on PyPI.

hf download sunandbean/flux2-klein-p100a --include "code/*" "media/*" --local-dir flux2-klein-p100a && cd flux2-klein-p100a
pip install -e code/                  # the Python API only
pip install -e "code/[server]"        # also the HTTP server (fastapi, uvicorn, pydantic)

from PIL import Image
from tt_flux2_klein import KleinTT, open_device, close_device

dev = open_device()
try:
    model = KleinTT(dev)              # downloads black-forest-labs/FLUX.2-klein-4B@e7b7dc2 into your HF cache

    # text to image
    image, timing = model.generate("A red fox standing in fresh snow at dawn, wildlife photography",
                                   width=1024, height=1024, seed=42)

    # editing / reference-guided generation: pass 1-3 images
    edited, timing = model.generate("Replace the cup of coffee with a tall glass of orange juice.",
                                    width=512, height=512, seed=3,
                                    images=[Image.open("media/croissant512.png")])
    edited.save("output.png")
finally:
    close_device(dev)

Tasks	Text to image, editing, and reference-guided generation — the same generate call, with or without images.
Size	width / height 512–1920 px, 8 px steps, at most 1920×1088 pixels.
Steps	Fixed at 4. klein is distilled: guidance_scale is 1.0 and there is no CFG, so each step is one DiT call.
References	1–3 images, passed at their original aspect ratio (reduced to ≤1024² area and center-cropped to a multiple of 16), not squeezed onto a square canvas.
Output	(PIL.Image, Timing). Timing.values holds text_encoder_s, reference_encode_s, context_s, dit_s, vae_decode_s.
Weights	black-forest-labs/FLUX.2-klein-4B pinned to e7b7dc27f91deacad38e78976d1f2b499d76a294. KLEIN_SNAPSHOT overrides.

Reference encodings are cached (last 4 images), so reusing a source image across prompts skips the host VAE encode — about 2.9 s at 1024².

Serving (HTTP)

uvicorn tt_flux2_klein.server:app --host 0.0.0.0 --port 20000
curl -s localhost:20000/predict -H 'Content-Type: application/json' \
  -d '{"prompt":"A red fox standing in fresh snow at dawn","width":1024,"height":1024,"seed":42}'

• POST /predict: prompt, optional images (1–3 base64 PNG/JPEG/WebP), width (1024), height (1024), seed (42), num_steps (4).
• GET /health, GET /info.
• Contract errors — more than 3 images, a step count other than 4, a size outside 512–1920 or off the 8-px grid — return 422 without touching the card.

Demo

1024×1024, seed 1234, 4 steps, generated on the p100a.

A cozy cafe storefront with a wooden sign that says "OPEN DAILY", morning light	A bowl of ramen with a soft-boiled egg, scallions and chashu, overhead food photography

Performance

Warm, batch 1, 4 steps, measured 2026-09-26 on an idle p100a: Tensix (worker) dispatch, 1 command queue, 11×10 compute grid.

Text to image, 1024×1024 — ten different prompts through the HTTP server:

End to end	2.80–3.43 s
DiT (4 steps)	1.33–1.36 s
VAE decode (on the card)	1.23–1.28 s
Text encoder	0.08–0.09 s
Context	0.001–0.004 s

Editing and reference-guided generation (operational measurement, same day):

1024² with a new reference	~5.5 s, of which the host VAE reference encode is 2.9 s
1024² reusing the same reference (cached)	~2.2 s
512² with a reference	~1.9 s
1920×1080 text to image, continuous	9.2 s (20.2 s for the first image at that size)
1080×1920 text to image, first at that size	29.2 s

Model load	22.3 s (81.3 s for a full service start)
Device DRAM after load / after warmup	12.9 / 13.0 GB of 29.9 GB

From the published container

The same image anyone can tt-model pull, reached over its own HTTP server on an idle card (2026-10-09). Three requests at 1024², seed 42:

	cold	warm	warm
End to end	38.32 s	2.93 s	2.95 s
Text encoder	6.32 s	0.074 s	0.075 s
DiT (4 steps)	12.06 s	1.323 s	1.327 s
VAE decode	19.07 s	1.212 s	1.219 s

All three produced a byte-identical PNG (sha256 14cb3542fffc8b81).

Against an RTX 5070 Ti

The GPU runs the diffusers Flux2KleinPipeline at bf16 with the same prompt, size, seed and 4-step schedule.

	p100a	RTX 5070 Ti
End to end, 1024²	2.95 s	9.0–10.3 s
Per DiT step	0.331 s	0.46 s
Peak device memory	12.9 GB of 29.9	14.7 GB of 15.5

klein at bf16 does not fit in 15.5 GB either, so the GPU run offloads stage by stage and the end-to-end figure includes that traffic. The per-step figure does not: 1.4× per DiT step.

Accuracy

Against diffusers 0.37.1 on CPU (klein_device_check.py):

Stage	edit 512²	848×624	1024²
Text encoder, valid rows	0.9999993	0.9999994	0.9999994
Text encoder, all 512 rows	0.99917	0.99960	0.99958
Reference tokens	0.99993	—	—
DiT first step	0.99967	0.99955	0.99935
End-to-end latents, reference text encoder	0.9991	0.9688	0.9942
End-to-end latents, on-card text encoder	0.9792	0.7809	0.9601
End-to-end decoded image, on-card text encoder	0.9868	0.8944	0.9804

Read the last three rows together. Every individual stage is at 0.999 or better, but klein runs only 4 steps with no CFG, so a small deviation in the conditioning is not averaged away — it steers the whole trajectory. Swapping in the CPU reference text encoder recovers most of the gap, which locates the sensitivity in the text encoder rather than in the DiT. 848×624 is the worst case; the same shape is also the worst case for Z-Image-Turbo on this card.

The images remain correct in composition and quality. See Caveats for what this does and does not mean.

Known failure mode

Across 24 editing and reference-guided cases run during the port, the failures that showed up were over-application of a constrained instruction: "make only the teapot green" turned the cup green as well, and "make only the sofa blue" shifted the cat's fur colour. Whether that is upstream model behaviour or something this port introduced was not isolated.

The full porting notes — upstream behaviour reproduced, the single-chip design, and why the 848×624 number is implementation variance rather than a bug — are in PORTING_NOTES.md.

Caveats

• Single card only. Mesh 1×1 with Tensix (worker) dispatch; no multi-chip scaling.
• 4 steps only. klein is distilled; other step counts are rejected.
• End-to-end PCC against the CPU reference is modest at some sizes (0.78 latents at 848×624). This is the 4-step, no-CFG schedule amplifying a text-encoder deviation of 6e-7 per valid row, not a broken stage — every stage measures 0.999+. Treat the port as visually equivalent, not numerically identical.
• Editing instructions can over-apply. "Only the teapot" became the teapot and the cup. Check the result before relying on a constrained edit.
• The text encoder is bit-identical to Qwen/Qwen3-4B (398 of 398 tensors), so that checkpoint can be substituted with KLEIN_TEXT_ENCODER.
• Upstream feeds all 512 padded text rows into the DiT with no mask, and uses layers 9, 18 and 27 of the text encoder concatenated to [512, 7680]; layers 28–36 are never used. This port reproduces that rather than "fixing" it.
• Modulation in klein is global per step, not per block — a detail that is easy to get wrong when porting from other DiTs.

Licensing

• Weights: black-forest-labs/FLUX.2-klein-4B, Apache-2.0, revision e7b7dc27f91deacad38e78976d1f2b499d76a294. Not redistributed here.
• Port code (code/): Apache-2.0 (SPDX headers on the modules), from SunandBean/tt-flux2-klein.
• code/tt_flux2_klein/vendor_tt_dit/: the tt_dit VAE from tt-metal at 01d6e7b (Apache-2.0), with non-square slicing and GroupNorm grid fixes.
• code/tt_flux2_klein/common.py: model-independent host helpers shared with the tt-z-image-turbo port (Apache-2.0).
• Derived from the changh95/qwen-image-2.1-p150 single-chip port (Apache-2.0).
• code/THIRD-PARTY.md lists every directory under code/ that came from somewhere else, with its upstream commit and where its licence text sits in this repo.

Run it as a container

Everything below is generated by tt-model package from the manifest this image was built from.

tt-model pull  sunandbean/flux2-klein-p100a --with-weights
tt-model serve sunandbean/flux2-klein-p100a

pull --with-weights downloads the Docker image and the black-forest-labs/FLUX.2-klein-4B weights at e7b7dc27f91deacad38e78976d1f2b499d76a294 (into your HF cache; they are not in the image). serve starts the model's own HTTP server on port 20000 (or the next free port, if that one is busy); the first start compiles kernels for your device, which takes several minutes, and the server is ready when it logs Application startup complete.

Request shape

POST /predict takes {"prompt": str, "images": [b64], "width": int, "height": int, "seed": int} and returns {"image": base64 PNG, "width", "height", "model", "license", "seed", "num_steps", "reference_count", "timing_ms"}. images is optional: none for text-to-image, up to three for an edit or a reference. klein is itself distilled to 4 steps; any other step count is a 422. A constrained edit can over-apply -- see Caveats on the model card.

Provenance

The exact sources the image was built from — the model code under code/ is byte-identical to what is inside the image (pyproject.toml, PYTHON.md and examples/ are packaging files this repo adds, and are not in the image):

component	built from
tt-metal	58648d44443223041148125843ca7c7d4780463e
code/ digest	8491c45396dbafe7 (sha256, first 16 hex digits)
built	2026-10-09T17:48:15+00:00 by tt-model 0.1.0