| """AuroraIMG2-50M-Preview: custom diffusers pipeline (load with trust_remote_code=True). |
| |
| 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 |
|
|
|
|
| |
| |
| |
| 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) |
|
|
|
|
| |
| |
| |
| def _unpatchify_flux2(z): |
| 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, |
| } |
|
|
|
|
| |
| |
| |
| 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 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) |
|
|
|
|
| |
| |
| |
| 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) |
|
|
| |
| 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 _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 |
| max_len = int(cfg.max_tokens) |
| cleaned = [" ".join(str(t).split()) for t in texts] |
| |
| |
| |
| 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) |
| 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)) |
| |
| 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) |
|
|
| |
| @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) |
