File size: 7,234 Bytes

28db2ee

---

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).