| """LiquidInbetween: a small learned in-betweener for photos. |
| |
| The architecture is RIFE's IFNet (Huang et al., ECCV 2022, MIT licence, |
| github.com/hzwer/ECCV2022-RIFE), extended with a timestep input so one pass |
| gives the frame at any t in (0, 1), not only the middle. The extra channel |
| starts at zero, so training begins exactly at RIFE's behaviour at t = 0.5. |
| """ |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| def conv(cin, cout, k=3, s=1, p=1): |
| return nn.Sequential(nn.Conv2d(cin, cout, k, s, p, bias=True), nn.PReLU(cout)) |
|
|
|
|
| def base_grid(n, h, w, device, dtype): |
| gx = torch.linspace(-1.0, 1.0, w, device=device, dtype=dtype).view(1, 1, 1, w).expand(n, -1, h, -1) |
| gy = torch.linspace(-1.0, 1.0, h, device=device, dtype=dtype).view(1, 1, h, 1).expand(n, -1, -1, w) |
| return torch.cat([gx, gy], 1) |
|
|
|
|
| def warp(img, flow): |
| """Backward warp: each output pixel reads `img` at its position + `flow` (pixels).""" |
| n, _, h, w = flow.shape |
| f = torch.cat([flow[:, 0:1] / ((w - 1.0) / 2.0), flow[:, 1:2] / ((h - 1.0) / 2.0)], 1) |
| |
| g = (base_grid(n, h, w, flow.device, flow.dtype) + f).clamp(-1.0, 1.0).permute(0, 2, 3, 1) |
| return F.grid_sample(img, g, mode="bilinear", padding_mode="zeros", align_corners=True) |
|
|
|
|
| class IFBlock(nn.Module): |
| def __init__(self, in_planes, c=90): |
| super().__init__() |
| self.conv0 = nn.Sequential(conv(in_planes, c // 2, 3, 2, 1), conv(c // 2, c, 3, 2, 1)) |
| self.convblock0 = nn.Sequential(conv(c, c), conv(c, c)) |
| self.convblock1 = nn.Sequential(conv(c, c), conv(c, c)) |
| self.convblock2 = nn.Sequential(conv(c, c), conv(c, c)) |
| self.convblock3 = nn.Sequential(conv(c, c), conv(c, c)) |
| self.conv1 = nn.Sequential(nn.ConvTranspose2d(c, c // 2, 4, 2, 1), nn.PReLU(c // 2), nn.ConvTranspose2d(c // 2, 4, 4, 2, 1)) |
| self.conv2 = nn.Sequential(nn.ConvTranspose2d(c, c // 2, 4, 2, 1), nn.PReLU(c // 2), nn.ConvTranspose2d(c // 2, 1, 4, 2, 1)) |
|
|
| def forward(self, x, flow, scale): |
| x = F.interpolate(x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False) |
| flow = F.interpolate(flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False) / scale |
| feat = self.conv0(torch.cat((x, flow), 1)) |
| feat = self.convblock0(feat) + feat |
| feat = self.convblock1(feat) + feat |
| feat = self.convblock2(feat) + feat |
| feat = self.convblock3(feat) + feat |
| flow = F.interpolate(self.conv1(feat), scale_factor=scale, mode="bilinear", align_corners=False) * scale |
| mask = F.interpolate(self.conv2(feat), scale_factor=scale, mode="bilinear", align_corners=False) |
| return flow, mask |
|
|
|
|
| class IFNet(nn.Module): |
| """RIFE's IFNet. With `timestep`, every block also sees t (and the |
| mirrored pass 1 − t), so the frame at any t comes out in one pass.""" |
|
|
| def __init__(self, timestep=False, c=90): |
| super().__init__() |
| self.timestep = timestep |
| extra = 1 if timestep else 0 |
| self.block0 = IFBlock(7 + 4 + extra, c) |
| self.block1 = IFBlock(7 + 4 + extra, c) |
| self.block2 = IFBlock(7 + 4 + extra, c) |
|
|
| def forward(self, img0, img1, t=None, all_scales=False): |
| n, _, h, w = img0.shape |
| flow = torch.zeros(n, 4, h, w, device=img0.device, dtype=img0.dtype) |
| mask = torch.zeros(n, 1, h, w, device=img0.device, dtype=img0.dtype) |
| if self.timestep: |
| tt = (t if t is not None else torch.full((n, 1, 1, 1), 0.5, device=img0.device, dtype=img0.dtype)).expand(n, 1, h, w) |
| w0, w1 = img0, img1 |
| outputs = [] |
| for block, scale in ((self.block0, 4), (self.block1, 2), (self.block2, 1)): |
| a = [w0, w1, mask] + ([tt] if self.timestep else []) |
| b = [w1, w0, -mask] + ([1 - tt] if self.timestep else []) |
| f0, m0 = block(torch.cat(a, 1), flow, scale) |
| f1, m1 = block(torch.cat(b, 1), torch.cat((flow[:, 2:4], flow[:, :2]), 1), scale) |
| flow = flow + (f0 + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2 |
| mask = mask + (m0 - m1) / 2 |
| w0 = warp(img0, flow[:, :2]) |
| w1 = warp(img1, flow[:, 2:4]) |
| if all_scales: |
| m = torch.sigmoid(mask) |
| outputs.append(w0 * m + w1 * (1 - m)) |
| m = torch.sigmoid(mask) |
| out = w0 * m + w1 * (1 - m) |
| return outputs if all_scales else out |
|
|
|
|
| def load_teacher(path): |
| """RIFE's ECCV 2022 weights (safetensors) into IFNet().""" |
| from safetensors.torch import load_file |
| net = IFNet(timestep=False) |
| state = load_file(path) |
| state = {k.replace("module.", ""): v for k, v in state.items() if not k.startswith("block_tea")} |
| missing, unexpected = net.load_state_dict(state, strict=False) |
| assert not missing, missing |
| return net |
|
|
|
|
| def student_from_teacher(teacher): |
| """IFNet(timestep=True) starting exactly as the teacher: the new t channel's |
| weights are zero.""" |
| student = IFNet(timestep=True) |
| sd = student.state_dict() |
| for k, v in teacher.state_dict().items(): |
| if k.endswith("conv0.0.0.weight"): |
| new = torch.zeros_like(sd[k]) |
| |
| new[:, :7] = v[:, :7] |
| new[:, 8:] = v[:, 7:] |
| sd[k] = new |
| else: |
| sd[k] = v |
| student.load_state_dict(sd) |
| return student |
|
|
