iris-image-gen

IRIS: Iterative Refinement Image Synthesizer

A mobile-first image generation architecture designed from recent research (2025-2026). 17/17 tests pass.

Colab Quick Start (Free Tier T4 — Just Run It)

One script, real dataset, real training, ~20 minutes total.

1. Open Google Colab
2. Set runtime to T4 GPU (Runtime → Change runtime type → T4 GPU)
3. Create a new cell and paste:

!wget -q https://huggingface.co/asdf98/iris-image-gen/resolve/main/colab_train_iris.py
%run colab_train_iris.py

What happens:

• Installs deps (~30s)
• Downloads IRIS code + DC-AE encoder + text encoder (~2min)
• Encodes 833 Pokemon images to latents (~2min on T4)
• Trains IRIS-Small (40M params) for 3000 steps (~15min on T4)
• Generates sample images and plots loss curve
• Saves checkpoint to ./iris_checkpoints/

Colab free tier (2025): T4 GPU (16GB VRAM), ~12.7GB RAM, PyTorch 2.5+

VRAM Budget (Colab T4, 16 GB)

Encoders are freed before training starts → plenty of headroom.

Configs for Different Hardware

Dependencies (All pip-installable, no special builds)

torch>=2.0            # preinstalled in Colab
diffusers>=0.32.0     # for AutoencoderDC
sentence-transformers  # for text encoding (all-MiniLM-L6-v2, 87MB)
datasets              # for HF dataset loading
accelerate            # diffusers dependency
huggingface_hub       # for downloading IRIS code

No flash-attn, no triton, no apex, no custom CUDA kernels. Pure PyTorch.

---

Model Variants

Architecture

IRIS combines five innovations:

1. PDE-SSM spatial mixing — O(N log N) Fourier-domain PDE, native 2D. PDE-SSM-DiT
2. Weight-shared refinement — 6 blocks × R iterations, same weights. GRN
3. Structured latent canvas — DC-AE with channel masking. DC-AE 1.5
4. Tiny decoder — 0.1M params PixelShuffle. SnapGen
5. MQA + 2D RoPE + QK-RMSNorm — Mobile-optimized. SnapGen++

Python API

from iris import IRIS, get_model_config, flow_matching_loss, euler_sample
import torch

# Create model (with text projection for 384-dim MiniLM embeddings)
model = IRIS(**get_model_config("iris-small"), text_dim=384)

# Training step
z_0 = torch.randn(4, 32, 16, 16) * 2.5  # DC-AE latents
text_emb = torch.randn(4, 1, 384)        # MiniLM text embeddings
losses = flow_matching_loss(model, z_0, text_emb, num_iterations=3)
losses["loss"].backward()

# Sampling
model.eval()
noise = torch.randn(1, 32, 16, 16)
with torch.no_grad():
    z_pred = euler_sample(model, noise, text_emb[:1], num_steps=20, num_iterations=3)
    image = model.decode_latent(z_pred)  # (1, 3, 512, 512)

Files

colab_train_iris.py      # <-- ONE-CLICK COLAB NOTEBOOK
iris/
  __init__.py            # Public API
  model.py               # IRIS model (Patchify, Unpatchify, TinyDecoder, IRIS)
  core.py                # RefinementCore (weight-shared block loop)
  pde_ssm.py             # SpectralConv2d, TokenDifferential, PDESSMBlock
  blocks.py              # MQA, RoPE2D, UIB-FFN, TimestepEmbed, IterationEmbed
  flow_matching.py        # Rectified flow loss, Euler sampler
  configs.py             # 5 model configurations (tiny → large)
  train.py               # Training utilities (dataset, scheduler)
  train_production.py    # CLI training script
  test_all.py            # 17-test suite

License

MIT