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