Upload 9 files

roop/processors/Enhance_CodeFormer.py

@@ -0,0 +1,70 @@
+from typing import Any, List, Callable
+import cv2 
+import numpy as np
+import onnxruntime
+import roop.globals
+
+from roop.typing import Face, Frame, FaceSet
+from roop.utilities import resolve_relative_path
+
+
+# THREAD_LOCK = threading.Lock()
+
+
+class Enhance_CodeFormer():
+    model_codeformer = None
+    devicename = None
+
+    processorname = 'codeformer'
+    type = 'enhance'
+    
+
+    def Initialize(self, devicename:str):
+        if self.model_codeformer is None:
+            # replace Mac mps with cpu for the moment
+            devicename = devicename.replace('mps', 'cpu')
+            self.devicename = devicename
+            model_path = resolve_relative_path('../models/CodeFormer/CodeFormerv0.1.onnx')
+            self.model_codeformer = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
+            self.model_inputs = self.model_codeformer.get_inputs()
+            model_outputs = self.model_codeformer.get_outputs()
+            self.io_binding = self.model_codeformer.io_binding()           
+            self.io_binding.bind_cpu_input(self.model_inputs[1].name, np.array([0.5]))
+            self.io_binding.bind_output(model_outputs[0].name, self.devicename)
+
+
+    def Run(self, source_faceset: FaceSet, target_face: Face, temp_frame: Frame) -> Frame:
+        input_size = temp_frame.shape[1]
+        # preprocess
+        temp_frame = cv2.resize(temp_frame, (512, 512), cv2.INTER_CUBIC)
+        temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB)
+        temp_frame = temp_frame.astype('float32') / 255.0
+        temp_frame = (temp_frame - 0.5) / 0.5
+        temp_frame = np.expand_dims(temp_frame, axis=0).transpose(0, 3, 1, 2)
+        
+        self.io_binding.bind_cpu_input(self.model_inputs[0].name, temp_frame.astype(np.float32))
+        self.model_codeformer.run_with_iobinding(self.io_binding)
+        ort_outs = self.io_binding.copy_outputs_to_cpu()
+        result = ort_outs[0][0]
+        del ort_outs
+        
+        # post-process
+        result = result.transpose((1, 2, 0))
+
+        un_min = -1.0
+        un_max = 1.0
+        result = np.clip(result, un_min, un_max)
+        result = (result - un_min) / (un_max - un_min)
+
+        result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
+        result = (result * 255.0).round()
+        scale_factor = int(result.shape[1] / input_size)       
+        return result.astype(np.uint8), scale_factor
+
+
+    def Release(self):
+        del self.model_codeformer
+        self.model_codeformer = None
+        del self.io_binding
+        self.io_binding = None
+

roop/processors/Enhance_DMDNet.py

@@ -0,0 +1,893 @@
+from typing import Any, List, Callable
+import cv2 
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.utils.spectral_norm as SpectralNorm
+import threading
+from torchvision.ops import roi_align
+
+from math import sqrt
+
+from torchvision.transforms.functional import normalize
+
+from roop.typing import Face, Frame, FaceSet
+
+
+THREAD_LOCK_DMDNET = threading.Lock()
+
+
+class Enhance_DMDNet():
+
+    model_dmdnet = None
+    torchdevice = None
+
+    processorname = 'dmdnet'
+    type = 'enhance'
+
+
+    def Initialize(self, devicename):
+        if self.model_dmdnet is None:
+            self.model_dmdnet = self.create(devicename)
+            
+
+    # temp_frame already cropped+aligned, bbox not
+    def Run(self, source_faceset: FaceSet, target_face: Face, temp_frame: Frame) -> Frame:
+        input_size = temp_frame.shape[1]
+
+        result = self.enhance_face(source_faceset, temp_frame, target_face)
+        scale_factor = int(result.shape[1] / input_size)       
+        return result.astype(np.uint8), scale_factor
+
+
+    def Release(self):
+        self.model_gfpgan = None
+
+
+    # https://stackoverflow.com/a/67174339
+    def landmarks106_to_68(self, pt106):
+        map106to68=[1,10,12,14,16,3,5,7,0,23,21,19,32,30,28,26,17,
+                        43,48,49,51,50,
+                        102,103,104,105,101,
+                        72,73,74,86,78,79,80,85,84,
+                        35,41,42,39,37,36,
+                        89,95,96,93,91,90,
+                        52,64,63,71,67,68,61,58,59,53,56,55,65,66,62,70,69,57,60,54
+                        ]
+      
+        pt68 = []
+        for i in range(68):
+            index = map106to68[i]
+            pt68.append(pt106[index])
+        return pt68
+
+        
+
+
+    def check_bbox(self, imgs, boxes):
+        boxes = boxes.view(-1, 4, 4)
+        colors = [(0, 255, 0), (0, 255, 0), (255, 255, 0), (255, 0, 0)]
+        i = 0
+        for img, box in zip(imgs, boxes):
+            img = (img + 1)/2 * 255
+            img2 = img.permute(1, 2, 0).float().cpu().flip(2).numpy().copy()
+            for idx, point in enumerate(box):
+                cv2.rectangle(img2, (int(point[0]), int(point[1])), (int(point[2]), int(point[3])), color=colors[idx], thickness=2)
+            cv2.imwrite('dmdnet_{:02d}.png'.format(i), img2)
+            i += 1
+
+
+    def trans_points2d(self, pts, M):
+        new_pts = np.zeros(shape=pts.shape, dtype=np.float32)
+        for i in range(pts.shape[0]):
+            pt = pts[i]
+            new_pt = np.array([pt[0], pt[1], 1.0], dtype=np.float32)
+            new_pt = np.dot(M, new_pt)
+            new_pts[i] = new_pt[0:2]
+
+        return new_pts
+
+
+    def enhance_face(self, ref_faceset: FaceSet, temp_frame, face: Face):
+        # preprocess
+        start_x, start_y, end_x, end_y = map(int, face['bbox'])
+        lm106 = face.landmark_2d_106
+        lq_landmarks = np.asarray(self.landmarks106_to_68(lm106))
+
+        if temp_frame.shape[0] != 512 or temp_frame.shape[1] != 512:
+            # scale to 512x512
+            scale_factor = 512 / temp_frame.shape[1]
+
+            M = face.matrix * scale_factor
+
+            lq_landmarks = self.trans_points2d(lq_landmarks, M)
+            temp_frame = cv2.resize(temp_frame, (512,512), interpolation = cv2.INTER_AREA)
+
+        if temp_frame.ndim == 2:
+            temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_GRAY2RGB)  # GGG
+        # else:
+        #     temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB)  # RGB
+
+        lq = read_img_tensor(temp_frame)
+
+        LQLocs = get_component_location(lq_landmarks)
+        # self.check_bbox(lq, LQLocs.unsqueeze(0))
+
+        # specific, change 1000 to 1 to activate
+        if len(ref_faceset.faces) > 1:
+            SpecificImgs = []
+            SpecificLocs = []
+            for i,face in enumerate(ref_faceset.faces):
+                lm106 = face.landmark_2d_106
+                lq_landmarks = np.asarray(self.landmarks106_to_68(lm106))
+                ref_image = ref_faceset.ref_images[i]
+                if ref_image.shape[0] != 512 or ref_image.shape[1] != 512:
+                    # scale to 512x512
+                    scale_factor = 512 / ref_image.shape[1]
+
+                    M = face.matrix * scale_factor
+
+                    lq_landmarks = self.trans_points2d(lq_landmarks, M)
+                    ref_image = cv2.resize(ref_image, (512,512), interpolation = cv2.INTER_AREA)
+
+                if ref_image.ndim == 2:
+                    temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_GRAY2RGB)  # GGG
+                # else:
+                #     temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB)  # RGB
+
+                ref_tensor = read_img_tensor(ref_image)
+                ref_locs = get_component_location(lq_landmarks)
+                # self.check_bbox(ref_tensor, ref_locs.unsqueeze(0))
+
+                SpecificImgs.append(ref_tensor)
+                SpecificLocs.append(ref_locs.unsqueeze(0))
+
+            SpecificImgs = torch.cat(SpecificImgs, dim=0)
+            SpecificLocs = torch.cat(SpecificLocs, dim=0)
+            # check_bbox(SpecificImgs, SpecificLocs)
+            SpMem256, SpMem128, SpMem64 = self.model_dmdnet.generate_specific_dictionary(sp_imgs = SpecificImgs.to(self.torchdevice), sp_locs = SpecificLocs)
+            SpMem256Para = {}
+            SpMem128Para = {}
+            SpMem64Para = {}
+            for k, v in SpMem256.items():
+                SpMem256Para[k] = v
+            for k, v in SpMem128.items():
+                SpMem128Para[k] = v
+            for k, v in SpMem64.items():
+                SpMem64Para[k] = v
+        else:
+            # generic
+            SpMem256Para, SpMem128Para, SpMem64Para = None, None, None
+
+        with torch.no_grad():
+            with THREAD_LOCK_DMDNET:
+                try:
+                    GenericResult, SpecificResult = self.model_dmdnet(lq = lq.to(self.torchdevice), loc = LQLocs.unsqueeze(0), sp_256 = SpMem256Para, sp_128 = SpMem128Para, sp_64 = SpMem64Para)
+                except Exception as e:
+                    print(f'Error {e} there may be something wrong with the detected component locations.')
+                    return temp_frame
+        
+        if SpecificResult is not None:
+            save_specific = SpecificResult * 0.5 + 0.5
+            save_specific = save_specific.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR
+            save_specific = np.clip(save_specific.float().cpu().numpy(), 0, 1) * 255.0
+            temp_frame =  save_specific.astype("uint8")
+            if False:
+                save_generic = GenericResult * 0.5 + 0.5
+                save_generic = save_generic.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR
+                save_generic = np.clip(save_generic.float().cpu().numpy(), 0, 1) * 255.0
+                check_lq = lq * 0.5 + 0.5
+                check_lq = check_lq.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR
+                check_lq = np.clip(check_lq.float().cpu().numpy(), 0, 1) * 255.0
+                cv2.imwrite('dmdnet_comparison.png', cv2.cvtColor(np.hstack((check_lq, save_generic, save_specific)),cv2.COLOR_RGB2BGR))
+        else:
+            save_generic = GenericResult * 0.5 + 0.5
+            save_generic = save_generic.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR
+            save_generic = np.clip(save_generic.float().cpu().numpy(), 0, 1) * 255.0
+            temp_frame =  save_generic.astype("uint8")
+        temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_RGB2BGR)  # RGB
+        return temp_frame
+
+    
+
+    def create(self, devicename):
+        self.torchdevice = torch.device(devicename)
+        model_dmdnet = DMDNet().to(self.torchdevice)
+        weights = torch.load('./models/DMDNet.pth') 
+        model_dmdnet.load_state_dict(weights, strict=True)
+
+        model_dmdnet.eval()
+        num_params = 0
+        for param in model_dmdnet.parameters():
+            num_params += param.numel()
+        return model_dmdnet
+
+    # print('{:>8s} : {}'.format('Using device', device))
+    # print('{:>8s} : {:.2f}M'.format('Model params', num_params/1e6))
+
+
+
+def read_img_tensor(Img=None): #rgb -1~1 
+    Img = Img.transpose((2, 0, 1))/255.0
+    Img = torch.from_numpy(Img).float()
+    normalize(Img, [0.5,0.5,0.5], [0.5,0.5,0.5], inplace=True)
+    ImgTensor = Img.unsqueeze(0)
+    return ImgTensor
+
+
+def get_component_location(Landmarks, re_read=False):
+    if re_read:
+        ReadLandmark = []
+        with open(Landmarks,'r') as f:
+            for line in f:
+                tmp = [float(i) for i in line.split(' ') if i != '\n']
+                ReadLandmark.append(tmp)
+        ReadLandmark = np.array(ReadLandmark) #
+        Landmarks = np.reshape(ReadLandmark, [-1, 2]) # 68*2
+    Map_LE_B = list(np.hstack((range(17,22), range(36,42))))
+    Map_RE_B = list(np.hstack((range(22,27), range(42,48))))
+    Map_LE = list(range(36,42))
+    Map_RE = list(range(42,48))
+    Map_NO = list(range(29,36))
+    Map_MO = list(range(48,68))
+
+    Landmarks[Landmarks>504]=504
+    Landmarks[Landmarks<8]=8
+    
+    #left eye
+    Mean_LE = np.mean(Landmarks[Map_LE],0)
+    L_LE1 = Mean_LE[1] - np.min(Landmarks[Map_LE_B,1])
+    L_LE1 = L_LE1 * 1.3
+    L_LE2 = L_LE1 / 1.9
+    L_LE_xy = L_LE1 + L_LE2
+    L_LE_lt = [L_LE_xy/2, L_LE1]
+    L_LE_rb = [L_LE_xy/2, L_LE2]
+    Location_LE = np.hstack((Mean_LE - L_LE_lt + 1, Mean_LE + L_LE_rb)).astype(int)
+
+    #right eye
+    Mean_RE = np.mean(Landmarks[Map_RE],0)
+    L_RE1 = Mean_RE[1] - np.min(Landmarks[Map_RE_B,1])
+    L_RE1 = L_RE1 * 1.3
+    L_RE2 = L_RE1 / 1.9
+    L_RE_xy = L_RE1 + L_RE2
+    L_RE_lt = [L_RE_xy/2, L_RE1]
+    L_RE_rb = [L_RE_xy/2, L_RE2]
+    Location_RE = np.hstack((Mean_RE - L_RE_lt + 1, Mean_RE + L_RE_rb)).astype(int)
+
+    #nose
+    Mean_NO = np.mean(Landmarks[Map_NO],0)
+    L_NO1 =( np.max([Mean_NO[0] - Landmarks[31][0], Landmarks[35][0] - Mean_NO[0]])) * 1.25
+    L_NO2 = (Landmarks[33][1] - Mean_NO[1]) * 1.1
+    L_NO_xy = L_NO1 * 2
+    L_NO_lt = [L_NO_xy/2, L_NO_xy - L_NO2]
+    L_NO_rb = [L_NO_xy/2, L_NO2]
+    Location_NO = np.hstack((Mean_NO - L_NO_lt + 1, Mean_NO + L_NO_rb)).astype(int)
+    
+    #mouth
+    Mean_MO = np.mean(Landmarks[Map_MO],0)
+    L_MO = np.max((np.max(np.max(Landmarks[Map_MO],0) - np.min(Landmarks[Map_MO],0))/2,16)) * 1.1
+    MO_O = Mean_MO - L_MO + 1
+    MO_T = Mean_MO + L_MO
+    MO_T[MO_T>510]=510
+    Location_MO = np.hstack((MO_O, MO_T)).astype(int)
+    return torch.cat([torch.FloatTensor(Location_LE).unsqueeze(0), torch.FloatTensor(Location_RE).unsqueeze(0), torch.FloatTensor(Location_NO).unsqueeze(0), torch.FloatTensor(Location_MO).unsqueeze(0)], dim=0)
+
+
+
+
+def calc_mean_std_4D(feat, eps=1e-5):
+    # eps is a small value added to the variance to avoid divide-by-zero.
+    size = feat.size()
+    assert (len(size) == 4)
+    N, C = size[:2]
+    feat_var = feat.view(N, C, -1).var(dim=2) + eps
+    feat_std = feat_var.sqrt().view(N, C, 1, 1)
+    feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
+    return feat_mean, feat_std
+
+def adaptive_instance_normalization_4D(content_feat, style_feat): # content_feat is ref feature, style is degradate feature
+    size = content_feat.size()
+    style_mean, style_std = calc_mean_std_4D(style_feat)
+    content_mean, content_std = calc_mean_std_4D(content_feat)
+    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
+    return normalized_feat * style_std.expand(size) + style_mean.expand(size)
+
+
+def convU(in_channels, out_channels,conv_layer, norm_layer, kernel_size=3, stride=1,dilation=1, bias=True):
+    return nn.Sequential(
+        SpectralNorm(conv_layer(in_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=((kernel_size-1)//2)*dilation, bias=bias)),
+        nn.LeakyReLU(0.2),
+        SpectralNorm(conv_layer(out_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=((kernel_size-1)//2)*dilation, bias=bias)),
+    )
+    
+
+class MSDilateBlock(nn.Module):
+    def __init__(self, in_channels,conv_layer=nn.Conv2d, norm_layer=nn.BatchNorm2d, kernel_size=3, dilation=[1,1,1,1], bias=True):
+        super(MSDilateBlock, self).__init__()
+        self.conv1 =  convU(in_channels, in_channels,conv_layer, norm_layer, kernel_size,dilation=dilation[0], bias=bias)
+        self.conv2 =  convU(in_channels, in_channels,conv_layer, norm_layer, kernel_size,dilation=dilation[1], bias=bias)
+        self.conv3 =  convU(in_channels, in_channels,conv_layer, norm_layer, kernel_size,dilation=dilation[2], bias=bias)
+        self.conv4 =  convU(in_channels, in_channels,conv_layer, norm_layer, kernel_size,dilation=dilation[3], bias=bias)
+        self.convi =  SpectralNorm(conv_layer(in_channels*4, in_channels, kernel_size=kernel_size, stride=1, padding=(kernel_size-1)//2, bias=bias))
+    def forward(self, x):
+        conv1 = self.conv1(x)
+        conv2 = self.conv2(x)
+        conv3 = self.conv3(x)
+        conv4 = self.conv4(x)
+        cat  = torch.cat([conv1, conv2, conv3, conv4], 1)
+        out = self.convi(cat) + x
+        return out
+
+
+class AdaptiveInstanceNorm(nn.Module):
+    def __init__(self, in_channel):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(in_channel)
+
+    def forward(self, input, style):
+        style_mean, style_std = calc_mean_std_4D(style)
+        out = self.norm(input)
+        size = input.size()
+        out = style_std.expand(size) * out + style_mean.expand(size)
+        return out
+
+class NoiseInjection(nn.Module):
+    def __init__(self, channel):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(1, channel, 1, 1))
+    def forward(self, image, noise):
+        if noise is None:
+            b, c, h, w = image.shape
+            noise = image.new_empty(b, 1, h, w).normal_()
+        return image + self.weight * noise
+
+class StyledUpBlock(nn.Module):
+    def __init__(self, in_channel, out_channel, kernel_size=3, padding=1,upsample=False, noise_inject=False):
+        super().__init__()
+
+        self.noise_inject = noise_inject
+        if upsample:
+            self.conv1 = nn.Sequential(
+                nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
+                SpectralNorm(nn.Conv2d(in_channel, out_channel, kernel_size, padding=padding)),
+                nn.LeakyReLU(0.2),
+            )
+        else:
+            self.conv1 = nn.Sequential(
+                SpectralNorm(nn.Conv2d(in_channel, out_channel, kernel_size, padding=padding)),
+                nn.LeakyReLU(0.2),
+                SpectralNorm(nn.Conv2d(out_channel, out_channel, kernel_size, padding=padding)),
+            )
+        self.convup = nn.Sequential(
+                nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
+                SpectralNorm(nn.Conv2d(out_channel, out_channel, kernel_size, padding=padding)),
+                nn.LeakyReLU(0.2),
+                SpectralNorm(nn.Conv2d(out_channel, out_channel, kernel_size, padding=padding)),
+            )
+        if self.noise_inject:
+            self.noise1 = NoiseInjection(out_channel)
+
+        self.lrelu1 = nn.LeakyReLU(0.2)
+
+        self.ScaleModel1 = nn.Sequential(
+            SpectralNorm(nn.Conv2d(in_channel,out_channel,3, 1, 1)),
+            nn.LeakyReLU(0.2),
+            SpectralNorm(nn.Conv2d(out_channel, out_channel, 3, 1, 1))
+        )
+        self.ShiftModel1 = nn.Sequential(
+            SpectralNorm(nn.Conv2d(in_channel,out_channel,3, 1, 1)),
+            nn.LeakyReLU(0.2),
+            SpectralNorm(nn.Conv2d(out_channel, out_channel, 3, 1, 1)),
+        )
+       
+    def forward(self, input, style):
+        out = self.conv1(input)
+        out = self.lrelu1(out)
+        Shift1 = self.ShiftModel1(style)
+        Scale1 = self.ScaleModel1(style)
+        out = out * Scale1 + Shift1
+        if self.noise_inject:
+            out = self.noise1(out, noise=None)
+        outup = self.convup(out)
+        return outup
+
+
+####################################################################
+###############Face Dictionary Generator
+####################################################################
+def AttentionBlock(in_channel):
+    return nn.Sequential(
+        SpectralNorm(nn.Conv2d(in_channel, in_channel, 3, 1, 1)),
+        nn.LeakyReLU(0.2),
+        SpectralNorm(nn.Conv2d(in_channel, in_channel, 3, 1, 1)),
+    )
+
+class DilateResBlock(nn.Module):
+    def __init__(self, dim, dilation=[5,3] ):
+        super(DilateResBlock, self).__init__()
+        self.Res = nn.Sequential(
+            SpectralNorm(nn.Conv2d(dim, dim, 3, 1, ((3-1)//2)*dilation[0], dilation[0])),
+            nn.LeakyReLU(0.2),
+            SpectralNorm(nn.Conv2d(dim, dim, 3, 1, ((3-1)//2)*dilation[1], dilation[1])),
+        )
+    def forward(self, x):
+        out = x + self.Res(x)
+        return out
+
+
+class KeyValue(nn.Module):
+    def __init__(self, indim, keydim, valdim):
+        super(KeyValue, self).__init__()
+        self.Key = nn.Sequential(
+            SpectralNorm(nn.Conv2d(indim, keydim, kernel_size=(3,3), padding=(1,1), stride=1)),
+            nn.LeakyReLU(0.2),
+            SpectralNorm(nn.Conv2d(keydim, keydim, kernel_size=(3,3), padding=(1,1), stride=1)),
+        )
+        self.Value = nn.Sequential(
+            SpectralNorm(nn.Conv2d(indim, valdim, kernel_size=(3,3), padding=(1,1), stride=1)),
+            nn.LeakyReLU(0.2),
+            SpectralNorm(nn.Conv2d(valdim, valdim, kernel_size=(3,3), padding=(1,1), stride=1)),
+        )
+    def forward(self, x):  
+        return self.Key(x), self.Value(x)
+
+class MaskAttention(nn.Module):
+    def __init__(self, indim):
+        super(MaskAttention, self).__init__()
+        self.conv1 = nn.Sequential(
+            SpectralNorm(nn.Conv2d(indim, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)),
+            nn.LeakyReLU(0.2),
+            SpectralNorm(nn.Conv2d(indim//3, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)),
+        )
+        self.conv2 = nn.Sequential(
+            SpectralNorm(nn.Conv2d(indim, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)),
+            nn.LeakyReLU(0.2),
+            SpectralNorm(nn.Conv2d(indim//3, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)),
+        )
+        self.conv3 = nn.Sequential(
+            SpectralNorm(nn.Conv2d(indim, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)),
+            nn.LeakyReLU(0.2),
+            SpectralNorm(nn.Conv2d(indim//3, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)),
+        )
+        self.convCat = nn.Sequential(
+            SpectralNorm(nn.Conv2d(indim//3 * 3, indim, kernel_size=(3,3), padding=(1,1), stride=1)),
+            nn.LeakyReLU(0.2),
+            SpectralNorm(nn.Conv2d(indim, indim, kernel_size=(3,3), padding=(1,1), stride=1)),
+        )
+    def forward(self, x, y, z):
+        c1 = self.conv1(x)
+        c2 = self.conv2(y)
+        c3 = self.conv3(z)
+        return self.convCat(torch.cat([c1,c2,c3], dim=1))
+
+class Query(nn.Module):
+    def __init__(self, indim, quedim):
+        super(Query, self).__init__()
+        self.Query = nn.Sequential(
+            SpectralNorm(nn.Conv2d(indim, quedim, kernel_size=(3,3), padding=(1,1), stride=1)),
+            nn.LeakyReLU(0.2),
+            SpectralNorm(nn.Conv2d(quedim, quedim, kernel_size=(3,3), padding=(1,1), stride=1)),
+        )
+    def forward(self, x):
+        return self.Query(x)
+
+def roi_align_self(input, location, target_size):
+    test = (target_size.item(),target_size.item())
+    return torch.cat([F.interpolate(input[i:i+1,:,location[i,1]:location[i,3],location[i,0]:location[i,2]],test,mode='bilinear',align_corners=False) for i in range(input.size(0))],0)
+
+class FeatureExtractor(nn.Module):
+    def __init__(self, ngf = 64, key_scale = 4):#
+        super().__init__()
+
+        self.key_scale = 4
+        self.part_sizes = np.array([80,80,50,110]) #
+        self.feature_sizes = np.array([256,128,64]) # 
+
+        self.conv1 = nn.Sequential(
+                SpectralNorm(nn.Conv2d(3, ngf, 3, 2, 1)),
+                nn.LeakyReLU(0.2),
+                SpectralNorm(nn.Conv2d(ngf, ngf, 3, 1, 1)),
+            )
+        self.conv2 = nn.Sequential(
+            SpectralNorm(nn.Conv2d(ngf, ngf, 3, 1, 1)),
+            nn.LeakyReLU(0.2),
+            SpectralNorm(nn.Conv2d(ngf, ngf, 3, 1, 1))
+        )
+        self.res1 = DilateResBlock(ngf, [5,3])
+        self.res2 = DilateResBlock(ngf, [5,3])
+
+        
+        self.conv3 = nn.Sequential(
+            SpectralNorm(nn.Conv2d(ngf, ngf*2, 3, 2, 1)),
+            nn.LeakyReLU(0.2),
+            SpectralNorm(nn.Conv2d(ngf*2, ngf*2, 3, 1, 1)),
+            )
+        self.conv4 = nn.Sequential(
+            SpectralNorm(nn.Conv2d(ngf*2, ngf*2, 3, 1, 1)),
+            nn.LeakyReLU(0.2),
+            SpectralNorm(nn.Conv2d(ngf*2, ngf*2, 3, 1, 1))
+        )
+        self.res3 = DilateResBlock(ngf*2, [3,1])
+        self.res4 = DilateResBlock(ngf*2, [3,1])
+
+        self.conv5 = nn.Sequential(
+            SpectralNorm(nn.Conv2d(ngf*2, ngf*4, 3, 2, 1)),
+            nn.LeakyReLU(0.2),
+            SpectralNorm(nn.Conv2d(ngf*4, ngf*4, 3, 1, 1)),
+        )
+        self.conv6 = nn.Sequential(
+            SpectralNorm(nn.Conv2d(ngf*4, ngf*4, 3, 1, 1)),
+            nn.LeakyReLU(0.2),
+            SpectralNorm(nn.Conv2d(ngf*4, ngf*4, 3, 1, 1))
+        )
+        self.res5 = DilateResBlock(ngf*4, [1,1])
+        self.res6 = DilateResBlock(ngf*4, [1,1])
+
+        self.LE_256_Q = Query(ngf, ngf // self.key_scale)
+        self.RE_256_Q = Query(ngf, ngf // self.key_scale)
+        self.MO_256_Q = Query(ngf, ngf // self.key_scale)
+        self.LE_128_Q = Query(ngf * 2, ngf * 2 // self.key_scale)
+        self.RE_128_Q = Query(ngf * 2, ngf * 2 // self.key_scale)
+        self.MO_128_Q = Query(ngf * 2, ngf * 2 // self.key_scale)
+        self.LE_64_Q = Query(ngf * 4, ngf * 4 // self.key_scale)
+        self.RE_64_Q = Query(ngf * 4, ngf * 4 // self.key_scale)
+        self.MO_64_Q = Query(ngf * 4, ngf * 4 // self.key_scale)
+
+
+    def forward(self, img, locs):
+        le_location = locs[:,0,:].int().cpu().numpy()
+        re_location = locs[:,1,:].int().cpu().numpy()
+        no_location = locs[:,2,:].int().cpu().numpy()
+        mo_location = locs[:,3,:].int().cpu().numpy()
+        
+
+        f1_0 = self.conv1(img) 
+        f1_1 = self.res1(f1_0)
+        f2_0 = self.conv2(f1_1)
+        f2_1 = self.res2(f2_0)
+
+        f3_0 = self.conv3(f2_1) 
+        f3_1 = self.res3(f3_0)
+        f4_0 = self.conv4(f3_1)
+        f4_1 = self.res4(f4_0)
+
+        f5_0 = self.conv5(f4_1) 
+        f5_1 = self.res5(f5_0)
+        f6_0 = self.conv6(f5_1)
+        f6_1 = self.res6(f6_0)
+
+
+        ####ROI Align
+        le_part_256 = roi_align_self(f2_1.clone(), le_location//2, self.part_sizes[0]//2)
+        re_part_256 = roi_align_self(f2_1.clone(), re_location//2, self.part_sizes[1]//2)
+        mo_part_256 = roi_align_self(f2_1.clone(), mo_location//2, self.part_sizes[3]//2)
+
+        le_part_128 = roi_align_self(f4_1.clone(), le_location//4, self.part_sizes[0]//4)
+        re_part_128 = roi_align_self(f4_1.clone(), re_location//4, self.part_sizes[1]//4)
+        mo_part_128 = roi_align_self(f4_1.clone(), mo_location//4, self.part_sizes[3]//4)
+
+        le_part_64 = roi_align_self(f6_1.clone(), le_location//8, self.part_sizes[0]//8)
+        re_part_64 = roi_align_self(f6_1.clone(), re_location//8, self.part_sizes[1]//8)
+        mo_part_64 = roi_align_self(f6_1.clone(), mo_location//8, self.part_sizes[3]//8)
+
+
+        le_256_q = self.LE_256_Q(le_part_256)
+        re_256_q = self.RE_256_Q(re_part_256)
+        mo_256_q = self.MO_256_Q(mo_part_256)
+
+        le_128_q = self.LE_128_Q(le_part_128)
+        re_128_q = self.RE_128_Q(re_part_128)
+        mo_128_q = self.MO_128_Q(mo_part_128)
+
+        le_64_q = self.LE_64_Q(le_part_64)
+        re_64_q = self.RE_64_Q(re_part_64)
+        mo_64_q = self.MO_64_Q(mo_part_64)
+
+        return {'f256': f2_1, 'f128': f4_1, 'f64': f6_1,\
+            'le256': le_part_256, 're256': re_part_256, 'mo256': mo_part_256, \
+            'le128': le_part_128, 're128': re_part_128, 'mo128': mo_part_128, \
+            'le64': le_part_64, 're64': re_part_64, 'mo64': mo_part_64, \
+            'le_256_q': le_256_q, 're_256_q': re_256_q, 'mo_256_q': mo_256_q,\
+            'le_128_q': le_128_q, 're_128_q': re_128_q, 'mo_128_q': mo_128_q,\
+            'le_64_q': le_64_q, 're_64_q': re_64_q, 'mo_64_q': mo_64_q}
+
+
+class DMDNet(nn.Module):
+    def __init__(self, ngf = 64, banks_num = 128):
+        super().__init__()
+        self.part_sizes = np.array([80,80,50,110]) # size for 512
+        self.feature_sizes = np.array([256,128,64]) # size for 512
+
+        self.banks_num = banks_num
+        self.key_scale = 4
+
+        self.E_lq = FeatureExtractor(key_scale = self.key_scale)
+        self.E_hq = FeatureExtractor(key_scale = self.key_scale)
+
+        self.LE_256_KV = KeyValue(ngf, ngf // self.key_scale, ngf)
+        self.RE_256_KV = KeyValue(ngf, ngf // self.key_scale, ngf)
+        self.MO_256_KV = KeyValue(ngf, ngf // self.key_scale, ngf)
+
+        self.LE_128_KV = KeyValue(ngf * 2 , ngf * 2 // self.key_scale, ngf * 2)
+        self.RE_128_KV = KeyValue(ngf * 2 , ngf * 2 // self.key_scale, ngf * 2)
+        self.MO_128_KV = KeyValue(ngf * 2 , ngf * 2 // self.key_scale, ngf * 2)
+
+        self.LE_64_KV = KeyValue(ngf * 4 , ngf * 4 // self.key_scale, ngf * 4)
+        self.RE_64_KV = KeyValue(ngf * 4 , ngf * 4 // self.key_scale, ngf * 4)
+        self.MO_64_KV = KeyValue(ngf * 4 , ngf * 4 // self.key_scale, ngf * 4)
+
+
+        self.LE_256_Attention = AttentionBlock(64)
+        self.RE_256_Attention = AttentionBlock(64)
+        self.MO_256_Attention = AttentionBlock(64)
+
+        self.LE_128_Attention = AttentionBlock(128)
+        self.RE_128_Attention = AttentionBlock(128)
+        self.MO_128_Attention = AttentionBlock(128)
+
+        self.LE_64_Attention = AttentionBlock(256)
+        self.RE_64_Attention = AttentionBlock(256)
+        self.MO_64_Attention = AttentionBlock(256)
+
+        self.LE_256_Mask = MaskAttention(64)
+        self.RE_256_Mask = MaskAttention(64)
+        self.MO_256_Mask = MaskAttention(64)
+
+        self.LE_128_Mask = MaskAttention(128)
+        self.RE_128_Mask = MaskAttention(128)
+        self.MO_128_Mask = MaskAttention(128)
+
+        self.LE_64_Mask = MaskAttention(256)
+        self.RE_64_Mask = MaskAttention(256)
+        self.MO_64_Mask = MaskAttention(256)
+
+        self.MSDilate = MSDilateBlock(ngf*4, dilation = [4,3,2,1])
+
+        self.up1 = StyledUpBlock(ngf*4, ngf*2, noise_inject=False) #
+        self.up2 = StyledUpBlock(ngf*2, ngf, noise_inject=False) #
+        self.up3 = StyledUpBlock(ngf, ngf, noise_inject=False) #
+        self.up4 = nn.Sequential( 
+            SpectralNorm(nn.Conv2d(ngf, ngf, 3, 1, 1)),
+            nn.LeakyReLU(0.2),
+            UpResBlock(ngf),
+            UpResBlock(ngf),
+            SpectralNorm(nn.Conv2d(ngf, 3, kernel_size=3, stride=1, padding=1)),
+            nn.Tanh()
+        )
+ 
+        # define generic memory, revise register_buffer to register_parameter for backward update
+        self.register_buffer('le_256_mem_key', torch.randn(128,16,40,40))
+        self.register_buffer('re_256_mem_key', torch.randn(128,16,40,40))
+        self.register_buffer('mo_256_mem_key', torch.randn(128,16,55,55))
+        self.register_buffer('le_256_mem_value', torch.randn(128,64,40,40))
+        self.register_buffer('re_256_mem_value', torch.randn(128,64,40,40))
+        self.register_buffer('mo_256_mem_value', torch.randn(128,64,55,55))
+        
+
+        self.register_buffer('le_128_mem_key', torch.randn(128,32,20,20))
+        self.register_buffer('re_128_mem_key', torch.randn(128,32,20,20))
+        self.register_buffer('mo_128_mem_key', torch.randn(128,32,27,27))
+        self.register_buffer('le_128_mem_value', torch.randn(128,128,20,20))
+        self.register_buffer('re_128_mem_value', torch.randn(128,128,20,20))
+        self.register_buffer('mo_128_mem_value', torch.randn(128,128,27,27))
+
+        self.register_buffer('le_64_mem_key', torch.randn(128,64,10,10))
+        self.register_buffer('re_64_mem_key', torch.randn(128,64,10,10))
+        self.register_buffer('mo_64_mem_key', torch.randn(128,64,13,13))
+        self.register_buffer('le_64_mem_value', torch.randn(128,256,10,10))
+        self.register_buffer('re_64_mem_value', torch.randn(128,256,10,10))
+        self.register_buffer('mo_64_mem_value', torch.randn(128,256,13,13))
+
+    
+    def readMem(self, k, v, q):
+        sim = F.conv2d(q, k)
+        score = F.softmax(sim/sqrt(sim.size(1)), dim=1) #B * S * 1 * 1 6*128
+        sb,sn,sw,sh = score.size()
+        s_m = score.view(sb, -1).unsqueeze(1)#2*1*M
+        vb,vn,vw,vh = v.size()
+        v_in = v.view(vb, -1).repeat(sb,1,1)#2*M*(c*w*h)
+        mem_out = torch.bmm(s_m, v_in).squeeze(1).view(sb, vn, vw,vh)
+        max_inds = torch.argmax(score, dim=1).squeeze()
+        return mem_out, max_inds
+    
+
+    def memorize(self, img, locs):
+        fs = self.E_hq(img, locs)
+        LE256_key, LE256_value = self.LE_256_KV(fs['le256'])
+        RE256_key, RE256_value = self.RE_256_KV(fs['re256'])
+        MO256_key, MO256_value = self.MO_256_KV(fs['mo256'])
+
+        LE128_key, LE128_value = self.LE_128_KV(fs['le128'])
+        RE128_key, RE128_value = self.RE_128_KV(fs['re128'])
+        MO128_key, MO128_value = self.MO_128_KV(fs['mo128'])
+
+        LE64_key, LE64_value = self.LE_64_KV(fs['le64'])
+        RE64_key, RE64_value = self.RE_64_KV(fs['re64'])
+        MO64_key, MO64_value = self.MO_64_KV(fs['mo64'])
+
+        Mem256 = {'LE256Key': LE256_key, 'LE256Value': LE256_value, 'RE256Key': RE256_key, 'RE256Value': RE256_value,'MO256Key': MO256_key, 'MO256Value': MO256_value}
+        Mem128 = {'LE128Key': LE128_key, 'LE128Value': LE128_value, 'RE128Key': RE128_key, 'RE128Value': RE128_value,'MO128Key': MO128_key, 'MO128Value': MO128_value}
+        Mem64 = {'LE64Key': LE64_key, 'LE64Value': LE64_value, 'RE64Key': RE64_key, 'RE64Value': RE64_value,'MO64Key': MO64_key, 'MO64Value': MO64_value}
+ 
+        FS256 = {'LE256F':fs['le256'], 'RE256F':fs['re256'], 'MO256F':fs['mo256']}
+        FS128 = {'LE128F':fs['le128'], 'RE128F':fs['re128'], 'MO128F':fs['mo128']}
+        FS64 = {'LE64F':fs['le64'], 'RE64F':fs['re64'], 'MO64F':fs['mo64']}
+        
+        return Mem256, Mem128, Mem64
+
+    def enhancer(self, fs_in, sp_256=None, sp_128=None, sp_64=None):
+        le_256_q = fs_in['le_256_q']
+        re_256_q = fs_in['re_256_q']
+        mo_256_q = fs_in['mo_256_q']
+
+        le_128_q = fs_in['le_128_q']
+        re_128_q = fs_in['re_128_q']
+        mo_128_q = fs_in['mo_128_q']
+
+        le_64_q = fs_in['le_64_q']
+        re_64_q = fs_in['re_64_q']
+        mo_64_q = fs_in['mo_64_q']
+
+        
+        ####for 256
+        le_256_mem_g, le_256_inds = self.readMem(self.le_256_mem_key, self.le_256_mem_value, le_256_q)
+        re_256_mem_g, re_256_inds = self.readMem(self.re_256_mem_key, self.re_256_mem_value, re_256_q)
+        mo_256_mem_g, mo_256_inds = self.readMem(self.mo_256_mem_key, self.mo_256_mem_value, mo_256_q)
+
+        le_128_mem_g, le_128_inds = self.readMem(self.le_128_mem_key, self.le_128_mem_value, le_128_q)
+        re_128_mem_g, re_128_inds = self.readMem(self.re_128_mem_key, self.re_128_mem_value, re_128_q)
+        mo_128_mem_g, mo_128_inds = self.readMem(self.mo_128_mem_key, self.mo_128_mem_value, mo_128_q)
+
+        le_64_mem_g, le_64_inds = self.readMem(self.le_64_mem_key, self.le_64_mem_value, le_64_q)
+        re_64_mem_g, re_64_inds = self.readMem(self.re_64_mem_key, self.re_64_mem_value, re_64_q)
+        mo_64_mem_g, mo_64_inds = self.readMem(self.mo_64_mem_key, self.mo_64_mem_value, mo_64_q)
+
+        if sp_256 is not None and sp_128 is not None and sp_64 is not None:
+            le_256_mem_s, _ = self.readMem(sp_256['LE256Key'], sp_256['LE256Value'], le_256_q)
+            re_256_mem_s, _ = self.readMem(sp_256['RE256Key'], sp_256['RE256Value'], re_256_q)
+            mo_256_mem_s, _ = self.readMem(sp_256['MO256Key'], sp_256['MO256Value'], mo_256_q)
+            le_256_mask = self.LE_256_Mask(fs_in['le256'],le_256_mem_s,le_256_mem_g)
+            le_256_mem = le_256_mask*le_256_mem_s + (1-le_256_mask)*le_256_mem_g
+            re_256_mask = self.RE_256_Mask(fs_in['re256'],re_256_mem_s,re_256_mem_g)
+            re_256_mem = re_256_mask*re_256_mem_s + (1-re_256_mask)*re_256_mem_g
+            mo_256_mask = self.MO_256_Mask(fs_in['mo256'],mo_256_mem_s,mo_256_mem_g)
+            mo_256_mem = mo_256_mask*mo_256_mem_s + (1-mo_256_mask)*mo_256_mem_g
+
+            le_128_mem_s, _ = self.readMem(sp_128['LE128Key'], sp_128['LE128Value'], le_128_q)
+            re_128_mem_s, _ = self.readMem(sp_128['RE128Key'], sp_128['RE128Value'], re_128_q)
+            mo_128_mem_s, _ = self.readMem(sp_128['MO128Key'], sp_128['MO128Value'], mo_128_q)
+            le_128_mask = self.LE_128_Mask(fs_in['le128'],le_128_mem_s,le_128_mem_g)
+            le_128_mem = le_128_mask*le_128_mem_s + (1-le_128_mask)*le_128_mem_g
+            re_128_mask = self.RE_128_Mask(fs_in['re128'],re_128_mem_s,re_128_mem_g)
+            re_128_mem = re_128_mask*re_128_mem_s + (1-re_128_mask)*re_128_mem_g
+            mo_128_mask = self.MO_128_Mask(fs_in['mo128'],mo_128_mem_s,mo_128_mem_g)
+            mo_128_mem = mo_128_mask*mo_128_mem_s + (1-mo_128_mask)*mo_128_mem_g
+
+            le_64_mem_s, _ = self.readMem(sp_64['LE64Key'], sp_64['LE64Value'], le_64_q)
+            re_64_mem_s, _ = self.readMem(sp_64['RE64Key'], sp_64['RE64Value'], re_64_q)
+            mo_64_mem_s, _ = self.readMem(sp_64['MO64Key'], sp_64['MO64Value'], mo_64_q)
+            le_64_mask = self.LE_64_Mask(fs_in['le64'],le_64_mem_s,le_64_mem_g)
+            le_64_mem = le_64_mask*le_64_mem_s + (1-le_64_mask)*le_64_mem_g
+            re_64_mask = self.RE_64_Mask(fs_in['re64'],re_64_mem_s,re_64_mem_g)
+            re_64_mem = re_64_mask*re_64_mem_s + (1-re_64_mask)*re_64_mem_g
+            mo_64_mask = self.MO_64_Mask(fs_in['mo64'],mo_64_mem_s,mo_64_mem_g)
+            mo_64_mem = mo_64_mask*mo_64_mem_s + (1-mo_64_mask)*mo_64_mem_g
+        else:
+            le_256_mem = le_256_mem_g
+            re_256_mem = re_256_mem_g
+            mo_256_mem = mo_256_mem_g
+            le_128_mem = le_128_mem_g
+            re_128_mem = re_128_mem_g
+            mo_128_mem = mo_128_mem_g
+            le_64_mem = le_64_mem_g
+            re_64_mem = re_64_mem_g
+            mo_64_mem = mo_64_mem_g
+
+        le_256_mem_norm = adaptive_instance_normalization_4D(le_256_mem, fs_in['le256'])
+        re_256_mem_norm = adaptive_instance_normalization_4D(re_256_mem, fs_in['re256'])
+        mo_256_mem_norm = adaptive_instance_normalization_4D(mo_256_mem, fs_in['mo256'])
+        
+        ####for 128
+        le_128_mem_norm = adaptive_instance_normalization_4D(le_128_mem, fs_in['le128'])
+        re_128_mem_norm = adaptive_instance_normalization_4D(re_128_mem, fs_in['re128'])
+        mo_128_mem_norm = adaptive_instance_normalization_4D(mo_128_mem, fs_in['mo128'])
+        
+        ####for 64
+        le_64_mem_norm = adaptive_instance_normalization_4D(le_64_mem, fs_in['le64'])
+        re_64_mem_norm = adaptive_instance_normalization_4D(re_64_mem, fs_in['re64'])
+        mo_64_mem_norm = adaptive_instance_normalization_4D(mo_64_mem, fs_in['mo64'])
+    
+
+        EnMem256 = {'LE256Norm': le_256_mem_norm, 'RE256Norm': re_256_mem_norm, 'MO256Norm': mo_256_mem_norm}
+        EnMem128 = {'LE128Norm': le_128_mem_norm, 'RE128Norm': re_128_mem_norm, 'MO128Norm': mo_128_mem_norm}
+        EnMem64 = {'LE64Norm': le_64_mem_norm, 'RE64Norm': re_64_mem_norm, 'MO64Norm': mo_64_mem_norm}
+        Ind256 = {'LE': le_256_inds, 'RE': re_256_inds, 'MO': mo_256_inds}
+        Ind128 = {'LE': le_128_inds, 'RE': re_128_inds, 'MO': mo_128_inds}
+        Ind64 = {'LE': le_64_inds, 'RE': re_64_inds, 'MO': mo_64_inds}
+        return EnMem256, EnMem128, EnMem64, Ind256, Ind128, Ind64
+
+    def reconstruct(self, fs_in, locs, memstar):
+        le_256_mem_norm, re_256_mem_norm, mo_256_mem_norm = memstar[0]['LE256Norm'], memstar[0]['RE256Norm'], memstar[0]['MO256Norm']
+        le_128_mem_norm, re_128_mem_norm, mo_128_mem_norm = memstar[1]['LE128Norm'], memstar[1]['RE128Norm'], memstar[1]['MO128Norm']
+        le_64_mem_norm, re_64_mem_norm, mo_64_mem_norm = memstar[2]['LE64Norm'], memstar[2]['RE64Norm'], memstar[2]['MO64Norm']
+
+        le_256_final = self.LE_256_Attention(le_256_mem_norm - fs_in['le256']) * le_256_mem_norm + fs_in['le256']
+        re_256_final = self.RE_256_Attention(re_256_mem_norm - fs_in['re256']) * re_256_mem_norm + fs_in['re256']
+        mo_256_final = self.MO_256_Attention(mo_256_mem_norm - fs_in['mo256']) * mo_256_mem_norm + fs_in['mo256']
+        
+        le_128_final = self.LE_128_Attention(le_128_mem_norm - fs_in['le128']) * le_128_mem_norm + fs_in['le128']
+        re_128_final = self.RE_128_Attention(re_128_mem_norm - fs_in['re128']) * re_128_mem_norm + fs_in['re128']
+        mo_128_final = self.MO_128_Attention(mo_128_mem_norm - fs_in['mo128']) * mo_128_mem_norm + fs_in['mo128']
+        
+        le_64_final = self.LE_64_Attention(le_64_mem_norm - fs_in['le64']) * le_64_mem_norm + fs_in['le64']
+        re_64_final = self.RE_64_Attention(re_64_mem_norm - fs_in['re64']) * re_64_mem_norm + fs_in['re64']
+        mo_64_final = self.MO_64_Attention(mo_64_mem_norm - fs_in['mo64']) * mo_64_mem_norm + fs_in['mo64']
+
+
+        le_location = locs[:,0,:]
+        re_location = locs[:,1,:]
+        mo_location = locs[:,3,:]
+
+        # Somehow with latest Torch it doesn't like numpy wrappers anymore
+        
+        # le_location = le_location.cpu().int().numpy()
+        # re_location = re_location.cpu().int().numpy()
+        # mo_location = mo_location.cpu().int().numpy()
+        le_location = le_location.cpu().int()
+        re_location = re_location.cpu().int()
+        mo_location = mo_location.cpu().int()
+
+        up_in_256 = fs_in['f256'].clone()# * 0
+        up_in_128 = fs_in['f128'].clone()# * 0
+        up_in_64 = fs_in['f64'].clone()# * 0
+
+        for i in range(fs_in['f256'].size(0)):
+            up_in_256[i:i+1,:,le_location[i,1]//2:le_location[i,3]//2,le_location[i,0]//2:le_location[i,2]//2] = F.interpolate(le_256_final[i:i+1,:,:,:].clone(), (le_location[i,3]//2-le_location[i,1]//2,le_location[i,2]//2-le_location[i,0]//2),mode='bilinear',align_corners=False)
+            up_in_256[i:i+1,:,re_location[i,1]//2:re_location[i,3]//2,re_location[i,0]//2:re_location[i,2]//2] = F.interpolate(re_256_final[i:i+1,:,:,:].clone(), (re_location[i,3]//2-re_location[i,1]//2,re_location[i,2]//2-re_location[i,0]//2),mode='bilinear',align_corners=False)
+            up_in_256[i:i+1,:,mo_location[i,1]//2:mo_location[i,3]//2,mo_location[i,0]//2:mo_location[i,2]//2] = F.interpolate(mo_256_final[i:i+1,:,:,:].clone(), (mo_location[i,3]//2-mo_location[i,1]//2,mo_location[i,2]//2-mo_location[i,0]//2),mode='bilinear',align_corners=False)
+            
+            up_in_128[i:i+1,:,le_location[i,1]//4:le_location[i,3]//4,le_location[i,0]//4:le_location[i,2]//4] = F.interpolate(le_128_final[i:i+1,:,:,:].clone(), (le_location[i,3]//4-le_location[i,1]//4,le_location[i,2]//4-le_location[i,0]//4),mode='bilinear',align_corners=False)
+            up_in_128[i:i+1,:,re_location[i,1]//4:re_location[i,3]//4,re_location[i,0]//4:re_location[i,2]//4] = F.interpolate(re_128_final[i:i+1,:,:,:].clone(), (re_location[i,3]//4-re_location[i,1]//4,re_location[i,2]//4-re_location[i,0]//4),mode='bilinear',align_corners=False)
+            up_in_128[i:i+1,:,mo_location[i,1]//4:mo_location[i,3]//4,mo_location[i,0]//4:mo_location[i,2]//4] = F.interpolate(mo_128_final[i:i+1,:,:,:].clone(), (mo_location[i,3]//4-mo_location[i,1]//4,mo_location[i,2]//4-mo_location[i,0]//4),mode='bilinear',align_corners=False)
+
+            up_in_64[i:i+1,:,le_location[i,1]//8:le_location[i,3]//8,le_location[i,0]//8:le_location[i,2]//8] = F.interpolate(le_64_final[i:i+1,:,:,:].clone(), (le_location[i,3]//8-le_location[i,1]//8,le_location[i,2]//8-le_location[i,0]//8),mode='bilinear',align_corners=False)
+            up_in_64[i:i+1,:,re_location[i,1]//8:re_location[i,3]//8,re_location[i,0]//8:re_location[i,2]//8] = F.interpolate(re_64_final[i:i+1,:,:,:].clone(), (re_location[i,3]//8-re_location[i,1]//8,re_location[i,2]//8-re_location[i,0]//8),mode='bilinear',align_corners=False)
+            up_in_64[i:i+1,:,mo_location[i,1]//8:mo_location[i,3]//8,mo_location[i,0]//8:mo_location[i,2]//8] = F.interpolate(mo_64_final[i:i+1,:,:,:].clone(), (mo_location[i,3]//8-mo_location[i,1]//8,mo_location[i,2]//8-mo_location[i,0]//8),mode='bilinear',align_corners=False)
+        
+        ms_in_64 = self.MSDilate(fs_in['f64'].clone())
+        fea_up1 = self.up1(ms_in_64, up_in_64)
+        fea_up2 = self.up2(fea_up1, up_in_128) #
+        fea_up3 = self.up3(fea_up2, up_in_256) #
+        output = self.up4(fea_up3) #
+        return output
+
+    def generate_specific_dictionary(self, sp_imgs=None, sp_locs=None):
+        return self.memorize(sp_imgs, sp_locs)
+
+    def forward(self, lq=None, loc=None, sp_256 = None, sp_128 = None, sp_64 = None):
+        try:
+            fs_in = self.E_lq(lq, loc) # low quality images
+        except Exception as e:
+            print(e)
+
+        GeMemNorm256, GeMemNorm128, GeMemNorm64, Ind256, Ind128, Ind64 = self.enhancer(fs_in)
+        GeOut = self.reconstruct(fs_in, loc, memstar = [GeMemNorm256, GeMemNorm128, GeMemNorm64])
+        if sp_256 is not None and sp_128 is not None and sp_64 is not None:
+            GSMemNorm256, GSMemNorm128, GSMemNorm64, _, _, _ = self.enhancer(fs_in, sp_256, sp_128, sp_64)
+            GSOut = self.reconstruct(fs_in, loc, memstar = [GSMemNorm256, GSMemNorm128, GSMemNorm64])
+        else:
+            GSOut = None
+        return GeOut, GSOut
+
+class UpResBlock(nn.Module):
+    def __init__(self, dim, conv_layer = nn.Conv2d, norm_layer = nn.BatchNorm2d):
+        super(UpResBlock, self).__init__()
+        self.Model = nn.Sequential(
+            SpectralNorm(conv_layer(dim, dim, 3, 1, 1)),
+            nn.LeakyReLU(0.2),
+            SpectralNorm(conv_layer(dim, dim, 3, 1, 1)),
+        )
+    def forward(self, x):
+        out = x + self.Model(x)
+        return out

roop/processors/Enhance_GFPGAN.py

@@ -0,0 +1,72 @@
+from typing import Any, List, Callable
+import cv2 
+import numpy as np
+import onnxruntime
+import roop.globals
+
+from roop.typing import Face, Frame, FaceSet
+from roop.utilities import resolve_relative_path
+
+
+# THREAD_LOCK = threading.Lock()
+
+
+class Enhance_GFPGAN():
+
+    model_gfpgan = None
+    name = None
+    devicename = None
+
+    processorname = 'gfpgan'
+    type = 'enhance'
+
+
+    def Initialize(self, devicename):
+        if self.model_gfpgan is None:
+            model_path = resolve_relative_path('../models/GFPGANv1.4.onnx')
+            self.model_gfpgan = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
+            # replace Mac mps with cpu for the moment
+            devicename = devicename.replace('mps', 'cpu')
+            self.devicename = devicename
+
+        self.name = self.model_gfpgan.get_inputs()[0].name
+
+    def Run(self, source_faceset: FaceSet, target_face: Face, temp_frame: Frame) -> Frame:
+        # preprocess
+        input_size = temp_frame.shape[1]
+        temp_frame = cv2.resize(temp_frame, (512, 512), cv2.INTER_CUBIC)
+
+        temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB)
+        temp_frame = temp_frame.astype('float32') / 255.0
+        temp_frame = (temp_frame - 0.5) / 0.5
+        temp_frame = np.expand_dims(temp_frame, axis=0).transpose(0, 3, 1, 2)
+
+        io_binding = self.model_gfpgan.io_binding()           
+        io_binding.bind_cpu_input("input", temp_frame)
+        io_binding.bind_output("1288", self.devicename)
+        self.model_gfpgan.run_with_iobinding(io_binding)
+        ort_outs = io_binding.copy_outputs_to_cpu()
+        result = ort_outs[0][0]
+
+        # post-process
+        result = np.clip(result, -1, 1)
+        result = (result + 1) / 2
+        result = result.transpose(1, 2, 0) * 255.0
+        result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
+        scale_factor = int(result.shape[1] / input_size)       
+        return result.astype(np.uint8), scale_factor
+
+
+    def Release(self):
+        self.model_gfpgan = None
+
+
+
+
+
+
+
+
+
+
+

roop/processors/Enhance_GPEN.py

@@ -0,0 +1,58 @@
+from typing import Any, List, Callable
+import cv2 
+import numpy as np
+import onnxruntime
+import roop.globals
+
+from roop.typing import Face, Frame, FaceSet
+from roop.utilities import resolve_relative_path
+
+
+class Enhance_GPEN():
+
+    model_gpen = None
+    name = None
+    devicename = None
+
+    processorname = 'gpen'
+    type = 'enhance'
+
+
+    def Initialize(self, devicename):
+        if self.model_gpen is None:
+            model_path = resolve_relative_path('../models/GPEN-BFR-512.onnx')
+            self.model_gpen = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
+            # replace Mac mps with cpu for the moment
+            devicename = devicename.replace('mps', 'cpu')
+            self.devicename = devicename
+
+        self.name = self.model_gpen.get_inputs()[0].name
+
+    def Run(self, source_faceset: FaceSet, target_face: Face, temp_frame: Frame) -> Frame:
+        # preprocess
+        input_size = temp_frame.shape[1]
+        temp_frame = cv2.resize(temp_frame, (512, 512), cv2.INTER_CUBIC)
+
+        temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB)
+        temp_frame = temp_frame.astype('float32') / 255.0
+        temp_frame = (temp_frame - 0.5) / 0.5
+        temp_frame = np.expand_dims(temp_frame, axis=0).transpose(0, 3, 1, 2)
+
+        io_binding = self.model_gpen.io_binding()           
+        io_binding.bind_cpu_input("input", temp_frame)
+        io_binding.bind_output("output", self.devicename)
+        self.model_gpen.run_with_iobinding(io_binding)
+        ort_outs = io_binding.copy_outputs_to_cpu()
+        result = ort_outs[0][0]
+
+        # post-process
+        result = np.clip(result, -1, 1)
+        result = (result + 1) / 2
+        result = result.transpose(1, 2, 0) * 255.0
+        result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
+        scale_factor = int(result.shape[1] / input_size)       
+        return result.astype(np.uint8), scale_factor
+
+
+    def Release(self):
+        self.model_gpen = None

roop/processors/Enhance_RestoreFormerPPlus.py

@@ -0,0 +1,59 @@
+from typing import Any, List, Callable
+import cv2 
+import numpy as np
+import onnxruntime
+import roop.globals
+
+from roop.typing import Face, Frame, FaceSet
+from roop.utilities import resolve_relative_path
+
+class Enhance_RestoreFormerPPlus():
+    model_restoreformerpplus = None
+    devicename = None
+    name = None
+
+    processorname = 'restoreformer++'
+    type = 'enhance'
+    
+
+    def Initialize(self, devicename:str):
+        if self.model_restoreformerpplus is None:
+            # replace Mac mps with cpu for the moment
+            devicename = devicename.replace('mps', 'cpu')
+            self.devicename = devicename
+            model_path = resolve_relative_path('../models/restoreformer_plus_plus.onnx')
+            self.model_restoreformerpplus = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
+            self.model_inputs = self.model_restoreformerpplus.get_inputs()
+            model_outputs = self.model_restoreformerpplus.get_outputs()
+            self.io_binding = self.model_restoreformerpplus.io_binding()
+            self.io_binding.bind_output(model_outputs[0].name, self.devicename)
+
+    def Run(self, source_faceset: FaceSet, target_face: Face, temp_frame: Frame) -> Frame:
+        # preprocess
+        input_size = temp_frame.shape[1]
+        temp_frame = cv2.resize(temp_frame, (512, 512), cv2.INTER_CUBIC)
+        temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB)
+        temp_frame = temp_frame.astype('float32') / 255.0
+        temp_frame = (temp_frame - 0.5) / 0.5
+        temp_frame = np.expand_dims(temp_frame, axis=0).transpose(0, 3, 1, 2)
+        
+        self.io_binding.bind_cpu_input(self.model_inputs[0].name, temp_frame) # .astype(np.float32)
+        self.model_restoreformerpplus.run_with_iobinding(self.io_binding)
+        ort_outs = self.io_binding.copy_outputs_to_cpu()
+        result = ort_outs[0][0]
+        del ort_outs 
+        
+        result = np.clip(result, -1, 1)
+        result = (result + 1) / 2
+        result = result.transpose(1, 2, 0) * 255.0
+        result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
+        scale_factor = int(result.shape[1] / input_size)       
+        return result.astype(np.uint8), scale_factor
+
+
+    def Release(self):
+        del self.model_restoreformerpplus
+        self.model_restoreformerpplus = None
+        del self.io_binding
+        self.io_binding = None
+

roop/processors/FaceSwapInsightFace.py

@@ -0,0 +1,66 @@
+import roop.globals
+import cv2
+import numpy as np
+import onnx
+import onnxruntime
+
+from roop.typing import Face, Frame
+from roop.utilities import resolve_relative_path
+
+
+
+class FaceSwapInsightFace():
+    model_swap_insightface = None
+
+
+    processorname = 'faceswap'
+    type = 'swap'
+
+
+    def Initialize(self, devicename):
+        if self.model_swap_insightface is None:
+            model_path = resolve_relative_path('../models/inswapper_128.onnx')
+            graph = onnx.load(model_path).graph
+            self.emap = onnx.numpy_helper.to_array(graph.initializer[-1])
+            devicename = devicename.replace('mps', 'cpu')
+            self.devicename = devicename
+            self.input_mean = 0.0
+            self.input_std = 255.0
+            #cuda_options = {"arena_extend_strategy": "kSameAsRequested", 'cudnn_conv_algo_search': 'DEFAULT'}            
+            sess_options = onnxruntime.SessionOptions()
+            sess_options.enable_cpu_mem_arena = False            
+            self.model_swap_insightface = onnxruntime.InferenceSession(model_path, sess_options, providers=roop.globals.execution_providers)
+            # replace Mac mps with cpu for the moment
+
+
+    
+    def Run(self, source_face: Face, target_face: Face, temp_frame: Frame) -> Frame:
+        blob = cv2.dnn.blobFromImage(temp_frame, 1.0 / self.input_std, (128, 128),
+                                      (self.input_mean, self.input_mean, self.input_mean), swapRB=True)
+        latent = source_face.normed_embedding.reshape((1,-1))
+        latent = np.dot(latent, self.emap)
+        latent /= np.linalg.norm(latent)
+        io_binding = self.model_swap_insightface.io_binding()           
+        io_binding.bind_cpu_input("target", blob)
+        io_binding.bind_cpu_input("source", latent)
+        io_binding.bind_output("output", self.devicename)
+        self.model_swap_insightface.run_with_iobinding(io_binding)
+        ort_outs = io_binding.copy_outputs_to_cpu()[0]
+        img_fake = ort_outs.transpose((0,2,3,1))[0]
+        return np.clip(255 * img_fake, 0, 255).astype(np.uint8)[:,:,::-1]
+
+
+        img_fake, M = self.model_swap_insightface.get(temp_frame, target_face, source_face, paste_back=False)
+    #    target_face.matrix = M
+    #    return img_fake 
+
+
+    def Release(self):
+        del self.model_swap_insightface
+        self.model_swap_insightface = None
+
+
+                
+
+
+

roop/processors/Mask_Clip2Seg.py

@@ -0,0 +1,89 @@
+import cv2
+import numpy as np
+import torch
+import threading
+from torchvision import transforms
+from clip.clipseg import CLIPDensePredT
+import numpy as np
+
+from roop.typing import Frame
+
+THREAD_LOCK_CLIP = threading.Lock()
+
+
+class Mask_Clip2Seg():
+
+    model_clip = None
+
+    processorname = 'clip2seg'
+    type = 'mask'
+
+
+    def Initialize(self, devicename):
+        if self.model_clip is None:
+            self.model_clip = CLIPDensePredT(version='ViT-B/16', reduce_dim=64, complex_trans_conv=True)
+            self.model_clip.eval();
+            self.model_clip.load_state_dict(torch.load('models/CLIP/rd64-uni-refined.pth', map_location=torch.device('cpu')), strict=False)
+
+        device = torch.device(devicename)
+        self.model_clip.to(device)
+
+
+    def Run(self, img1, keywords:str) -> Frame:
+        if keywords is None or len(keywords) < 1 or img1 is None:
+            return img1
+        
+        source_image_small = cv2.resize(img1, (256,256))
+        
+        img_mask = np.full((source_image_small.shape[0],source_image_small.shape[1]), 0, dtype=np.float32)
+        mask_border = 1
+        l = 0
+        t = 0
+        r = 1
+        b = 1
+        
+        mask_blur = 5
+        clip_blur = 5
+        
+        img_mask = cv2.rectangle(img_mask, (mask_border+int(l), mask_border+int(t)), 
+                                (256 - mask_border-int(r), 256-mask_border-int(b)), (255, 255, 255), -1)    
+        img_mask = cv2.GaussianBlur(img_mask, (mask_blur*2+1,mask_blur*2+1), 0)    
+        img_mask /= 255
+
+        
+        input_image = source_image_small
+
+        transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            transforms.Resize((256, 256)),
+        ])
+        img = transform(input_image).unsqueeze(0)
+
+        thresh = 0.5
+        prompts = keywords.split(',')
+        with THREAD_LOCK_CLIP:
+            with torch.no_grad():
+                preds = self.model_clip(img.repeat(len(prompts),1,1,1), prompts)[0]
+        clip_mask = torch.sigmoid(preds[0][0])
+        for i in range(len(prompts)-1):
+            clip_mask += torch.sigmoid(preds[i+1][0])
+           
+        clip_mask = clip_mask.data.cpu().numpy()
+        np.clip(clip_mask, 0, 1)
+        
+        clip_mask[clip_mask>thresh] = 1.0
+        clip_mask[clip_mask<=thresh] = 0.0
+        kernel = np.ones((5, 5), np.float32)
+        clip_mask = cv2.dilate(clip_mask, kernel, iterations=1)
+        clip_mask = cv2.GaussianBlur(clip_mask, (clip_blur*2+1,clip_blur*2+1), 0)
+       
+        img_mask *= clip_mask
+        img_mask[img_mask<0.0] = 0.0
+        return img_mask
+       
+
+
+    def Release(self):
+        self.model_clip = None
+

roop/processors/Mask_XSeg.py

@@ -0,0 +1,54 @@
+import numpy as np
+import cv2
+import onnxruntime
+import threading
+import roop.globals
+
+from roop.typing import Frame
+from roop.utilities import resolve_relative_path
+
+THREAD_LOCK_CLIP = threading.Lock()
+
+
+class Mask_XSeg():
+
+    model_xseg = None
+
+    processorname = 'xseg'
+    type = 'mask'
+
+
+    def Initialize(self, devicename):
+        if self.model_xseg is None:
+            model_path = resolve_relative_path('../models/xseg.onnx')
+            onnxruntime.set_default_logger_severity(3)
+            self.model_xseg = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
+            self.model_inputs = self.model_xseg.get_inputs()
+            self.model_outputs = self.model_xseg.get_outputs()
+
+            # replace Mac mps with cpu for the moment
+            devicename = devicename.replace('mps', 'cpu')
+            self.devicename = devicename
+
+
+    def Run(self, img1, keywords:str) -> Frame:
+        temp_frame = cv2.resize(img1, (256, 256), cv2.INTER_CUBIC)
+        temp_frame = temp_frame.astype('float32') / 255.0
+        temp_frame = temp_frame[None, ...]
+        io_binding = self.model_xseg.io_binding()           
+        io_binding.bind_cpu_input(self.model_inputs[0].name, temp_frame)
+        io_binding.bind_output(self.model_outputs[0].name, self.devicename)
+        self.model_xseg.run_with_iobinding(io_binding)
+        ort_outs = io_binding.copy_outputs_to_cpu()
+        result = ort_outs[0][0]
+        result = np.clip(result, 0, 1.0)
+        result[result < 0.1] = 0
+        # invert values to mask areas to keep
+        result = 1.0 - result
+        return result       
+
+
+    def Release(self):
+        self.model_xseg = None
+
+