Card: whole-process peak memory and --low-memory (pjit-mlx 0.1.1)

4e983d0 verified 4 days ago

3.68 kB

license: apache-2.0
library_name: mlx
pipeline_tag: text-to-image
base_model: Linum-AI/pjit-21k
tags:
  - mlx
  - apple-silicon
  - text-to-image
  - diffusion
  - pixel-space
  - pjit-21k

P-JiT 21K, MLX (8bit)

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: 8-bit (compact: 3.6 GB peak for the whole process with --low-memory, 5.0 GB without). It is as clean as bf16 at 100% and usually gives the same picture for a seed (14 of 15 gallery tests, with more fine-detail drift), but it can fork on sensitive prompts.

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	79 s	139 s
Peak memory, whole process	5.0 GB; 3.6 GB with `--low-memory`	14.9 GB
Download	6.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 8bit --low-memory --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 8bit downloads this repo on first use. To keep a local copy, run hf download VisualInference/P-JiT-21K-MLX-8bit --local-dir P-JiT-21K-MLX-8bit and pass --dit P-JiT-21K-MLX-8bit/dit --text-encoder P-JiT-21K-MLX-8bit/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, with more fine-detail drift (median 29.6 dB), but sensitive prompts can fork to a different image of the same quality.

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

This conversion and the runner: Apache-2.0.