Upload model/archs/unet.py with huggingface_hub

model/archs/unet.py

@@ -1,53 +1,53 @@
-'''
-Codes are from:
-https://github.com/jaxony/unet-pytorch/blob/master/model.py
-'''
-
-import torch
-import torch.nn as nn
-from diffusers import UNet2DModel
-import einops
-class UNetPP(nn.Module):
-    '''
-        Wrapper for UNet in diffusers
-    '''
-    def __init__(self, in_channels):
-        super(UNetPP, self).__init__()
-        self.in_channels = in_channels
-        self.unet = UNet2DModel(
-                sample_size=[256, 256*3],
-                in_channels=in_channels,
-                out_channels=32,
-                layers_per_block=2,
-                block_out_channels=(64, 128, 128, 128*2, 128*2, 128*4, 128*4),
-                down_block_types=(
-                    "DownBlock2D",
-                    "DownBlock2D",
-                    "DownBlock2D",
-                    "AttnDownBlock2D",
-                    "AttnDownBlock2D",
-                    "AttnDownBlock2D",
-                    "DownBlock2D",
-                ),
-                up_block_types=(
-                    "UpBlock2D",
-                    "AttnUpBlock2D",
-                    "AttnUpBlock2D",
-                    "AttnUpBlock2D",
-                    "UpBlock2D",
-                    "UpBlock2D",
-                    "UpBlock2D",
-                ),
-            )
-             
-        # self.unet.enable_xformers_memory_efficient_attention()    
-        if in_channels > 12:
-            self.learned_plane = torch.nn.parameter.Parameter(torch.zeros([1,in_channels-12,256,256*3]))
-
-    def forward(self, x, t=256):
-        learned_plane = self.learned_plane
-        if x.shape[1] < self.in_channels:
-            learned_plane = einops.repeat(learned_plane, '1 C H W -> B C H W', B=x.shape[0]).to(x.device)
-            x = torch.cat([x, learned_plane], dim = 1)
-        return self.unet(x, t).sample
-
+'''
+Codes are from:
+https://github.com/jaxony/unet-pytorch/blob/master/model.py
+'''
+
+import torch
+import torch.nn as nn
+from diffusers import UNet2DModel
+import einops
+class UNetPP(nn.Module):
+    '''
+        Wrapper for UNet in diffusers
+    '''
+    def __init__(self, in_channels):
+        super(UNetPP, self).__init__()
+        self.in_channels = in_channels
+        self.unet = UNet2DModel(
+                sample_size=[256, 256*3],
+                in_channels=in_channels,
+                out_channels=32,
+                layers_per_block=2,
+                block_out_channels=(64, 128, 128, 128*2, 128*2, 128*4, 128*4),
+                down_block_types=(
+                    "DownBlock2D",
+                    "DownBlock2D",
+                    "DownBlock2D",
+                    "AttnDownBlock2D",
+                    "AttnDownBlock2D",
+                    "AttnDownBlock2D",
+                    "DownBlock2D",
+                ),
+                up_block_types=(
+                    "UpBlock2D",
+                    "AttnUpBlock2D",
+                    "AttnUpBlock2D",
+                    "AttnUpBlock2D",
+                    "UpBlock2D",
+                    "UpBlock2D",
+                    "UpBlock2D",
+                ),
+            )
+             
+        self.unet.enable_xformers_memory_efficient_attention()    
+        if in_channels > 12:
+            self.learned_plane = torch.nn.parameter.Parameter(torch.zeros([1,in_channels-12,256,256*3]))
+
+    def forward(self, x, t=256):
+        learned_plane = self.learned_plane
+        if x.shape[1] < self.in_channels:
+            learned_plane = einops.repeat(learned_plane, '1 C H W -> B C H W', B=x.shape[0]).to(x.device)
+            x = torch.cat([x, learned_plane], dim = 1)
+        return self.unet(x, t).sample
+