Create modeling_flow_match.py

modeling_flow_match.py

@@ -0,0 +1,412 @@
+"""
+FlowMatchRelay model — HuggingFace compatible.
+
+Usage:
+    from transformers import AutoModel
+    model = AutoModel.from_pretrained(
+        "AbstractPhil/geolip-diffusion-proto",
+        trust_remote_code=True
+    )
+
+    # Generate samples
+    samples = model.sample(n_samples=8, class_label=3)  # 8 cats
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from transformers import PreTrainedModel
+from .configuration_flow_match import FlowMatchRelayConfig
+
+
+# ══════════════════════════════════════════════════════════════════
+# CONSTELLATION RELAY
+# ══════════════════════════════════════════════════════════════════
+
+class ConstellationRelay(nn.Module):
+    """
+    Geometric regulator for feature maps.
+    Fixed anchors on S^(d-1), multi-phase stroboscope triangulation,
+    gated residual correction.
+    """
+    def __init__(self, channels, patch_dim=16, n_anchors=16, n_phases=3,
+                 pw_hidden=32, gate_init=-3.0, mode='channel'):
+        super().__init__()
+        assert channels % patch_dim == 0
+        self.channels = channels
+        self.patch_dim = patch_dim
+        self.n_patches = channels // patch_dim
+        self.n_anchors = n_anchors
+        self.n_phases = n_phases
+        self.mode = mode
+
+        P, A, d = self.n_patches, n_anchors, patch_dim
+
+        home = torch.empty(P, A, d)
+        nn.init.xavier_normal_(home.view(P * A, d))
+        home = F.normalize(home.view(P, A, d), dim=-1)
+        self.register_buffer('home', home)
+        self.anchors = nn.Parameter(home.clone())
+
+        tri_dim = n_phases * A
+        self.pw_w1 = nn.Parameter(torch.empty(P, tri_dim, pw_hidden))
+        self.pw_b1 = nn.Parameter(torch.zeros(1, P, pw_hidden))
+        self.pw_w2 = nn.Parameter(torch.empty(P, pw_hidden, d))
+        self.pw_b2 = nn.Parameter(torch.zeros(1, P, d))
+        for p in range(P):
+            nn.init.xavier_normal_(self.pw_w1.data[p])
+            nn.init.xavier_normal_(self.pw_w2.data[p])
+        self.pw_norm = nn.LayerNorm(d)
+        self.gates = nn.Parameter(torch.full((P,), gate_init))
+        self.norm = nn.LayerNorm(channels)
+
+    def drift(self):
+        h, c = F.normalize(self.home, dim=-1), F.normalize(self.anchors, dim=-1)
+        return torch.acos((h * c).sum(-1).clamp(-1 + 1e-7, 1 - 1e-7))
+
+    def at_phase(self, t):
+        h, c = F.normalize(self.home, dim=-1), F.normalize(self.anchors, dim=-1)
+        omega = self.drift().unsqueeze(-1)
+        so = omega.sin().clamp(min=1e-7)
+        return torch.sin((1-t)*omega)/so * h + torch.sin(t*omega)/so * c
+
+    def _relay_core(self, x_flat):
+        N, C = x_flat.shape
+        P, A, d = self.n_patches, self.n_anchors, self.patch_dim
+        x_n = self.norm(x_flat)
+        patches = x_n.reshape(N, P, d)
+        patches_n = F.normalize(patches, dim=-1)
+        phases = torch.linspace(0, 1, self.n_phases, device=x_flat.device).tolist()
+        tris = []
+        for t in phases:
+            at = F.normalize(self.at_phase(t), dim=-1)
+            tris.append(1.0 - torch.einsum('npd,pad->npa', patches_n, at))
+        tri = torch.cat(tris, dim=-1)
+        h = F.gelu(torch.einsum('npt,pth->nph', tri, self.pw_w1) + self.pw_b1)
+        pw = self.pw_norm(torch.einsum('nph,phd->npd', h, self.pw_w2) + self.pw_b2)
+        g = self.gates.sigmoid().unsqueeze(0).unsqueeze(-1)
+        blended = g * pw + (1-g) * patches
+        return x_flat + blended.reshape(N, C)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        if self.mode == 'channel':
+            pooled = x.mean(dim=(-2, -1))
+            relayed = self._relay_core(pooled)
+            scale = (relayed / (pooled + 1e-8)).unsqueeze(-1).unsqueeze(-1)
+            return x * scale.clamp(-3, 3)
+        else:
+            x_flat = x.permute(0, 2, 3, 1).reshape(B * H * W, C)
+            out = self._relay_core(x_flat)
+            return out.reshape(B, H, W, C).permute(0, 3, 1, 2)
+
+
+# ══════════════════════════════════════════════════════════════════
+# BUILDING BLOCKS
+# ══════════════════════════════════════════════════════════════════
+
+class SinusoidalPosEmb(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, t):
+        half = self.dim // 2
+        emb = math.log(10000) / (half - 1)
+        emb = torch.exp(torch.arange(half, device=t.device, dtype=t.dtype) * -emb)
+        emb = t.unsqueeze(-1) * emb.unsqueeze(0)
+        return torch.cat([emb.sin(), emb.cos()], dim=-1)
+
+
+class AdaGroupNorm(nn.Module):
+    def __init__(self, channels, cond_dim, n_groups=8):
+        super().__init__()
+        self.gn = nn.GroupNorm(min(n_groups, channels), channels, affine=False)
+        self.proj = nn.Linear(cond_dim, channels * 2)
+        nn.init.zeros_(self.proj.weight)
+        nn.init.zeros_(self.proj.bias)
+
+    def forward(self, x, cond):
+        x = self.gn(x)
+        scale, shift = self.proj(cond).unsqueeze(-1).unsqueeze(-1).chunk(2, dim=1)
+        return x * (1 + scale) + shift
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, channels, cond_dim, use_relay=False,
+                 relay_patch_dim=16, relay_n_anchors=16, relay_n_phases=3,
+                 relay_pw_hidden=32, relay_gate_init=-3.0, relay_mode='channel'):
+        super().__init__()
+        self.dw_conv = nn.Conv2d(channels, channels, 7, padding=3, groups=channels)
+        self.norm = AdaGroupNorm(channels, cond_dim)
+        self.pw1 = nn.Conv2d(channels, channels * 4, 1)
+        self.pw2 = nn.Conv2d(channels * 4, channels, 1)
+        self.act = nn.GELU()
+        self.relay = ConstellationRelay(
+            channels,
+            patch_dim=min(relay_patch_dim, channels),
+            n_anchors=min(relay_n_anchors, channels),
+            n_phases=relay_n_phases,
+            pw_hidden=relay_pw_hidden,
+            gate_init=relay_gate_init,
+            mode=relay_mode) if use_relay else None
+
+    def forward(self, x, cond):
+        residual = x
+        x = self.dw_conv(x)
+        x = self.norm(x, cond)
+        x = self.pw1(x)
+        x = self.act(x)
+        x = self.pw2(x)
+        x = residual + x
+        if self.relay is not None:
+            x = self.relay(x)
+        return x
+
+
+class SelfAttnBlock(nn.Module):
+    def __init__(self, channels, n_heads=4):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = channels // n_heads
+        self.norm = nn.GroupNorm(8, channels)
+        self.qkv = nn.Conv2d(channels, channels * 3, 1)
+        self.out = nn.Conv2d(channels, channels, 1)
+        nn.init.zeros_(self.out.weight)
+        nn.init.zeros_(self.out.bias)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        residual = x
+        x = self.norm(x)
+        qkv = self.qkv(x).reshape(B, 3, self.n_heads, self.head_dim, H * W)
+        q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2]
+        attn = F.scaled_dot_product_attention(q, k, v)
+        out = attn.reshape(B, C, H, W)
+        return residual + self.out(out)
+
+
+class Downsample(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.conv = nn.Conv2d(channels, channels, 3, stride=2, padding=1)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class Upsample(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.conv = nn.Conv2d(channels, channels, 3, padding=1)
+
+    def forward(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        return self.conv(x)
+
+
+# ══════════════════════════════════════════════════════════════════
+# FLOW MATCHING UNET
+# ══════════════════════════════════════════════════════════════════
+
+class FlowMatchUNet(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        in_channels = config.in_channels
+        base_channels = config.base_channels
+        channel_mults = config.channel_mults
+        n_classes = config.n_classes
+        cond_dim = config.cond_dim
+        use_relay = config.use_relay
+        self.channel_mults = channel_mults
+
+        # Relay kwargs
+        rk = dict(
+            relay_patch_dim=config.relay_patch_dim,
+            relay_n_anchors=config.relay_n_anchors,
+            relay_n_phases=config.relay_n_phases,
+            relay_pw_hidden=config.relay_pw_hidden,
+            relay_gate_init=config.relay_gate_init,
+            relay_mode=config.relay_mode,
+        )
+
+        self.time_emb = nn.Sequential(
+            SinusoidalPosEmb(cond_dim),
+            nn.Linear(cond_dim, cond_dim), nn.GELU(),
+            nn.Linear(cond_dim, cond_dim))
+        self.class_emb = nn.Embedding(n_classes, cond_dim)
+        self.in_conv = nn.Conv2d(in_channels, base_channels, 3, padding=1)
+
+        # Encoder
+        self.enc = nn.ModuleList()
+        self.enc_down = nn.ModuleList()
+        ch_in = base_channels
+        enc_channels = [base_channels]
+
+        for i, mult in enumerate(channel_mults):
+            ch_out = base_channels * mult
+            self.enc.append(nn.ModuleList([
+                ConvBlock(ch_in, cond_dim) if ch_in == ch_out
+                else nn.Sequential(nn.Conv2d(ch_in, ch_out, 1),
+                                   ConvBlock(ch_out, cond_dim)),
+                ConvBlock(ch_out, cond_dim),
+            ]))
+            ch_in = ch_out
+            enc_channels.append(ch_out)
+            if i < len(channel_mults) - 1:
+                self.enc_down.append(Downsample(ch_out))
+
+        # Middle
+        mid_ch = ch_in
+        self.mid_block1 = ConvBlock(mid_ch, cond_dim, use_relay=use_relay, **rk)
+        self.mid_attn = SelfAttnBlock(mid_ch, n_heads=4)
+        self.mid_block2 = ConvBlock(mid_ch, cond_dim, use_relay=use_relay, **rk)
+
+        # Decoder
+        self.dec_up = nn.ModuleList()
+        self.dec_skip_proj = nn.ModuleList()
+        self.dec = nn.ModuleList()
+
+        for i in range(len(channel_mults) - 1, -1, -1):
+            ch_out = base_channels * channel_mults[i]
+            skip_ch = enc_channels.pop()
+            self.dec_skip_proj.append(nn.Conv2d(ch_in + skip_ch, ch_out, 1))
+            self.dec.append(nn.ModuleList([
+                ConvBlock(ch_out, cond_dim),
+                ConvBlock(ch_out, cond_dim),
+            ]))
+            ch_in = ch_out
+            if i > 0:
+                self.dec_up.append(Upsample(ch_out))
+
+        self.out_norm = nn.GroupNorm(8, ch_in)
+        self.out_conv = nn.Conv2d(ch_in, in_channels, 3, padding=1)
+        nn.init.zeros_(self.out_conv.weight)
+        nn.init.zeros_(self.out_conv.bias)
+
+    def forward(self, x, t, class_labels):
+        cond = self.time_emb(t) + self.class_emb(class_labels)
+        h = self.in_conv(x)
+        skips = [h]
+
+        for i in range(len(self.channel_mults)):
+            for block in self.enc[i]:
+                if isinstance(block, ConvBlock):
+                    h = block(h, cond)
+                elif isinstance(block, nn.Sequential):
+                    h = block[0](h)
+                    h = block[1](h, cond)
+                else:
+                    h = block(h)
+            skips.append(h)
+            if i < len(self.enc_down):
+                h = self.enc_down[i](h)
+
+        h = self.mid_block1(h, cond)
+        h = self.mid_attn(h)
+        h = self.mid_block2(h, cond)
+
+        for i in range(len(self.channel_mults)):
+            skip = skips.pop()
+            if i > 0:
+                h = self.dec_up[i - 1](h)
+            h = torch.cat([h, skip], dim=1)
+            h = self.dec_skip_proj[i](h)
+            for block in self.dec[i]:
+                h = block(h, cond)
+
+        h = self.out_norm(h)
+        h = F.silu(h)
+        return self.out_conv(h)
+
+
+# ══════════════════════════════════════════════════════════════════
+# HUGGINGFACE PRETRAINED MODEL WRAPPER
+# ══════════════════════════════════════════════════════════════════
+
+class FlowMatchRelayModel(PreTrainedModel):
+    """
+    HuggingFace-compatible wrapper for flow matching with constellation relay.
+
+    Load:
+        model = AutoModel.from_pretrained(
+            "AbstractPhil/geolip-diffusion-proto", trust_remote_code=True)
+
+    Generate:
+        images = model.sample(n_samples=8, class_label=3)
+    """
+    config_class = FlowMatchRelayConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.unet = FlowMatchUNet(config)
+
+    def forward(self, x, t, class_labels):
+        """
+        Predict velocity field for flow matching.
+
+        Args:
+            x: (B, 3, H, W) noisy images
+            t: (B,) timesteps in [0, 1]
+            class_labels: (B,) integer class labels
+
+        Returns:
+            v_pred: (B, 3, H, W) predicted velocity
+        """
+        return self.unet(x, t, class_labels)
+
+    @torch.no_grad()
+    def sample(self, n_samples=8, n_steps=None, class_label=None, device=None):
+        """
+        Generate images via Euler ODE integration.
+
+        Args:
+            n_samples: number of images to generate
+            n_steps: ODE integration steps (default from config)
+            class_label: optional class conditioning (0-9 for CIFAR-10)
+            device: target device
+
+        Returns:
+            images: (n_samples, 3, 32, 32) in [0, 1]
+        """
+        if device is None:
+            device = next(self.parameters()).device
+        if n_steps is None:
+            n_steps = self.config.n_sample_steps
+
+        self.eval()
+        x = torch.randn(n_samples, self.config.in_channels,
+                        self.config.image_size, self.config.image_size,
+                        device=device)
+
+        if class_label is not None:
+            labels = torch.full((n_samples,), class_label,
+                               dtype=torch.long, device=device)
+        else:
+            labels = torch.randint(0, self.config.n_classes,
+                                  (n_samples,), device=device)
+
+        dt = 1.0 / n_steps
+        for step in range(n_steps):
+            t_val = 1.0 - step * dt
+            t = torch.full((n_samples,), t_val, device=device)
+            v = self.unet(x, t, labels)
+            x = x - v * dt
+
+        # [-1, 1] → [0, 1]
+        return (x.clamp(-1, 1) + 1) / 2
+
+    def get_relay_diagnostics(self):
+        """Report constellation relay drift and gate values."""
+        diagnostics = {}
+        for name, module in self.named_modules():
+            if isinstance(module, ConstellationRelay):
+                drift = module.drift().mean().item()
+                gate = module.gates.sigmoid().mean().item()
+                diagnostics[name] = {
+                    'drift_rad': drift,
+                    'drift_deg': math.degrees(drift),
+                    'gate': gate,
+                }
+        return diagnostics