Ashrafb commited on
Commit
ae4009d
·
verified ·
1 Parent(s): aa1f42d

Update vtoonify_model.py

Browse files
Files changed (1) hide show
  1. vtoonify_model.py +6 -4
vtoonify_model.py CHANGED
@@ -195,7 +195,7 @@ class Model():
195
  return np.zeros((256, 256, 3), np.uint8), 'Oops, something wrong with the style type. Please go to Step 1 and load model again.'
196
 
197
  try:
198
- with torch.no_grad():
199
  if self.color_transfer:
200
  s_w = exstyle
201
  else:
@@ -212,9 +212,11 @@ class Model():
212
  logging.info(f"Output from VToonify shape: {y_tilde.shape}")
213
  print('*** Toonify %dx%d image with style of %s' % (y_tilde.shape[2], y_tilde.shape[3], style_type))
214
 
215
- 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)
216
-
217
-
 
 
218
  def tensor2cv2(self, img):
219
  """Convert a tensor image to OpenCV format."""
220
  tmp = ((img.cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8).copy()
 
195
  return np.zeros((256, 256, 3), np.uint8), 'Oops, something wrong with the style type. Please go to Step 1 and load model again.'
196
 
197
  try:
198
+ with torch.no_grad():
199
  if self.color_transfer:
200
  s_w = exstyle
201
  else:
 
212
  logging.info(f"Output from VToonify shape: {y_tilde.shape}")
213
  print('*** Toonify %dx%d image with style of %s' % (y_tilde.shape[2], y_tilde.shape[3], style_type))
214
 
215
+ 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)
216
+ except Exception as e:
217
+ logging.error(f"Error during model execution: {e}")
218
+ return np.zeros((256, 256, 3), np.uint8), f"Error during processing: {str(e)}"
219
+
220
  def tensor2cv2(self, img):
221
  """Convert a tensor image to OpenCV format."""
222
  tmp = ((img.cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8).copy()