Create model.py

model.py

@@ -0,0 +1,60 @@
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+class TimeEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.proj = nn.Sequential(
+            nn.Linear(1, dim),
+            nn.SiLU(),
+            nn.Linear(dim, dim)
+        )
+        
+    def forward(self, t):
+        return self.proj(t)
+
+class Conv3DBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, time_dim):
+        super().__init__()
+        self.time_mlp = nn.Linear(time_dim, out_ch)
+        self.conv = nn.Conv3d(in_ch, out_ch, kernel_size=3, padding=1)
+        self.norm = nn.BatchNorm3d(out_ch)
+        
+    def forward(self, x, t_emb):
+        t_emb = self.time_mlp(t_emb).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+        return F.silu(self.norm(self.conv(x) + t_emb))
+
+class UNet3D(nn.Module):
+    def __init__(self, in_ch=3, out_ch=3, text_dim=768):
+        super().__init__()
+        self.time_embed = TimeEmbedding(256)
+        self.text_proj = nn.Linear(text_dim, 256)
+        
+        # Downsample
+        self.down1 = Conv3DBlock(in_ch, 64, 256)
+        self.down2 = Conv3DBlock(64, 128, 256)
+        self.down3 = Conv3DBlock(128, 256, 256)
+        
+        # Upsample
+        self.up1 = Conv3DBlock(256 + 128, 128, 256)
+        self.up2 = Conv3DBlock(128 + 64, 64, 256)
+        self.up3 = nn.Conv3d(64, out_ch, kernel_size=3, padding=1)
+        
+    def forward(self, x, t, text_emb):
+        t_emb = self.time_embed(t)
+        text_emb = self.text_proj(text_emb)
+        c_emb = t_emb + text_emb
+        
+        # Downsample
+        x1 = self.down1(x, c_emb)
+        x2 = self.down2(F.max_pool3d(x1, 2), c_emb)
+        x3 = self.down3(F.max_pool3d(x2, 2), c_emb)
+        
+        # Upsample
+        x = F.interpolate(x3, scale_factor=2)
+        x = self.up1(torch.cat([x, x2], dim=1), c_emb)
+        x = F.interpolate(x, scale_factor=2)
+        x = self.up2(torch.cat([x, x1], dim=1), c_emb)
+        x = self.up3(x)
+        return x