Surjo-Image-Preview

Text-To-Image • 108M Parameters • Looped Latent DiT, Flow Matching

Samples

---

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

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

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

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.

---

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 for proving the tiny-T2I lane, and for the inspiration.
• Thanks to the FLUX-Reason-6M team for the training dataset, and to QLNI for the precomputed latents.
• Thanks to Black Forest Labs for the FLUX.2 VAE.
• Thanks to the authors of i1 (configuration), Looped-DiT (looping), XSA (exclusive self-attention), and PixelModel-v6 (training recipe).