VisualInference/P-JiT-21K-MLX-bf16

P-JiT 21K, MLX (bf16)

This is an MLX export of Linum-AI/pjit-21k, Linum's 2.18B pixel-space text-to-image DiT. It comes with its Qwen3.5-4B caption encoder, trimmed to the 28 blocks the model reads. The runner is pjit-mlx.

Tier: bf16 (reference precision, seed parity with Linum's gallery).

Rendered on a Mac mini from Linum's gallery prompts and seeds; Linum's H100 originals are at linum.ai/field-notes/pjit-21k.

M4 Pro, one 1152×864 image, 50 steps	This tier	PyTorch on MPS
Total time	72 s	139 s
Peak memory, whole process	6.7 GB; 5.6 GB with --low-memory	14.9 GB
Download	11.9 GB	18.0 GB

Peak memory is the process's peak physical footprint as macOS reports it, the same meter in both columns. --low-memory evaluates the DiT block by block and caps MLX's buffer cache; the pixels are identical.

pip install git+https://github.com/felix-ab/pjit-mlx
pjit-mlx --tier bf16 --seeds 4 --aspect 4:3 \
  --prompt "Bright midday sun lights a white man in his thirties in a head-and-shoulders portrait, squinting slightly with a relaxed half smile. His faded blue chambray shirt is open at the collar, and a field of tall grass blurs behind him."

--tier bf16 downloads this repo on first use. To keep a local copy, run hf download VisualInference/P-JiT-21K-MLX-bf16 --local-dir P-JiT-21K-MLX-bf16 and pass --dit P-JiT-21K-MLX-bf16/dit --text-encoder P-JiT-21K-MLX-bf16/text_encoder instead.

Files

• dit/: the DiT (model.safetensors) with Linum's config.json, plus a pjit_mlx block that gives the format, dtype, quantization and source revision.
• text_encoder/: Qwen3.5-4B blocks 0-27 (model.safetensors), the bf16 token embedding (embed_tokens.safetensors, read row by row from a memory map), and the tokenizer.

How exact it is

• In fp32 the MLX DiT and text encoder match the PyTorch reference to a relative error of about 3e-6.
• A full 50-step fp32 run matches the reference sampler on 99.5% of pixels, with the rest within 1/255.
• In bf16 they sit at the reference's own bf16 rounding noise.
• The bf16 tier reproduces Linum's gallery images from their published prompt and seed: 15 tested, all the same picture, median PSNR 35 dB against the AVIF originals. It does this by reproducing the H100 SXM noise draw exactly.
• The 8-bit tier shifts the first prediction a little more than bf16 does. Most seeds still give the gallery picture, but sensitive prompts can fork.

Full numbers and the scripts behind them are in the runner's README.

Limits

These are the model's own limits (see Linum's card):

• It was trained on dense 50-80 word captions.
• It cannot render text.
• It struggles with full-body figures and crowds.

License

• P-JiT 21K weights: Apache-2.0, © Linum Inc.
• Qwen3.5-4B: Apache-2.0, © Alibaba Cloud.
• This conversion and the runner: Apache-2.0.