Update vtoonify_model.py

vtoonify_model.py

@@ -92,7 +92,7 @@ class Model():
         else:
             self.color_transfer = False
         if style_type not in self.style_types.keys():
-            return torch.zeros(1,18,512).to(self.device), 'Oops, wrong Style Type. Please select a valid model.'
+            return None, 'Oops, wrong Style Type. Please select a valid model.'
         model_path, ind = self.style_types[style_type]
         style_path = os.path.join('models',os.path.dirname(model_path),'exstyle_code.npy')
         self.vtoonify.load_state_dict(torch.load(huggingface_hub.hf_hub_download(MODEL_REPO,'models/'+model_path), 
@@ -106,7 +106,7 @@ class Model():
     def detect_and_align(self, frame, top, bottom, left, right, return_para=False):
         message = 'Error: no face detected! Please retry or change the photo.'
         paras = get_video_crop_parameter(frame, self.landmarkpredictor, [left, right, top, bottom])
-        instyle = torch.zeros(1,18,512).to(self.device)
+        instyle = None
         h, w, scale = 0, 0, 0
         if paras is not None:
             h,w,top,bottom,left,right,scale = paras
@@ -136,16 +136,18 @@ class Model():
     #@torch.inference_mode()
     def detect_and_align_image(self, image: str, top: int, bottom: int, left: int, right: int
                               ) -> tuple[np.ndarray, torch.Tensor, str]:
-        
+        if image is None:
+            return np.zeros((256,256,3), np.uint8), None, 'Error: fail to load empty file.'
         frame = cv2.imread(image)
         if frame is None:
-            return np.zeros((256,256,3), np.uint8), torch.zeros(1,18,512).to(self.device), 'Error: fail to load the image.'       
+            return np.zeros((256,256,3), np.uint8), None, 'Error: fail to load the image.'       
         frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
         return self.detect_and_align(frame, top, bottom, left, right)
     
     def detect_and_align_video(self, video: str, top: int, bottom: int, left: int, right: int
                               ) -> tuple[np.ndarray, torch.Tensor, str]:
-
+        if video is None:
+            return np.zeros((256,256,3), np.uint8), None, 'Error: fail to load empty file.'
         video_cap = cv2.VideoCapture(video)
         if video_cap.get(7) == 0:
             video_cap.release()
@@ -157,7 +159,9 @@ class Model():
     
     def detect_and_align_full_video(self, video: str, top: int, bottom: int, left: int, right: int) -> tuple[str, torch.Tensor, str]:
         message = 'Error: no face detected! Please retry or change the video.'
-        instyle = torch.zeros(1,18,512).to(self.device)
+        instyle = None
+        if video is None:
+            return 'default.mp4', instyle, 'Error: fail to load empty file.'
         video_cap = cv2.VideoCapture(video)
         if video_cap.get(7) == 0:
             video_cap.release()
@@ -212,6 +216,8 @@ class Model():
         return ((y_tilde[0].cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8), 'Successfully toonify the image'
     
     def video_tooniy(self, aligned_video: str, instyle: torch.Tensor, exstyle: torch.Tensor, style_degree: float) -> tuple[str, str]:
+        if aligned_video is None:
+            return 'output.mp4', 'Opps, something wrong with the input. Please go to Step 2 and Rescale Video again.'         
         video_cap = cv2.VideoCapture(aligned_video)
         if instyle is None or aligned_face is None or video_cap.get(7) == 0:
             video_cap.release()