Upload iris/blocks.py

iris/blocks.py

@@ -0,0 +1,164 @@
+"""
+Core building blocks for IRIS: attention, FFN, cross-attention, embeddings.
+
+Design principles:
+- MQA (Multi-Query Attention) everywhere — shared K,V across heads
+- UIB-FFN (Universal Inverted Bottleneck) — depthwise separable, expansion=2
+- QK-RMSNorm for training stability (from SANA-Sprint)
+- 2D RoPE for spatial position encoding
+- Timestep addition (not AdaLN) — saves params (from HTH)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Optional
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        rms = torch.sqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
+        return (x.float() / rms * self.weight.float()).to(x.dtype)
+
+
+class RotaryEmbedding2D(nn.Module):
+    def __init__(self, dim: int, max_size: int = 64):
+        super().__init__()
+        self.dim = dim
+        half_dim = dim // 4
+        inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half_dim, dtype=torch.float32) / half_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def _build_cache(self, H, W, device, dtype):
+        h_pos = torch.arange(H, device=device, dtype=torch.float32)
+        w_pos = torch.arange(W, device=device, dtype=torch.float32)
+        inv = self.inv_freq.to(device)
+        h_freqs = torch.outer(h_pos, inv)[:, None, :].expand(H, W, -1)
+        w_freqs = torch.outer(w_pos, inv)[None, :, :].expand(H, W, -1)
+        freqs = torch.cat([h_freqs, w_freqs], dim=-1).reshape(H * W, -1)
+        return freqs.cos().to(dtype), freqs.sin().to(dtype)
+
+    def forward(self, x, H, W):
+        N = H * W
+        cos_c, sin_c = self._build_cache(H, W, x.device, x.dtype)
+        if x.dim() == 4:
+            cos_c = cos_c[None, None, :N, :]
+            sin_c = sin_c[None, None, :N, :]
+        else:
+            cos_c = cos_c[None, :N, :]
+            sin_c = sin_c[None, :N, :]
+        d = cos_c.shape[-1]
+        x1, x2, xr = x[..., :d], x[..., d:2*d], x[..., 2*d:]
+        return torch.cat([x1*cos_c - x2*sin_c, x1*sin_c + x2*cos_c, xr], dim=-1)
+
+
+class MultiQueryCrossAttention(nn.Module):
+    def __init__(self, dim, num_heads=4, qk_norm=True):
+        super().__init__()
+        assert dim % num_heads == 0
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, self.head_dim, bias=False)
+        self.v_proj = nn.Linear(dim, self.head_dim, bias=False)
+        self.out_proj = nn.Linear(dim, dim, bias=False)
+        self.q_norm = RMSNorm(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = RMSNorm(self.head_dim) if qk_norm else nn.Identity()
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x, context):
+        B, N, D = x.shape
+        S = context.shape[1]
+        residual = x
+        x = self.norm(x)
+        q = self.q_proj(x).view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(context).view(B, S, 1, self.head_dim).transpose(1, 2)
+        v = self.v_proj(context).view(B, S, 1, self.head_dim).transpose(1, 2)
+        q, k = self.q_norm(q), self.k_norm(k)
+        k = k.expand(-1, self.num_heads, -1, -1)
+        v = v.expand(-1, self.num_heads, -1, -1)
+        attn = F.scaled_dot_product_attention(q, k, v, scale=1.0/math.sqrt(self.head_dim))
+        return residual + self.out_proj(attn.transpose(1, 2).reshape(B, N, D))
+
+
+class MultiQuerySelfAttention(nn.Module):
+    def __init__(self, dim, num_heads=4, qk_norm=True):
+        super().__init__()
+        assert dim % num_heads == 0
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, self.head_dim, bias=False)
+        self.v_proj = nn.Linear(dim, self.head_dim, bias=False)
+        self.out_proj = nn.Linear(dim, dim, bias=False)
+        self.q_norm = RMSNorm(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = RMSNorm(self.head_dim) if qk_norm else nn.Identity()
+        self.norm = nn.LayerNorm(dim)
+        self.rope = RotaryEmbedding2D(self.head_dim)
+
+    def forward(self, x, H, W):
+        B, N, D = x.shape
+        residual = x
+        x = self.norm(x)
+        q = self.q_proj(x).view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, N, 1, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, N, 1, self.head_dim).transpose(1, 2)
+        q, k = self.q_norm(q), self.k_norm(k)
+        q, k = self.rope(q, H, W), self.rope(k, H, W)
+        k = k.expand(-1, self.num_heads, -1, -1)
+        v = v.expand(-1, self.num_heads, -1, -1)
+        attn = F.scaled_dot_product_attention(q, k, v, scale=1.0/math.sqrt(self.head_dim))
+        return residual + self.out_proj(attn.transpose(1, 2).reshape(B, N, D))
+
+
+class UIBFFN(nn.Module):
+    def __init__(self, dim, expansion=2, spatial_size=4):
+        super().__init__()
+        hidden = dim * expansion
+        self.norm = nn.LayerNorm(dim)
+        self.pw_up = nn.Linear(dim, hidden, bias=False)
+        self.gate = nn.Linear(dim, hidden, bias=False)
+        self.dw_conv = nn.Conv2d(hidden, hidden, 3, padding=1, groups=hidden, bias=True)
+        self.pw_down = nn.Linear(hidden, dim, bias=False)
+
+    def forward(self, x, H, W):
+        B, N, D = x.shape
+        residual = x
+        x = self.norm(x)
+        h = self.pw_up(x)
+        g = F.silu(self.gate(x))
+        h_2d = h.view(B, H, W, -1).permute(0, 3, 1, 2)
+        h = self.dw_conv(h_2d).permute(0, 2, 3, 1).reshape(B, N, -1)
+        return residual + self.pw_down(h * g)
+
+
+class TimestepEmbedding(nn.Module):
+    def __init__(self, dim, max_period=10000):
+        super().__init__()
+        self.dim = dim
+        self.max_period = max_period
+        self.mlp = nn.Sequential(nn.Linear(dim, dim*4), nn.SiLU(), nn.Linear(dim*4, dim))
+
+    def forward(self, t):
+        half = self.dim // 2
+        freqs = torch.exp(-math.log(self.max_period) * torch.arange(half, device=t.device, dtype=torch.float32) / half)
+        args = t[:, None].float() * freqs[None, :]
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if self.dim % 2:
+            emb = F.pad(emb, (0, 1))
+        return self.mlp(emb.to(t.dtype))
+
+
+class IterationEmbedding(nn.Module):
+    def __init__(self, dim, max_iterations=8):
+        super().__init__()
+        self.embed = nn.Embedding(max_iterations, dim)
+
+    def forward(self, iter_idx, batch_size, device):
+        return self.embed(torch.full((batch_size,), iter_idx, device=device, dtype=torch.long))