File size: 1,988 Bytes
b0da3df |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 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 |