hugging-apps/echo-memory
0
1import torch2import torch.nn as nn3import torch.nn.functional as F4import numpy as np5from PIL import Image6 7 8def warp(tenInput, tenFlow, device):9 backwarp_tenGrid = {}10 k = (str(tenFlow.device), str(tenFlow.size()))11 if k not in backwarp_tenGrid:12 tenHorizontal = torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=device).view(13 1, 1, 1, tenFlow.shape[3]).expand(tenFlow.shape[0], -1, tenFlow.shape[2], -1)14 tenVertical = torch.linspace(-1.0, 1.0, tenFlow.shape[2], device=device).view(15 1, 1, tenFlow.shape[2], 1).expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3])16 backwarp_tenGrid[k] = torch.cat(17 [tenHorizontal, tenVertical], 1).to(device)18 19 tenFlow = torch.cat([tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0),20 tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0)], 1)21 22 g = (backwarp_tenGrid[k] + tenFlow).permute(0, 2, 3, 1)23 return torch.nn.functional.grid_sample(input=tenInput, grid=g, mode='bilinear', padding_mode='border', align_corners=True)24 25 26def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):27 return nn.Sequential(28 nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,29 padding=padding, dilation=dilation, bias=True), 30 nn.PReLU(out_planes)31 )32 33 34class IFBlock(nn.Module):35 def __init__(self, in_planes, c=64):36 super(IFBlock, self).__init__()37 self.conv0 = nn.Sequential(conv(in_planes, c//2, 3, 2, 1), conv(c//2, c, 3, 2, 1),)38 self.convblock0 = nn.Sequential(conv(c, c), conv(c, c))39 self.convblock1 = nn.Sequential(conv(c, c), conv(c, c))40 self.convblock2 = nn.Sequential(conv(c, c), conv(c, c))41 self.convblock3 = nn.Sequential(conv(c, c), conv(c, c))42 self.conv1 = nn.Sequential(nn.ConvTranspose2d(c, c//2, 4, 2, 1), nn.PReLU(c//2), nn.ConvTranspose2d(c//2, 4, 4, 2, 1))43 self.conv2 = nn.Sequential(nn.ConvTranspose2d(c, c//2, 4, 2, 1), nn.PReLU(c//2), nn.ConvTranspose2d(c//2, 1, 4, 2, 1))44 45 def forward(self, x, flow, scale=1):46 x = F.interpolate(x, scale_factor= 1. / scale, mode="bilinear", align_corners=False, recompute_scale_factor=False)47 flow = F.interpolate(flow, scale_factor= 1. / scale, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 1. / scale48 feat = self.conv0(torch.cat((x, flow), 1))49 feat = self.convblock0(feat) + feat50 feat = self.convblock1(feat) + feat51 feat = self.convblock2(feat) + feat52 feat = self.convblock3(feat) + feat 53 flow = self.conv1(feat)54 mask = self.conv2(feat)55 flow = F.interpolate(flow, scale_factor=scale, mode="bilinear", align_corners=False, recompute_scale_factor=False) * scale56 mask = F.interpolate(mask, scale_factor=scale, mode="bilinear", align_corners=False, recompute_scale_factor=False)57 return flow, mask58 59 60class IFNet(nn.Module):61 def __init__(self, **kwargs):62 super(IFNet, self).__init__()63 self.block0 = IFBlock(7+4, c=90)64 self.block1 = IFBlock(7+4, c=90)65 self.block2 = IFBlock(7+4, c=90)66 self.block_tea = IFBlock(10+4, c=90)67 68 def forward(self, x, scale_list=[4, 2, 1], training=False):69 if training == False:70 channel = x.shape[1] // 271 img0 = x[:, :channel]72 img1 = x[:, channel:]73 flow_list = []74 merged = []75 mask_list = []76 warped_img0 = img077 warped_img1 = img178 flow = (x[:, :4]).detach() * 079 mask = (x[:, :1]).detach() * 080 block = [self.block0, self.block1, self.block2]81 for i in range(3):82 f0, m0 = block[i](torch.cat((warped_img0[:, :3], warped_img1[:, :3], mask), 1), flow, scale=scale_list[i])83 f1, m1 = block[i](torch.cat((warped_img1[:, :3], warped_img0[:, :3], -mask), 1), torch.cat((flow[:, 2:4], flow[:, :2]), 1), scale=scale_list[i])84 flow = flow + (f0 + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 285 mask = mask + (m0 + (-m1)) / 286 mask_list.append(mask)87 flow_list.append(flow)88 warped_img0 = warp(img0, flow[:, :2], device=x.device)89 warped_img1 = warp(img1, flow[:, 2:4], device=x.device)90 merged.append((warped_img0, warped_img1))91 '''92 c0 = self.contextnet(img0, flow[:, :2])93 c1 = self.contextnet(img1, flow[:, 2:4])94 tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)95 res = tmp[:, 1:4] * 2 - 196 '''97 for i in range(3):98 mask_list[i] = torch.sigmoid(mask_list[i])99 merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i]) 100 return flow_list, mask_list[2], merged101 102 @staticmethod103 def state_dict_converter():104 return IFNetStateDictConverter()105 106 107class IFNetStateDictConverter:108 def __init__(self):109 pass110 111 def from_diffusers(self, state_dict):112 state_dict_ = {k.replace("module.", ""): v for k, v in state_dict.items()}113 return state_dict_114 115 def from_civitai(self, state_dict):116 return self.from_diffusers(state_dict), {"upcast_to_float32": True}117 118 119class RIFEInterpolater:120 def __init__(self, model, device="cuda"):121 self.model = model122 self.device = device123 # IFNet only does not support float16124 self.torch_dtype = torch.float32125 126 @staticmethod127 def from_model_manager(model_manager):128 return RIFEInterpolater(model_manager.fetch_model("rife"), device=model_manager.device)129 130 def process_image(self, image):131 width, height = image.size132 if width % 32 != 0 or height % 32 != 0:133 width = (width + 31) // 32134 height = (height + 31) // 32135 image = image.resize((width, height))136 image = torch.Tensor(np.array(image, dtype=np.float32)[:, :, [2,1,0]] / 255).permute(2, 0, 1)137 return image138 139 def process_images(self, images):140 images = [self.process_image(image) for image in images]141 images = torch.stack(images)142 return images143 144 def decode_images(self, images):145 images = (images[:, [2,1,0]].permute(0, 2, 3, 1) * 255).clip(0, 255).numpy().astype(np.uint8)146 images = [Image.fromarray(image) for image in images]147 return images148 149 def add_interpolated_images(self, images, interpolated_images):150 output_images = []151 for image, interpolated_image in zip(images, interpolated_images):152 output_images.append(image)153 output_images.append(interpolated_image)154 output_images.append(images[-1])155 return output_images156 157 158 @torch.no_grad()159 def interpolate_(self, images, scale=1.0):160 input_tensor = self.process_images(images)161 input_tensor = torch.cat((input_tensor[:-1], input_tensor[1:]), dim=1)162 input_tensor = input_tensor.to(device=self.device, dtype=self.torch_dtype)163 flow, mask, merged = self.model(input_tensor, [4/scale, 2/scale, 1/scale])164 output_images = self.decode_images(merged[2].cpu())165 if output_images[0].size != images[0].size:166 output_images = [image.resize(images[0].size) for image in output_images]167 return output_images168 169 170 @torch.no_grad()171 def interpolate(self, images, scale=1.0, batch_size=4, num_iter=1, progress_bar=lambda x:x):172 # Preprocess173 processed_images = self.process_images(images)174 175 for iter in range(num_iter):176 # Input177 input_tensor = torch.cat((processed_images[:-1], processed_images[1:]), dim=1)178 179 # Interpolate180 output_tensor = []181 for batch_id in progress_bar(range(0, input_tensor.shape[0], batch_size)):182 batch_id_ = min(batch_id + batch_size, input_tensor.shape[0])183 batch_input_tensor = input_tensor[batch_id: batch_id_]184 batch_input_tensor = batch_input_tensor.to(device=self.device, dtype=self.torch_dtype)185 flow, mask, merged = self.model(batch_input_tensor, [4/scale, 2/scale, 1/scale])186 output_tensor.append(merged[2].cpu())187 188 # Output189 output_tensor = torch.concat(output_tensor, dim=0).clip(0, 1)190 processed_images = self.add_interpolated_images(processed_images, output_tensor)191 processed_images = torch.stack(processed_images)192 193 # To images194 output_images = self.decode_images(processed_images)195 if output_images[0].size != images[0].size:196 output_images = [image.resize(images[0].size) for image in output_images]197 return output_images198 199 200class RIFESmoother(RIFEInterpolater):201 def __init__(self, model, device="cuda"):202 super(RIFESmoother, self).__init__(model, device=device)203 204 @staticmethod205 def from_model_manager(model_manager):206 return RIFEInterpolater(model_manager.fetch_model("rife"), device=model_manager.device)207 208 def process_tensors(self, input_tensor, scale=1.0, batch_size=4):209 output_tensor = []210 for batch_id in range(0, input_tensor.shape[0], batch_size):211 batch_id_ = min(batch_id + batch_size, input_tensor.shape[0])212 batch_input_tensor = input_tensor[batch_id: batch_id_]213 batch_input_tensor = batch_input_tensor.to(device=self.device, dtype=self.torch_dtype)214 flow, mask, merged = self.model(batch_input_tensor, [4/scale, 2/scale, 1/scale])215 output_tensor.append(merged[2].cpu())216 output_tensor = torch.concat(output_tensor, dim=0)217 return output_tensor218 219 @torch.no_grad()220 def __call__(self, rendered_frames, scale=1.0, batch_size=4, num_iter=1, **kwargs):221 # Preprocess222 processed_images = self.process_images(rendered_frames)223 224 for iter in range(num_iter):225 # Input226 input_tensor = torch.cat((processed_images[:-2], processed_images[2:]), dim=1)227 228 # Interpolate229 output_tensor = self.process_tensors(input_tensor, scale=scale, batch_size=batch_size)230 231 # Blend232 input_tensor = torch.cat((processed_images[1:-1], output_tensor), dim=1)233 output_tensor = self.process_tensors(input_tensor, scale=scale, batch_size=batch_size)234 235 # Add to frames236 processed_images[1:-1] = output_tensor237 238 # To images239 output_images = self.decode_images(processed_images)240 if output_images[0].size != rendered_frames[0].size:241 output_images = [image.resize(rendered_frames[0].size) for image in output_images]242 return output_images243