PKUWilliamYang commited on
Commit
dbf6cac
1 Parent(s): ae52247

Update vtoonify_model.py

Browse files
Files changed (1) hide show
  1. vtoonify_model.py +2 -1
vtoonify_model.py CHANGED
@@ -206,7 +206,8 @@ class Model():
206
  int(video_cap.get(4)*4)))
207
 
208
  batch_frames = []
209
- batch_size = 4
 
210
  with torch.no_grad():
211
  if self.color_transfer:
212
  s_w = exstyle
 
206
  int(video_cap.get(4)*4)))
207
 
208
  batch_frames = []
209
+ batch_size = max(1, int(round(4 * 256 * 256 / video_cap.get(3) / video_cap.get(4))))
210
+ print('Using batch size of %d on %d frames'%(batch_size, num))
211
  with torch.no_grad():
212
  if self.color_transfer:
213
  s_w = exstyle