Upload iris/core.py

iris/core.py

@@ -0,0 +1,66 @@
+"""IRIS Refinement Core: Weight-shared denoising backbone."""
+
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+
+from .pde_ssm import PDESSMBlock
+from .blocks import (MultiQueryCrossAttention, MultiQuerySelfAttention, UIBFFN, TimestepEmbedding, IterationEmbedding)
+
+
+class IRISBlock(nn.Module):
+    def __init__(self, dim, num_heads=4, spatial_size=4, use_attn=False, ffn_expansion=2):
+        super().__init__()
+        self.cross_attn = MultiQueryCrossAttention(dim, num_heads=num_heads)
+        if use_attn:
+            self.spatial_mixer = MultiQuerySelfAttention(dim, num_heads=num_heads)
+        else:
+            self.spatial_mixer = PDESSMBlock(dim, spatial_size=spatial_size)
+        self.ffn = UIBFFN(dim, expansion=ffn_expansion, spatial_size=spatial_size)
+        self.use_attn = use_attn
+
+    def forward(self, x, context, H, W):
+        x = self.cross_attn(x, context)
+        x = self.spatial_mixer(x, H, W)
+        x = self.ffn(x, H, W)
+        return x
+
+
+class RefinementCore(nn.Module):
+    """6 blocks (5 PDE-SSM + 1 Self-Attention), applied R times with same weights."""
+
+    def __init__(self, dim=512, num_blocks=6, num_heads=4, spatial_size=4, max_iterations=8, ffn_expansion=2, gradient_checkpointing=True):
+        super().__init__()
+        self.dim = dim
+        self.num_blocks = num_blocks
+        self.max_iterations = max_iterations
+        self.gradient_checkpointing = gradient_checkpointing
+
+        self.blocks = nn.ModuleList()
+        for i in range(num_blocks):
+            use_attn = (i == num_blocks - 1)
+            self.blocks.append(IRISBlock(dim=dim, num_heads=num_heads, spatial_size=spatial_size, use_attn=use_attn, ffn_expansion=ffn_expansion))
+
+        self.timestep_embed = TimestepEmbedding(dim)
+        self.iter_embed = IterationEmbedding(dim, max_iterations=max_iterations)
+        self.final_norm = nn.LayerNorm(dim)
+
+    def _single_iteration(self, x, context, t_emb, iter_emb, H, W):
+        cond = (t_emb + iter_emb).unsqueeze(1)
+        x = x + cond
+        for block in self.blocks:
+            x = block(x, context, H, W)
+        return x
+
+    def forward(self, x, context, t, H, W, num_iterations=4):
+        B = x.shape[0]
+        t_emb = self.timestep_embed(t)
+
+        for r in range(num_iterations):
+            iter_emb = self.iter_embed(r, B, x.device)
+            if self.gradient_checkpointing and self.training:
+                x = checkpoint(self._single_iteration, x, context, t_emb, iter_emb, H, W, use_reentrant=False)
+            else:
+                x = self._single_iteration(x, context, t_emb, iter_emb, H, W)
+
+        return self.final_norm(x)