diff --git a/modules/rife/__init__.py b/modules/rife/__init__.py index b6c712e0c..e26c7aafc 100644 --- a/modules/rife/__init__.py +++ b/modules/rife/__init__.py @@ -16,16 +16,21 @@ from modules import devices, shared, paths from modules.logger import log -model_url = 'https://github.com/vladmandic/rife/raw/main/model/flownet-v46.pkl' +# Practical-RIFE v4.25 weights (MIT). Default URL points at the upstream HolyWu mirror; +# can be swapped to a self-hosted mirror without any other code change. +model_url = 'https://github.com/HolyWu/vs-rife/releases/download/model/flownet_v4.25.pkl' model: RifeModel = None -def load(model_path: str = 'rife/flownet-v46.pkl'): +def load(model_path: str = 'rife/flownet_v4.25.pkl'): global model # pylint: disable=global-statement if model is None: from modules import modelloader model_dir = os.path.join(paths.models_path, 'RIFE') - model_path = modelloader.load_file_from_url(url=model_url, model_dir=model_dir, file_name='flownet-v46.pkl') + model_path = modelloader.load_file_from_url(url=model_url, model_dir=model_dir, file_name='flownet_v4.25.pkl') + legacy_path = os.path.join(model_dir, 'flownet-v46.pkl') + if os.path.exists(legacy_path): + log.info(f'Video interpolate: legacy v3.9 weights at "{legacy_path}" are no longer used and can be deleted') log.debug(f'Video interpolate: model="{model_path}"') model = RifeModel() model.load_model(model_path, -1) @@ -144,8 +149,8 @@ def interpolate_nchw(images: list, count: int = 2, scale: float = 1.0): I1 = f_pad(frame.unsqueeze(0)) for i in range(count-1): output = model.inference(I0, I1, (i+1) * 1. / (count), scale) - interpolated.append(output) - interpolated.append(I1) + interpolated.append(output[:, :, :h, :w]) + interpolated.append(I1[:, :, :h, :w]) pbar.update(1) t1 = time.time() diff --git a/modules/rife/model_ifnet.py b/modules/rife/model_ifnet.py index df32b59a1..06ab8ef2c 100644 --- a/modules/rife/model_ifnet.py +++ b/modules/rife/model_ifnet.py @@ -1,49 +1,60 @@ -import os -import sys import torch import torch.nn as nn import torch.nn.functional as F -sys.path.append(os.path.dirname(__file__)) -from warplayer import warp # pylint: disable=wrong-import-position - - -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +from modules.rife.warplayer import warp def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): return nn.Sequential( - nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, - padding=padding, dilation=dilation, bias=True), - nn.LeakyReLU(0.2, True) + nn.Conv2d( + in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, bias=True + ), + nn.LeakyReLU(0.2, True), ) -def conv_bn(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): - return nn.Sequential( - nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, - padding=padding, dilation=dilation, bias=False), - nn.BatchNorm2d(out_planes), - nn.LeakyReLU(0.2, True) - ) + +class Head(nn.Module): + def __init__(self): + super(Head, self).__init__() + self.cnn0 = nn.Conv2d(3, 16, 3, 2, 1) + self.cnn1 = nn.Conv2d(16, 16, 3, 1, 1) + self.cnn2 = nn.Conv2d(16, 16, 3, 1, 1) + self.cnn3 = nn.ConvTranspose2d(16, 4, 4, 2, 1) + self.relu = nn.LeakyReLU(0.2, True) + + def forward(self, x, feat=False): + x = x.clamp(0.0, 1.0) + x0 = self.cnn0(x) + x = self.relu(x0) + x1 = self.cnn1(x) + x = self.relu(x1) + x2 = self.cnn2(x) + x = self.relu(x2) + x3 = self.cnn3(x) + if feat: + return [x0, x1, x2, x3] + return x3 + class ResConv(nn.Module): def __init__(self, c, dilation=1): - super().__init__() - self.conv = nn.Conv2d(c, c, 3, 1, dilation, dilation=dilation, groups=1\ -) + super(ResConv, self).__init__() + self.conv = nn.Conv2d(c, c, 3, 1, dilation, dilation=dilation, groups=1) self.beta = nn.Parameter(torch.ones((1, c, 1, 1)), requires_grad=True) self.relu = nn.LeakyReLU(0.2, True) def forward(self, x): return self.relu(self.conv(x) * self.beta + x) + class IFBlock(nn.Module): def __init__(self, in_planes, c=64): - super().__init__() + super(IFBlock, self).__init__() self.conv0 = nn.Sequential( - conv(in_planes, c//2, 3, 2, 1), - conv(c//2, c, 3, 2, 1), - ) + conv(in_planes, c // 2, 3, 2, 1), + conv(c // 2, c, 3, 2, 1), + ) self.convblock = nn.Sequential( ResConv(c), ResConv(c), @@ -54,45 +65,39 @@ class IFBlock(nn.Module): ResConv(c), ResConv(c), ) - self.lastconv = nn.Sequential( - nn.ConvTranspose2d(c, 4*6, 4, 2, 1), - nn.PixelShuffle(2) - ) + self.lastconv = nn.Sequential(nn.ConvTranspose2d(c, 4 * 13, 4, 2, 1), nn.PixelShuffle(2)) def forward(self, x, flow=None, scale=1): - x = F.interpolate(x, scale_factor= 1. / scale, mode="bilinear", align_corners=False) + x = F.interpolate(x, scale_factor=1.0 / scale, mode="bilinear") if flow is not None: - flow = F.interpolate(flow, scale_factor= 1. / scale, mode="bilinear", align_corners=False) * 1. / scale + flow = F.interpolate(flow, scale_factor=1.0 / scale, mode="bilinear") / scale x = torch.cat((x, flow), 1) feat = self.conv0(x) feat = self.convblock(feat) tmp = self.lastconv(feat) - tmp = F.interpolate(tmp, scale_factor=scale, mode="bilinear", align_corners=False) + tmp = F.interpolate(tmp, scale_factor=scale, mode="bilinear") flow = tmp[:, :4] * scale mask = tmp[:, 4:5] - return flow, mask + feat = tmp[:, 5:] + return flow, mask, feat + class IFNet(nn.Module): - def __init__(self): - super().__init__() - self.block0 = IFBlock(7, c=192) - self.block1 = IFBlock(8+4, c=128) - self.block2 = IFBlock(8+4, c=96) - self.block3 = IFBlock(8+4, c=64) - # self.contextnet = Contextnet() - # self.unet = Unet() + def __init__(self, scale=1, ensemble=False): + super(IFNet, self).__init__() + self.block0 = IFBlock(7 + 8, c=192) + self.block1 = IFBlock(8 + 4 + 8 + 8, c=128) + self.block2 = IFBlock(8 + 4 + 8 + 8, c=96) + self.block3 = IFBlock(8 + 4 + 8 + 8, c=64) + self.block4 = IFBlock(8 + 4 + 8 + 8, c=32) + self.encode = Head() + self.scale_list = [16 / scale, 8 / scale, 4 / scale, 2 / scale, 1 / scale] + if ensemble: + raise ValueError("rife: ensemble is not supported in v4.25") - def forward( self, x, timestep=0.5, scale_list=None, training=False, fastmode=True, ensemble=False): # pylint: disable=dangerous-default-value, unused-argument - if scale_list is None: - scale_list = [8, 4, 2, 1] - if training is False: - channel = x.shape[1] // 2 - img0 = x[:, :channel] - img1 = x[:, channel:] - if not torch.is_tensor(timestep): - timestep = (x[:, :1].clone() * 0 + 1) * timestep - else: - timestep = timestep.repeat(1, 1, img0.shape[2], img0.shape[3]) + def forward(self, img0, img1, timestep, tenFlow_div, backwarp_tenGrid, f0, f1): + img0 = img0.clamp(0.0, 1.0) + img1 = img1.clamp(0.0, 1.0) flow_list = [] merged = [] mask_list = [] @@ -100,28 +105,26 @@ class IFNet(nn.Module): warped_img1 = img1 flow = None mask = None - # loss_cons = 0 - block = [self.block0, self.block1, self.block2, self.block3] - for i in range(4): + block = [self.block0, self.block1, self.block2, self.block3, self.block4] + for i in range(5): if flow is None: - flow, mask = block[i](torch.cat((img0[:, :3], img1[:, :3], timestep), 1), None, scale=scale_list[i]) - if ensemble: - f1, m1 = block[i](torch.cat((img1[:, :3], img0[:, :3], 1-timestep), 1), None, scale=scale_list[i]) - flow = (flow + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2 - mask = (mask + (-m1)) / 2 + flow, mask, feat = block[i]( + torch.cat((img0, img1, f0, f1, timestep), 1), None, scale=self.scale_list[i] + ) else: - f0, m0 = block[i](torch.cat((warped_img0[:, :3], warped_img1[:, :3], timestep, mask), 1), flow, scale=scale_list[i]) - if ensemble: - f1, m1 = block[i](torch.cat((warped_img1[:, :3], warped_img0[:, :3], 1-timestep, -mask), 1), torch.cat((flow[:, 2:4], flow[:, :2]), 1), scale=scale_list[i]) # pylint: disable=invalid-unary-operand-type - f0 = (f0 + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2 - m0 = (m0 + (-m1)) / 2 - flow = flow + f0 - mask = mask + m0 + wf0 = warp(f0, flow[:, :2], tenFlow_div, backwarp_tenGrid) + wf1 = warp(f1, flow[:, 2:4], tenFlow_div, backwarp_tenGrid) + fd, m0, feat = block[i]( + torch.cat((warped_img0, warped_img1, wf0, wf1, timestep, mask, feat), 1), + flow, + scale=self.scale_list[i], + ) + mask = m0 + flow = flow + fd mask_list.append(mask) flow_list.append(flow) - warped_img0 = warp(img0, flow[:, :2]) - warped_img1 = warp(img1, flow[:, 2:4]) + warped_img0 = warp(img0, flow[:, :2], tenFlow_div, backwarp_tenGrid) + warped_img1 = warp(img1, flow[:, 2:4], tenFlow_div, backwarp_tenGrid) merged.append((warped_img0, warped_img1)) - mask_list[3] = torch.sigmoid(mask_list[3]) - merged[3] = merged[3][0] * mask_list[3] + merged[3][1] * (1 - mask_list[3]) - return flow_list, mask_list[3], merged + mask = torch.sigmoid(mask) + return warped_img0 * mask + warped_img1 * (1 - mask) diff --git a/modules/rife/model_rife.py b/modules/rife/model_rife.py index 52d359923..ee3aa1ecf 100644 --- a/modules/rife/model_rife.py +++ b/modules/rife/model_rife.py @@ -1,8 +1,5 @@ import torch -from torch.optim import AdamW -from torch.nn.parallel import DistributedDataParallel as DDP from modules.rife.model_ifnet import IFNet -from modules.rife.loss import EPE, SOBEL from modules import devices @@ -10,12 +7,11 @@ class RifeModel: def __init__(self, local_rank=-1): self.flownet = IFNet() self.device() - self.optimG = AdamW(self.flownet.parameters(), lr=1e-6, weight_decay=1e-4) - self.epe = EPE() - self.version = 3.9 - # self.vgg = VGGPerceptualLoss().to(device) - self.sobel = SOBEL() + self.version = 4.25 + self.tenFlow_div_cache = {} + self.backwarp_tenGrid_cache = {} if local_rank != -1: + from torch.nn.parallel import DistributedDataParallel as DDP self.flownet = DDP(self.flownet, device_ids=[local_rank], output_device=local_rank) def train(self): @@ -26,7 +22,7 @@ class RifeModel: def device(self): self.flownet.to(devices.device) - self.flownet.to(devices.dtype) + self.flownet.to(torch.float32) # bfloat16 produces visible checkerboard artifacts at the new IFNet's depth def load_model(self, model_file, rank=0): def convert(param): @@ -44,36 +40,30 @@ class RifeModel: if rank == 0: torch.save(self.flownet.state_dict(), model_file) - def inference(self, img0, img1, timestep=0.5, scale=1.0): - imgs = torch.cat((img0, img1), 1) - scale_list = [8/scale, 4/scale, 2/scale, 1/scale] - _flow, _mask, merged = self.flownet(imgs, timestep, scale_list) - return merged[3] + def grid_for(self, h, w, device, dtype): + key = (h, w, str(device), str(dtype)) + grid = self.backwarp_tenGrid_cache.get(key) + if grid is None: + tenHorizontal = torch.linspace(-1.0, 1.0, w, device=device, dtype=dtype).view(1, 1, 1, w).expand(1, -1, h, -1) + tenVertical = torch.linspace(-1.0, 1.0, h, device=device, dtype=dtype).view(1, 1, h, 1).expand(1, -1, -1, w) + grid = torch.cat([tenHorizontal, tenVertical], 1) + self.backwarp_tenGrid_cache[key] = grid + div = self.tenFlow_div_cache.get(key) + if div is None: + div = torch.tensor([(w - 1.0) / 2.0, (h - 1.0) / 2.0], device=device, dtype=dtype) + self.tenFlow_div_cache[key] = div + return grid, div - def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None): # pylint: disable=unused-argument - for param_group in self.optimG.param_groups: - param_group['lr'] = learning_rate - # img0 = imgs[:, :3] - # img1 = imgs[:, 3:] - if training: - self.train() - else: - self.eval() - scale = [8, 4, 2, 1] - flow, mask, merged = self.flownet(torch.cat((imgs, gt), 1), scale=scale, training=training) - loss_l1 = (merged[3] - gt).abs().mean() - loss_smooth = self.sobel(flow[3], flow[3]*0).mean() - # loss_vgg = self.vgg(merged[2], gt) - if training: - self.optimG.zero_grad() - loss_G = loss_l1 + loss_smooth * 0.1 - loss_G.backward() - self.optimG.step() - # else: - # flow_teacher = flow[2] - return merged[3], { - 'mask': mask, - 'flow': flow[3][:, :2], - 'loss_l1': loss_l1, - 'loss_smooth': loss_smooth, - } + def inference(self, img0, img1, timestep=0.5, scale=1.0): + in_dtype = img0.dtype + img0 = img0.float() + img1 = img1.float() + n, _c, h, w = img0.shape + device = img0.device + backwarp_tenGrid, tenFlow_div = self.grid_for(h, w, device, torch.float32) + self.flownet.scale_list = [16 / scale, 8 / scale, 4 / scale, 2 / scale, 1 / scale] + f0 = self.flownet.encode(img0) + f1 = self.flownet.encode(img1) + timestep_t = torch.full((n, 1, h, w), timestep, device=device, dtype=torch.float32) + out = self.flownet(img0, img1, timestep_t, tenFlow_div, backwarp_tenGrid, f0, f1) + return out.to(in_dtype) diff --git a/modules/rife/warplayer.py b/modules/rife/warplayer.py index abcd6aa5c..a1427e5ba 100644 --- a/modules/rife/warplayer.py +++ b/modules/rife/warplayer.py @@ -1,17 +1,12 @@ import torch -from modules import devices +import torch.nn.functional as F -backwarp_tenGrid = {} +def warp(tenInput, tenFlow, tenFlow_div, backwarp_tenGrid): + dtype = tenInput.dtype + tenInput = tenInput.to(torch.float) + tenFlow = tenFlow.to(torch.float) - -def warp(tenInput, tenFlow): - k = (str(tenFlow.device), str(tenFlow.size())) - if k not in backwarp_tenGrid: - tenHorizontal = torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=devices.device).view(1, 1, 1, tenFlow.shape[3]).expand(tenFlow.shape[0], -1, tenFlow.shape[2], -1) - tenVertical = torch.linspace(-1.0, 1.0, tenFlow.shape[2], device=devices.device).view(1, 1, tenFlow.shape[2], 1).expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3]) - backwarp_tenGrid[k] = torch.cat([tenHorizontal, tenVertical], 1).to(devices.device) - tenFlow = torch.cat([tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0), - tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0)], 1) - grid = (backwarp_tenGrid[k] + tenFlow).permute(0, 2, 3, 1).to(devices.dtype) - return torch.nn.functional.grid_sample(input=tenInput, grid=grid, mode='bilinear', padding_mode='border', align_corners=True) + tenFlow = torch.cat([tenFlow[:, 0:1] / tenFlow_div[0], tenFlow[:, 1:2] / tenFlow_div[1]], 1) + g = (backwarp_tenGrid + tenFlow).permute(0, 2, 3, 1) + return F.grid_sample(input=tenInput, grid=g, mode="bilinear", padding_mode="border", align_corners=True).to(dtype)