Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-07 13:30:46 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(rife): upgrade vendored RIFE to Practical-RIFE v4.25

Browse Source 
- Vendor IFNet_HDv3 v4.25 (5 IFBlocks, Head encoder, feat channel) and
  v4 warplayer with explicit (tenFlow_div, backwarp_tenGrid) signature
- Rewrite RifeModel.inference for the new forward signature with
  per-(H,W,device,dtype) caching of tenFlow_div and backwarp_tenGrid
- Force fp32 inference: bf16 produced visible checkerboard at the new
  IFNet's depth (was hidden by v3.9's shallower architecture)
- Crop padded frames in interpolate_nchw before output (was missing,
  produced gray bar on non-128-aligned inputs)
- Drop training scaffolding (AdamW, EPE/SOBEL, update method)
- Log obsolete legacy v3.9 weights file on first v4.25 load instead of
  silently deleting user data
- Default download URL is HolyWu vs-rife mirror (MIT, byte-identical
  upstream weights); swap to project-hosted URL before merge
This commit is contained in:
CalamitousFelicitousness
2026-04-25 21:14:18 +01:00
parent ce92c574c7
commit dc4c58d0cb
4 changed files with 123 additions and 130 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/rife/__init__.py
+10 -5
View File
@@ -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()
modules/rife/model_ifnet.py
+74 -71
View File
@@ -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)
modules/rife/model_rife.py
+31 -41
View File
@@ -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)
modules/rife/warplayer.py
+8 -13
View File
@@ -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)