Update vtoonify_model.py

vtoonify_model.py

@@ -193,27 +193,28 @@ class Model():
             exstyle, _ = self.load_model(style_type)
         if exstyle is None:
             return np.zeros((256, 256, 3), np.uint8), 'Oops, 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]
-
-            x = self.transform(aligned_face).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], 
+            
+        try:
+           with torch.no_grad():
+               if self.color_transfer:
+                   s_w = exstyle
+               else:
+                   s_w = instyle.clone()
+                   s_w[:, :7] = exstyle[:, :7]
+
+               x = self.transform(aligned_face).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()
-            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)
-            logging.info(f"Output from VToonify shape: {y_tilde.shape}")
-            print('*** Toonify %dx%d image with style of %s'%(y_tilde.shape[2], y_tilde.shape[3], style_type))
+               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)
+               logging.info(f"Output from VToonify shape: {y_tilde.shape}")
+               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()