---
license: apache-2.0
pipeline_tag: text-to-image
library_name: diffusers
tags:
- small
- dit
- flow-matching
- custom_code
datasets:
- QLNI/FLUX-Reason-6M-flux2-latents
language:
- en
---

<h1 align="center">Surjo-Image-Preview</h1>

<p align="center">
  Text-To-Image • 108M Parameters • Looped Latent DiT, Flow Matching
</p>

## Samples

<table>
<tr>
<td><img src="samples/sample_car.png" alt="car"></td>
<td><img src="samples/sample_vase.png" alt="vase"></td>
</tr>
<tr>
<td><img src="samples/sample_strawberries.png" alt="strawberries"></td>
<td><img src="samples/sample_sea.png" alt="sea"></td>
</tr>
</table>

---

## Model

### About the model

A tiny looped diffusion transformer (DiT) for 256px text-to-image, trained from scratch with rectified flow matching.

- Pipeline: text-to-image
- Parameter Count: 108.3M
- Encoder: frozen `LiquidAI/LFM2.5-Encoder-230M` (1024-dim, 128 tokens)
- VAE: FLUX.2 small-decoder mirror (`SurjoLabs/FLUX.2-small-decoder-vae`)
- Image resolution: 256²
- Latents: post-BN packed `(128, 16, 16)`, 256 tokens
- Weights shipped are the EMA artifact (best-equivalent)

### Model config

- `DIM`: 512
- `DEPTH`: 18 (7 pre + 4 looped ×2 + 7 post, effective 22)
- `HEADS`: 8 × 64, QK-RMSNorm
- `POS`: 2D-RoPE
- `MLP`: SwiGLU
- `TEXT_DIM`: 1024 (RMSNorm + projector)
- `REPA_DIM`: 384 (DINOv2-S, HASTE schedule, off after 70%)
- `LOOP_ITERS`: 2
- Custom architecture, loading needs `trust_remote_code=True`

---

## Training

### Dataset

~5.89M rows of prebuilt FLUX.2 latents + first-caption-per-row (`QLNI/FLUX-Reason-6M-flux2-latents`), 10 epochs (58.9M samples).

### Recipe

AdamW (0.9/0.99, no decay, clip 1.0), LR 2e-4 to ~2e-5 cosine (5k warmup; decay horizon overshot by the warmup length, so it lands at ~2.35e-5 instead of the 2e-5 floor, cosmetic, disclosed), fp32 masters / bf16 compute, EMA 0.9999. 57,510 steps, global batch 1024, final val NLL **0.30116**.

---

## How to run the model

```bash
pip install "transformers==5.13.1" "diffusers==0.39.0" accelerate safetensors pillow
```

```python
import torch
from diffusers import AutoModel, AutoencoderKLFlux2
from transformers import AutoTokenizer, AutoModelForMaskedLM

REPO = "SurjoLabs/Surjo-Image-Preview"
device = "cuda"
dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16

model = AutoModel.from_pretrained(REPO, trust_remote_code=True,
                                  torch_dtype=dtype).to(device).eval()
tok = AutoTokenizer.from_pretrained("LiquidAI/LFM2.5-Encoder-230M",
                                    trust_remote_code=True)
tenc = AutoModelForMaskedLM.from_pretrained(
    "LiquidAI/LFM2.5-Encoder-230M", trust_remote_code=True,
    dtype=dtype).to(device).eval().lfm2
vae = AutoencoderKLFlux2.from_pretrained(
    "SurjoLabs/FLUX.2-small-decoder-vae", subfolder="vae",
    torch_dtype=dtype).to(device).eval()
mean, var, eps = (vae.bn.running_mean.float(), vae.bn.running_var.float(),
                  float(vae.config.batch_norm_eps))

prompt = "a cozy cabin in a snowy forest at dusk, warm light in the windows"
ids = tok([prompt], padding="max_length", truncation=True, max_length=128,
          return_tensors="pt")
ctx = tenc(**{k: v.to(device) for k, v in ids.items()}).last_hidden_state.to(dtype)
uids = tok([""], padding="max_length", truncation=True, max_length=128,
           return_tensors="pt")
uctx = tenc(**{k: v.to(device) for k, v in uids.items()}).last_hidden_state.to(dtype)
cmask = ids["attention_mask"].to(device=device, dtype=dtype)
umask = uids["attention_mask"].to(device=device, dtype=dtype)
amp = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16

steps, cfg_scale, dt = 50, 3.0, 1.0 / 50
gen = torch.Generator(device=device).manual_seed(0)
z = torch.randn(1, 128, 16, 16, device=device, generator=gen, dtype=dtype)
with torch.no_grad():
    for i in range(steps):
        t = torch.full((1,), i * dt, device=device)
        with torch.autocast(device_type="cuda", dtype=amp):
            out = model(torch.cat([z, z]), torch.cat([t, t]),
                        torch.cat([uctx, ctx]),
                        ctx_mask=torch.cat([umask, cmask]),
                        deepsup=False, loop_iters=2)
        v_u, v_c = out.float().chunk(2)
        z = (z + dt * (v_u + cfg_scale * (v_c - v_u))).to(dtype)
    std = torch.sqrt(var.view(1, -1, 1, 1) + eps)
    lat = (z.float() * std.to(z.device)
           + mean.view(1, -1, 1, 1).to(z.device)).to(dtype)
    img = vae.decode(lat.view(1, 32, 2, 2, 16, 16
                              ).permute(0, 1, 4, 2, 5, 3).reshape(1, 32, 32, 32)).sample
    img = (img.clamp(-1, 1) + 1) / 2

from PIL import Image
Image.fromarray(
    ((img[0].permute(1, 2, 0).float().cpu().numpy()) * 255).astype("uint8")
).save("out.png")
```

### Recommended settings for sampling

- `--steps`: 50
- `--cfg`: 3.0
- `loop_iters`: 2 (model default)
- fp16 on pre-sm80 GPUs (e.g. T4), bf16 otherwise

---

## Evaluation

GenEval2 ([Kamath et al., 2025](https://arxiv.org/abs/2512.16853)): 800 prompts, Soft-TIFA with Qwen3-VL-8B-Instruct, 1 image per prompt (256px, 50 steps, CFG 3.0, seed 42). Headline is Soft-TIFA GM.

| metric | Surjo-Image-Preview | Supra2-IMG |
|---|---|---|
| Soft-TIFA GM | **8.07** | 6.80 |
| Soft-TIFA AM | **42.30** | 36.38 |
| object | **61.03** | 46.77 |
| attribute | **45.80** | 43.56 |
| count | **22.30** | 22.14 |
| position | 28.91 | **31.07** |
| verb | **2.03** | 1.51 |

Takes objects/attributes/count; cedes position. Both models floor on verbs (~2%) and struggle counting (~22%): small-model limits, not a gap.

---

## Future

These are the final weights for this preview. Surjo-Image (full release) will bring more data, refined architecture, and higher optimization.

---

![Agni](agni-logo-dark.svg)

Trained using Agni

---

## Limitations

- Renders legible text poorly; faces and hands degrade out-of-distribution.
- English-only captions; 256px single-scale; small-model weak draws on complex multi-subject scenes.
- Web-scale source biases inherited.

---

## AI Usage

SurjoLabs has never trained a Text2Image model prior. AI was used heavily compared to previous models to write code. Research and designing the model was done by a human.

---

## Acknowledgments

- Thanks to [SupraLabs](https://huggingface.co/SupraLabs/Supra2-IMG) for proving the tiny-T2I lane, and for the inspiration.
- Thanks to the [FLUX-Reason-6M](https://huggingface.co/datasets/LucasFang/FLUX-Reason-6M) team for the training dataset, and to [QLNI](https://huggingface.co/datasets/QLNI/FLUX-Reason-6M-flux2-latents) for the precomputed latents.
- Thanks to Black Forest Labs for the [FLUX.2 VAE](https://huggingface.co/black-forest-labs/FLUX.2-small-decoder).
- Thanks to the authors of [i1](https://arxiv.org/abs/2606.11289) (configuration), [Looped-DiT](https://arxiv.org/abs/2609.40305) (looping), [XSA](https://arxiv.org/abs/2603.09078) (exclusive self-attention), and [PixelModel-v6](https://huggingface.co/bench-labs/PixelModel-v6) (training recipe).