Upload 2 files

FlashScape.safetensors

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

network_diffusion_unet.py

@@ -5,11 +5,12 @@ from torch.utils.checkpoint import checkpoint
 
 
 class SinusoidalEmbedding(nn.Module):
-    def __init__(self, embedding_dim=128, scaling=1000):
+    def __init__(self, embedding_dim=128, base=1000, scaling=1000):
         super().__init__()
         self.embedding_dim = embedding_dim
         half_dim = embedding_dim // 2
-        freqs = torch.exp(-math.log(10000) * torch.arange(0, half_dim) / half_dim)
+        freqs = torch.exp(-math.log(base) * torch.arange(0, half_dim) / half_dim)
+        # at base 1000, max-range = +=500pi = -1571 to 1571
         self.scaling = nn.parameter.Buffer(torch.tensor(scaling))
         self.freqs = nn.parameter.Buffer(freqs)
 
@@ -27,8 +28,9 @@ class SinusoidalPositionalEmbedding2D(nn.Module):
         assert embedding_dim % 2 == 0, "embedding_dim must be even"
         self.embedding_dim = embedding_dim
         half_dim = self.embedding_dim // 2
-        div_term = torch.exp(torch.arange(0, half_dim, 2, dtype=torch.float32) * (-math.log(10000.0) / half_dim))
-        self.div_term = nn.parameter.Buffer(div_term)
+        div_term = torch.exp(torch.arange(0, half_dim, 2) * (-math.log(100.0) / half_dim))
+        # Since our grid size is small, 100 should be enough
+        self.div_term = nn.parameter.Buffer(div_term.to(torch.float32))
 
     def forward(self, height, width):
         """Generate embeddings for a grid of size (height, width)."""
@@ -124,83 +126,85 @@ class ImageLinearAttention(nn.Module):
 class ResConvBlock(nn.Module):
     def __init__(self, channels, time_dim):
         super().__init__()
-        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1, bias=False, padding_mode='replicate')
-        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1, padding_mode='replicate')
+        self.first_conv = nn.Conv2d(channels, channels, 3, padding=1, bias=False, padding_mode='replicate')
+        self.second_conv = nn.Conv2d(channels, channels, 3, padding=1, padding_mode='replicate')
         self.gn1 = nn.GroupNorm(8, channels, affine=True)
         self.gn2 = nn.GroupNorm(8, channels, affine=False)
-        self.time_affine = nn.Linear(time_dim, channels * 2)
+        self.embd_affine = nn.Linear(time_dim, channels * 2)
         self.act = nn.LeakyReLU(inplace=True)
 
     def forward(self, x, t_emb):
         # Get affine parameters from time embedding
-        affine_params = self.time_affine(t_emb)
+        affine_params = self.embd_affine(t_emb)
         scale, shift = affine_params.chunk(2, dim=1)
 
         # First convolution path
-        h = self.conv1(self.act(self.gn1(x)))
+        h = self.first_conv(self.act(self.gn1(x)))
 
         # Second convolution path with adaptive normalization
         h = self.gn2(h)
         h = h * (1 + scale[:, :, None, None]) + shift[:, :, None, None]
-        h = self.conv2(self.act(h))
+        h = self.second_conv(self.act(h))
 
         return x + h
 
 
 class DiTLayer(nn.Module):
-    def __init__(self, d_model, nhead, dim_feedforward=1024):
+    def __init__(self, d_model, embd_dim, nhead, dim_feedforward=1024):
         super().__init__()
-        self.norm1 = nn.LayerNorm(d_model)
-        self.attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
-        self.norm2 = nn.LayerNorm(d_model)
+        self.norm1 = nn.LayerNorm(d_model, elementwise_affine=False)
+        self.attn = nn.MultiheadAttention(d_model, nhead, batch_first=False)
+        self.norm2 = nn.LayerNorm(d_model, elementwise_affine=False)
+        self.embd_affine = nn.Linear(embd_dim, 6*d_model)
         self.ffn = nn.Sequential(
             nn.Linear(d_model, dim_feedforward),
             nn.LeakyReLU(0.2, inplace=True),
             nn.Linear(dim_feedforward, d_model),
         )
 
-    def forward(self, src):
+    def forward(self, x, embd):
+        affine_params = self.embd_affine(embd)
+        scale1, scale2, shift1, shift2, alpha1, alpha2 = affine_params.chunk(6, dim=1)
+
         # Self-attention block
-        attn_output, _ = self.attn(self.norm1(src), self.norm1(src), self.norm1(src))
-        src = src + attn_output
+        x = self.norm1(x)
+        x = x * (1 + scale1[None, :, :]) + shift1[None, :, :]
+
+        attn_output, _ = self.attn(x, x, x)
+        x = x + attn_output * alpha1[None, :, :]
 
         # Feedforward block
-        ffn_output = self.ffn(self.norm2(src))
-        src = src + ffn_output
-        return src
+        x = self.norm2(x)
+        x = x * (1 + scale2[None, :, :]) + shift2[None, :, :]
+        ffn_output = self.ffn(x)
+        x = x + ffn_output * alpha2[None, :, :]
+        return x
 
 
 class DiTBlock(nn.Module):
-    def __init__(self, channels, patch_size, hidden_size, nhead, num_layers=2):
+    def __init__(self, channels, embd_dim, patch_size, nhead, num_layers):
         super().__init__()
         self.patch_size = patch_size
         self.patchify = nn.Unfold(kernel_size=patch_size, stride=patch_size)
-        self.patch_embedding_in = nn.Linear(channels * patch_size**2, hidden_size)
+        hidden_size = channels * patch_size**2
         self.pos_embd = SinusoidalPositionalEmbedding2D(hidden_size)
-        self.waterlevel_embd = SinusoidalEmbedding(hidden_size, 10)
-        self.patch_embedding_out = nn.Linear(hidden_size, channels * patch_size**2)
         self.dit_layers = nn.ModuleList([
-            DiTLayer(hidden_size, nhead, 2*hidden_size)
+            DiTLayer(hidden_size, embd_dim, nhead, 2*hidden_size)
             for _ in range(num_layers)
         ])
-        self.norm = nn.GroupNorm(8, channels)
 
-    def forward(self, src, water_level):
-        B, C, H, W = src.shape
+    def forward(self, x, embd):
+        B, C, H, W = x.shape
         H_p, W_p = H // self.patch_size, W // self.patch_size
-        x = self.norm(src)
-        x = self.patchify(x).permute(0, 2, 1)
-        x = self.patch_embedding_in(x)
+        x = self.patchify(x).permute(0, 2, 1)  # [B, num_patches, d_main]
         pos_embd = self.pos_embd(H_p, W_p).to(dtype=x.dtype)
         x = x + pos_embd.unsqueeze(0)
-        water_level_cls = self.waterlevel_embd(water_level).unsqueeze(1)
-        x = torch.cat((x, water_level_cls), dim=1)
+        x = x.permute(1, 0, 2) # [num_patches, B, d_main)
         for dit_layer in self.dit_layers:
-            x = dit_layer(x)
-        x = self.patch_embedding_out(x).permute(0, 2, 1)
-        x = x[:, :, :-1]
+            x = dit_layer(x, embd)
+        x = x.permute(1, 2, 0)  # [B, d_main, num_patches]
         x = nn.functional.fold(x, (H, W), (self.patch_size, self.patch_size), stride=(self.patch_size, self.patch_size))
-        return src + x
+        return x
 
 
 class UpBlock(nn.Module):
@@ -220,16 +224,22 @@ class UpBlock(nn.Module):
         return x
 
 class UpBlockWithDit(nn.Module):
-    def __init__(self, in_ch, out_ch, patch_size, hidden_size, nhead, time_dim, cat):
+    def __init__(self, in_ch, mid_ch, out_ch, patch_size, nhead, time_dim, layers, cat):
         super().__init__()
         self.res = ResConvBlock(in_ch, time_dim)
-        self.dit = DiTBlock(in_ch, patch_size, hidden_size, nhead, 4)
+        self.down_map = nn.Conv2d(in_ch, mid_ch, kernel_size=1, bias=False)
+        self.down_norm = nn.GroupNorm(4, mid_ch)
+        self.dit = DiTBlock(mid_ch, time_dim, patch_size, nhead, layers)
+        self.up_map = nn.Conv2d(mid_ch, in_ch, kernel_size=1)
         self.up = nn.ConvTranspose2d(in_ch, out_ch, 4, stride=2, padding=1)
         self.cat = cat
 
-    def forward(self, x, t_emb, water_level, skip=None):
-        x = self.res(x, t_emb)
-        x = self.dit(x, water_level)
+    def forward(self, x, embd, skip=None):
+        x = self.res(x, embd)
+        h = self.down_norm(self.down_map(x))
+        h = self.dit(h, embd)
+        h = self.up_map(h)
+        x = x + h
         x = self.up(x)
         if self.cat:
             x = torch.cat([x, skip], dim=1)
@@ -313,7 +323,9 @@ class ConditionalUNet(nn.Module):
 class ConditionalUNetDiT(nn.Module):
     def __init__(self, base_ch=8, embd_dim=16):
         super().__init__()
-        self.time_embd = SinusoidalEmbedding(embd_dim, 1000)
+        self.time_embd = SinusoidalEmbedding(embd_dim, scaling=1000)
+        self.waterlevel_embd = SinusoidalEmbedding(embd_dim, scaling=1)
+        embd_dim *= 2
 
         # Input channels = noisy height (1) + ridge map (1) + lake map (1)
         self.expand = nn.Conv2d(3, base_ch, 3, padding=1, padding_mode='replicate')
@@ -321,30 +333,41 @@ class ConditionalUNetDiT(nn.Module):
 
         self.down0 = nn.Conv2d(base_ch, base_ch * 2, 4, stride=2, padding=1, padding_mode='replicate') # 1024->512
         self.enc_1 = ResConvBlock(base_ch * 2, embd_dim)
-        self.enc_1_dit = DiTBlock(base_ch * 2, 16, 1024, 8, 4)
+        #self.enc_1_dit = DiTBlock(base_ch * 2, 16, 1024, 8, 4)
 
         self.down1 = nn.Conv2d(base_ch * 2, base_ch * 4, 4, stride=2, padding=1, padding_mode='replicate') # 512->256
 
-        self.up1 = UpBlockWithDit(base_ch * 4, base_ch * 2, 8, 1024, 8, embd_dim, False) # 256->512
-        self.up0 = UpBlockWithDit(base_ch * 2, base_ch, 16, 1024, 8, embd_dim, True) # 512->1024
+        self.up1 = UpBlockWithDit(base_ch * 4, base_ch, base_ch * 2, 8, 8, embd_dim, 6, False) # 256->512
+        self.up0 = UpBlockWithDit(base_ch * 2, base_ch//2, base_ch, 16, 16, embd_dim, 3, True) # 512->1024
         self.out = ResConvBlock(base_ch * 2, embd_dim)
         self.final = nn.Conv2d(base_ch * 2, 1, 1)
 
+    def initialize(self):
+        for name, m in self.named_modules():
+            if isinstance(m, nn.Linear) and ('embd_affine' in name or 'water_level_affine' in name):
+                m.weight.data.zero_()
+                m.bias.data.zero_()
+            if isinstance(m, nn.Conv2d) and 'second_conv' in name:
+                m.weight.data.zero_()
+                m.bias.data.zero_()
+
     def forward(self, x, ridge_map, basin_map, water_level, t):
         t_embed = self.time_embd(t).to(x.dtype)
+        waterlevel_embd = self.waterlevel_embd(water_level).to(x.dtype)
+        embeds = torch.cat([t_embed, waterlevel_embd], dim=1)
         # x: noisy height map, ridge_map: binary edges, basin_map: binary basins, water_level: the estimate sea level
         h0 = torch.cat([x, ridge_map, basin_map], dim=1)  # concat condition
         # encode
-        h0 = checkpoint(run_block, self.expand, h0, use_reentrant=False) if self.training else self.expand(h0)
-        h0 = checkpoint(run_block, self.enc_0, h0, t_embed, use_reentrant=False) if self.training else self.enc_0(h0, t_embed)
-        h1 = checkpoint(run_block, self.down0, h0, use_reentrant=False) if self.training else self.down0(h0)
-        h1 = checkpoint(run_block, self.enc_1, h1, t_embed, use_reentrant=False) if self.training else self.enc_1(h1, t_embed)
-        h1 = checkpoint(run_block, self.enc_1_dit, h1, water_level, use_reentrant=False) if self.training else self.enc_1_dit(h1, water_level) # 512x512
-        h2 = checkpoint(run_block, self.down1, h1, use_reentrant=False) if self.training else self.down1(h1)  # 256x256
+        h0 = self.expand(h0)
+        h0 = self.enc_0(h0, embeds)
+        h1 = self.down0(h0)
+        h1 = self.enc_1(h1, embeds) # 512x512
+        #h1 = checkpoint(run_block, self.enc_1_dit, h1, water_level, use_reentrant=False) if self.training else self.enc_1_dit(h1, water_level)
+        h2 = self.down1(h1)  # 256x256
         # decode with skip connections
-        out = checkpoint(run_block, self.up1, h2, t_embed, water_level, h1, use_reentrant=False) if self.training else self.up1(h2, t_embed, water_level, h1)  # 512x512
-        out = checkpoint(run_block, self.up0, out, t_embed, water_level, h0, use_reentrant=False) if self.training else self.up0(out, t_embed, water_level, h0)  # 1024x1024
-        out = checkpoint(run_block, self.out, out, t_embed, use_reentrant=False) if self.training else self.out(out, t_embed)
+        out = self.up1(h2, embeds, h1)  # 512x512
+        out = self.up0(out, embeds, h0)  # 1024x1024
+        out = self.out(out, embeds)
         out = self.final(out)
         return out  # predicted noise for diffusion loss