Upload unet.py with huggingface_hub

unet.py

@@ -1,12 +1,9 @@
-"""E3Diff UNet Architecture - exact copy from original with fixed imports."""
-
 import math
 import torch
 from torch import nn
 import torch.nn.functional as F
 from inspect import isfunction
-from softpool import soft_pool2d, SoftPool2d
-
+import numpy as np
 
 def exists(x):
     return x is not None
@@ -17,7 +14,7 @@ def default(val, d):
         return val
     return d() if isfunction(d) else d
 
-
+# PositionalEncoding Source： https://github.com/lmnt-com/wavegrad/blob/master/src/wavegrad/model.py
 class PositionalEncoding(nn.Module):
     def __init__(self, dim):
         super().__init__()
@@ -77,6 +74,9 @@ class Downsample(nn.Module):
         return self.conv(x)
 
 
+# building block modules
+
+
 class Block(nn.Module):
     def __init__(self, dim, dim_out, groups=32, dropout=0, stride=1):
         super().__init__()
@@ -96,7 +96,7 @@ class ResnetBlock(nn.Module):
         super().__init__()
         self.noise_func = FeatureWiseAffine(
             noise_level_emb_dim, dim_out, use_affine_level)
-        self.c_func = nn.Conv2d(dim_out, dim_out, 1)
+        self.c_func =  nn.Conv2d(dim_out, dim_out, 1)
 
         self.block1 = Block(dim, dim_out, groups=norm_groups)
         self.block2 = Block(dim_out, dim_out, groups=norm_groups, dropout=dropout)
@@ -104,17 +104,22 @@ class ResnetBlock(nn.Module):
             dim, dim_out, 1) if dim != dim_out else nn.Identity()
 
     def forward(self, x, time_emb, c):
+        # b, c, h, w = x.shape
         h = self.block1(x)
         h = self.noise_func(h, time_emb)
         h = self.block2(h)
+
         h = self.c_func(c) + h
         return h + self.res_conv(x)
 
+ 
 
 class SelfAttention(nn.Module):
     def __init__(self, in_channel, n_head=1, norm_groups=32):
         super().__init__()
+
         self.n_head = n_head
+
         self.norm = nn.GroupNorm(norm_groups, in_channel)
         self.qkv = nn.Conv2d(in_channel, in_channel * 3, 1, bias=False)
         self.out = nn.Conv2d(in_channel, in_channel, 1)
@@ -126,7 +131,7 @@ class SelfAttention(nn.Module):
 
         norm = self.norm(input)
         qkv = self.qkv(norm).view(batch, n_head, head_dim * 3, height, width)
-        query, key, value = qkv.chunk(3, dim=2)
+        query, key, value = qkv.chunk(3, dim=2)  # bhdyx
 
         attn = torch.einsum(
             "bnchw, bncyx -> bnhwyx", query, key
@@ -140,6 +145,10 @@ class SelfAttention(nn.Module):
 
         return out + input
 
+ 
+
+
+
 
 class ResnetBlocWithAttn(nn.Module):
     def __init__(self, dim, dim_out, *, noise_level_emb_dim=None, norm_groups=32, dropout=0, with_attn=False, size=256):
@@ -151,76 +160,12 @@ class ResnetBlocWithAttn(nn.Module):
             self.attn = SelfAttention(dim_out, norm_groups=norm_groups)
 
     def forward(self, x, time_emb, c, t=0, save_flag=False, file_i=0):
-        x = self.res_block(x, time_emb, c)
-        if self.with_attn:
+        x = self.res_block(x, time_emb, c)   # resblock(x + self.noise_func(noise_embed)) + con1_1(c)
+        if(self.with_attn):
             x = self.attn(x, t=t, save_flag=save_flag, file_num=file_i)
         return x
 
 
-class ResBlock_normal(nn.Module):
-    def __init__(self, dim, dim_out, dropout=0, norm_groups=32):
-        super().__init__()
-        self.block1 = Block(dim, dim_out, groups=norm_groups)
-        self.block2 = Block(dim_out, dim_out, groups=norm_groups, dropout=dropout)
-        self.res_conv = nn.Conv2d(
-            dim, dim_out, 1) if dim != dim_out else nn.Identity()
-
-    def forward(self, x):
-        b, c, h, w = x.shape
-        h = self.block1(x)
-        h = self.block2(h)
-        return h + self.res_conv(x)
-
-
-class CPEN(nn.Module):
-    """Condition Pyramid Encoder Network - EXACT architecture from E3Diff."""
-    def __init__(self, inchannel=1):
-        super(CPEN, self).__init__()
-        self.pool = SoftPool2d(kernel_size=(2, 2), stride=(2, 2))
-
-        self.E1 = nn.Sequential(
-            nn.Conv2d(inchannel, 64, kernel_size=3, padding=1),
-            Swish()
-        )
-
-        self.E2 = nn.Sequential(
-            ResBlock_normal(64, 128, dropout=0, norm_groups=16),
-            ResBlock_normal(128, 128, dropout=0, norm_groups=16),
-        )
-
-        self.E3 = nn.Sequential(
-            ResBlock_normal(128, 256, dropout=0, norm_groups=16),
-            ResBlock_normal(256, 256, dropout=0, norm_groups=16),
-        )
-
-        self.E4 = nn.Sequential(
-            ResBlock_normal(256, 512, dropout=0, norm_groups=16),
-            ResBlock_normal(512, 512, dropout=0, norm_groups=16),
-        )
-
-        self.E5 = nn.Sequential(
-            ResBlock_normal(512, 512, dropout=0, norm_groups=16),
-            ResBlock_normal(512, 1024, dropout=0, norm_groups=16),
-        )
-
-    def forward(self, x):
-        x1 = self.E1(x)        # 256x256, 64ch
-
-        x2 = self.pool(x1)     # 128x128
-        x2 = self.E2(x2)       # 128x128, 128ch
-
-        x3 = self.pool(x2)     # 64x64
-        x3 = self.E3(x3)       # 64x64, 256ch
-
-        x4 = self.pool(x3)     # 32x32
-        x4 = self.E4(x4)       # 32x32, 512ch
-
-        x5 = self.pool(x4)     # 16x16
-        x5 = self.E5(x5)       # 16x16, 1024ch
-
-        return x1, x2, x3, x4, x5
-
-
 class UNet(nn.Module):
     def __init__(
         self,
@@ -229,7 +174,7 @@ class UNet(nn.Module):
         inner_channel=32,
         norm_groups=32,
         channel_mults=(1, 2, 4, 8, 8),
-        attn_res=(8,),
+        attn_res=(8),
         res_blocks=3,
         dropout=0,
         with_noise_level_emb=True,
@@ -251,35 +196,37 @@ class UNet(nn.Module):
             noise_level_channel = None
             self.noise_level_mlp = None
 
+ 
+
+
         self.res_blocks = res_blocks
         num_mults = len(channel_mults)
         self.num_mults = num_mults
         pre_channel = inner_channel
         feat_channels = [pre_channel]
         now_res = image_size
-        
-        downs = [nn.Conv2d(in_channel, inner_channel, kernel_size=3, padding=1)]
+        downs = [nn.Conv2d(in_channel, inner_channel,
+                           kernel_size=3, padding=1)]
         for ind in range(num_mults):
             is_last = (ind == num_mults - 1)
             use_attn = (now_res in attn_res)
             channel_mult = inner_channel * channel_mults[ind]
             for _ in range(0, res_blocks):
                 downs.append(ResnetBlocWithAttn(
-                    pre_channel, channel_mult, noise_level_emb_dim=noise_level_channel, 
-                    norm_groups=norm_groups, dropout=dropout, with_attn=use_attn, size=now_res))
+                    pre_channel, channel_mult, noise_level_emb_dim=noise_level_channel, norm_groups=norm_groups, dropout=dropout, with_attn=use_attn,size=now_res))
                 feat_channels.append(channel_mult)
                 pre_channel = channel_mult
             if not is_last:
                 downs.append(Downsample(pre_channel))
                 feat_channels.append(pre_channel)
-                now_res = now_res // 2
+                now_res = now_res//2
         self.downs = nn.ModuleList(downs)
 
         self.mid = nn.ModuleList([
-            ResnetBlocWithAttn(pre_channel, pre_channel, noise_level_emb_dim=noise_level_channel, 
-                               norm_groups=norm_groups, dropout=dropout, with_attn=True, size=now_res),
-            ResnetBlocWithAttn(pre_channel, pre_channel, noise_level_emb_dim=noise_level_channel, 
-                               norm_groups=norm_groups, dropout=dropout, with_attn=False, size=now_res)
+            ResnetBlocWithAttn(pre_channel, pre_channel, noise_level_emb_dim=noise_level_channel, norm_groups=norm_groups,
+                               dropout=dropout, with_attn=True,size=now_res),
+            ResnetBlocWithAttn(pre_channel, pre_channel, noise_level_emb_dim=noise_level_channel, norm_groups=norm_groups,
+                               dropout=dropout, with_attn=False,size=now_res)
         ])
 
         ups = []
@@ -287,66 +234,164 @@ class UNet(nn.Module):
             is_last = (ind < 1)
             use_attn = (now_res in attn_res)
             channel_mult = inner_channel * channel_mults[ind]
-            for _ in range(0, res_blocks + 1):
+            for _ in range(0, res_blocks+1):
                 ups.append(ResnetBlocWithAttn(
-                    pre_channel + feat_channels.pop(), channel_mult, noise_level_emb_dim=noise_level_channel, 
-                    norm_groups=norm_groups, dropout=dropout, with_attn=use_attn, size=now_res))
+                    pre_channel+feat_channels.pop(), channel_mult, noise_level_emb_dim=noise_level_channel, norm_groups=norm_groups,
+                        dropout=dropout, with_attn=use_attn, size=now_res))
                 pre_channel = channel_mult
             if not is_last:
                 ups.append(Upsample(pre_channel))
-                now_res = now_res * 2
+                now_res = now_res*2
+
         self.ups = nn.ModuleList(ups)
 
         self.final_conv = Block(pre_channel, default(out_channel, in_channel), groups=norm_groups)
+
         
-        self.condition = CPEN(inchannel=condition_ch)
+        self.condition = CPEN(inchannel = condition_ch)  # canny+sar
         self.condition_ch = condition_ch
+        # self.c_func2 = nn.Linear(128, 128)   #128  256  512  1024
+        self.mi = 0
+
 
+
+
+        
     def forward(self, x, time, img_s1=None, class_label=None, return_condition=False, t_ori=0):
-        condition = x[:, :self.condition_ch, ...].clone()
+        # x torch.cat([x_in['SR'], x_noisy], dim=1)
+        condition = x[:, :self.condition_ch, ...].clone()  
         x = x[:, self.condition_ch:, ...]
         
+
         c1, c2, c3, c4, c5 = self.condition(condition)
         c_base = [c1, c2, c3, c4, c5]
         
+
+ 
+
+        
         c = []
         for i in range(len(c_base)):
             for _ in range(self.res_blocks):
-                c.append(c_base[i])
+                c.append(c_base[i])       
 
-        t = self.noise_level_mlp(time) if exists(self.noise_level_mlp) else None
+        t = self.noise_level_mlp(time) if exists(
+            self.noise_level_mlp) else None
 
+ 
+        
         feats = []
-        i = 0
+        i=0
         for layer in self.downs:
             if isinstance(layer, ResnetBlocWithAttn):
+                
                 x = layer(x, t, c[i])
-                i += 1
+                # print(x.shape)
+                i+=1
             else:
                 x = layer(x)
+
             feats.append(x)
+            
+            
 
         for layer in self.mid:
             if isinstance(layer, ResnetBlocWithAttn):
                 x = layer(x, t, c5)
+                # print(x.shape)
             else:
                 x = layer(x)
+            
 
+        
         c_base = [c5, c4, c3, c2, c1]
         c = []
         for i in range(len(c_base)):
-            for _ in range(self.res_blocks + 1):
-                c.append(c_base[i])
-        
+            for _ in range(self.res_blocks+1):
+                c.append(c_base[i])   
         i = 0
         for layer in self.ups:
             if isinstance(layer, ResnetBlocWithAttn):
+                # print(x.shape)
                 x = layer(torch.cat((x, feats.pop()), dim=1), t, c[i])
-                i += 1
+                # print(x.shape)
+                i+=1
             else:
                 x = layer(x)
-
+            
         if not return_condition:
             return self.final_conv(x)
         else:
             return self.final_conv(x), [c1, c2, c3, c4, c5]
+
+ 
+
+class ResBlock_normal(nn.Module):
+    def __init__(self, dim, dim_out, dropout=0, norm_groups=32):
+        super().__init__()
+        
+        self.block1 = Block(dim, dim_out, groups=norm_groups)
+        self.block2 = Block(dim_out, dim_out, groups=norm_groups, dropout=dropout)
+        self.res_conv = nn.Conv2d(
+            dim, dim_out, 1) if dim != dim_out else nn.Identity()
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        h = self.block1(x)
+        h = self.block2(h)
+        return h + self.res_conv(x)
+
+
+from SoftPool import soft_pool2d, SoftPool2d
+class CPEN(nn.Module):
+    def __init__(self, inchannel = 1):
+        super(CPEN, self).__init__()
+        self.pool = SoftPool2d(kernel_size=(2,2), stride=(2,2))
+        # self.scale=scale
+        # if scale == 2:
+
+        self.E1= nn.Sequential(nn.Conv2d(inchannel, 64, kernel_size=3, padding=1),
+                Swish())
+        
+         
+
+        self.E2=nn.Sequential(
+            ResBlock_normal(64, 128, dropout=0, norm_groups=16), 
+            ResBlock_normal(128, 128, dropout=0, norm_groups=16),             
+            )
+
+        self.E3=nn.Sequential(
+            ResBlock_normal(128, 256, dropout=0, norm_groups=16), 
+            ResBlock_normal(256, 256, dropout=0, norm_groups=16),             
+            )
+
+        self.E4=nn.Sequential(
+            ResBlock_normal(256, 512, dropout=0, norm_groups=16), 
+            ResBlock_normal(512, 512, dropout=0, norm_groups=16),             
+            )
+
+        self.E5=nn.Sequential(
+            ResBlock_normal(512, 512, dropout=0, norm_groups=16), 
+            ResBlock_normal(512, 1024, dropout=0, norm_groups=16),             
+            )
+
+ 
+
+    def forward(self, x):
+
+        x1 = self.E1(x)
+
+        x2 = self.pool(x1)
+        x2 = self.E2(x2)
+
+        x3 = self.pool(x2)
+        x3 = self.E3(x3)
+
+
+        x4 = self.pool(x3)
+        x4 = self.E4(x4)
+
+        x5 = self.pool(x4)
+        x5 = self.E5(x5)
+
+        return x1, x2, x3, x4, x5