| 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) |
| |
| |
| self.down1 = Conv3DBlock(in_ch, 64, 256) |
| self.down2 = Conv3DBlock(64, 128, 256) |
| self.down3 = Conv3DBlock(128, 256, 256) |
| |
| |
| 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 |
| |
| |
| 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) |
| |
| |
| 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 |
