Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* 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
+samples/grid.png filter=lfs diff=lfs merge=lfs -text

model_index.json

@@ -0,0 +1,20 @@
+{
+  "_class_name": "AuroraIMG2Pipeline",
+  "_diffusers_version": "0.42.1",
+  "text_encoder": [
+    "pipeline",
+    "AuroraLMTextEncoder"
+  ],
+  "tokenizer": [
+    "transformers",
+    "PreTrainedTokenizerFast"
+  ],
+  "transformer": [
+    "pipeline",
+    "AuroraDenoiser"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLFlux2"
+  ]
+}

pipeline.py

@@ -0,0 +1,595 @@
+"""AuroraIMG2-style text-to-image pipeline (load with trust_remote_code=True).
+
+Based on the AuroraIMG2-50M-Preview pipeline (Apache-2.0), extended with a second text encoder type:
+  * AuroraTextEncoder    -> the released from-scratch encoder (unchanged)
+  * AuroraLMTextEncoder  -> a frozen pretrained causal LM (SlayerLab/Slayer149 GPT, Apache-2.0) + a small bidirectional
+                            refiner; its LM weights are stored inside text_encoder/ so the repo is self-contained
+Generated by trenowanie_generatora_obrazkow/export.py.
+
+Components (all stored in the repo as separate sub-folders):
+  * tokenizer     -> from-scratch byte-level BPE (8192 tokens), loaded with transformers
+  * text_encoder  -> AuroraTextEncoder, ~12.7M params, trained from scratch (defined below)
+  * transformer   -> AuroraDenoiser, ~50M params flow-matching DiT, trained from scratch (defined below)
+  * vae           -> pretrained FLUX.2 small decoder (AutoencoderKLFlux2, Apache-2.0), frozen
+
+Sampling is plain Euler on t = linspace(1, 0, steps + 1) with classifier-free guidance, i.e. exactly
+what was used for the reported benchmark numbers.
+
+NOTE: diffusers' custom-pipeline loader expects exactly ONE DiffusionPipeline subclass in this file.
+"""
+import contextlib
+import math
+from typing import List, Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers import ConfigMixin, DiffusionPipeline, ModelMixin
+from diffusers.configuration_utils import register_to_config
+from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+
+# =============================================================================
+# Building blocks (identical to the training code)
+# =============================================================================
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps).to(x.dtype) * self.weight
+
+
+def rotate_half_pairs(x: torch.Tensor) -> torch.Tensor:
+    even = x[..., 0::2]
+    odd = x[..., 1::2]
+    return torch.stack((-odd, even), dim=-1).flatten(-2)
+
+
+def _angles(positions: torch.Tensor, dim: int) -> torch.Tensor:
+    inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    return torch.outer(positions.float(), inv_freq)
+
+
+def build_rope_1d(n: int, d: int):
+    ang = _angles(torch.arange(n), d)
+    return ang.cos().repeat_interleave(2, dim=-1)[None, None], ang.sin().repeat_interleave(2, dim=-1)[None, None]
+
+
+def build_rope_2d(side: int, d: int):
+    """First half of the head dim encodes the row, second half the column. Tokens are row-major."""
+    coords = torch.arange(side)
+    rows, cols = torch.meshgrid(coords, coords, indexing="ij")
+    ang = torch.cat([_angles(rows.reshape(-1), d // 2), _angles(cols.reshape(-1), d // 2)], dim=-1)
+    return ang.cos().repeat_interleave(2, dim=-1)[None, None], ang.sin().repeat_interleave(2, dim=-1)[None, None]
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    return x * cos.to(x.dtype) + rotate_half_pairs(x) * sin.to(x.dtype)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim: int, hidden: int):
+        super().__init__()
+        self.in_proj = nn.Linear(dim, hidden * 2, bias=False)
+        self.out_proj = nn.Linear(hidden, dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        a, b = self.in_proj(x).chunk(2, dim=-1)
+        return self.out_proj(F.silu(a) * b)
+
+
+class TextAttention(nn.Module):
+    def __init__(self, dim: int, heads: int, kv_heads: int, head_dim: int):
+        super().__init__()
+        self.heads, self.kv_heads, self.head_dim = heads, kv_heads, head_dim
+        self.qkv = nn.Linear(dim, (heads + 2 * kv_heads) * head_dim, bias=False)
+        self.out = nn.Linear(heads * head_dim, dim, bias=False)
+
+    def forward(self, x, valid_mask, cos, sin):
+        b, t, _ = x.shape
+        q, k, v = self.qkv(x).split(
+            [self.heads * self.head_dim, self.kv_heads * self.head_dim, self.kv_heads * self.head_dim], dim=-1)
+        q = q.view(b, t, self.heads, self.head_dim).transpose(1, 2)
+        k = k.view(b, t, self.kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(b, t, self.kv_heads, self.head_dim).transpose(1, 2)
+        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+        rep = self.heads // self.kv_heads
+        k, v = k.repeat_interleave(rep, dim=1), v.repeat_interleave(rep, dim=1)
+        h = F.scaled_dot_product_attention(q, k, v, attn_mask=valid_mask[:, None, None, :].bool())
+        return self.out(h.transpose(1, 2).reshape(b, t, -1))
+
+
+class TextBlock(nn.Module):
+    def __init__(self, dim: int, heads: int, kv_heads: int, head_dim: int, mlp_hidden: int):
+        super().__init__()
+        self.norm1 = RMSNorm(dim)
+        self.attn = TextAttention(dim, heads, kv_heads, head_dim)
+        self.norm2 = RMSNorm(dim)
+        self.mlp = SwiGLU(dim, mlp_hidden)
+
+    def forward(self, x, valid_mask, cos, sin):
+        x = x + self.attn(self.norm1(x), valid_mask, cos, sin)
+        return x + self.mlp(self.norm2(x))
+
+
+class SelfAttention(nn.Module):
+    """Global self-attention over the 16x16 latent tokens with 2D RoPE and QK-norm."""
+
+    def __init__(self, dim: int, heads: int):
+        super().__init__()
+        self.heads = heads
+        self.head_dim = dim // heads
+        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.out = nn.Linear(dim, dim, bias=False)
+        self.q_norm = RMSNorm(self.head_dim)
+        self.k_norm = RMSNorm(self.head_dim)
+
+    def forward(self, x, cos, sin):
+        b, t, c = x.shape
+        q, k, v = self.qkv(x).view(b, t, 3, self.heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0)
+        q, k = apply_rope(self.q_norm(q), cos, sin), apply_rope(self.k_norm(k), cos, sin)
+        h = F.scaled_dot_product_attention(q, k, v)
+        return self.out(h.transpose(1, 2).reshape(b, t, c))
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, dim: int, heads: int, text_dim: int):
+        super().__init__()
+        self.heads = heads
+        self.head_dim = dim // heads
+        self.q = nn.Linear(dim, dim, bias=False)
+        self.kv = nn.Linear(text_dim, 2 * dim, bias=False)
+        self.out = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x, text, text_valid):
+        b, t, c = x.shape
+        n = text.shape[1]
+        q = self.q(x).view(b, t, self.heads, self.head_dim).transpose(1, 2)
+        k, v = self.kv(text).chunk(2, dim=-1)
+        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)
+        h = F.scaled_dot_product_attention(q, k, v, attn_mask=text_valid[:, None, None, :].bool())
+        return self.out(h.transpose(1, 2).reshape(b, t, c))
+
+
+class DenoiserBlock(nn.Module):
+    """DiT block with PixArt-style adaLN-single: a shared time->(6*dim) projection plus a per-block table."""
+
+    def __init__(self, dim: int, heads: int, text_dim: int, mlp_hidden: int):
+        super().__init__()
+        self.norm1 = RMSNorm(dim)
+        self.attn = SelfAttention(dim, heads)
+        self.norm2 = RMSNorm(dim)
+        self.cross = CrossAttention(dim, heads, text_dim)
+        self.norm3 = RMSNorm(dim)
+        self.mlp = SwiGLU(dim, mlp_hidden)
+        self.mod_table = nn.Parameter(torch.zeros(6, dim))
+
+    def forward(self, x, t6, text, text_valid, cos, sin):
+        shift1, scale1, gate1, shift2, scale2, gate2 = (t6 + self.mod_table[None]).to(x.dtype).unbind(1)
+        h = self.norm1(x) * (1 + scale1[:, None]) + shift1[:, None]
+        x = x + gate1[:, None] * self.attn(h, cos, sin)
+        x = x + self.cross(self.norm2(x), text, text_valid)
+        h = self.norm3(x) * (1 + scale2[:, None]) + shift2[:, None]
+        return x + gate2[:, None] * self.mlp(h)
+
+
+def timestep_embedding(t: torch.Tensor, dim: int = 256) -> torch.Tensor:
+    t = t.float() * 1000.0
+    half = dim // 2
+    freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device).float() / max(half - 1, 1))
+    angles = t[:, None] * freqs[None, :]
+    return torch.cat([angles.sin(), angles.cos()], dim=-1)
+
+
+# =============================================================================
+# Frozen-LM text encoder (GoLLeM / Slayer GPT, inference architecture, Apache-2.0)
+# =============================================================================
+class LMRMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(d))
+        self.eps = eps
+
+    def forward(self, x):
+        return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).to(x.dtype) * self.weight
+
+
+def lm_apply_rope(x, base=100000.0):
+    """Parameter-free RoPE on [B, H, T, D], interleaved convention (even/odd channel pairs)."""
+    _, _, T, dim = x.shape
+    pos = torch.arange(T, device=x.device, dtype=torch.float32)
+    freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=x.device, dtype=torch.float32) / dim))
+    ang = torch.outer(pos, freq)
+    cos, sin = ang.cos().to(x.dtype)[None, None], ang.sin().to(x.dtype)[None, None]
+    even, odd = x[..., ::2], x[..., 1::2]
+    return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
+
+
+class LMSwiGLU(nn.Module):
+    def __init__(self, d, hidden):
+        super().__init__()
+        self.gate = nn.Linear(d, hidden, bias=False)
+        self.up = nn.Linear(d, hidden, bias=False)
+        self.down = nn.Linear(hidden, d, bias=False)
+
+    def forward(self, x):
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class LMBlock(nn.Module):
+    def __init__(self, d, nh, ffn_mult, eps, rope_theta, value_residual, qk_norm, is_first=False):
+        super().__init__()
+        self.ln1, self.ln2 = LMRMSNorm(d, eps), LMRMSNorm(d, eps)
+        self.qkv = nn.Linear(d, 3 * d)
+        self.proj = nn.Linear(d, d)
+        self.mlp = LMSwiGLU(d, int(round(ffn_mult * d)))
+        self.nh, self.d, self.rope_theta = nh, d, rope_theta
+        self.value_residual, self.qk_norm, self.is_first = value_residual, qk_norm, is_first
+        if value_residual and not is_first:
+            self.vr_lambda = nn.Parameter(torch.zeros(1))
+        if qk_norm:
+            self.q_norm, self.k_norm = LMRMSNorm(d // nh, eps), LMRMSNorm(d // nh, eps)
+
+    def forward(self, x, v0=None):
+        B, T, D = x.size()
+        q, k, v = self.qkv(self.ln1(x)).split(self.d, dim=2)
+        hd = D // self.nh
+        q, k, v = (t.view(B, T, self.nh, hd).transpose(1, 2) for t in (q, k, v))
+        if self.qk_norm:
+            q, k = self.q_norm(q), self.k_norm(k)
+        q, k = lm_apply_rope(q, self.rope_theta), lm_apply_rope(k, self.rope_theta)
+        if self.value_residual:
+            if self.is_first:
+                v0 = v
+            else:
+                v = v + self.vr_lambda * v0
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        x = x + self.proj(y.transpose(1, 2).contiguous().view(B, T, D))
+        return x + self.mlp(self.ln2(x)), v0
+
+
+class LMTrunk(nn.Module):
+    """Token embedding + the first n blocks of the LM (no final norm, no head)."""
+
+    def __init__(self, vocab, n_blocks, d, nh, ffn_mult, eps, rope_theta, value_residual, qk_norm):
+        super().__init__()
+        self.tok = nn.Embedding(vocab, d)
+        self.blocks = nn.ModuleList([LMBlock(d, nh, ffn_mult, eps, rope_theta, value_residual, qk_norm, i == 0)
+                                     for i in range(n_blocks)])
+
+    def forward(self, idx):
+        x, v0 = self.tok(idx), None
+        for b in self.blocks:
+            x, v0 = b(x, v0)
+        return x
+
+
+# =============================================================================
+# FLUX.2 latent packing layouts (used to turn the model's (128,16,16) latents into the decoder's input)
+# =============================================================================
+def _unpatchify_flux2(z):  # official FLUX.2 pipeline layout: channel = c*4 + dh*2 + dw
+    b, c4, h, w = z.shape
+    return z.reshape(b, c4 // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c4 // 4, h * 2, w * 2)
+
+
+def _unpatchify_swap_subpixel(z):
+    b, c4, h, w = z.shape
+    return z.reshape(b, c4 // 4, 2, 2, h, w).permute(0, 1, 4, 3, 5, 2).reshape(b, c4 // 4, h * 2, w * 2)
+
+
+def _unpatchify_subpixel_major(z):
+    b, c4, h, w = z.shape
+    return z.reshape(b, 2, 2, c4 // 4, h, w).permute(0, 3, 4, 1, 5, 2).reshape(b, c4 // 4, h * 2, w * 2)
+
+
+def _unpatchify_subpixel_major_swapped(z):
+    b, c4, h, w = z.shape
+    return z.reshape(b, 2, 2, c4 // 4, h, w).permute(0, 3, 4, 2, 5, 1).reshape(b, c4 // 4, h * 2, w * 2)
+
+
+UNPATCHIFY_LAYOUTS = {
+    "flux2_pipeline": _unpatchify_flux2,
+    "swap_subpixel": _unpatchify_swap_subpixel,
+    "subpixel_major": _unpatchify_subpixel_major,
+    "subpixel_major_swapped": _unpatchify_subpixel_major_swapped,
+}
+
+
+# =============================================================================
+# Diffusers-style models
+# =============================================================================
+class AuroraTextEncoder(ModelMixin, ConfigMixin):
+    """Small from-scratch text encoder: GQA transformer with 1D RoPE, SwiGLU and RMSNorm."""
+
+    @register_to_config
+    def __init__(
+        self,
+        vocab_size: int = 8192,
+        dim: int = 384,
+        depth: int = 6,
+        heads: int = 6,
+        kv_heads: int = 2,
+        head_dim: int = 64,
+        mlp_hidden: int = 1024,
+        max_tokens: int = 64,
+        pad_id: int = 0,
+        bos_id: int = 2,
+        eos_id: int = 3,
+        null_id: int = 4,
+        mask_id: int = 5,
+    ):
+        super().__init__()
+        self.token_embedding = nn.Embedding(vocab_size, dim)
+        self.blocks = nn.ModuleList([TextBlock(dim, heads, kv_heads, head_dim, mlp_hidden) for _ in range(depth)])
+        self.norm = RMSNorm(dim)
+        cos, sin = build_rope_1d(max_tokens, head_dim)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+    def forward(self, ids: torch.Tensor):
+        valid = ids.ne(self.config.pad_id)
+        h = self.token_embedding(ids)
+        for block in self.blocks:
+            h = block(h, valid, self.rope_cos, self.rope_sin)
+        return self.norm(h), valid
+
+
+class AuroraLMTextEncoder(ModelMixin, ConfigMixin):
+    """Frozen causal LM trunk -> RMSNorm (+ learned null prompt) -> bidirectional refiner blocks -> denoiser text tokens.
+    Captions: [BOS] ids[:max_tokens-1] [PAD]...; the CFG null prompt is [NULL] [PAD]... (ids outside the LM vocab)."""
+
+    @register_to_config
+    def __init__(
+        self,
+        kind: str = "lm",
+        vocab_size: int = 12288,
+        lm_blocks: int = 18,  # LM blocks used (the trained run takes the stream before the last of 19)
+        lm_dim: int = 768,
+        lm_heads: int = 12,
+        lm_ffn_mult: float = 2.8125,
+        lm_norm_eps: float = 1e-6,
+        lm_rope_theta: float = 100000.0,
+        lm_value_residual: bool = True,
+        lm_qk_norm: bool = True,
+        dim: int = 768,
+        depth: int = 2,
+        heads: int = 12,
+        kv_heads: int = 4,
+        head_dim: int = 64,
+        mlp_hidden: int = 2048,
+        max_tokens: int = 96,
+        bos_id: int = 12285,
+        pad_id: int = 65535,
+        null_id: int = 65534,
+    ):
+        super().__init__()
+        self.lm = LMTrunk(vocab_size, lm_blocks, lm_dim, lm_heads, lm_ffn_mult, lm_norm_eps, lm_rope_theta,
+                          lm_value_residual, lm_qk_norm)
+        self.in_norm = RMSNorm(lm_dim)
+        self.null = nn.Parameter(torch.zeros(1, 1, lm_dim))
+        self.in_proj = nn.Linear(lm_dim, dim, bias=False) if dim != lm_dim else nn.Identity()
+        self.blocks = nn.ModuleList([TextBlock(dim, heads, kv_heads, head_dim, mlp_hidden) for _ in range(depth)])
+        self.norm = RMSNorm(dim)
+        cos, sin = build_rope_1d(max_tokens, head_dim)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+    def forward(self, ids: torch.Tensor):
+        cfg = self.config
+        valid = ids.ne(cfg.pad_id)
+        is_null = ids[:, 0].eq(cfg.null_id)
+        h = self.lm(ids.masked_fill(~valid | ids.eq(cfg.null_id), 0))  # causal + right padding: pads never seen
+        h = self.in_norm(h)
+        h = torch.where(is_null[:, None, None], self.null.to(h.dtype).expand_as(h), h)
+        h = self.in_proj(h)
+        for block in self.blocks:
+            h = block(h, valid, self.rope_cos, self.rope_sin)
+        return self.norm(h), valid
+
+
+class AuroraDenoiser(ModelMixin, ConfigMixin):
+    """Flow-matching DiT: velocity predictor on (B, 128, 16, 16) FLUX.2 latents, one latent pixel per token."""
+
+    @register_to_config
+    def __init__(
+        self,
+        latent_channels: int = 128,
+        latent_size: int = 16,
+        dim: int = 640,
+        depth: int = 8,
+        heads: int = 10,
+        mlp_hidden: int = 1536,
+        text_dim: int = 384,
+        latent_layout: str = "flux2_pipeline",
+    ):
+        super().__init__()
+        self.patch_in = nn.Linear(latent_channels, dim, bias=False)
+        self.time_mlp = nn.Sequential(nn.Linear(256, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.ada = nn.Linear(dim, 6 * dim)
+        self.blocks = nn.ModuleList([DenoiserBlock(dim, heads, text_dim, mlp_hidden) for _ in range(depth)])
+        self.norm = RMSNorm(dim)
+        self.patch_out = nn.Linear(dim, latent_channels, bias=False)
+        cos, sin = build_rope_2d(latent_size, dim // heads)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+    def forward(self, z_t: torch.Tensor, t: torch.Tensor, text: torch.Tensor, text_valid: torch.Tensor):
+        b, c, h, w = z_t.shape
+        dim = self.config.dim
+        x = self.patch_in(z_t.flatten(2).transpose(1, 2))
+        t_emb = self.time_mlp(timestep_embedding(t, 256))
+        t6 = self.ada(F.silu(t_emb)).view(b, 6, dim)
+        for block in self.blocks:
+            x = block(x, t6, text, text_valid, self.rope_cos, self.rope_sin)
+        return self.patch_out(self.norm(x)).transpose(1, 2).reshape(b, c, h, w)
+
+
+# =============================================================================
+# Pipeline
+# =============================================================================
+class AuroraIMG2Pipeline(DiffusionPipeline):
+    """Text-to-image pipeline (256x256). Usage:
+
+        pipe = DiffusionPipeline.from_pretrained(repo, custom_pipeline=repo, trust_remote_code=True)
+        image = pipe("a red fox in the snow", num_inference_steps=50, guidance_scale=3.0).images[0]
+    """
+
+    def __init__(self, tokenizer, text_encoder, transformer, vae):
+        super().__init__()
+        self.register_modules(tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae)
+
+    # ---- helpers -------------------------------------------------------------------------------
+    def _autocast(self):
+        device = self._execution_device
+        if device.type == "cuda":
+            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+            return torch.autocast(device_type="cuda", dtype=dtype)
+        return contextlib.nullcontext()
+
+    def _is_lm(self) -> bool:
+        return getattr(self.text_encoder.config, "kind", "") == "lm"
+
+    def _tokenize(self, texts: List[str]) -> torch.Tensor:
+        """[BOS] + BPE ids (max 62) + [EOS], padded with [PAD] to max_tokens. Same as in training."""
+        cfg = self.text_encoder.config
+        if self._is_lm():  # [BOS] ids[:max-1] [PAD]..., no leading space (prepare_data.py --encoder slayer149)
+            cleaned = [" ".join(str(t).split()) for t in texts]
+            encoded = self.tokenizer(cleaned, add_special_tokens=False)["input_ids"]
+            rows = torch.full((len(cleaned), int(cfg.max_tokens)), int(cfg.pad_id), dtype=torch.long)
+            for i, ids in enumerate(encoded):
+                ids = list(ids)[: int(cfg.max_tokens) - 1]
+                rows[i, 0] = int(cfg.bos_id)
+                if ids:
+                    rows[i, 1:1 + len(ids)] = torch.tensor(ids, dtype=torch.long)
+            return rows
+        max_len = int(cfg.max_tokens)
+        cleaned = [" ".join(str(t).split()) for t in texts]
+        # The training tokenizer uses ByteLevel(add_prefix_space=True), i.e. every caption was tokenised with a
+        # leading space. The transformers wrapper can drop that setting when the tokenizer is saved/reloaded, so the
+        # space is added explicitly here (a no-op if the backend already adds it).
+        cleaned = [(" " + t) if t else t for t in cleaned]
+        encoded = self.tokenizer(cleaned, add_special_tokens=False)["input_ids"]
+        rows = torch.full((len(cleaned), max_len), int(cfg.pad_id), dtype=torch.long)
+        for i, ids in enumerate(encoded):
+            ids = list(ids)[: max_len - 2]
+            rows[i, 0] = int(cfg.bos_id)
+            if ids:
+                rows[i, 1:1 + len(ids)] = torch.tensor(ids, dtype=torch.long)
+            rows[i, 1 + len(ids)] = int(cfg.eos_id)
+        return rows
+
+    def _null_ids(self, batch: int) -> torch.Tensor:
+        """The learned empty prompt used for classifier-free guidance: [BOS] [NULL] [EOS] [PAD]..."""
+        cfg = self.text_encoder.config
+        row = torch.full((int(cfg.max_tokens),), int(cfg.pad_id), dtype=torch.long)
+        if self._is_lm():
+            row[0] = int(cfg.null_id)
+            return row[None, :].repeat(batch, 1)
+        row[0], row[1], row[2] = int(cfg.bos_id), int(cfg.null_id), int(cfg.eos_id)
+        return row[None, :].repeat(batch, 1)
+
+    def encode_prompt(self, prompts: List[str], negative_prompts: Optional[List[str]], do_cfg: bool, device):
+        ids = self._tokenize(prompts).to(device)
+        with self._autocast():
+            cond, cond_valid = self.text_encoder(ids)
+            if not do_cfg:
+                return cond, cond_valid
+            if negative_prompts is None:
+                uncond_ids = self._null_ids(len(prompts))
+            else:
+                uncond_ids = self._tokenize(negative_prompts)
+            uncond, uncond_valid = self.text_encoder(uncond_ids.to(device))
+        # unconditional first, conditional second (the order the model was evaluated with)
+        return torch.cat([uncond, cond], dim=0), torch.cat([uncond_valid, cond_valid], dim=0)
+
+    @torch.no_grad()
+    def decode_latents(self, latents: torch.Tensor, chunk_size: int = 32) -> torch.Tensor:
+        """Model latents (B,128,16,16) -> images in [-1, 1] via the FLUX.2 small decoder."""
+        vae = self.vae
+        eps = float(getattr(vae.config, "batch_norm_eps", 1e-4))
+        mean = vae.bn.running_mean.detach().float().view(1, -1, 1, 1).to(latents.device)
+        std = torch.sqrt(vae.bn.running_var.detach().float() + eps).view(1, -1, 1, 1).to(latents.device)
+        unpatchify = UNPATCHIFY_LAYOUTS[self.transformer.config.latent_layout]
+        outs = []
+        for i in range(0, latents.shape[0], chunk_size):
+            z = latents[i:i + chunk_size].float() * std + mean
+            z = unpatchify(z)
+            outs.append(vae.decode(z.to(vae.dtype)).sample.float().clamp(-1, 1))
+        return torch.cat(outs, dim=0)
+
+    # ---- main entry point ----------------------------------------------------------------------
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 3.0,
+        num_images_per_prompt: int = 1,
+        generator: Optional[torch.Generator] = None,
+        latents: Optional[torch.Tensor] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+    ):
+        """
+        Args:
+            prompt: text prompt(s). English; prompts longer than 62 BPE tokens are truncated.
+            negative_prompt: optional. None (default) uses the learned empty prompt the model was trained
+                with for classifier-free guidance. Arbitrary negative prompts were not part of training.
+            num_inference_steps: Euler steps. 50 = recommended (used for the benchmarks), 20 = fast.
+            guidance_scale: classifier-free guidance scale. 3.0 recommended.
+            output_type: "pil", "np", "pt" (tensor in [0,1]) or "latent" (raw model latents).
+        """
+        base = [prompt] if isinstance(prompt, str) else list(prompt)
+        prompts = [p for p in base for _ in range(num_images_per_prompt)]
+        negs = None
+        if negative_prompt is not None:
+            neg_base = [negative_prompt] * len(base) if isinstance(negative_prompt, str) else list(negative_prompt)
+            if len(neg_base) != len(base):
+                raise ValueError("negative_prompt must be a string or have the same length as prompt.")
+            negs = [n for n in neg_base for _ in range(num_images_per_prompt)]
+
+        device = self._execution_device
+        b = len(prompts)
+        do_cfg = guidance_scale != 1.0
+        context, context_valid = self.encode_prompt(prompts, negs, do_cfg, device)
+
+        tcfg = self.transformer.config
+        shape = (b, tcfg.latent_channels, tcfg.latent_size, tcfg.latent_size)
+        if latents is None:
+            z = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
+        else:
+            z = latents.to(device=device, dtype=torch.float32)
+
+        ts = torch.linspace(1.0, 0.0, num_inference_steps + 1, device=device)
+        for i in self.progress_bar(range(num_inference_steps)):
+            t0, t1 = ts[i], ts[i + 1]
+            with self._autocast():
+                if do_cfg:
+                    v = self.transformer(torch.cat([z, z], dim=0), t0.expand(2 * b), context, context_valid)
+                    v_uncond, v_cond = v.chunk(2, dim=0)
+                    velocity = v_uncond + guidance_scale * (v_cond - v_uncond)
+                else:
+                    velocity = self.transformer(z, t0.expand(b), context, context_valid)
+            z = z + (t1 - t0) * velocity.float()
+
+        if output_type == "latent":
+            images = z
+        else:
+            img = (self.decode_latents(z) + 1.0) / 2.0
+            if output_type == "pt":
+                images = img
+            else:
+                arr = img.cpu().permute(0, 2, 3, 1).float().numpy()
+                images = self.numpy_to_pil(arr) if output_type == "pil" else arr
+
+        self.maybe_free_model_hooks()
+        if not return_dict:
+            return (images,)
+        return ImagePipelineOutput(images=images)

text_encoder/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

text_encoder/NOTICE

@@ -0,0 +1,30 @@
+Slayer149
+
+Weights and architecture derive from SlayerLab/gollem-v5-ckpts
+revision 963440ce6ab4ada7e95da4c1faa28ebb00c082d4, published under Apache-2.0.
+modeling_gollem.py extracts the inference classes unchanged from the pinned
+GoLLeM r6 training implementation. The continuation, export and release scripts
+were prepared for this 149M run.
+
+The likelihood routines in glint_metrics.py are adapted from
+Glint-Research/Glint-1.3/benchmark.py, revision
+c3f99e246aa2c64382f9668dc864533617482d0a, whose repository declares MIT.
+Copyright remains with the original contributors. MIT license text:
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

text_encoder/config.json

@@ -0,0 +1,24 @@
+{
+  "_class_name": "AuroraLMTextEncoder",
+  "_diffusers_version": "0.42.1",
+  "kind": "lm",
+  "vocab_size": 12288,
+  "lm_blocks": 18,
+  "lm_dim": 768,
+  "lm_heads": 12,
+  "lm_ffn_mult": 2.8125,
+  "lm_norm_eps": 1e-06,
+  "lm_rope_theta": 100000.0,
+  "lm_value_residual": true,
+  "lm_qk_norm": true,
+  "dim": 768,
+  "depth": 2,
+  "heads": 12,
+  "kv_heads": 4,
+  "head_dim": 64,
+  "mlp_hidden": 2048,
+  "max_tokens": 96,
+  "bos_id": 12285,
+  "pad_id": 65535,
+  "null_id": 65534
+}

text_encoder/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6ebdabb5c5cb8a2d8244f3d6ef254f26506df358110216d526bee019d02e2b67
+size 616652324

text_encoder/pipeline.py

@@ -0,0 +1,595 @@
+"""AuroraIMG2-style text-to-image pipeline (load with trust_remote_code=True).
+
+Based on the AuroraIMG2-50M-Preview pipeline (Apache-2.0), extended with a second text encoder type:
+  * AuroraTextEncoder    -> the released from-scratch encoder (unchanged)
+  * AuroraLMTextEncoder  -> a frozen pretrained causal LM (SlayerLab/Slayer149 GPT, Apache-2.0) + a small bidirectional
+                            refiner; its LM weights are stored inside text_encoder/ so the repo is self-contained
+Generated by trenowanie_generatora_obrazkow/export.py.
+
+Components (all stored in the repo as separate sub-folders):
+  * tokenizer     -> from-scratch byte-level BPE (8192 tokens), loaded with transformers
+  * text_encoder  -> AuroraTextEncoder, ~12.7M params, trained from scratch (defined below)
+  * transformer   -> AuroraDenoiser, ~50M params flow-matching DiT, trained from scratch (defined below)
+  * vae           -> pretrained FLUX.2 small decoder (AutoencoderKLFlux2, Apache-2.0), frozen
+
+Sampling is plain Euler on t = linspace(1, 0, steps + 1) with classifier-free guidance, i.e. exactly
+what was used for the reported benchmark numbers.
+
+NOTE: diffusers' custom-pipeline loader expects exactly ONE DiffusionPipeline subclass in this file.
+"""
+import contextlib
+import math
+from typing import List, Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers import ConfigMixin, DiffusionPipeline, ModelMixin
+from diffusers.configuration_utils import register_to_config
+from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+
+# =============================================================================
+# Building blocks (identical to the training code)
+# =============================================================================
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps).to(x.dtype) * self.weight
+
+
+def rotate_half_pairs(x: torch.Tensor) -> torch.Tensor:
+    even = x[..., 0::2]
+    odd = x[..., 1::2]
+    return torch.stack((-odd, even), dim=-1).flatten(-2)
+
+
+def _angles(positions: torch.Tensor, dim: int) -> torch.Tensor:
+    inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    return torch.outer(positions.float(), inv_freq)
+
+
+def build_rope_1d(n: int, d: int):
+    ang = _angles(torch.arange(n), d)
+    return ang.cos().repeat_interleave(2, dim=-1)[None, None], ang.sin().repeat_interleave(2, dim=-1)[None, None]
+
+
+def build_rope_2d(side: int, d: int):
+    """First half of the head dim encodes the row, second half the column. Tokens are row-major."""
+    coords = torch.arange(side)
+    rows, cols = torch.meshgrid(coords, coords, indexing="ij")
+    ang = torch.cat([_angles(rows.reshape(-1), d // 2), _angles(cols.reshape(-1), d // 2)], dim=-1)
+    return ang.cos().repeat_interleave(2, dim=-1)[None, None], ang.sin().repeat_interleave(2, dim=-1)[None, None]
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    return x * cos.to(x.dtype) + rotate_half_pairs(x) * sin.to(x.dtype)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim: int, hidden: int):
+        super().__init__()
+        self.in_proj = nn.Linear(dim, hidden * 2, bias=False)
+        self.out_proj = nn.Linear(hidden, dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        a, b = self.in_proj(x).chunk(2, dim=-1)
+        return self.out_proj(F.silu(a) * b)
+
+
+class TextAttention(nn.Module):
+    def __init__(self, dim: int, heads: int, kv_heads: int, head_dim: int):
+        super().__init__()
+        self.heads, self.kv_heads, self.head_dim = heads, kv_heads, head_dim
+        self.qkv = nn.Linear(dim, (heads + 2 * kv_heads) * head_dim, bias=False)
+        self.out = nn.Linear(heads * head_dim, dim, bias=False)
+
+    def forward(self, x, valid_mask, cos, sin):
+        b, t, _ = x.shape
+        q, k, v = self.qkv(x).split(
+            [self.heads * self.head_dim, self.kv_heads * self.head_dim, self.kv_heads * self.head_dim], dim=-1)
+        q = q.view(b, t, self.heads, self.head_dim).transpose(1, 2)
+        k = k.view(b, t, self.kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(b, t, self.kv_heads, self.head_dim).transpose(1, 2)
+        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+        rep = self.heads // self.kv_heads
+        k, v = k.repeat_interleave(rep, dim=1), v.repeat_interleave(rep, dim=1)
+        h = F.scaled_dot_product_attention(q, k, v, attn_mask=valid_mask[:, None, None, :].bool())
+        return self.out(h.transpose(1, 2).reshape(b, t, -1))
+
+
+class TextBlock(nn.Module):
+    def __init__(self, dim: int, heads: int, kv_heads: int, head_dim: int, mlp_hidden: int):
+        super().__init__()
+        self.norm1 = RMSNorm(dim)
+        self.attn = TextAttention(dim, heads, kv_heads, head_dim)
+        self.norm2 = RMSNorm(dim)
+        self.mlp = SwiGLU(dim, mlp_hidden)
+
+    def forward(self, x, valid_mask, cos, sin):
+        x = x + self.attn(self.norm1(x), valid_mask, cos, sin)
+        return x + self.mlp(self.norm2(x))
+
+
+class SelfAttention(nn.Module):
+    """Global self-attention over the 16x16 latent tokens with 2D RoPE and QK-norm."""
+
+    def __init__(self, dim: int, heads: int):
+        super().__init__()
+        self.heads = heads
+        self.head_dim = dim // heads
+        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.out = nn.Linear(dim, dim, bias=False)
+        self.q_norm = RMSNorm(self.head_dim)
+        self.k_norm = RMSNorm(self.head_dim)
+
+    def forward(self, x, cos, sin):
+        b, t, c = x.shape
+        q, k, v = self.qkv(x).view(b, t, 3, self.heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0)
+        q, k = apply_rope(self.q_norm(q), cos, sin), apply_rope(self.k_norm(k), cos, sin)
+        h = F.scaled_dot_product_attention(q, k, v)
+        return self.out(h.transpose(1, 2).reshape(b, t, c))
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, dim: int, heads: int, text_dim: int):
+        super().__init__()
+        self.heads = heads
+        self.head_dim = dim // heads
+        self.q = nn.Linear(dim, dim, bias=False)
+        self.kv = nn.Linear(text_dim, 2 * dim, bias=False)
+        self.out = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x, text, text_valid):
+        b, t, c = x.shape
+        n = text.shape[1]
+        q = self.q(x).view(b, t, self.heads, self.head_dim).transpose(1, 2)
+        k, v = self.kv(text).chunk(2, dim=-1)
+        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)
+        h = F.scaled_dot_product_attention(q, k, v, attn_mask=text_valid[:, None, None, :].bool())
+        return self.out(h.transpose(1, 2).reshape(b, t, c))
+
+
+class DenoiserBlock(nn.Module):
+    """DiT block with PixArt-style adaLN-single: a shared time->(6*dim) projection plus a per-block table."""
+
+    def __init__(self, dim: int, heads: int, text_dim: int, mlp_hidden: int):
+        super().__init__()
+        self.norm1 = RMSNorm(dim)
+        self.attn = SelfAttention(dim, heads)
+        self.norm2 = RMSNorm(dim)
+        self.cross = CrossAttention(dim, heads, text_dim)
+        self.norm3 = RMSNorm(dim)
+        self.mlp = SwiGLU(dim, mlp_hidden)
+        self.mod_table = nn.Parameter(torch.zeros(6, dim))
+
+    def forward(self, x, t6, text, text_valid, cos, sin):
+        shift1, scale1, gate1, shift2, scale2, gate2 = (t6 + self.mod_table[None]).to(x.dtype).unbind(1)
+        h = self.norm1(x) * (1 + scale1[:, None]) + shift1[:, None]
+        x = x + gate1[:, None] * self.attn(h, cos, sin)
+        x = x + self.cross(self.norm2(x), text, text_valid)
+        h = self.norm3(x) * (1 + scale2[:, None]) + shift2[:, None]
+        return x + gate2[:, None] * self.mlp(h)
+
+
+def timestep_embedding(t: torch.Tensor, dim: int = 256) -> torch.Tensor:
+    t = t.float() * 1000.0
+    half = dim // 2
+    freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device).float() / max(half - 1, 1))
+    angles = t[:, None] * freqs[None, :]
+    return torch.cat([angles.sin(), angles.cos()], dim=-1)
+
+
+# =============================================================================
+# Frozen-LM text encoder (GoLLeM / Slayer GPT, inference architecture, Apache-2.0)
+# =============================================================================
+class LMRMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(d))
+        self.eps = eps
+
+    def forward(self, x):
+        return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).to(x.dtype) * self.weight
+
+
+def lm_apply_rope(x, base=100000.0):
+    """Parameter-free RoPE on [B, H, T, D], interleaved convention (even/odd channel pairs)."""
+    _, _, T, dim = x.shape
+    pos = torch.arange(T, device=x.device, dtype=torch.float32)
+    freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=x.device, dtype=torch.float32) / dim))
+    ang = torch.outer(pos, freq)
+    cos, sin = ang.cos().to(x.dtype)[None, None], ang.sin().to(x.dtype)[None, None]
+    even, odd = x[..., ::2], x[..., 1::2]
+    return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
+
+
+class LMSwiGLU(nn.Module):
+    def __init__(self, d, hidden):
+        super().__init__()
+        self.gate = nn.Linear(d, hidden, bias=False)
+        self.up = nn.Linear(d, hidden, bias=False)
+        self.down = nn.Linear(hidden, d, bias=False)
+
+    def forward(self, x):
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class LMBlock(nn.Module):
+    def __init__(self, d, nh, ffn_mult, eps, rope_theta, value_residual, qk_norm, is_first=False):
+        super().__init__()
+        self.ln1, self.ln2 = LMRMSNorm(d, eps), LMRMSNorm(d, eps)
+        self.qkv = nn.Linear(d, 3 * d)
+        self.proj = nn.Linear(d, d)
+        self.mlp = LMSwiGLU(d, int(round(ffn_mult * d)))
+        self.nh, self.d, self.rope_theta = nh, d, rope_theta
+        self.value_residual, self.qk_norm, self.is_first = value_residual, qk_norm, is_first
+        if value_residual and not is_first:
+            self.vr_lambda = nn.Parameter(torch.zeros(1))
+        if qk_norm:
+            self.q_norm, self.k_norm = LMRMSNorm(d // nh, eps), LMRMSNorm(d // nh, eps)
+
+    def forward(self, x, v0=None):
+        B, T, D = x.size()
+        q, k, v = self.qkv(self.ln1(x)).split(self.d, dim=2)
+        hd = D // self.nh
+        q, k, v = (t.view(B, T, self.nh, hd).transpose(1, 2) for t in (q, k, v))
+        if self.qk_norm:
+            q, k = self.q_norm(q), self.k_norm(k)
+        q, k = lm_apply_rope(q, self.rope_theta), lm_apply_rope(k, self.rope_theta)
+        if self.value_residual:
+            if self.is_first:
+                v0 = v
+            else:
+                v = v + self.vr_lambda * v0
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        x = x + self.proj(y.transpose(1, 2).contiguous().view(B, T, D))
+        return x + self.mlp(self.ln2(x)), v0
+
+
+class LMTrunk(nn.Module):
+    """Token embedding + the first n blocks of the LM (no final norm, no head)."""
+
+    def __init__(self, vocab, n_blocks, d, nh, ffn_mult, eps, rope_theta, value_residual, qk_norm):
+        super().__init__()
+        self.tok = nn.Embedding(vocab, d)
+        self.blocks = nn.ModuleList([LMBlock(d, nh, ffn_mult, eps, rope_theta, value_residual, qk_norm, i == 0)
+                                     for i in range(n_blocks)])
+
+    def forward(self, idx):
+        x, v0 = self.tok(idx), None
+        for b in self.blocks:
+            x, v0 = b(x, v0)
+        return x
+
+
+# =============================================================================
+# FLUX.2 latent packing layouts (used to turn the model's (128,16,16) latents into the decoder's input)
+# =============================================================================
+def _unpatchify_flux2(z):  # official FLUX.2 pipeline layout: channel = c*4 + dh*2 + dw
+    b, c4, h, w = z.shape
+    return z.reshape(b, c4 // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c4 // 4, h * 2, w * 2)
+
+
+def _unpatchify_swap_subpixel(z):
+    b, c4, h, w = z.shape
+    return z.reshape(b, c4 // 4, 2, 2, h, w).permute(0, 1, 4, 3, 5, 2).reshape(b, c4 // 4, h * 2, w * 2)
+
+
+def _unpatchify_subpixel_major(z):
+    b, c4, h, w = z.shape
+    return z.reshape(b, 2, 2, c4 // 4, h, w).permute(0, 3, 4, 1, 5, 2).reshape(b, c4 // 4, h * 2, w * 2)
+
+
+def _unpatchify_subpixel_major_swapped(z):
+    b, c4, h, w = z.shape
+    return z.reshape(b, 2, 2, c4 // 4, h, w).permute(0, 3, 4, 2, 5, 1).reshape(b, c4 // 4, h * 2, w * 2)
+
+
+UNPATCHIFY_LAYOUTS = {
+    "flux2_pipeline": _unpatchify_flux2,
+    "swap_subpixel": _unpatchify_swap_subpixel,
+    "subpixel_major": _unpatchify_subpixel_major,
+    "subpixel_major_swapped": _unpatchify_subpixel_major_swapped,
+}
+
+
+# =============================================================================
+# Diffusers-style models
+# =============================================================================
+class AuroraTextEncoder(ModelMixin, ConfigMixin):
+    """Small from-scratch text encoder: GQA transformer with 1D RoPE, SwiGLU and RMSNorm."""
+
+    @register_to_config
+    def __init__(
+        self,
+        vocab_size: int = 8192,
+        dim: int = 384,
+        depth: int = 6,
+        heads: int = 6,
+        kv_heads: int = 2,
+        head_dim: int = 64,
+        mlp_hidden: int = 1024,
+        max_tokens: int = 64,
+        pad_id: int = 0,
+        bos_id: int = 2,
+        eos_id: int = 3,
+        null_id: int = 4,
+        mask_id: int = 5,
+    ):
+        super().__init__()
+        self.token_embedding = nn.Embedding(vocab_size, dim)
+        self.blocks = nn.ModuleList([TextBlock(dim, heads, kv_heads, head_dim, mlp_hidden) for _ in range(depth)])
+        self.norm = RMSNorm(dim)
+        cos, sin = build_rope_1d(max_tokens, head_dim)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+    def forward(self, ids: torch.Tensor):
+        valid = ids.ne(self.config.pad_id)
+        h = self.token_embedding(ids)
+        for block in self.blocks:
+            h = block(h, valid, self.rope_cos, self.rope_sin)
+        return self.norm(h), valid
+
+
+class AuroraLMTextEncoder(ModelMixin, ConfigMixin):
+    """Frozen causal LM trunk -> RMSNorm (+ learned null prompt) -> bidirectional refiner blocks -> denoiser text tokens.
+    Captions: [BOS] ids[:max_tokens-1] [PAD]...; the CFG null prompt is [NULL] [PAD]... (ids outside the LM vocab)."""
+
+    @register_to_config
+    def __init__(
+        self,
+        kind: str = "lm",
+        vocab_size: int = 12288,
+        lm_blocks: int = 18,  # LM blocks used (the trained run takes the stream before the last of 19)
+        lm_dim: int = 768,
+        lm_heads: int = 12,
+        lm_ffn_mult: float = 2.8125,
+        lm_norm_eps: float = 1e-6,
+        lm_rope_theta: float = 100000.0,
+        lm_value_residual: bool = True,
+        lm_qk_norm: bool = True,
+        dim: int = 768,
+        depth: int = 2,
+        heads: int = 12,
+        kv_heads: int = 4,
+        head_dim: int = 64,
+        mlp_hidden: int = 2048,
+        max_tokens: int = 96,
+        bos_id: int = 12285,
+        pad_id: int = 65535,
+        null_id: int = 65534,
+    ):
+        super().__init__()
+        self.lm = LMTrunk(vocab_size, lm_blocks, lm_dim, lm_heads, lm_ffn_mult, lm_norm_eps, lm_rope_theta,
+                          lm_value_residual, lm_qk_norm)
+        self.in_norm = RMSNorm(lm_dim)
+        self.null = nn.Parameter(torch.zeros(1, 1, lm_dim))
+        self.in_proj = nn.Linear(lm_dim, dim, bias=False) if dim != lm_dim else nn.Identity()
+        self.blocks = nn.ModuleList([TextBlock(dim, heads, kv_heads, head_dim, mlp_hidden) for _ in range(depth)])
+        self.norm = RMSNorm(dim)
+        cos, sin = build_rope_1d(max_tokens, head_dim)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+    def forward(self, ids: torch.Tensor):
+        cfg = self.config
+        valid = ids.ne(cfg.pad_id)
+        is_null = ids[:, 0].eq(cfg.null_id)
+        h = self.lm(ids.masked_fill(~valid | ids.eq(cfg.null_id), 0))  # causal + right padding: pads never seen
+        h = self.in_norm(h)
+        h = torch.where(is_null[:, None, None], self.null.to(h.dtype).expand_as(h), h)
+        h = self.in_proj(h)
+        for block in self.blocks:
+            h = block(h, valid, self.rope_cos, self.rope_sin)
+        return self.norm(h), valid
+
+
+class AuroraDenoiser(ModelMixin, ConfigMixin):
+    """Flow-matching DiT: velocity predictor on (B, 128, 16, 16) FLUX.2 latents, one latent pixel per token."""
+
+    @register_to_config
+    def __init__(
+        self,
+        latent_channels: int = 128,
+        latent_size: int = 16,
+        dim: int = 640,
+        depth: int = 8,
+        heads: int = 10,
+        mlp_hidden: int = 1536,
+        text_dim: int = 384,
+        latent_layout: str = "flux2_pipeline",
+    ):
+        super().__init__()
+        self.patch_in = nn.Linear(latent_channels, dim, bias=False)
+        self.time_mlp = nn.Sequential(nn.Linear(256, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.ada = nn.Linear(dim, 6 * dim)
+        self.blocks = nn.ModuleList([DenoiserBlock(dim, heads, text_dim, mlp_hidden) for _ in range(depth)])
+        self.norm = RMSNorm(dim)
+        self.patch_out = nn.Linear(dim, latent_channels, bias=False)
+        cos, sin = build_rope_2d(latent_size, dim // heads)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+    def forward(self, z_t: torch.Tensor, t: torch.Tensor, text: torch.Tensor, text_valid: torch.Tensor):
+        b, c, h, w = z_t.shape
+        dim = self.config.dim
+        x = self.patch_in(z_t.flatten(2).transpose(1, 2))
+        t_emb = self.time_mlp(timestep_embedding(t, 256))
+        t6 = self.ada(F.silu(t_emb)).view(b, 6, dim)
+        for block in self.blocks:
+            x = block(x, t6, text, text_valid, self.rope_cos, self.rope_sin)
+        return self.patch_out(self.norm(x)).transpose(1, 2).reshape(b, c, h, w)
+
+
+# =============================================================================
+# Pipeline
+# =============================================================================
+class AuroraIMG2Pipeline(DiffusionPipeline):
+    """Text-to-image pipeline (256x256). Usage:
+
+        pipe = DiffusionPipeline.from_pretrained(repo, custom_pipeline=repo, trust_remote_code=True)
+        image = pipe("a red fox in the snow", num_inference_steps=50, guidance_scale=3.0).images[0]
+    """
+
+    def __init__(self, tokenizer, text_encoder, transformer, vae):
+        super().__init__()
+        self.register_modules(tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae)
+
+    # ---- helpers -------------------------------------------------------------------------------
+    def _autocast(self):
+        device = self._execution_device
+        if device.type == "cuda":
+            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+            return torch.autocast(device_type="cuda", dtype=dtype)
+        return contextlib.nullcontext()
+
+    def _is_lm(self) -> bool:
+        return getattr(self.text_encoder.config, "kind", "") == "lm"
+
+    def _tokenize(self, texts: List[str]) -> torch.Tensor:
+        """[BOS] + BPE ids (max 62) + [EOS], padded with [PAD] to max_tokens. Same as in training."""
+        cfg = self.text_encoder.config
+        if self._is_lm():  # [BOS] ids[:max-1] [PAD]..., no leading space (prepare_data.py --encoder slayer149)
+            cleaned = [" ".join(str(t).split()) for t in texts]
+            encoded = self.tokenizer(cleaned, add_special_tokens=False)["input_ids"]
+            rows = torch.full((len(cleaned), int(cfg.max_tokens)), int(cfg.pad_id), dtype=torch.long)
+            for i, ids in enumerate(encoded):
+                ids = list(ids)[: int(cfg.max_tokens) - 1]
+                rows[i, 0] = int(cfg.bos_id)
+                if ids:
+                    rows[i, 1:1 + len(ids)] = torch.tensor(ids, dtype=torch.long)
+            return rows
+        max_len = int(cfg.max_tokens)
+        cleaned = [" ".join(str(t).split()) for t in texts]
+        # The training tokenizer uses ByteLevel(add_prefix_space=True), i.e. every caption was tokenised with a
+        # leading space. The transformers wrapper can drop that setting when the tokenizer is saved/reloaded, so the
+        # space is added explicitly here (a no-op if the backend already adds it).
+        cleaned = [(" " + t) if t else t for t in cleaned]
+        encoded = self.tokenizer(cleaned, add_special_tokens=False)["input_ids"]
+        rows = torch.full((len(cleaned), max_len), int(cfg.pad_id), dtype=torch.long)
+        for i, ids in enumerate(encoded):
+            ids = list(ids)[: max_len - 2]
+            rows[i, 0] = int(cfg.bos_id)
+            if ids:
+                rows[i, 1:1 + len(ids)] = torch.tensor(ids, dtype=torch.long)
+            rows[i, 1 + len(ids)] = int(cfg.eos_id)
+        return rows
+
+    def _null_ids(self, batch: int) -> torch.Tensor:
+        """The learned empty prompt used for classifier-free guidance: [BOS] [NULL] [EOS] [PAD]..."""
+        cfg = self.text_encoder.config
+        row = torch.full((int(cfg.max_tokens),), int(cfg.pad_id), dtype=torch.long)
+        if self._is_lm():
+            row[0] = int(cfg.null_id)
+            return row[None, :].repeat(batch, 1)
+        row[0], row[1], row[2] = int(cfg.bos_id), int(cfg.null_id), int(cfg.eos_id)
+        return row[None, :].repeat(batch, 1)
+
+    def encode_prompt(self, prompts: List[str], negative_prompts: Optional[List[str]], do_cfg: bool, device):
+        ids = self._tokenize(prompts).to(device)
+        with self._autocast():
+            cond, cond_valid = self.text_encoder(ids)
+            if not do_cfg:
+                return cond, cond_valid
+            if negative_prompts is None:
+                uncond_ids = self._null_ids(len(prompts))
+            else:
+                uncond_ids = self._tokenize(negative_prompts)
+            uncond, uncond_valid = self.text_encoder(uncond_ids.to(device))
+        # unconditional first, conditional second (the order the model was evaluated with)
+        return torch.cat([uncond, cond], dim=0), torch.cat([uncond_valid, cond_valid], dim=0)
+
+    @torch.no_grad()
+    def decode_latents(self, latents: torch.Tensor, chunk_size: int = 32) -> torch.Tensor:
+        """Model latents (B,128,16,16) -> images in [-1, 1] via the FLUX.2 small decoder."""
+        vae = self.vae
+        eps = float(getattr(vae.config, "batch_norm_eps", 1e-4))
+        mean = vae.bn.running_mean.detach().float().view(1, -1, 1, 1).to(latents.device)
+        std = torch.sqrt(vae.bn.running_var.detach().float() + eps).view(1, -1, 1, 1).to(latents.device)
+        unpatchify = UNPATCHIFY_LAYOUTS[self.transformer.config.latent_layout]
+        outs = []
+        for i in range(0, latents.shape[0], chunk_size):
+            z = latents[i:i + chunk_size].float() * std + mean
+            z = unpatchify(z)
+            outs.append(vae.decode(z.to(vae.dtype)).sample.float().clamp(-1, 1))
+        return torch.cat(outs, dim=0)
+
+    # ---- main entry point ----------------------------------------------------------------------
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 3.0,
+        num_images_per_prompt: int = 1,
+        generator: Optional[torch.Generator] = None,
+        latents: Optional[torch.Tensor] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+    ):
+        """
+        Args:
+            prompt: text prompt(s). English; prompts longer than 62 BPE tokens are truncated.
+            negative_prompt: optional. None (default) uses the learned empty prompt the model was trained
+                with for classifier-free guidance. Arbitrary negative prompts were not part of training.
+            num_inference_steps: Euler steps. 50 = recommended (used for the benchmarks), 20 = fast.
+            guidance_scale: classifier-free guidance scale. 3.0 recommended.
+            output_type: "pil", "np", "pt" (tensor in [0,1]) or "latent" (raw model latents).
+        """
+        base = [prompt] if isinstance(prompt, str) else list(prompt)
+        prompts = [p for p in base for _ in range(num_images_per_prompt)]
+        negs = None
+        if negative_prompt is not None:
+            neg_base = [negative_prompt] * len(base) if isinstance(negative_prompt, str) else list(negative_prompt)
+            if len(neg_base) != len(base):
+                raise ValueError("negative_prompt must be a string or have the same length as prompt.")
+            negs = [n for n in neg_base for _ in range(num_images_per_prompt)]
+
+        device = self._execution_device
+        b = len(prompts)
+        do_cfg = guidance_scale != 1.0
+        context, context_valid = self.encode_prompt(prompts, negs, do_cfg, device)
+
+        tcfg = self.transformer.config
+        shape = (b, tcfg.latent_channels, tcfg.latent_size, tcfg.latent_size)
+        if latents is None:
+            z = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
+        else:
+            z = latents.to(device=device, dtype=torch.float32)
+
+        ts = torch.linspace(1.0, 0.0, num_inference_steps + 1, device=device)
+        for i in self.progress_bar(range(num_inference_steps)):
+            t0, t1 = ts[i], ts[i + 1]
+            with self._autocast():
+                if do_cfg:
+                    v = self.transformer(torch.cat([z, z], dim=0), t0.expand(2 * b), context, context_valid)
+                    v_uncond, v_cond = v.chunk(2, dim=0)
+                    velocity = v_uncond + guidance_scale * (v_cond - v_uncond)
+                else:
+                    velocity = self.transformer(z, t0.expand(b), context, context_valid)
+            z = z + (t1 - t0) * velocity.float()
+
+        if output_type == "latent":
+            images = z
+        else:
+            img = (self.decode_latents(z) + 1.0) / 2.0
+            if output_type == "pt":
+                images = img
+            else:
+                arr = img.cpu().permute(0, 2, 3, 1).float().numpy()
+                images = self.numpy_to_pil(arr) if output_type == "pil" else arr
+
+        self.maybe_free_model_hooks()
+        if not return_dict:
+            return (images,)
+        return ImagePipelineOutput(images=images)

tokenizer/special_tokens_map.json

@@ -0,0 +1,5 @@
+{
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>"
+}

tokenizer/tokenizer_config.json

@@ -0,0 +1,7 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "model_max_length": 1024,
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>"
+}

transformer/config.json

@@ -0,0 +1,12 @@
+{
+  "_class_name": "AuroraDenoiser",
+  "_diffusers_version": "0.42.1",
+  "depth": 16,
+  "dim": 1024,
+  "heads": 16,
+  "latent_channels": 128,
+  "latent_layout": "flux2_pipeline",
+  "latent_size": 16,
+  "mlp_hidden": 2432,
+  "text_dim": 768
+}

transformer/diffusion_pytorch_model.safetensors

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

transformer/pipeline.py

@@ -0,0 +1,595 @@
+"""AuroraIMG2-style text-to-image pipeline (load with trust_remote_code=True).
+
+Based on the AuroraIMG2-50M-Preview pipeline (Apache-2.0), extended with a second text encoder type:
+  * AuroraTextEncoder    -> the released from-scratch encoder (unchanged)
+  * AuroraLMTextEncoder  -> a frozen pretrained causal LM (SlayerLab/Slayer149 GPT, Apache-2.0) + a small bidirectional
+                            refiner; its LM weights are stored inside text_encoder/ so the repo is self-contained
+Generated by trenowanie_generatora_obrazkow/export.py.
+
+Components (all stored in the repo as separate sub-folders):
+  * tokenizer     -> from-scratch byte-level BPE (8192 tokens), loaded with transformers
+  * text_encoder  -> AuroraTextEncoder, ~12.7M params, trained from scratch (defined below)
+  * transformer   -> AuroraDenoiser, ~50M params flow-matching DiT, trained from scratch (defined below)
+  * vae           -> pretrained FLUX.2 small decoder (AutoencoderKLFlux2, Apache-2.0), frozen
+
+Sampling is plain Euler on t = linspace(1, 0, steps + 1) with classifier-free guidance, i.e. exactly
+what was used for the reported benchmark numbers.
+
+NOTE: diffusers' custom-pipeline loader expects exactly ONE DiffusionPipeline subclass in this file.
+"""
+import contextlib
+import math
+from typing import List, Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers import ConfigMixin, DiffusionPipeline, ModelMixin
+from diffusers.configuration_utils import register_to_config
+from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+
+# =============================================================================
+# Building blocks (identical to the training code)
+# =============================================================================
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps).to(x.dtype) * self.weight
+
+
+def rotate_half_pairs(x: torch.Tensor) -> torch.Tensor:
+    even = x[..., 0::2]
+    odd = x[..., 1::2]
+    return torch.stack((-odd, even), dim=-1).flatten(-2)
+
+
+def _angles(positions: torch.Tensor, dim: int) -> torch.Tensor:
+    inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    return torch.outer(positions.float(), inv_freq)
+
+
+def build_rope_1d(n: int, d: int):
+    ang = _angles(torch.arange(n), d)
+    return ang.cos().repeat_interleave(2, dim=-1)[None, None], ang.sin().repeat_interleave(2, dim=-1)[None, None]
+
+
+def build_rope_2d(side: int, d: int):
+    """First half of the head dim encodes the row, second half the column. Tokens are row-major."""
+    coords = torch.arange(side)
+    rows, cols = torch.meshgrid(coords, coords, indexing="ij")
+    ang = torch.cat([_angles(rows.reshape(-1), d // 2), _angles(cols.reshape(-1), d // 2)], dim=-1)
+    return ang.cos().repeat_interleave(2, dim=-1)[None, None], ang.sin().repeat_interleave(2, dim=-1)[None, None]
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    return x * cos.to(x.dtype) + rotate_half_pairs(x) * sin.to(x.dtype)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim: int, hidden: int):
+        super().__init__()
+        self.in_proj = nn.Linear(dim, hidden * 2, bias=False)
+        self.out_proj = nn.Linear(hidden, dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        a, b = self.in_proj(x).chunk(2, dim=-1)
+        return self.out_proj(F.silu(a) * b)
+
+
+class TextAttention(nn.Module):
+    def __init__(self, dim: int, heads: int, kv_heads: int, head_dim: int):
+        super().__init__()
+        self.heads, self.kv_heads, self.head_dim = heads, kv_heads, head_dim
+        self.qkv = nn.Linear(dim, (heads + 2 * kv_heads) * head_dim, bias=False)
+        self.out = nn.Linear(heads * head_dim, dim, bias=False)
+
+    def forward(self, x, valid_mask, cos, sin):
+        b, t, _ = x.shape
+        q, k, v = self.qkv(x).split(
+            [self.heads * self.head_dim, self.kv_heads * self.head_dim, self.kv_heads * self.head_dim], dim=-1)
+        q = q.view(b, t, self.heads, self.head_dim).transpose(1, 2)
+        k = k.view(b, t, self.kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(b, t, self.kv_heads, self.head_dim).transpose(1, 2)
+        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+        rep = self.heads // self.kv_heads
+        k, v = k.repeat_interleave(rep, dim=1), v.repeat_interleave(rep, dim=1)
+        h = F.scaled_dot_product_attention(q, k, v, attn_mask=valid_mask[:, None, None, :].bool())
+        return self.out(h.transpose(1, 2).reshape(b, t, -1))
+
+
+class TextBlock(nn.Module):
+    def __init__(self, dim: int, heads: int, kv_heads: int, head_dim: int, mlp_hidden: int):
+        super().__init__()
+        self.norm1 = RMSNorm(dim)
+        self.attn = TextAttention(dim, heads, kv_heads, head_dim)
+        self.norm2 = RMSNorm(dim)
+        self.mlp = SwiGLU(dim, mlp_hidden)
+
+    def forward(self, x, valid_mask, cos, sin):
+        x = x + self.attn(self.norm1(x), valid_mask, cos, sin)
+        return x + self.mlp(self.norm2(x))
+
+
+class SelfAttention(nn.Module):
+    """Global self-attention over the 16x16 latent tokens with 2D RoPE and QK-norm."""
+
+    def __init__(self, dim: int, heads: int):
+        super().__init__()
+        self.heads = heads
+        self.head_dim = dim // heads
+        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.out = nn.Linear(dim, dim, bias=False)
+        self.q_norm = RMSNorm(self.head_dim)
+        self.k_norm = RMSNorm(self.head_dim)
+
+    def forward(self, x, cos, sin):
+        b, t, c = x.shape
+        q, k, v = self.qkv(x).view(b, t, 3, self.heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0)
+        q, k = apply_rope(self.q_norm(q), cos, sin), apply_rope(self.k_norm(k), cos, sin)
+        h = F.scaled_dot_product_attention(q, k, v)
+        return self.out(h.transpose(1, 2).reshape(b, t, c))
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, dim: int, heads: int, text_dim: int):
+        super().__init__()
+        self.heads = heads
+        self.head_dim = dim // heads
+        self.q = nn.Linear(dim, dim, bias=False)
+        self.kv = nn.Linear(text_dim, 2 * dim, bias=False)
+        self.out = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x, text, text_valid):
+        b, t, c = x.shape
+        n = text.shape[1]
+        q = self.q(x).view(b, t, self.heads, self.head_dim).transpose(1, 2)
+        k, v = self.kv(text).chunk(2, dim=-1)
+        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)
+        h = F.scaled_dot_product_attention(q, k, v, attn_mask=text_valid[:, None, None, :].bool())
+        return self.out(h.transpose(1, 2).reshape(b, t, c))
+
+
+class DenoiserBlock(nn.Module):
+    """DiT block with PixArt-style adaLN-single: a shared time->(6*dim) projection plus a per-block table."""
+
+    def __init__(self, dim: int, heads: int, text_dim: int, mlp_hidden: int):
+        super().__init__()
+        self.norm1 = RMSNorm(dim)
+        self.attn = SelfAttention(dim, heads)
+        self.norm2 = RMSNorm(dim)
+        self.cross = CrossAttention(dim, heads, text_dim)
+        self.norm3 = RMSNorm(dim)
+        self.mlp = SwiGLU(dim, mlp_hidden)
+        self.mod_table = nn.Parameter(torch.zeros(6, dim))
+
+    def forward(self, x, t6, text, text_valid, cos, sin):
+        shift1, scale1, gate1, shift2, scale2, gate2 = (t6 + self.mod_table[None]).to(x.dtype).unbind(1)
+        h = self.norm1(x) * (1 + scale1[:, None]) + shift1[:, None]
+        x = x + gate1[:, None] * self.attn(h, cos, sin)
+        x = x + self.cross(self.norm2(x), text, text_valid)
+        h = self.norm3(x) * (1 + scale2[:, None]) + shift2[:, None]
+        return x + gate2[:, None] * self.mlp(h)
+
+
+def timestep_embedding(t: torch.Tensor, dim: int = 256) -> torch.Tensor:
+    t = t.float() * 1000.0
+    half = dim // 2
+    freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device).float() / max(half - 1, 1))
+    angles = t[:, None] * freqs[None, :]
+    return torch.cat([angles.sin(), angles.cos()], dim=-1)
+
+
+# =============================================================================
+# Frozen-LM text encoder (GoLLeM / Slayer GPT, inference architecture, Apache-2.0)
+# =============================================================================
+class LMRMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(d))
+        self.eps = eps
+
+    def forward(self, x):
+        return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).to(x.dtype) * self.weight
+
+
+def lm_apply_rope(x, base=100000.0):
+    """Parameter-free RoPE on [B, H, T, D], interleaved convention (even/odd channel pairs)."""
+    _, _, T, dim = x.shape
+    pos = torch.arange(T, device=x.device, dtype=torch.float32)
+    freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=x.device, dtype=torch.float32) / dim))
+    ang = torch.outer(pos, freq)
+    cos, sin = ang.cos().to(x.dtype)[None, None], ang.sin().to(x.dtype)[None, None]
+    even, odd = x[..., ::2], x[..., 1::2]
+    return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
+
+
+class LMSwiGLU(nn.Module):
+    def __init__(self, d, hidden):
+        super().__init__()
+        self.gate = nn.Linear(d, hidden, bias=False)
+        self.up = nn.Linear(d, hidden, bias=False)
+        self.down = nn.Linear(hidden, d, bias=False)
+
+    def forward(self, x):
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class LMBlock(nn.Module):
+    def __init__(self, d, nh, ffn_mult, eps, rope_theta, value_residual, qk_norm, is_first=False):
+        super().__init__()
+        self.ln1, self.ln2 = LMRMSNorm(d, eps), LMRMSNorm(d, eps)
+        self.qkv = nn.Linear(d, 3 * d)
+        self.proj = nn.Linear(d, d)
+        self.mlp = LMSwiGLU(d, int(round(ffn_mult * d)))
+        self.nh, self.d, self.rope_theta = nh, d, rope_theta
+        self.value_residual, self.qk_norm, self.is_first = value_residual, qk_norm, is_first
+        if value_residual and not is_first:
+            self.vr_lambda = nn.Parameter(torch.zeros(1))
+        if qk_norm:
+            self.q_norm, self.k_norm = LMRMSNorm(d // nh, eps), LMRMSNorm(d // nh, eps)
+
+    def forward(self, x, v0=None):
+        B, T, D = x.size()
+        q, k, v = self.qkv(self.ln1(x)).split(self.d, dim=2)
+        hd = D // self.nh
+        q, k, v = (t.view(B, T, self.nh, hd).transpose(1, 2) for t in (q, k, v))
+        if self.qk_norm:
+            q, k = self.q_norm(q), self.k_norm(k)
+        q, k = lm_apply_rope(q, self.rope_theta), lm_apply_rope(k, self.rope_theta)
+        if self.value_residual:
+            if self.is_first:
+                v0 = v
+            else:
+                v = v + self.vr_lambda * v0
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        x = x + self.proj(y.transpose(1, 2).contiguous().view(B, T, D))
+        return x + self.mlp(self.ln2(x)), v0
+
+
+class LMTrunk(nn.Module):
+    """Token embedding + the first n blocks of the LM (no final norm, no head)."""
+
+    def __init__(self, vocab, n_blocks, d, nh, ffn_mult, eps, rope_theta, value_residual, qk_norm):
+        super().__init__()
+        self.tok = nn.Embedding(vocab, d)
+        self.blocks = nn.ModuleList([LMBlock(d, nh, ffn_mult, eps, rope_theta, value_residual, qk_norm, i == 0)
+                                     for i in range(n_blocks)])
+
+    def forward(self, idx):
+        x, v0 = self.tok(idx), None
+        for b in self.blocks:
+            x, v0 = b(x, v0)
+        return x
+
+
+# =============================================================================
+# FLUX.2 latent packing layouts (used to turn the model's (128,16,16) latents into the decoder's input)
+# =============================================================================
+def _unpatchify_flux2(z):  # official FLUX.2 pipeline layout: channel = c*4 + dh*2 + dw
+    b, c4, h, w = z.shape
+    return z.reshape(b, c4 // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c4 // 4, h * 2, w * 2)
+
+
+def _unpatchify_swap_subpixel(z):
+    b, c4, h, w = z.shape
+    return z.reshape(b, c4 // 4, 2, 2, h, w).permute(0, 1, 4, 3, 5, 2).reshape(b, c4 // 4, h * 2, w * 2)
+
+
+def _unpatchify_subpixel_major(z):
+    b, c4, h, w = z.shape
+    return z.reshape(b, 2, 2, c4 // 4, h, w).permute(0, 3, 4, 1, 5, 2).reshape(b, c4 // 4, h * 2, w * 2)
+
+
+def _unpatchify_subpixel_major_swapped(z):
+    b, c4, h, w = z.shape
+    return z.reshape(b, 2, 2, c4 // 4, h, w).permute(0, 3, 4, 2, 5, 1).reshape(b, c4 // 4, h * 2, w * 2)
+
+
+UNPATCHIFY_LAYOUTS = {
+    "flux2_pipeline": _unpatchify_flux2,
+    "swap_subpixel": _unpatchify_swap_subpixel,
+    "subpixel_major": _unpatchify_subpixel_major,
+    "subpixel_major_swapped": _unpatchify_subpixel_major_swapped,
+}
+
+
+# =============================================================================
+# Diffusers-style models
+# =============================================================================
+class AuroraTextEncoder(ModelMixin, ConfigMixin):
+    """Small from-scratch text encoder: GQA transformer with 1D RoPE, SwiGLU and RMSNorm."""
+
+    @register_to_config
+    def __init__(
+        self,
+        vocab_size: int = 8192,
+        dim: int = 384,
+        depth: int = 6,
+        heads: int = 6,
+        kv_heads: int = 2,
+        head_dim: int = 64,
+        mlp_hidden: int = 1024,
+        max_tokens: int = 64,
+        pad_id: int = 0,
+        bos_id: int = 2,
+        eos_id: int = 3,
+        null_id: int = 4,
+        mask_id: int = 5,
+    ):
+        super().__init__()
+        self.token_embedding = nn.Embedding(vocab_size, dim)
+        self.blocks = nn.ModuleList([TextBlock(dim, heads, kv_heads, head_dim, mlp_hidden) for _ in range(depth)])
+        self.norm = RMSNorm(dim)
+        cos, sin = build_rope_1d(max_tokens, head_dim)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+    def forward(self, ids: torch.Tensor):
+        valid = ids.ne(self.config.pad_id)
+        h = self.token_embedding(ids)
+        for block in self.blocks:
+            h = block(h, valid, self.rope_cos, self.rope_sin)
+        return self.norm(h), valid
+
+
+class AuroraLMTextEncoder(ModelMixin, ConfigMixin):
+    """Frozen causal LM trunk -> RMSNorm (+ learned null prompt) -> bidirectional refiner blocks -> denoiser text tokens.
+    Captions: [BOS] ids[:max_tokens-1] [PAD]...; the CFG null prompt is [NULL] [PAD]... (ids outside the LM vocab)."""
+
+    @register_to_config
+    def __init__(
+        self,
+        kind: str = "lm",
+        vocab_size: int = 12288,
+        lm_blocks: int = 18,  # LM blocks used (the trained run takes the stream before the last of 19)
+        lm_dim: int = 768,
+        lm_heads: int = 12,
+        lm_ffn_mult: float = 2.8125,
+        lm_norm_eps: float = 1e-6,
+        lm_rope_theta: float = 100000.0,
+        lm_value_residual: bool = True,
+        lm_qk_norm: bool = True,
+        dim: int = 768,
+        depth: int = 2,
+        heads: int = 12,
+        kv_heads: int = 4,
+        head_dim: int = 64,
+        mlp_hidden: int = 2048,
+        max_tokens: int = 96,
+        bos_id: int = 12285,
+        pad_id: int = 65535,
+        null_id: int = 65534,
+    ):
+        super().__init__()
+        self.lm = LMTrunk(vocab_size, lm_blocks, lm_dim, lm_heads, lm_ffn_mult, lm_norm_eps, lm_rope_theta,
+                          lm_value_residual, lm_qk_norm)
+        self.in_norm = RMSNorm(lm_dim)
+        self.null = nn.Parameter(torch.zeros(1, 1, lm_dim))
+        self.in_proj = nn.Linear(lm_dim, dim, bias=False) if dim != lm_dim else nn.Identity()
+        self.blocks = nn.ModuleList([TextBlock(dim, heads, kv_heads, head_dim, mlp_hidden) for _ in range(depth)])
+        self.norm = RMSNorm(dim)
+        cos, sin = build_rope_1d(max_tokens, head_dim)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+    def forward(self, ids: torch.Tensor):
+        cfg = self.config
+        valid = ids.ne(cfg.pad_id)
+        is_null = ids[:, 0].eq(cfg.null_id)
+        h = self.lm(ids.masked_fill(~valid | ids.eq(cfg.null_id), 0))  # causal + right padding: pads never seen
+        h = self.in_norm(h)
+        h = torch.where(is_null[:, None, None], self.null.to(h.dtype).expand_as(h), h)
+        h = self.in_proj(h)
+        for block in self.blocks:
+            h = block(h, valid, self.rope_cos, self.rope_sin)
+        return self.norm(h), valid
+
+
+class AuroraDenoiser(ModelMixin, ConfigMixin):
+    """Flow-matching DiT: velocity predictor on (B, 128, 16, 16) FLUX.2 latents, one latent pixel per token."""
+
+    @register_to_config
+    def __init__(
+        self,
+        latent_channels: int = 128,
+        latent_size: int = 16,
+        dim: int = 640,
+        depth: int = 8,
+        heads: int = 10,
+        mlp_hidden: int = 1536,
+        text_dim: int = 384,
+        latent_layout: str = "flux2_pipeline",
+    ):
+        super().__init__()
+        self.patch_in = nn.Linear(latent_channels, dim, bias=False)
+        self.time_mlp = nn.Sequential(nn.Linear(256, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.ada = nn.Linear(dim, 6 * dim)
+        self.blocks = nn.ModuleList([DenoiserBlock(dim, heads, text_dim, mlp_hidden) for _ in range(depth)])
+        self.norm = RMSNorm(dim)
+        self.patch_out = nn.Linear(dim, latent_channels, bias=False)
+        cos, sin = build_rope_2d(latent_size, dim // heads)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+    def forward(self, z_t: torch.Tensor, t: torch.Tensor, text: torch.Tensor, text_valid: torch.Tensor):
+        b, c, h, w = z_t.shape
+        dim = self.config.dim
+        x = self.patch_in(z_t.flatten(2).transpose(1, 2))
+        t_emb = self.time_mlp(timestep_embedding(t, 256))
+        t6 = self.ada(F.silu(t_emb)).view(b, 6, dim)
+        for block in self.blocks:
+            x = block(x, t6, text, text_valid, self.rope_cos, self.rope_sin)
+        return self.patch_out(self.norm(x)).transpose(1, 2).reshape(b, c, h, w)
+
+
+# =============================================================================
+# Pipeline
+# =============================================================================
+class AuroraIMG2Pipeline(DiffusionPipeline):
+    """Text-to-image pipeline (256x256). Usage:
+
+        pipe = DiffusionPipeline.from_pretrained(repo, custom_pipeline=repo, trust_remote_code=True)
+        image = pipe("a red fox in the snow", num_inference_steps=50, guidance_scale=3.0).images[0]
+    """
+
+    def __init__(self, tokenizer, text_encoder, transformer, vae):
+        super().__init__()
+        self.register_modules(tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae)
+
+    # ---- helpers -------------------------------------------------------------------------------
+    def _autocast(self):
+        device = self._execution_device
+        if device.type == "cuda":
+            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+            return torch.autocast(device_type="cuda", dtype=dtype)
+        return contextlib.nullcontext()
+
+    def _is_lm(self) -> bool:
+        return getattr(self.text_encoder.config, "kind", "") == "lm"
+
+    def _tokenize(self, texts: List[str]) -> torch.Tensor:
+        """[BOS] + BPE ids (max 62) + [EOS], padded with [PAD] to max_tokens. Same as in training."""
+        cfg = self.text_encoder.config
+        if self._is_lm():  # [BOS] ids[:max-1] [PAD]..., no leading space (prepare_data.py --encoder slayer149)
+            cleaned = [" ".join(str(t).split()) for t in texts]
+            encoded = self.tokenizer(cleaned, add_special_tokens=False)["input_ids"]
+            rows = torch.full((len(cleaned), int(cfg.max_tokens)), int(cfg.pad_id), dtype=torch.long)
+            for i, ids in enumerate(encoded):
+                ids = list(ids)[: int(cfg.max_tokens) - 1]
+                rows[i, 0] = int(cfg.bos_id)
+                if ids:
+                    rows[i, 1:1 + len(ids)] = torch.tensor(ids, dtype=torch.long)
+            return rows
+        max_len = int(cfg.max_tokens)
+        cleaned = [" ".join(str(t).split()) for t in texts]
+        # The training tokenizer uses ByteLevel(add_prefix_space=True), i.e. every caption was tokenised with a
+        # leading space. The transformers wrapper can drop that setting when the tokenizer is saved/reloaded, so the
+        # space is added explicitly here (a no-op if the backend already adds it).
+        cleaned = [(" " + t) if t else t for t in cleaned]
+        encoded = self.tokenizer(cleaned, add_special_tokens=False)["input_ids"]
+        rows = torch.full((len(cleaned), max_len), int(cfg.pad_id), dtype=torch.long)
+        for i, ids in enumerate(encoded):
+            ids = list(ids)[: max_len - 2]
+            rows[i, 0] = int(cfg.bos_id)
+            if ids:
+                rows[i, 1:1 + len(ids)] = torch.tensor(ids, dtype=torch.long)
+            rows[i, 1 + len(ids)] = int(cfg.eos_id)
+        return rows
+
+    def _null_ids(self, batch: int) -> torch.Tensor:
+        """The learned empty prompt used for classifier-free guidance: [BOS] [NULL] [EOS] [PAD]..."""
+        cfg = self.text_encoder.config
+        row = torch.full((int(cfg.max_tokens),), int(cfg.pad_id), dtype=torch.long)
+        if self._is_lm():
+            row[0] = int(cfg.null_id)
+            return row[None, :].repeat(batch, 1)
+        row[0], row[1], row[2] = int(cfg.bos_id), int(cfg.null_id), int(cfg.eos_id)
+        return row[None, :].repeat(batch, 1)
+
+    def encode_prompt(self, prompts: List[str], negative_prompts: Optional[List[str]], do_cfg: bool, device):
+        ids = self._tokenize(prompts).to(device)
+        with self._autocast():
+            cond, cond_valid = self.text_encoder(ids)
+            if not do_cfg:
+                return cond, cond_valid
+            if negative_prompts is None:
+                uncond_ids = self._null_ids(len(prompts))
+            else:
+                uncond_ids = self._tokenize(negative_prompts)
+            uncond, uncond_valid = self.text_encoder(uncond_ids.to(device))
+        # unconditional first, conditional second (the order the model was evaluated with)
+        return torch.cat([uncond, cond], dim=0), torch.cat([uncond_valid, cond_valid], dim=0)
+
+    @torch.no_grad()
+    def decode_latents(self, latents: torch.Tensor, chunk_size: int = 32) -> torch.Tensor:
+        """Model latents (B,128,16,16) -> images in [-1, 1] via the FLUX.2 small decoder."""
+        vae = self.vae
+        eps = float(getattr(vae.config, "batch_norm_eps", 1e-4))
+        mean = vae.bn.running_mean.detach().float().view(1, -1, 1, 1).to(latents.device)
+        std = torch.sqrt(vae.bn.running_var.detach().float() + eps).view(1, -1, 1, 1).to(latents.device)
+        unpatchify = UNPATCHIFY_LAYOUTS[self.transformer.config.latent_layout]
+        outs = []
+        for i in range(0, latents.shape[0], chunk_size):
+            z = latents[i:i + chunk_size].float() * std + mean
+            z = unpatchify(z)
+            outs.append(vae.decode(z.to(vae.dtype)).sample.float().clamp(-1, 1))
+        return torch.cat(outs, dim=0)
+
+    # ---- main entry point ----------------------------------------------------------------------
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 3.0,
+        num_images_per_prompt: int = 1,
+        generator: Optional[torch.Generator] = None,
+        latents: Optional[torch.Tensor] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+    ):
+        """
+        Args:
+            prompt: text prompt(s). English; prompts longer than 62 BPE tokens are truncated.
+            negative_prompt: optional. None (default) uses the learned empty prompt the model was trained
+                with for classifier-free guidance. Arbitrary negative prompts were not part of training.
+            num_inference_steps: Euler steps. 50 = recommended (used for the benchmarks), 20 = fast.
+            guidance_scale: classifier-free guidance scale. 3.0 recommended.
+            output_type: "pil", "np", "pt" (tensor in [0,1]) or "latent" (raw model latents).
+        """
+        base = [prompt] if isinstance(prompt, str) else list(prompt)
+        prompts = [p for p in base for _ in range(num_images_per_prompt)]
+        negs = None
+        if negative_prompt is not None:
+            neg_base = [negative_prompt] * len(base) if isinstance(negative_prompt, str) else list(negative_prompt)
+            if len(neg_base) != len(base):
+                raise ValueError("negative_prompt must be a string or have the same length as prompt.")
+            negs = [n for n in neg_base for _ in range(num_images_per_prompt)]
+
+        device = self._execution_device
+        b = len(prompts)
+        do_cfg = guidance_scale != 1.0
+        context, context_valid = self.encode_prompt(prompts, negs, do_cfg, device)
+
+        tcfg = self.transformer.config
+        shape = (b, tcfg.latent_channels, tcfg.latent_size, tcfg.latent_size)
+        if latents is None:
+            z = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
+        else:
+            z = latents.to(device=device, dtype=torch.float32)
+
+        ts = torch.linspace(1.0, 0.0, num_inference_steps + 1, device=device)
+        for i in self.progress_bar(range(num_inference_steps)):
+            t0, t1 = ts[i], ts[i + 1]
+            with self._autocast():
+                if do_cfg:
+                    v = self.transformer(torch.cat([z, z], dim=0), t0.expand(2 * b), context, context_valid)
+                    v_uncond, v_cond = v.chunk(2, dim=0)
+                    velocity = v_uncond + guidance_scale * (v_cond - v_uncond)
+                else:
+                    velocity = self.transformer(z, t0.expand(b), context, context_valid)
+            z = z + (t1 - t0) * velocity.float()
+
+        if output_type == "latent":
+            images = z
+        else:
+            img = (self.decode_latents(z) + 1.0) / 2.0
+            if output_type == "pt":
+                images = img
+            else:
+                arr = img.cpu().permute(0, 2, 3, 1).float().numpy()
+                images = self.numpy_to_pil(arr) if output_type == "pil" else arr
+
+        self.maybe_free_model_hooks()
+        if not return_dict:
+            return (images,)
+        return ImagePipelineOutput(images=images)

vae/config.json

@@ -0,0 +1,46 @@
+{
+  "_class_name": "AutoencoderKLFlux2",
+  "_diffusers_version": "0.42.1",
+  "_name_or_path": "black-forest-labs/FLUX.2-small-decoder",
+  "act_fn": "silu",
+  "batch_norm_eps": 0.0001,
+  "batch_norm_momentum": 0.1,
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "decoder_block_out_channels": [
+    96,
+    192,
+    384,
+    384
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 32,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "patch_size": [
+    2,
+    2
+  ],
+  "sample_size": 1024,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}

vae/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:377e8be52f1971978b0af970f2c7db9e40e4e0c9e2da0785ec5d0de461e51f04
+size 249521372