diff --git a/CHANGELOG.md b/CHANGELOG.md index a83202e94..a5b138854 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ # Change Log for SD.Next -## Update for 2023-11-29 +## Update for 2023-11-30 - **Diffusers** - **HDR latent control**, based on [article](https://huggingface.co/blog/TimothyAlexisVass/explaining-the-sdxl-latent-space#long-prompts-at-high-guidance-scales-becoming-possible) @@ -38,13 +38,13 @@ - applies to any model that supports video generation, e.g. AnimateDiff and StableVideoDiffusion - support for GIF and MP4 - output folder for videos is in *settings -> image paths -> video* -- **Model merge** - - add **SD-XL ReBasin** support, thanks @AI-Casanova - **General** - - further UI optimizations for **mobile devices**, thanks @iDeNoh + - **model merge** add **SD-XL ReBasin** support, thanks @AI-Casanova + - further UI optimizations for **mobile devices**, thanks @iDeNoh - log level defaults to info for console and debug for log file - better prompt display in process tab - increase maximum lora cache values + - fix extra networks sorting - fix controlnet compatibility issues in original backend - fix img2img/inpaint paste params - fix save text file for manually saved images diff --git a/modules/processing.py b/modules/processing.py index 1131fad85..01f1dd3af 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -122,7 +122,6 @@ class StableDiffusionProcessing: The first set of paramaters: sd_models -> do_not_reload_embeddings represent the minimum required to create a StableDiffusionProcessing """ def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt: str = "", styles: List[str] = None, seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1, seed_enable_extras: bool = True, sampler_name: str = None, latent_sampler: str = None, batch_size: int = 1, n_iter: int = 1, steps: int = 50, cfg_scale: float = 7.0, image_cfg_scale: float = None, clip_skip: int = 1, width: int = 512, height: int = 512, full_quality: bool = True, restore_faces: bool = False, tiling: bool = False, do_not_save_samples: bool = False, do_not_save_grid: bool = False, extra_generation_params: Dict[Any, Any] = None, overlay_images: Any = None, negative_prompt: str = None, eta: float = None, do_not_reload_embeddings: bool = False, denoising_strength: float = 0, diffusers_guidance_rescale: float = 0.7, hdr_clamp: bool = False, hdr_boundary: float = 4.0, hdr_threshold: float = 3.5, hdr_center: bool = False, hdr_channel_shift: float = 0.8, hdr_full_shift: float = 0.8, hdr_maximize: bool = False, hdr_max_center: float = 0.6, hdr_max_boundry: float = 1.0, override_settings: Dict[str, Any] = None, override_settings_restore_afterwards: bool = True, sampler_index: int = None, script_args: list = None): # pylint: disable=unused-argument - self.outpath_samples: str = outpath_samples self.outpath_grids: str = outpath_grids self.prompt: str = prompt @@ -166,8 +165,8 @@ class StableDiffusionProcessing: self.disable_extra_networks = False self.token_merging_ratio = 0 self.token_merging_ratio_hr = 0 - self.scripts = None - self.script_args = script_args or [] + # self.scripts = modules.scripts.ScriptRunner() # set via property + # self.script_args = script_args or [] # set via property self.per_script_args = {} self.all_prompts = None self.all_negative_prompts = None @@ -1070,6 +1069,8 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing): self.refiner_prompt = refiner_prompt self.refiner_negative = refiner_negative self.sampler = None + self.scripts = None + self.script_args = [] def init(self, all_prompts, all_seeds, all_subseeds): if shared.backend == shared.Backend.DIFFUSERS: @@ -1232,6 +1233,8 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing): self.is_batch = False self.scale_by = 1.0 self.sampler = None + self.scripts = None + self.script_args = [] def init(self, all_prompts, all_seeds, all_subseeds): if shared.backend == shared.Backend.DIFFUSERS and self.image_mask is not None: diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 995c84183..2049fa881 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -82,7 +82,6 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro raise AssertionError('Interrupted...') time.sleep(0.1) if kwargs.get('latents', None) is None: - print('HERE NO') return kwargs kwargs = correction_callback(p, timestep, kwargs) shared.state.current_latent = kwargs['latents'] diff --git a/modules/rife/loss.py b/modules/rife/loss.py new file mode 100644 index 000000000..993f319fb --- /dev/null +++ b/modules/rife/loss.py @@ -0,0 +1,122 @@ +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision.models as models +from modules import devices + + +class EPE(nn.Module): + def __init__(self): + super(EPE, self).__init__() + + def forward(self, flow, gt, loss_mask): + loss_map = (flow - gt.detach()) ** 2 + loss_map = (loss_map.sum(1, True) + 1e-6) ** 0.5 + return loss_map * loss_mask + + +class Ternary(nn.Module): + def __init__(self): + super(Ternary, self).__init__() + patch_size = 7 + out_channels = patch_size * patch_size + self.w = np.eye(out_channels).reshape( + (patch_size, patch_size, 1, out_channels)) + self.w = np.transpose(self.w, (3, 2, 0, 1)) + self.w = torch.tensor(self.w).float().to(devices.device) + + def transform(self, img): + patches = F.conv2d(img, self.w, padding=3, bias=None) + transf = patches - img + transf_norm = transf / torch.sqrt(0.81 + transf**2) + return transf_norm + + def rgb2gray(self, rgb): + r, g, b = rgb[:, 0:1, :, :], rgb[:, 1:2, :, :], rgb[:, 2:3, :, :] + gray = 0.2989 * r + 0.5870 * g + 0.1140 * b + return gray + + def hamming(self, t1, t2): + dist = (t1 - t2) ** 2 + dist_norm = torch.mean(dist / (0.1 + dist), 1, True) + return dist_norm + + def valid_mask(self, t, padding): + n, _, h, w = t.size() + inner = torch.ones(n, 1, h - 2 * padding, w - 2 * padding).type_as(t) + mask = F.pad(inner, [padding] * 4) + return mask + + def forward(self, img0, img1): + img0 = self.transform(self.rgb2gray(img0)) + img1 = self.transform(self.rgb2gray(img1)) + return self.hamming(img0, img1) * self.valid_mask(img0, 1) + + +class SOBEL(nn.Module): + def __init__(self): + super(SOBEL, self).__init__() + self.kernelX = torch.tensor([ + [1, 0, -1], + [2, 0, -2], + [1, 0, -1], + ]).float() + self.kernelY = self.kernelX.clone().T + self.kernelX = self.kernelX.unsqueeze(0).unsqueeze(0).to(devices.device) + self.kernelY = self.kernelY.unsqueeze(0).unsqueeze(0).to(devices.device) + + def forward(self, pred, gt): + N, C, H, W = pred.shape[0], pred.shape[1], pred.shape[2], pred.shape[3] + img_stack = torch.cat( + [pred.reshape(N*C, 1, H, W), gt.reshape(N*C, 1, H, W)], 0) + sobel_stack_x = F.conv2d(img_stack, self.kernelX, padding=1) + sobel_stack_y = F.conv2d(img_stack, self.kernelY, padding=1) + pred_X, gt_X = sobel_stack_x[:N*C], sobel_stack_x[N*C:] + pred_Y, gt_Y = sobel_stack_y[:N*C], sobel_stack_y[N*C:] + L1X, L1Y = torch.abs(pred_X-gt_X), torch.abs(pred_Y-gt_Y) + loss = L1X+L1Y + return loss + + +class MeanShift(nn.Conv2d): + def __init__(self, data_mean, data_std, data_range=1, norm=True): + c = len(data_mean) + super(MeanShift, self).__init__(c, c, kernel_size=1) + std = torch.Tensor(data_std) + self.weight.data = torch.eye(c).view(c, c, 1, 1) + if norm: + self.weight.data.div_(std.view(c, 1, 1, 1)) + self.bias.data = -1 * data_range * torch.Tensor(data_mean) + self.bias.data.div_(std) + else: + self.weight.data.mul_(std.view(c, 1, 1, 1)) + self.bias.data = data_range * torch.Tensor(data_mean) + self.requires_grad = False + + +class VGGPerceptualLoss(torch.nn.Module): + def __init__(self, rank=0): # pylint: disable=unused-argument + super(VGGPerceptualLoss, self).__init__() + pretrained = True + self.vgg_pretrained_features = models.vgg19( + pretrained=pretrained).features + self.normalize = MeanShift([0.485, 0.456, 0.406], [ + 0.229, 0.224, 0.225], norm=True).cuda() + for param in self.parameters(): + param.requires_grad = False + + def forward(self, X, Y, indices=None): + X = self.normalize(X) + Y = self.normalize(Y) + indices = [2, 7, 12, 21, 30] + weights = [1.0/2.6, 1.0/4.8, 1.0/3.7, 1.0/5.6, 10/1.5] + k = 0 + loss = 0 + for i in range(indices[-1]): + X = self.vgg_pretrained_features[i](X) + Y = self.vgg_pretrained_features[i](Y) + if i+1 in indices: + loss += weights[k] * (X - Y.detach()).abs().mean() * 0.1 + k += 1 + return loss diff --git a/modules/rife/model_ifnet.py b/modules/rife/model_ifnet.py new file mode 100644 index 000000000..91598bde8 --- /dev/null +++ b/modules/rife/model_ifnet.py @@ -0,0 +1,134 @@ +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") + + +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) + ) + +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 ResConv(nn.Module): + def __init__(self, c, dilation=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(IFBlock, self).__init__() + self.conv0 = nn.Sequential( + conv(in_planes, c//2, 3, 2, 1), + conv(c//2, c, 3, 2, 1), + ) + self.convblock = nn.Sequential( + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ) + self.lastconv = nn.Sequential( + nn.ConvTranspose2d(c, 4*6, 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) + if flow is not None: + flow = F.interpolate(flow, scale_factor= 1. / scale, mode="bilinear", align_corners=False) * 1. / 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) + flow = tmp[:, :4] * scale + mask = tmp[:, 4:5] + return flow, mask + +class IFNet(nn.Module): + def __init__(self): + super(IFNet, self).__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 forward( self, x, timestep=0.5, scale_list=[8, 4, 2, 1], training=False, fastmode=True, ensemble=False): + 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]) + flow_list = [] + merged = [] + mask_list = [] + warped_img0 = img0 + warped_img1 = img1 + flow = None + mask = None + # loss_cons = 0 + block = [self.block0, self.block1, self.block2, self.block3] + for i in range(4): + 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 + 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 + mask_list.append(mask) + flow_list.append(flow) + warped_img0 = warp(img0, flow[:, :2]) + warped_img1 = warp(img1, flow[:, 2:4]) + 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]) + if not fastmode: + print('contextnet is removed') + ''' + c0 = self.contextnet(img0, flow[:, :2]) + c1 = self.contextnet(img1, flow[:, 2:4]) + tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1) + res = tmp[:, :3] * 2 - 1 + merged[3] = torch.clamp(merged[3] + res, 0, 1) + ''' + return flow_list, mask_list[3], merged diff --git a/modules/rife/model_rife.py b/modules/rife/model_rife.py new file mode 100644 index 000000000..4084db6df --- /dev/null +++ b/modules/rife/model_rife.py @@ -0,0 +1,83 @@ +import torch +from torch.optim import AdamW +from torch.nn.parallel import DistributedDataParallel as DDP +from model_ifnet import IFNet +from loss import EPE, SOBEL +from modules import devices + + +class Model: + 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() + if local_rank != -1: + self.flownet = DDP(self.flownet, device_ids=[local_rank], output_device=local_rank) + + def train(self): + self.flownet.train() + + def eval(self): + self.flownet.eval() + + def device(self): + self.flownet.to(devices.device) + self.flownet.to(devices.dtype) + + def load_model(self, model_file, rank=0): + def convert(param): + if rank == -1: + return { + k.replace("module.", ""): v + for k, v in param.items() + if "module." in k + } + else: + return param + if rank <= 0: + if torch.cuda.is_available(): + self.flownet.load_state_dict(convert(torch.load(model_file)), False) + else: + self.flownet.load_state_dict(convert(torch.load(model_file, map_location='cpu')), False) + + def save_model(self, model_file, rank=0): + 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 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, + } diff --git a/modules/rife/refine.py b/modules/rife/refine.py new file mode 100644 index 000000000..fae3ddc5c --- /dev/null +++ b/modules/rife/refine.py @@ -0,0 +1,90 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from warplayer import warp + + +c = 16 + + +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) + ) + + +def conv_woact(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), + ) + + +def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1): # pylint: disable=unused-argument + return nn.Sequential( + torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1, bias=True), + nn.LeakyReLU(0.2, True) + ) + + +class Conv2(nn.Module): + def __init__(self, in_planes, out_planes, stride=2): + super(Conv2, self).__init__() + self.conv1 = conv(in_planes, out_planes, 3, stride, 1) + self.conv2 = conv(out_planes, out_planes, 3, 1, 1) + + def forward(self, x): + x = self.conv1(x) + x = self.conv2(x) + return x + + +class Contextnet(nn.Module): + def __init__(self): + super(Contextnet, self).__init__() + self.conv1 = Conv2(3, c) + self.conv2 = Conv2(c, 2*c) + self.conv3 = Conv2(2*c, 4*c) + self.conv4 = Conv2(4*c, 8*c) + + def forward(self, x, flow): + x = self.conv1(x) + flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5 + f1 = warp(x, flow) + x = self.conv2(x) + flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5 + f2 = warp(x, flow) + x = self.conv3(x) + flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5 + f3 = warp(x, flow) + x = self.conv4(x) + flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5 + f4 = warp(x, flow) + return [f1, f2, f3, f4] + + +class Unet(nn.Module): + def __init__(self): + super(Unet, self).__init__() + self.down0 = Conv2(17, 2*c) + self.down1 = Conv2(4*c, 4*c) + self.down2 = Conv2(8*c, 8*c) + self.down3 = Conv2(16*c, 16*c) + self.up0 = deconv(32*c, 8*c) + self.up1 = deconv(16*c, 4*c) + self.up2 = deconv(8*c, 2*c) + self.up3 = deconv(4*c, c) + self.conv = nn.Conv2d(c, 3, 3, 1, 1) + + def forward(self, img0, img1, warped_img0, warped_img1, mask, flow, c0, c1): + s0 = self.down0( + torch.cat((img0, img1, warped_img0, warped_img1, mask, flow), 1)) + s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1)) + s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1)) + s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1)) + x = self.up0(torch.cat((s3, c0[3], c1[3]), 1)) + x = self.up1(torch.cat((x, s2), 1)) + x = self.up2(torch.cat((x, s1), 1)) + x = self.up3(torch.cat((x, s0), 1)) + x = self.conv(x) + return torch.sigmoid(x) diff --git a/modules/rife/rife.py b/modules/rife/rife.py new file mode 100755 index 000000000..249b778cf --- /dev/null +++ b/modules/rife/rife.py @@ -0,0 +1,153 @@ +#!/bin/env python + +import _thread +import argparse +import os +import time +import tempfile +from queue import Queue +import filetype +import cv2 +import numpy as np +import torch +from torch.nn import functional as F +from tqdm.rich import tqdm +from ssim import ssim_matlab +from model_rife import Model +from modules import devices + + +model = None +count = 0 + + +def load(model_path: str = 'rife/flownet-v46.pkl'): + global model # pylint: disable=global-statement + model = Model() + model.load_model(model_path, -1) + model.eval() + model.device() + + +def interpolate(args): # pylint: disable=redefined-outer-name + print('start interpolate') + t0 = time.time() + if model is None: + load(args.model) + videogen = [] + if args.seq is None: + for f in os.listdir(args.input): + fn = os.path.join(args.input, f) + if os.path.isfile(fn) and filetype.is_image(fn): + videogen.append(fn) + else: + files = sorted(os.listdir(args.input)) + current = args.seq + for f in files: + seq = os.path.basename(f).split('-')[0] + if seq.isdigit() and int(seq) == current: + fn = os.path.join(args.input, f) + videogen.append(fn) + current += 1 + + videogen = sorted(videogen) + print(f'inputs: {len(videogen)} {[os.path.basename(f) for f in videogen]}') + # videogen.sort(key=lambda x:int(os.path.basename(x[:-4]))) + frame = cv2.imread(videogen[0], cv2.IMREAD_UNCHANGED)[:, :, ::-1].copy() + h, w, _ = frame.shape + if not os.path.exists(args.output): + os.mkdir(args.output) + + def write(output_dir, buffer): + global count # pylint: disable=global-statement + item = buffer.get() + while item is not None: + cv2.imwrite(f'{output_dir}/{count:0>6d}.jpg', item[:, :, ::-1]) + item = buffer.get() + count += 1 + + def execute(I0, I1, n): + if model.version >= 3.9: + res = [] + for i in range(n): + res.append(model.inference(I0, I1, (i+1) * 1. / (n+1), args.scale)) + return res + else: + middle = model.inference(I0, I1, args.scale) + if n == 1: + return [middle] + first_half = execute(I0, middle, n=n//2) + second_half = execute(middle, I1, n=n//2) + if n % 2: + return [*first_half, middle, *second_half] + else: + return [*first_half, *second_half] + + def pad(img): + return F.pad(img, padding).half() if args.fp16 else F.pad(img, padding) # pylint: disable=not-callable + + tmp = max(128, int(128 / args.scale)) + ph = ((h - 1) // tmp + 1) * tmp + pw = ((w - 1) // tmp + 1) * tmp + padding = (0, pw - w, 0, ph - h) + buffer = Queue(maxsize=8192) + _thread.start_new_thread(write, (args.output, buffer)) + + print(f'padded start: frames={args.buffer}') + for _i in range(args.buffer): # fill starting frames + buffer.put(frame) + + I1 = pad(torch.from_numpy(np.transpose(frame, (2,0,1))).to(devices.device, non_blocking=True).unsqueeze(0).float() / 255.) + with torch.no_grad(): + with tqdm(total=len(videogen), desc='interpolate', unit='frame') as pbar: + for f in videogen: + frame = cv2.imread(f, cv2.IMREAD_UNCHANGED)[:, :, ::-1].copy() + I0 = I1 + I1 = pad(torch.from_numpy(np.transpose(frame, (2,0,1))).to(devices.device, non_blocking=True).unsqueeze(0).float() / 255.) + I0_small = F.interpolate(I0, (32, 32), mode='bilinear', align_corners=False) + I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False) + ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3]) + if ssim > 0.99: # skip duplicate frames + continue + if ssim < args.change: + output = [] + for _i in range(args.buffer): # fill frames if change rate is above threshold + output.append(I0) + for _i in range(args.buffer): + output.append(I1) + else: + output = execute(I0, I1, args.multi-1) + for mid in output: + mid = (((mid[0] * 255.).byte().cpu().numpy().transpose(1, 2, 0))) + buffer.put(mid[:h, :w]) + buffer.put(frame) + pbar.update(1) + + print(f'padded end: frames={args.buffer}') + for _i in range(args.buffer): # fill ending frames + buffer.put(frame) + while not buffer.empty(): + time.sleep(0.5) + t1 = time.time() + print(f'end interpolate: input={len(videogen)} frames={count} time={round(t1 - t0, 2)}') + + +if __name__ == "__main__": + print('starting rife') + tmp_folder = os.path.join(tempfile.gettempdir(), f'rife-{time.strftime("%Y%m%d-%H%M%S")}') + parser = argparse.ArgumentParser(description='interpolate video frames using RIFE') + parser.add_argument('--model', type=str, default=os.path.abspath(os.path.join(os.path.dirname(__file__), 'model/flownet-v46.pkl')), help='path to model, default: %(default)s') + parser.add_argument('--input', type=str, required=True, default=None, help='input directory containing images, default: %(default)s') + parser.add_argument('--output', type=str, default=tmp_folder, help='output directory for interpolated images, default: %(default)s') + parser.add_argument('--scale', type=float, default=1.0, help='scale factor for interpolated images, default: %(default)s') + parser.add_argument('--multi', type=int, default=4, help='number of frames to interpolate between two input images, default: %(default)s') + parser.add_argument('--buffer', type=int, default=2, help='number of frames to buffer on scene change, default: %(default)s') + parser.add_argument('--change', type=float, default=0.3, help='scene change threshold (lower is more sensitive, default: %(default)s') + parser.add_argument('--fp16', action='store_true', help='use float16 precision instead of float32, default: %(default)s') + parser.add_argument('--fps', type=int, default=25, help='desired framerate, default: %(default)s') + parser.add_argument('--seq', type=int, default=None, help='image sequence start number, default: %(default)s') + parser.add_argument('--rm', action='store_true', help='remove interpolated images, default: %(default)s') + args = parser.parse_args() + print('args', args) + assert args.scale in [0.25, 0.5, 1.0, 2.0, 4.0] + interpolate(args) diff --git a/modules/rife/ssim.py b/modules/rife/ssim.py new file mode 100644 index 000000000..e2261ca7a --- /dev/null +++ b/modules/rife/ssim.py @@ -0,0 +1,174 @@ +from math import exp +import torch +import torch.nn.functional as F +from modules import devices + + +def gaussian(window_size, sigma): + gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)]) + return gauss/gauss.sum() + + +def create_window(window_size, channel=1): + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0).to(devices.device) + window = _2D_window.expand(channel, 1, window_size, window_size).contiguous() + return window + + +def create_window_3d(window_size, channel=1): + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()) + _3D_window = _2D_window.unsqueeze(2) @ (_1D_window.t()) + window = _3D_window.expand(1, channel, window_size, window_size, window_size).contiguous().to(devices.device) + return window + + +def ssim(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None): + # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh). + if val_range is None: + if torch.max(img1) > 128: + max_val = 255 + else: + max_val = 1 + + if torch.min(img1) < -0.5: + min_val = -1 + else: + min_val = 0 + L = max_val - min_val + else: + L = val_range + padd = 0 + (_, channel, height, width) = img1.size() + if window is None: + real_size = min(window_size, height, width) + window = create_window(real_size, channel=channel).to(img1.device) + # mu1 = F.conv2d(img1, window, padding=padd, groups=channel) + # mu2 = F.conv2d(img2, window, padding=padd, groups=channel) + mu1 = F.conv2d(F.pad(img1, (5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=channel) + mu2 = F.conv2d(F.pad(img2, (5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(F.pad(img1 * img1, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(F.pad(img2 * img2, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu2_sq + sigma12 = F.conv2d(F.pad(img1 * img2, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu1_mu2 + C1 = (0.01 * L) ** 2 + C2 = (0.03 * L) ** 2 + v1 = 2.0 * sigma12 + C2 + v2 = sigma1_sq + sigma2_sq + C2 + cs = torch.mean(v1 / v2) # contrast sensitivity + ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2) + if size_average: + ret = ssim_map.mean() + else: + ret = ssim_map.mean(1).mean(1).mean(1) + if full: + return ret, cs + return ret + + +def ssim_matlab(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None): + # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh). + if val_range is None: + if torch.max(img1) > 128: + max_val = 255 + else: + max_val = 1 + if torch.min(img1) < -0.5: + min_val = -1 + else: + min_val = 0 + L = max_val - min_val + else: + L = val_range + padd = 0 + (_, _, height, width) = img1.size() + if window is None: + real_size = min(window_size, height, width) + window = create_window_3d(real_size, channel=1).to(img1.device) + # Channel is set to 1 since we consider color images as volumetric images + img1 = img1.unsqueeze(1) + img2 = img2.unsqueeze(1) + mu1 = F.conv3d(F.pad(img1, (5, 5, 5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=1) + mu2 = F.conv3d(F.pad(img2, (5, 5, 5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=1) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv3d(F.pad(img1 * img1, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu1_sq + sigma2_sq = F.conv3d(F.pad(img2 * img2, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu2_sq + sigma12 = F.conv3d(F.pad(img1 * img2, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu1_mu2 + C1 = (0.01 * L) ** 2 + C2 = (0.03 * L) ** 2 + v1 = 2.0 * sigma12 + C2 + v2 = sigma1_sq + sigma2_sq + C2 + cs = torch.mean(v1 / v2) # contrast sensitivity + ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2) + if size_average: + ret = ssim_map.mean() + else: + ret = ssim_map.mean(1).mean(1).mean(1) + if full: + return ret, cs + return ret + + +def msssim(img1, img2, window_size=11, size_average=True, val_range=None, normalize=False): + local_device = img1.device + weights = torch.FloatTensor([0.0448, 0.2856, 0.3001, 0.2363, 0.1333]).to(local_device) + levels = weights.size()[0] + mssim = [] + mcs = [] + for _ in range(levels): + sim, cs = ssim(img1, img2, window_size=window_size, size_average=size_average, full=True, val_range=val_range) + mssim.append(sim) + mcs.append(cs) + img1 = F.avg_pool2d(img1, (2, 2)) + img2 = F.avg_pool2d(img2, (2, 2)) + mssim = torch.stack(mssim) + mcs = torch.stack(mcs) + # Normalize (to avoid NaNs during training unstable models, not compliant with original definition) + if normalize: + mssim = (mssim + 1) / 2 + mcs = (mcs + 1) / 2 + pow1 = mcs ** weights + pow2 = mssim ** weights + # From Matlab implementation https://ece.uwaterloo.ca/~z70wang/research/iwssim/ + output = torch.prod(pow1[:-1] * pow2[-1]) + return output + + +# Classes to re-use window +class SSIM(torch.nn.Module): + def __init__(self, window_size=11, size_average=True, val_range=None): + super(SSIM, self).__init__() + self.window_size = window_size + self.size_average = size_average + self.val_range = val_range + # Assume 3 channel for SSIM + self.channel = 3 + self.window = create_window(window_size, channel=self.channel) + + def forward(self, img1, img2): + (_, channel, _, _) = img1.size() + if channel == self.channel and self.window.dtype == img1.dtype: + window = self.window + else: + window = create_window(self.window_size, channel).to(img1.device).type(img1.dtype) + self.window = window + self.channel = channel + _ssim = ssim(img1, img2, window=window, window_size=self.window_size, size_average=self.size_average) + dssim = (1 - _ssim) / 2 + return dssim + + +class MSSSIM(torch.nn.Module): + def __init__(self, window_size=11, size_average=True, channel=3): + super(MSSSIM, self).__init__() + self.window_size = window_size + self.size_average = size_average + self.channel = channel + + def forward(self, img1, img2): + return msssim(img1, img2, window_size=self.window_size, size_average=self.size_average) diff --git a/modules/rife/warplayer.py b/modules/rife/warplayer.py new file mode 100644 index 000000000..8a5d6398d --- /dev/null +++ b/modules/rife/warplayer.py @@ -0,0 +1,17 @@ +import torch +from modules import devices + + +backwarp_tenGrid = {} + + +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) + g = (backwarp_tenGrid[k] + tenFlow).permute(0, 2, 3, 1) + return torch.nn.functional.grid_sample(input=tenInput, grid=g, mode='bilinear', padding_mode='border', align_corners=True)