Spaces:
Paused
Paused
| import torch | |
| import torch.nn as nn | |
| class Image2DResBlockWithTV(nn.Module): | |
| def __init__(self, dim, tdim, vdim): | |
| super().__init__() | |
| norm = lambda c: nn.GroupNorm(8, c) | |
| self.time_embed = nn.Conv2d(tdim, dim, 1, 1) | |
| self.view_embed = nn.Conv2d(vdim, dim, 1, 1) | |
| self.conv = nn.Sequential( | |
| norm(dim), | |
| nn.SiLU(True), | |
| nn.Conv2d(dim, dim, 3, 1, 1), | |
| norm(dim), | |
| nn.SiLU(True), | |
| nn.Conv2d(dim, dim, 3, 1, 1), | |
| ) | |
| def forward(self, x, t, v): | |
| return x+self.conv(x+self.time_embed(t)+self.view_embed(v)) | |
| class NoisyTargetViewEncoder(nn.Module): | |
| def __init__(self, time_embed_dim, viewpoint_dim, run_dim=16, output_dim=8): | |
| super().__init__() | |
| self.init_conv = nn.Conv2d(4, run_dim, 3, 1, 1) | |
| self.out_conv0 = Image2DResBlockWithTV(run_dim, time_embed_dim, viewpoint_dim) | |
| self.out_conv1 = Image2DResBlockWithTV(run_dim, time_embed_dim, viewpoint_dim) | |
| self.out_conv2 = Image2DResBlockWithTV(run_dim, time_embed_dim, viewpoint_dim) | |
| self.final_out = nn.Sequential( | |
| nn.GroupNorm(8, run_dim), | |
| nn.SiLU(True), | |
| nn.Conv2d(run_dim, output_dim, 3, 1, 1) | |
| ) | |
| def forward(self, x, t, v): | |
| B, DT = t.shape | |
| t = t.view(B, DT, 1, 1) | |
| B, DV = v.shape | |
| v = v.view(B, DV, 1, 1) | |
| x = self.init_conv(x) | |
| x = self.out_conv0(x, t, v) | |
| x = self.out_conv1(x, t, v) | |
| x = self.out_conv2(x, t, v) | |
| x = self.final_out(x) | |
| return x | |
| class SpatialUpTimeBlock(nn.Module): | |
| def __init__(self, x_in_dim, t_in_dim, out_dim): | |
| super().__init__() | |
| norm_act = lambda c: nn.GroupNorm(8, c) | |
| self.t_conv = nn.Conv3d(t_in_dim, x_in_dim, 1, 1) # 16 | |
| self.norm = norm_act(x_in_dim) | |
| self.silu = nn.SiLU(True) | |
| self.conv = nn.ConvTranspose3d(x_in_dim, out_dim, kernel_size=3, padding=1, output_padding=1, stride=2) | |
| def forward(self, x, t): | |
| x = x + self.t_conv(t) | |
| return self.conv(self.silu(self.norm(x))) | |
| class SpatialTimeBlock(nn.Module): | |
| def __init__(self, x_in_dim, t_in_dim, out_dim, stride): | |
| super().__init__() | |
| norm_act = lambda c: nn.GroupNorm(8, c) | |
| self.t_conv = nn.Conv3d(t_in_dim, x_in_dim, 1, 1) # 16 | |
| self.bn = norm_act(x_in_dim) | |
| self.silu = nn.SiLU(True) | |
| self.conv = nn.Conv3d(x_in_dim, out_dim, 3, stride=stride, padding=1) | |
| def forward(self, x, t): | |
| x = x + self.t_conv(t) | |
| return self.conv(self.silu(self.bn(x))) | |
| class SpatialTime3DNet(nn.Module): | |
| def __init__(self, time_dim=256, input_dim=128, dims=(32, 64, 128, 256)): | |
| super().__init__() | |
| d0, d1, d2, d3 = dims | |
| dt = time_dim | |
| self.init_conv = nn.Conv3d(input_dim, d0, 3, 1, 1) # 32 | |
| self.conv0 = SpatialTimeBlock(d0, dt, d0, stride=1) | |
| self.conv1 = SpatialTimeBlock(d0, dt, d1, stride=2) | |
| self.conv2_0 = SpatialTimeBlock(d1, dt, d1, stride=1) | |
| self.conv2_1 = SpatialTimeBlock(d1, dt, d1, stride=1) | |
| self.conv3 = SpatialTimeBlock(d1, dt, d2, stride=2) | |
| self.conv4_0 = SpatialTimeBlock(d2, dt, d2, stride=1) | |
| self.conv4_1 = SpatialTimeBlock(d2, dt, d2, stride=1) | |
| self.conv5 = SpatialTimeBlock(d2, dt, d3, stride=2) | |
| self.conv6_0 = SpatialTimeBlock(d3, dt, d3, stride=1) | |
| self.conv6_1 = SpatialTimeBlock(d3, dt, d3, stride=1) | |
| self.conv7 = SpatialUpTimeBlock(d3, dt, d2) | |
| self.conv8 = SpatialUpTimeBlock(d2, dt, d1) | |
| self.conv9 = SpatialUpTimeBlock(d1, dt, d0) | |
| def forward(self, x, t): | |
| B, C = t.shape | |
| t = t.view(B, C, 1, 1, 1) | |
| x = self.init_conv(x) | |
| conv0 = self.conv0(x, t) | |
| x = self.conv1(conv0, t) | |
| x = self.conv2_0(x, t) | |
| conv2 = self.conv2_1(x, t) | |
| x = self.conv3(conv2, t) | |
| x = self.conv4_0(x, t) | |
| conv4 = self.conv4_1(x, t) | |
| x = self.conv5(conv4, t) | |
| x = self.conv6_0(x, t) | |
| x = self.conv6_1(x, t) | |
| x = conv4 + self.conv7(x, t) | |
| x = conv2 + self.conv8(x, t) | |
| x = conv0 + self.conv9(x, t) | |
| return x | |
| class FrustumTVBlock(nn.Module): | |
| def __init__(self, x_dim, t_dim, v_dim, out_dim, stride): | |
| super().__init__() | |
| norm_act = lambda c: nn.GroupNorm(8, c) | |
| self.t_conv = nn.Conv3d(t_dim, x_dim, 1, 1) # 16 | |
| self.v_conv = nn.Conv3d(v_dim, x_dim, 1, 1) # 16 | |
| self.bn = norm_act(x_dim) | |
| self.silu = nn.SiLU(True) | |
| self.conv = nn.Conv3d(x_dim, out_dim, 3, stride=stride, padding=1) | |
| def forward(self, x, t, v): | |
| x = x + self.t_conv(t) + self.v_conv(v) | |
| return self.conv(self.silu(self.bn(x))) | |
| class FrustumTVUpBlock(nn.Module): | |
| def __init__(self, x_dim, t_dim, v_dim, out_dim): | |
| super().__init__() | |
| norm_act = lambda c: nn.GroupNorm(8, c) | |
| self.t_conv = nn.Conv3d(t_dim, x_dim, 1, 1) # 16 | |
| self.v_conv = nn.Conv3d(v_dim, x_dim, 1, 1) # 16 | |
| self.norm = norm_act(x_dim) | |
| self.silu = nn.SiLU(True) | |
| self.conv = nn.ConvTranspose3d(x_dim, out_dim, kernel_size=3, padding=1, output_padding=1, stride=2) | |
| def forward(self, x, t, v): | |
| x = x + self.t_conv(t) + self.v_conv(v) | |
| return self.conv(self.silu(self.norm(x))) | |
| class FrustumTV3DNet(nn.Module): | |
| def __init__(self, in_dim, t_dim, v_dim, dims=(32, 64, 128, 256)): | |
| super().__init__() | |
| self.conv0 = nn.Conv3d(in_dim, dims[0], 3, 1, 1) # 32 | |
| self.conv1 = FrustumTVBlock(dims[0], t_dim, v_dim, dims[1], 2) | |
| self.conv2 = FrustumTVBlock(dims[1], t_dim, v_dim, dims[1], 1) | |
| self.conv3 = FrustumTVBlock(dims[1], t_dim, v_dim, dims[2], 2) | |
| self.conv4 = FrustumTVBlock(dims[2], t_dim, v_dim, dims[2], 1) | |
| self.conv5 = FrustumTVBlock(dims[2], t_dim, v_dim, dims[3], 2) | |
| self.conv6 = FrustumTVBlock(dims[3], t_dim, v_dim, dims[3], 1) | |
| self.up0 = FrustumTVUpBlock(dims[3], t_dim, v_dim, dims[2]) | |
| self.up1 = FrustumTVUpBlock(dims[2], t_dim, v_dim, dims[1]) | |
| self.up2 = FrustumTVUpBlock(dims[1], t_dim, v_dim, dims[0]) | |
| def forward(self, x, t, v): | |
| B,DT = t.shape | |
| t = t.view(B,DT,1,1,1) | |
| B,DV = v.shape | |
| v = v.view(B,DV,1,1,1) | |
| b, _, d, h, w = x.shape | |
| x0 = self.conv0(x) | |
| x1 = self.conv2(self.conv1(x0, t, v), t, v) | |
| x2 = self.conv4(self.conv3(x1, t, v), t, v) | |
| x3 = self.conv6(self.conv5(x2, t, v), t, v) | |
| x2 = self.up0(x3, t, v) + x2 | |
| x1 = self.up1(x2, t, v) + x1 | |
| x0 = self.up2(x1, t, v) + x0 | |
| return {w: x0, w//2: x1, w//4: x2, w//8: x3} | |