Iris-3B

Pixel-space generation & general vision learner

Generative priors are a promising foundation for downstream vision tasks. In this project we explore pixel-space generative models as an alternative to vision foundation models such as DINOv2.

Iris-3B is a 3B-parameter diffusion transformer that generates directly in pixel space: no VAE, no latent space, every pixel is produced by the model, so nothing is lost to a lossy, texture-biased latent. It is not only an image generator: the same pixel-space prior, fine-tuned, estimates monocular depth and restores and upscales images — one backbone, three tasks. This repository contains the model, the training and fine-tuning code, and inference.

https://github.com/user-attachments/assets/2345a853-c673-4ddb-83b9-7e05c83198ed

Generations

Samples from the final checkpoint at native aspect ratios of about one megapixel (CFG 3, 100 steps, one fixed seed, no prompt rewriting).

More samples: project page.

One prior, many tasks

Iris-3B is a general vision learner, not only a text-to-image model: the same pixel-space backbone, fine-tuned with no architectural change, does dense vision tasks where detail matters. Both fine-tunes ship with the release (see Depth, image restoration and upscaling for usage).

Monocular depth — photo → relative depth, one forward pass:

Image restoration and upscaling — degraded input → restored image:

Interactive depth point clouds and before/after sliders: project page · demo.

Model

Install

pip install -e .            # Python >= 3.11, PyTorch >= 2.7.1
pip install -e ".[dev]"     # + pytest, ruff

Or with uv, which creates .venv and installs the package editable:

uv sync --extra dev         # add --extra ui for the demo
uv run pytest -q
uv run python scripts/sample.py --checkpoint iris-3b --prompt "..."

Inference

hf download speridlabs/iris-3b --local-dir iris-3b    # ~12 GB
python scripts/sample.py --checkpoint iris-3b \
    --prompt "a red fox sleeping in fresh snow, golden hour"

--checkpoint is a directory holding model.safetensors + config.yaml, or a training .pth (its EMA weights and embedded config are used). Defaults: 1024×1024, 100 DPM-Solver++ (order 2) steps, CFG 3. Use --txt-file for a prompt list, --steps, --seed, --negative-prompt to adjust. The device defaults to CUDA, then Apple MPS, then CPU. On CUDA and MPS the weights are held in BF16 (--weights-dtype fp32 keeps the FP32 release precision), which halves their memory at the same speed, since inference runs under BF16 autocast either way.

Faster sampling, training-free (each trades some fidelity for speed; compare against the defaults on your prompts):

• --steps 30: second-order DPM-Solver++ usually needs far fewer than the default 100 steps; cost is linear in steps.
• --cfg-interval 0 0.9: skip CFG on the noisiest steps, which then cost one forward instead of two (t = 1 is noise; with shift 4 about a third of the steps lie above 0.9).
• --cache-threshold 0.1: reuse the trunk's output between adjacent steps while the first block's input barely changes (first-block caching). A reused step costs ~15% of a full one; the pixel head always runs on the current image. 0 disables it; at most 3 steps in a row are reused.

On an M3 Max (40-core GPU) a 1024×1024 forward takes ~6.4 s with CFG and ~3.1 s without, so the defaults take ~11 min per image and --steps 30 --cfg-interval 0 0.9 ~2.6 min; with --cache-threshold 0.1 on top, one measured image reused the trunk on 18 of 30 steps and took 77 s.

Turbo LoRA (few-step distillation)

scripts/distill_turbo.py distills the model into a 4-step LoRA with the guidance baked in (DMD2: the student, the teacher and the DMD critic share one frozen copy of the weights; only two rank-64 adapters train, from prompts alone). Train at ~1 MP: the released checkpoint shows 16 px block artifacts at 512, which a 512 student would inherit.

python scripts/distill_turbo.py --checkpoint iris-3b --prompts prompts.txt \
    --out output/turbo --device cuda:0 --text-device cuda:1 --iters 3000
python scripts/sample.py --checkpoint iris-3b --lora output/turbo/latest/student.safetensors \
    --prompt "a red fox sleeping in fresh snow, golden hour"

--lora also takes a Hub id: --lora pentacoxian-dev/iris-3b-turbo-lora loads the published default adapter, ...:r1-1500 a named one (the repo's config.json manifest lists them).

The adapter is merged into the weights at load, so a step costs one plain forward (no CFG batch): 4 steps at 1024² take ~11–13 s per image on an M3 Max and ~0.9 s on an RTX 5090.

Speed demo

demo/speed.py is a Gradio page that times every mode on the same prompt and seed: Fast (the 4-step turbo LoRA, on its own merged bf16 copy so it runs at sample.py speed; --live-lora saves the 6 GB at ~10% per step), Balanced (the training-free 30-step settings) and Standard (the 100-step sampler). Models are warmed up at launch. It reports wall time, forward passes, time per forward, images per minute and the speed-up over the 100-step sampler (measured once that has run, otherwise estimated from measured forward times).

pip install -e ".[ui]"
python demo/speed.py --lora output/turbo/latest/student.safetensors   # http://127.0.0.1:7861

Macs with 16 GB

--preset mac16 fits a 16 GB Mac: the prompts are encoded in a short-lived child process (the text encoder and the image model never coexist), the image model is streamed onto the GPU tensor by tensor with the LoRA merged on the way, weight-only int8 is applied to the text encoder and the image trunk (--quantize none keeps bf16), and images are generated one at a time. Measured on an M3 Max with GPU allocations capped at a 16 GB Mac's ~10.7 GB working set: 6.9 GB peak footprint and ~11 s per 1024² turbo image (base M-series GPUs are slower). int8 kept text embeddings at cosine 0.9999 and images visually unchanged (28–38 dB PSNR against bf16 on the 4-step sampler).

python scripts/sample.py --checkpoint iris-3b --preset mac16 \
    --lora output/turbo/latest/student.safetensors --prompt "a red fox sleeping in fresh snow"

Turn a training checkpoint into a release directory (EMA weights only, FP32 by default, config reduced to model/text-encoder/flow settings):

python scripts/export_checkpoint.py output/run/checkpoints/latest.pth exported/iris-3b

Depth, image restoration and upscaling

The release also ships Iris-3B fine-tuned for monocular depth estimation and for image restoration and upscaling — the general-vision-learner side of the model — in the depth/ and upscaler/ folders of the same Hugging Face repo (~12 GB each). Both run the full model in a single forward pass with the empty prompt, so no text encoder is loaded.

python scripts/depth.py photo.jpg --out depth_out      # <name>.npy + colorized <name>.png
python scripts/upscale.py photo.jpg --out upscaled     # <name>_x4.png

• Depth returns relative log depth at the input resolution (−1 near, +1 far). It is affine-invariant, not metric: fit a scale and shift in log space to compare with ground truth. Inputs above 1024 px on the long side are downscaled for the model (--max-side 0 keeps native resolution).
• Restoration is a one-step restorer trained on 1024×1024 crops (4× of 256×256 inputs). The input is upsampled bicubically to 4×; outputs smaller than 1024 on the short side are run as one tile, larger ones as overlapping 1024×1024 tiles (stride 512) blended with a Gaussian window, followed by a wavelet colour fix against the bicubic upsample. By default inputs are first downscaled to at most 512 px short side / 1024 px long side, which caps the output at 2048×4096 (21 tiles); --no-budget keeps the input size.

Python API: iris3b.downstream.depth.DepthPredictor and iris3b.downstream.restoration.Restorer. Fine-tuning checkpoints are converted with scripts/export_downstream.py.

Demo

The Gradio app behind the Hugging Face Space (text-to-image, depth and upscaling tabs) is in demo/; it runs on ZeroGPU there and on any local GPU. All three models together take ~45 GB of GPU memory; with IRIS_OFFLOAD=1 they stay on the CPU and each request moves only the model it needs to the GPU (~20 GB peak), which is what the Space does:

pip install -e ".[ui]"
python demo/app.py          # IRIS_MODEL_REPO=<dir or repo id> to use other weights

Training

Data

Training reads tar shards indexed by a wids-meta.json. Each sample is an image plus a JSON with its caption, height and width. Build them from a folder of images with .txt/.json caption sidecars, or from a Hugging Face dataset:

python scripts/prepare_wids.py folder /path/to/images --out /data/my-set
python scripts/prepare_wids.py hf user/dataset --split train \
    --image-column image --caption-column text --out /data/my-set

Recipe

Iris-3B was trained in three resolution stages followed by supervised fine-tuning, each resuming from the previous one. The configs list only what differs from the defaults in src/iris3b/config.py:

Common to all stages: hybrid Muon — Muon (momentum 0.95, Nesterov, RMS-matched LR) for hidden weight matrices, AdamW (β₁ 0.9, β₂ 0.95) for embeddings, norms and biases — with weight decay 0, gradient clipping at 0.5 and a constant LR of 1e-4 after a 2K-step linear warmup unless noted. EMA decay 0.9999, 10% caption dropout for classifier-free guidance, logit-normal (0, 1) timesteps with the resolution-dependent shift above. From the 512 stage on, RoPE uses isotropic coordinates for non-square images. bf16, FSDP2 hybrid sharding, torch.compile.

Single node:

torchrun --nproc_per_node=8 scripts/train.py --config configs/iris3b/stage1_256.yaml \
    data.data_dirs=[/data/my-set] train.expected_global_batch=128

Multi-node (run on every node with its own NODE_RANK):

NNODES=8 NODE_RANK=0 MASTER_ADDR=node0 \
    scripts/launch_multinode.sh configs/iris3b/stage1_256.yaml data.data_dirs=[/data/my-set]

Any config key can be overridden as key=value after the config path.

Tests

pytest -q

Acknowledgements

The dual-level design (a patch-level transformer followed by a pixel-level transformer head) follows the architecture described in PixelDiT. This is an independent implementation; it contains no code from that project. We also build on DiT, REPA, Muon through Dion, Qwen3-VL and DINOv2.

License

The code in this repository and the model weights on Hugging Face are licensed under the Apache License 2.0; see also NOTICE.

Third-party components are not redistributed here and remain under their own licenses:

Citation

@article{licai2026iris,
  title   = {Iris-3B: Going Beyond the Latent with Pixel-Space
             Diffusion Training, Conversion and Fine-Tuning},
  author  = {Li Cai, Hanqiu and Garabito, Chema},
  journal = {arXiv preprint arXiv:2610.09450},
  year    = {2026},
  url     = {https://arxiv.org/abs/2610.09450}
}

---

Made with ❤️ by speridlabs.com