Super-squash branch 'main' using huggingface_hub

.gitattributes

@@ -0,0 +1,166 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+grid_00000500.png filter=lfs diff=lfs merge=lfs -text
+grid_00001000.png filter=lfs diff=lfs merge=lfs -text
+grid_00001500.png filter=lfs diff=lfs merge=lfs -text
+grid_00002000.png filter=lfs diff=lfs merge=lfs -text
+grid_00002500.png filter=lfs diff=lfs merge=lfs -text
+grid_00003000.png filter=lfs diff=lfs merge=lfs -text
+grid_00003500.png filter=lfs diff=lfs merge=lfs -text
+grid_00004000.png filter=lfs diff=lfs merge=lfs -text
+grid_00004500.png filter=lfs diff=lfs merge=lfs -text
+grid_00005000.png filter=lfs diff=lfs merge=lfs -text
+grid_00005500.png filter=lfs diff=lfs merge=lfs -text
+grid_00006000.png filter=lfs diff=lfs merge=lfs -text
+grid_00006500.png filter=lfs diff=lfs merge=lfs -text
+grid_00007000.png filter=lfs diff=lfs merge=lfs -text
+grid_00007500.png filter=lfs diff=lfs merge=lfs -text
+grid_00008000.png filter=lfs diff=lfs merge=lfs -text
+grid_00008500.png filter=lfs diff=lfs merge=lfs -text
+grid_00009000.png filter=lfs diff=lfs merge=lfs -text
+grid_00009500.png filter=lfs diff=lfs merge=lfs -text
+grid_00010000.png filter=lfs diff=lfs merge=lfs -text
+grid_00010500.png filter=lfs diff=lfs merge=lfs -text
+grid_00011000.png filter=lfs diff=lfs merge=lfs -text
+grid_00011500.png filter=lfs diff=lfs merge=lfs -text
+grid_00012000.png filter=lfs diff=lfs merge=lfs -text
+grid_00012500.png filter=lfs diff=lfs merge=lfs -text
+grid_00013000.png filter=lfs diff=lfs merge=lfs -text
+grid_00013500.png filter=lfs diff=lfs merge=lfs -text
+grid_00014000.png filter=lfs diff=lfs merge=lfs -text
+grid_00014500.png filter=lfs diff=lfs merge=lfs -text
+grid_00015000.png filter=lfs diff=lfs merge=lfs -text
+grid_00015500.png filter=lfs diff=lfs merge=lfs -text
+grid_00016000.png filter=lfs diff=lfs merge=lfs -text
+grid_00016500.png filter=lfs diff=lfs merge=lfs -text
+grid_00017000.png filter=lfs diff=lfs merge=lfs -text
+grid_00017500.png filter=lfs diff=lfs merge=lfs -text
+grid_00018000.png filter=lfs diff=lfs merge=lfs -text
+grid_00018500.png filter=lfs diff=lfs merge=lfs -text
+grid_00019000.png filter=lfs diff=lfs merge=lfs -text
+grid_00019500.png filter=lfs diff=lfs merge=lfs -text
+grid_00020000.png filter=lfs diff=lfs merge=lfs -text
+grid_00020500.png filter=lfs diff=lfs merge=lfs -text
+grid_00021000.png filter=lfs diff=lfs merge=lfs -text
+grid_00021500.png filter=lfs diff=lfs merge=lfs -text
+grid_00022000.png filter=lfs diff=lfs merge=lfs -text
+grid_00022500.png filter=lfs diff=lfs merge=lfs -text
+grid_00023000.png filter=lfs diff=lfs merge=lfs -text
+grid_00023500.png filter=lfs diff=lfs merge=lfs -text
+grid_00024000.png filter=lfs diff=lfs merge=lfs -text
+grid_00024500.png filter=lfs diff=lfs merge=lfs -text
+grid_00025000.png filter=lfs diff=lfs merge=lfs -text
+grid_00025500.png filter=lfs diff=lfs merge=lfs -text
+grid_00026000.png filter=lfs diff=lfs merge=lfs -text
+grid_00026500.png filter=lfs diff=lfs merge=lfs -text
+grid_00027000.png filter=lfs diff=lfs merge=lfs -text
+grid_00027500.png filter=lfs diff=lfs merge=lfs -text
+grid_00028000.png filter=lfs diff=lfs merge=lfs -text
+grid_00028500.png filter=lfs diff=lfs merge=lfs -text
+grid_00029000.png filter=lfs diff=lfs merge=lfs -text
+grid_00029500.png filter=lfs diff=lfs merge=lfs -text
+grid_00030000.png filter=lfs diff=lfs merge=lfs -text
+grid_00030500.png filter=lfs diff=lfs merge=lfs -text
+grid_00031000.png filter=lfs diff=lfs merge=lfs -text
+grid_00031500.png filter=lfs diff=lfs merge=lfs -text
+grid_00032000.png filter=lfs diff=lfs merge=lfs -text
+grid_00032500.png filter=lfs diff=lfs merge=lfs -text
+grid_00033000.png filter=lfs diff=lfs merge=lfs -text
+grid_00033500.png filter=lfs diff=lfs merge=lfs -text
+grid_00034000.png filter=lfs diff=lfs merge=lfs -text
+grid_00034500.png filter=lfs diff=lfs merge=lfs -text
+grid_00035000.png filter=lfs diff=lfs merge=lfs -text
+grid_00035500.png filter=lfs diff=lfs merge=lfs -text
+grid_00036000.png filter=lfs diff=lfs merge=lfs -text
+grid_00036500.png filter=lfs diff=lfs merge=lfs -text
+grid_00037000.png filter=lfs diff=lfs merge=lfs -text
+grid_00037500.png filter=lfs diff=lfs merge=lfs -text
+grid_00038000.png filter=lfs diff=lfs merge=lfs -text
+grid_00038500.png filter=lfs diff=lfs merge=lfs -text
+grid_00039000.png filter=lfs diff=lfs merge=lfs -text
+grid_00039500.png filter=lfs diff=lfs merge=lfs -text
+grid_00040000.png filter=lfs diff=lfs merge=lfs -text
+grid_00040500.png filter=lfs diff=lfs merge=lfs -text
+grid_00041000.png filter=lfs diff=lfs merge=lfs -text
+grid_00041500.png filter=lfs diff=lfs merge=lfs -text
+grid_00042000.png filter=lfs diff=lfs merge=lfs -text
+grid_00042500.png filter=lfs diff=lfs merge=lfs -text
+grid_00043000.png filter=lfs diff=lfs merge=lfs -text
+grid_00043500.png filter=lfs diff=lfs merge=lfs -text
+grid_00044000.png filter=lfs diff=lfs merge=lfs -text
+grid_00044500.png filter=lfs diff=lfs merge=lfs -text
+grid_00045000.png filter=lfs diff=lfs merge=lfs -text
+grid_00045500.png filter=lfs diff=lfs merge=lfs -text
+grid_00046000.png filter=lfs diff=lfs merge=lfs -text
+grid_00046500.png filter=lfs diff=lfs merge=lfs -text
+grid_00047000.png filter=lfs diff=lfs merge=lfs -text
+grid_00047500.png filter=lfs diff=lfs merge=lfs -text
+grid_00048000.png filter=lfs diff=lfs merge=lfs -text
+grid_00048500.png filter=lfs diff=lfs merge=lfs -text
+grid_00049000.png filter=lfs diff=lfs merge=lfs -text
+grid_00049500.png filter=lfs diff=lfs merge=lfs -text
+grid_00050000.png filter=lfs diff=lfs merge=lfs -text
+grid_00050500.png filter=lfs diff=lfs merge=lfs -text
+grid_00051000.png filter=lfs diff=lfs merge=lfs -text
+grid_00051500.png filter=lfs diff=lfs merge=lfs -text
+grid_00052000.png filter=lfs diff=lfs merge=lfs -text
+grid_00052500.png filter=lfs diff=lfs merge=lfs -text
+grid_00053000.png filter=lfs diff=lfs merge=lfs -text
+grid_00053500.png filter=lfs diff=lfs merge=lfs -text
+grid_00054000.png filter=lfs diff=lfs merge=lfs -text
+grid_00054500.png filter=lfs diff=lfs merge=lfs -text
+grid_00055000.png filter=lfs diff=lfs merge=lfs -text
+grid_00055500.png filter=lfs diff=lfs merge=lfs -text
+grid_00056000.png filter=lfs diff=lfs merge=lfs -text
+grid_00056500.png filter=lfs diff=lfs merge=lfs -text
+grid_00057000.png filter=lfs diff=lfs merge=lfs -text
+grid_00057500.png filter=lfs diff=lfs merge=lfs -text
+grid_00057510.png filter=lfs diff=lfs merge=lfs -text
+samples/card_0.png filter=lfs diff=lfs merge=lfs -text
+samples/card_1.png filter=lfs diff=lfs merge=lfs -text
+samples/sample_cabin.png filter=lfs diff=lfs merge=lfs -text
+samples/sample_ramen.png filter=lfs diff=lfs merge=lfs -text
+tmp_sheet_s7.png filter=lfs diff=lfs merge=lfs -text
+tmp_sheet_s21.png filter=lfs diff=lfs merge=lfs -text
+tmp_sheet_simple.png filter=lfs diff=lfs merge=lfs -text
+tmp_simple_3.png filter=lfs diff=lfs merge=lfs -text
+tmp_simple_5.png filter=lfs diff=lfs merge=lfs -text
+samples/sample_strawberries.png filter=lfs diff=lfs merge=lfs -text
+tmp_sheet_fix.png filter=lfs diff=lfs merge=lfs -text
+tmp_fix_2.png filter=lfs diff=lfs merge=lfs -text
+samples/sample_sea.png filter=lfs diff=lfs merge=lfs -text
+tmp_snippet_test.png filter=lfs diff=lfs merge=lfs -text
+tmp_snipcheck.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,200 @@
+---
+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).

config.json

@@ -0,0 +1,26 @@
+{
+  "_class_name": "SurjoImagePreviewDiT",
+  "auto_map": {
+    "AutoModel": "modeling_surjo_image_preview.SurjoImagePreviewDiT"
+  },
+  "deepsup": true,
+  "dim": 512,
+  "freq_dim": 256,
+  "heads": 8,
+  "in_ch": 128,
+  "loop_iters": 2,
+  "mlp_ratio": 4.0,
+  "model_type": "surjo_image_preview",
+  "n_tokens": 256,
+  "repa_dim": 384,
+  "rope_theta": 10000.0,
+  "split": [
+    7,
+    4,
+    7
+  ],
+  "t_scale": 1000.0,
+  "text_dim": 1024,
+  "transformers_version": "5.13.1",
+  "xsa": true
+}

configuration_surjo_image_preview.py

@@ -0,0 +1,57 @@
+# Copyright 2026 SurjoLabs. Apache 2.0.
+"""Surjo-Image-Preview model configuration (looped DiT + QK-RMSNorm + 2D-RoPE + SwiGLU)."""
+
+from typing import List, Optional
+
+from transformers import PretrainedConfig
+
+
+class SurjoImagePreviewConfig(PretrainedConfig):
+    r"""Configuration for the Surjo-Image-Preview looped diffusion transformer."""
+
+    model_type = "surjo_image_preview"
+
+    def __init__(
+        self,
+        dim: int = 512,
+        heads: int = 8,
+        mlp_ratio: float = 4.0,
+        split: Optional[List[int]] = None,
+        loop_iters: int = 4,
+        xsa: bool = True,
+        deepsup: bool = True,
+        n_tokens: int = 256,
+        in_ch: int = 128,
+        text_dim: int = 1024,
+        repa_dim: int = 384,
+        freq_dim: int = 256,
+        t_scale: float = 1000.0,
+        rope_theta: float = 10000.0,
+        **kwargs,
+    ):
+        self.dim = dim
+        self.heads = heads
+        self.mlp_ratio = mlp_ratio
+        self.split = list(split) if split is not None else [7, 4, 7]
+        self.loop_iters = loop_iters
+        self.xsa = xsa
+        self.deepsup = deepsup
+        self.n_tokens = n_tokens
+        self.in_ch = in_ch
+        self.text_dim = text_dim
+        self.repa_dim = repa_dim
+        self.freq_dim = freq_dim
+        self.t_scale = t_scale
+        self.rope_theta = rope_theta
+
+        if dim % heads != 0:
+            raise ValueError(f"dim {dim} must be divisible by heads {heads}")
+        if len(self.split) != 3 or sum(self.split) < 3:
+            raise ValueError(f"bad split {self.split}")
+        import math
+
+        grid = int(math.sqrt(n_tokens))
+        if grid * grid != n_tokens:
+            raise ValueError(f"n_tokens {n_tokens} must be a perfect square")
+
+        super().__init__(**kwargs)

diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:586c3911e0b2a56798fc117356c612bd2e3cfc5560884dee37276be6ef5c9b70
+size 433126528

modeling_surjo_image_preview.py

@@ -0,0 +1,394 @@
+# Copyright 2026 SurjoLabs. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+try:
+    from diffusers import ConfigMixin, ModelMixin
+    from diffusers.configuration_utils import register_to_config
+
+    HAS_DIFFUSERS = True
+except ImportError:  # pragma: no cover
+    import functools
+    import inspect
+    import json
+    import os
+
+    HAS_DIFFUSERS = False
+
+    class ConfigMixin:
+        def register_to_config(self, **kwargs):
+            self._internal_dict = dict(kwargs)
+
+        @property
+        def config(self):
+            return self._internal_dict
+
+    def register_to_config(init):
+        sig = inspect.signature(init)
+
+        @functools.wraps(init)
+        def wrapper(self, *args, **kwargs):
+            bound = sig.bind(self, *args, **kwargs)
+            bound.apply_defaults()
+            out = init(self, *args, **kwargs)
+            self.register_to_config(**{k: v for k, v in bound.arguments.items() if k != "self"})
+            return out
+
+        return wrapper
+
+    class ModelMixin(nn.Module, ConfigMixin):
+        config_name = "config.json"
+        weights_name = "diffusion_pytorch_model.safetensors"
+
+        def save_pretrained(self, save_directory, **kwargs):
+            os.makedirs(save_directory, exist_ok=True)
+            cfg = dict(self.config)
+            cfg["_class_name"] = type(self).__name__
+            cfg["_diffusers_version"] = "none"
+            tmp = os.path.join(save_directory, self.config_name + ".tmp")
+            with open(tmp, "w", encoding="utf-8") as f:
+                json.dump(cfg, f, indent=2, sort_keys=True)
+            os.replace(tmp, os.path.join(save_directory, self.config_name))
+            from safetensors.torch import save_file
+
+            clean = {k: v.detach().to("cpu", copy=True).contiguous() for k, v in self.state_dict().items()}
+            tmp = os.path.join(save_directory, self.weights_name + ".tmp")
+            save_file(clean, tmp)
+            os.replace(tmp, os.path.join(save_directory, self.weights_name))
+
+        @classmethod
+        def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+            with open(os.path.join(pretrained_model_name_or_path, cls.config_name), encoding="utf-8") as f:
+                cfg = json.load(f)
+            cfg = {k: v for k, v in cfg.items() if not k.startswith("_") and k != "torch_dtype"}
+            model = cls(**cfg)
+            from safetensors.torch import load_file
+
+            sd = load_file(os.path.join(pretrained_model_name_or_path, cls.weights_name))
+            model.load_state_dict(sd, strict=True)
+            return model
+
+
+def timestep_embedding(t: torch.Tensor, dim: int, max_period: float = 10000.0) -> torch.Tensor:
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period) * torch.arange(half, dtype=torch.float32, device=t.device) / max(1, half)
+    )
+    args = t.float().unsqueeze(-1) * freqs.unsqueeze(0)
+    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        emb = F.pad(emb, (0, 1))
+    return emb
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden: int, freq_dim: int = 256, scale: float = 1000.0):
+        super().__init__()
+        self.freq_dim = freq_dim
+        self.scale = scale
+        self.mlp = nn.Sequential(nn.Linear(freq_dim, hidden), nn.SiLU(), nn.Linear(hidden, hidden))
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        return self.mlp(timestep_embedding(self.scale * t, self.freq_dim))
+
+
+FP_CLAMP = 65504.0  # bf16 max; keeps autocast casts finite
+
+
+class RMSNorm1P(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) * (1.0 + self.weight)
+
+
+def rope_2d_tables(grid: int, head_half: int, theta: float = 10000.0) -> tuple[torch.Tensor, ...]:
+    n_freq = head_half // 2  # one freq per rotate-half pair within each axis-half
+    freqs = 1.0 / (theta ** (torch.arange(0, head_half, 2, dtype=torch.float32) / head_half))
+    pos = torch.arange(grid, dtype=torch.float32)
+    ang = torch.outer(pos, freqs)
+    return torch.cos(ang), torch.sin(ang), torch.cos(ang), torch.sin(ang)
+
+
+def apply_rope_2d(
+    t: torch.Tensor,
+    cos_h: torch.Tensor,
+    sin_h: torch.Tensor,
+    cos_w: torch.Tensor,
+    sin_w: torch.Tensor,
+    grid: int,
+) -> torch.Tensor:
+    b, h, _, d = t.shape
+    half = d // 2
+    t = t.view(b, h, grid, grid, d)
+    ta, tb = t[..., :half], t[..., half:]
+
+    def _rot(x, cos, sin):
+        x1, x2 = x[..., 0::2], x[..., 1::2]
+        o = torch.empty_like(x)
+        o[..., 0::2] = x1 * cos - x2 * sin
+        o[..., 1::2] = x1 * sin + x2 * cos
+        return o
+
+    ta = _rot(ta, cos_h[:, None, :], sin_h[:, None, :])
+    tb = _rot(tb, cos_w[None, :, :], sin_w[None, :, :])
+    return torch.cat([ta, tb], dim=-1).reshape(b, h, grid * grid, d)
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, heads: int, xsa: bool):
+        super().__init__()
+        self.heads = heads
+        self.head_dim = dim // heads
+        self.xsa = xsa
+        self.qk_norm_q = RMSNorm1P(self.head_dim)
+        self.qk_norm_k = RMSNorm1P(self.head_dim)
+        self.qkv = nn.Linear(dim, dim * 3)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor, rope=None) -> torch.Tensor:
+        b, n, _ = x.shape
+        q, k, v = self.qkv(x).chunk(3, dim=-1)
+        q = q.view(b, n, self.heads, self.head_dim).transpose(1, 2)
+        k = k.view(b, n, self.heads, self.head_dim).transpose(1, 2)
+        v = v.view(b, n, self.heads, self.head_dim).transpose(1, 2)
+        q = self.qk_norm_q(q)
+        k = self.qk_norm_k(k)
+        if rope is not None:
+            cos_h, sin_h, cos_w, sin_w, grid = rope
+            dt = q.dtype
+            q = apply_rope_2d(q, cos_h.to(dt), sin_h.to(dt), cos_w.to(dt), sin_w.to(dt), grid)
+            k = apply_rope_2d(k, cos_h.to(dt), sin_h.to(dt), cos_w.to(dt), sin_w.to(dt), grid)
+        o = F.scaled_dot_product_attention(q, k, v)
+        if self.xsa:
+            # XSA (arXiv:2603.09078): strip each token's own value-component so attention carries context only
+            v_hat = F.normalize(v.float(), dim=-1).to(o.dtype)
+            o = o - v_hat * (v_hat.float() * o.float()).sum(dim=-1, keepdim=True).to(o.dtype)
+        o = o.transpose(1, 2).reshape(b, n, -1)
+        return self.proj(o)
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, dim: int, heads: int):
+        super().__init__()
+        self.heads = heads
+        self.head_dim = dim // heads
+        self.qk_norm_q = RMSNorm1P(self.head_dim)
+        self.qk_norm_k = RMSNorm1P(self.head_dim)
+        self.q = nn.Linear(dim, dim)
+        self.kv = nn.Linear(dim, dim * 2)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        ctx: torch.Tensor,
+        ctx_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        b, n, _ = x.shape
+        m = ctx.shape[1]
+        q = self.qk_norm_q(self.q(x).view(b, n, self.heads, self.head_dim).transpose(1, 2))
+        kv = self.kv(ctx).view(b, m, 2, self.heads, self.head_dim)
+        k = self.qk_norm_k(kv[:, :, 0].transpose(1, 2))
+        v = kv[:, :, 1].transpose(1, 2)
+        bias = None
+        if ctx_mask is not None:
+            bias = (1.0 - ctx_mask.to(dtype=q.dtype, device=q.device)) * (-1e4)
+            bias = bias[:, None, None, :]
+        o = F.scaled_dot_product_attention(q, k, v, attn_mask=bias)
+        o = o.transpose(1, 2).reshape(b, n, -1)
+        return self.proj(o)
+
+
+class DiTBlock(nn.Module):
+    def __init__(self, dim: int, heads: int, mlp_ratio: float, xsa: bool):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False)
+        self.attn = SelfAttention(dim, heads, xsa=xsa)
+        self.norm_c = nn.LayerNorm(dim, elementwise_affine=True)
+        self.cross = CrossAttention(dim, heads)
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False)
+        hidden = max(8, (int(dim * mlp_ratio * 2 / 3) // 8) * 8)  # SwiGLU width at GELU-4x param parity
+        self.mlp_gate = nn.Linear(dim, hidden)
+        self.mlp_up = nn.Linear(dim, hidden)
+        self.mlp_down = nn.Linear(hidden, dim)
+        self.ada = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
+        nn.init.zeros_(self.ada[-1].weight)
+        nn.init.zeros_(self.ada[-1].bias)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        temb: torch.Tensor,
+        ctx: torch.Tensor,
+        ctx_mask: torch.Tensor | None = None,
+        rope=None,
+    ) -> torch.Tensor:
+        shift1, scale1, gate1, shift2, scale2, gate2 = self.ada(temb).chunk(6, dim=-1)
+        s1 = scale1.unsqueeze(1)
+        g1 = gate1.unsqueeze(1)
+        s2 = scale2.unsqueeze(1)
+        g2 = gate2.unsqueeze(1)
+        h = self.norm1(x)
+        x = x + g1 * self.attn(h * (1.0 + s1) + shift1.unsqueeze(1), rope=rope)
+        x = x + self.cross(self.norm_c(x), ctx, ctx_mask)
+        h = self.norm2(x)
+        h = h * (1.0 + s2) + shift2.unsqueeze(1)
+        x = x + g2 * self.mlp_down(F.silu(self.mlp_gate(h)) * self.mlp_up(h))
+        return x.clamp(-FP_CLAMP, FP_CLAMP)
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, dim: int, out_ch: int):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
+        self.linear = nn.Linear(dim, out_ch)
+        self.ada = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 2))
+        nn.init.zeros_(self.ada[-1].weight)
+        nn.init.zeros_(self.ada[-1].bias)
+
+    def forward(self, x: torch.Tensor, temb: torch.Tensor) -> torch.Tensor:
+        shift, scale = self.ada(temb).chunk(2, dim=-1)
+        h = self.norm(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        return self.linear(h)
+
+
+class SurjoImagePreviewDiT(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(
+        self,
+        dim: int = 512,
+        heads: int = 8,
+        mlp_ratio: float = 4.0,
+        split: tuple[int, int, int] = (7, 4, 7),
+        loop_iters: int = 4,
+        xsa: bool = True,
+        deepsup: bool = True,
+        n_tokens: int = 256,
+        in_ch: int = 128,
+        text_dim: int = 1024,
+        repa_dim: int = 384,  # DINOv2-S patch dim
+        freq_dim: int = 256,
+        t_scale: float = 1000.0,  # t arrives in [0,1]
+        rope_theta: float = 10000.0,
+    ):
+        super().__init__()
+        if dim % heads != 0:
+            raise ValueError(f"dim {dim} must be divisible by heads {heads}")
+        if len(split) != 3 or sum(split) < 3:
+            raise ValueError(f"bad split {split}")
+        self.dim = dim
+        self.split = tuple(split)
+        self.loop_iters = loop_iters
+        self.deepsup = deepsup
+        self.n_tokens = n_tokens
+        self.in_ch = in_ch
+        self.grid = int(math.sqrt(n_tokens))
+        if self.grid * self.grid != n_tokens:
+            raise ValueError(f"n_tokens {n_tokens} must be a perfect square")
+
+        self.x_embed = nn.Linear(in_ch, dim)
+        cos_h, sin_h, cos_w, sin_w = rope_2d_tables(self.grid, (dim // heads) // 2, theta=rope_theta)
+        self.register_buffer("rope_cos_h", cos_h)
+        self.register_buffer("rope_sin_h", sin_h)
+        self.register_buffer("rope_cos_w", cos_w)
+        self.register_buffer("rope_sin_w", sin_w)
+        self.t_embed = TimestepEmbedder(dim, freq_dim=freq_dim, scale=t_scale)
+        t_hidden = max(8, (text_dim * 2 // 8) * 8)
+        self.ctx_in = RMSNorm1P(text_dim)
+        self.ctx_mlp = nn.Sequential(
+            nn.Linear(text_dim, t_hidden), nn.GELU(approximate="tanh"), nn.Linear(t_hidden, dim)
+        )
+        self.repa_proj = nn.Linear(dim, repa_dim)
+
+        self.pre_blocks = nn.ModuleList(DiTBlock(dim, heads, mlp_ratio, xsa=False) for _ in range(split[0]))
+        self.loop_blocks = nn.ModuleList(DiTBlock(dim, heads, mlp_ratio, xsa=xsa) for _ in range(split[1]))
+        self.post_blocks = nn.ModuleList(DiTBlock(dim, heads, mlp_ratio, xsa=False) for _ in range(split[2]))
+        self.final = FinalLayer(dim, in_ch)
+
+    @property
+    def n_effective_depth(self) -> int:
+        return self.split[0] + self.split[1] * self.loop_iters + self.split[2]
+
+    def forward(
+        self,
+        z: torch.Tensor,
+        t: torch.Tensor,
+        ctx: torch.Tensor,
+        *,
+        ctx_mask: torch.Tensor | None = None,
+        deepsup: bool | None = None,
+        loop_iters: int | None = None,
+        capture_mid: bool = False,
+    ) -> torch.Tensor | list[torch.Tensor] | tuple:
+        deep = self.deepsup if deepsup is None else deepsup
+        k = self.loop_iters if loop_iters is None else loop_iters
+        if k < 1:
+            raise ValueError(f"loop_iters must be >= 1, got {k}")
+
+        b = z.shape[0]
+        x = z.flatten(2).transpose(1, 2)  # raster order = DINO patch order (REPA cosine assumes it)
+        if x.shape[1] != self.n_tokens:
+            raise ValueError(f"expected {self.n_tokens} tokens, got {x.shape[1]}")
+        x = self.x_embed(x)
+        temb = self.t_embed(t)
+        ctx = self.ctx_mlp(self.ctx_in(ctx))
+        rope = (self.rope_cos_h, self.rope_sin_h, self.rope_cos_w, self.rope_sin_w, self.grid)
+
+        for blk in self.pre_blocks:
+            x = blk(x, temb, ctx, ctx_mask, rope=rope)
+        mid = x if capture_mid else None
+
+        exits: list[torch.Tensor] = []
+        for _ in range(k):
+            for blk in self.loop_blocks:
+                x = blk(x, temb, ctx, ctx_mask, rope=rope)
+            if deep:
+                exits.append(x)
+
+        def through_post_and_final(h: torch.Tensor) -> torch.Tensor:
+            for blk in self.post_blocks:
+                h = blk(h, temb, ctx, ctx_mask, rope=rope)
+            out = self.final(h, temb)
+            return out.transpose(1, 2).reshape(b, self.in_ch, self.grid, self.grid)
+
+        if deep:
+            outs = [through_post_and_final(h) for h in exits]
+            return (outs, mid) if capture_mid else outs
+        out = through_post_and_final(x)
+        return (out, mid) if capture_mid else out
+
+    def num_params(self) -> int:
+        return sum(p.numel() for p in self.parameters())
+
+
+_ROPE_BUFFERS = ("rope_cos_h", "rope_sin_h", "rope_cos_w", "rope_sin_w")
+
+
+def load_ema_state(ema_model: nn.Module, ema_state: dict) -> None:
+    missing, unexpected = ema_model.load_state_dict(ema_state, strict=False)
+    bad_missing = [k for k in missing if k not in _ROPE_BUFFERS]
+    if bad_missing or unexpected:
+        raise RuntimeError(f"EMA load mismatch: missing={bad_missing} unexpected={unexpected}")