Update vtoonify_model.py

vtoonify_model.py

@@ -201,30 +201,26 @@ class Model():
 
         return 'input.mp4', instyle, 'Successfully rescale the video to (%d, %d)'%(bottom-top, right-left)
     
-    def image_toonify(self, aligned_face: np.ndarray, instyle: torch.Tensor, exstyle: torch.Tensor, style_degree: float, style_type: str) -> tuple[np.ndarray, str]:
-        #print(style_type + ' ' + self.style_name)
-        if instyle is None or aligned_face is None:
-            return np.zeros((256,256,3), np.uint8), 'Opps, something wrong with the input. Please go to Step 2 and Rescale Image/First Frame again.'
-        if self.style_name != style_type:
-            exstyle, _  = self.load_model(style_type)
-        if exstyle is None:
-            return np.zeros((256,256,3), np.uint8), 'Opps, something wrong with the style type. Please go to Step 1 and load model again.'
-        with torch.no_grad():
-            if self.color_transfer:
-                s_w = exstyle
-            else:
-                s_w = instyle.clone()
-                s_w[:,:7] = exstyle[:,:7]
+def image_toonify(self, aligned_face: np.ndarray, instyle: torch.Tensor, exstyle: torch.Tensor, style_degree: float, style_type: str) -> tuple[np.ndarray, str]:
+    if instyle is None or aligned_face is None:
+        return np.zeros((256, 256, 3), np.uint8), 'Opps, something wrong with the input. Please go to Step 2 and Rescale Image/First Frame again.'
+    if self.style_name != style_type:
+        exstyle, _ = self.load_model(style_type)
+    if exstyle is None:
+        return np.zeros((256, 256, 3), np.uint8), 'Opps, something wrong with the style type. Please go to Step 1 and load model again.'
+    with torch.no_grad():
+        s_w = instyle.clone()
+        s_w[:, :7] = exstyle[:, :7]
+
+        x = self.transform(aligned_face).unsqueeze(dim=0).to(self.device)
+        x_p = F.interpolate(self.parsingpredictor(2 * (F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)))[0],
+                            scale_factor=0.5, recompute_scale_factor=False).detach()
+        inputs = torch.cat((x, x_p / 16.), dim=1)
+        y_tilde = self.vtoonify(inputs, s_w.repeat(inputs.size(0), 1, 1), d_s=style_degree)
+        y_tilde = torch.clamp(y_tilde, -1, 1)
+    print('*** Toonify %dx%d image with style of %s' % (y_tilde.shape[2], y_tilde.shape[3], style_type))
+    return ((y_tilde[0].cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8), 'Successfully toonify the image with style of %s' % (self.style_name)
 
-            x = self.transform(aligned_face).unsqueeze(dim=0).to(self.device)
-            x_p = F.interpolate(self.parsingpredictor(2*(F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)))[0], 
-                                scale_factor=0.5, recompute_scale_factor=False).detach()
-            inputs = torch.cat((x, x_p/16.), dim=1)
-            y_tilde = self.vtoonify(inputs, s_w.repeat(inputs.size(0), 1, 1), d_s = style_degree)        
-            y_tilde = torch.clamp(y_tilde, -1, 1)
-        print('*** Toonify %dx%d image with style of %s'%(y_tilde.shape[2], y_tilde.shape[3], style_type))
-        return ((y_tilde[0].cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8), 'Successfully toonify the image with style of %s'%(self.style_name)
-    
     def video_tooniy(self, aligned_video: str, instyle: torch.Tensor, exstyle: torch.Tensor, style_degree: float, style_type: str) -> tuple[str, str]:
         #print(style_type + ' ' + self.style_name)
         if aligned_video is None: