Update vtoonify_model.py

vtoonify_model.py

@@ -186,14 +186,18 @@ class Model():
 
         return np.array(img)
 
+    
     def image_toonify(self, aligned_face: np.ndarray, instyle: torch.Tensor, exstyle: torch.Tensor, style_degree: float, style_type: str) -> tuple:
         if instyle is None or aligned_face is None:
+            logging.error("Invalid input: instyle or aligned_face is None.")
             return np.zeros((256, 256, 3), np.uint8), 'Oops, 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:
+            logging.error("Failed to load style model.")
             return np.zeros((256, 256, 3), np.uint8), 'Oops, something wrong with the style type. Please go to Step 1 and load model again.'
-            
+    
         try:
             with torch.no_grad():
                if self.color_transfer:
@@ -202,7 +206,9 @@ class Model():
                    s_w = instyle.clone()
                    s_w[:, :7] = exstyle[:, :7]
 
-               x = self.transform(aligned_face).unsqueeze(dim=0).to(self.device)
+            # Ensure the input is resized to 256x256
+               aligned_face_resized = cv2.resize(aligned_face, (256, 256))
+               x = self.transform(aligned_face_resized).unsqueeze(dim=0).to(self.device)
                logging.info(f"Input to VToonify shape: {x.shape}")
                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()
@@ -213,10 +219,10 @@ class Model():
                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)
+    
         except Exception as e:
             logging.error(f"Error during model execution: {e}")
             return np.zeros((256, 256, 3), np.uint8), f"Error during processing: {str(e)}"
-           
     def tensor2cv2(self, img):
         """Convert a tensor image to OpenCV format."""
         tmp = ((img.cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8).copy()