Update web-demos/hugging_face/inpainter/base_inpainter.py

web-demos/hugging_face/inpainter/base_inpainter.py

@@ -203,9 +203,13 @@ class ProInpainter:
 
 		size = frames[0].size
 		# The ouput size should be divided by 2 so that it can encoded by libx264
-		if ratio !=1.0: size = (int(ratio*size[0])//2*2, int(ratio*size[1])//2*2)
+		size = (int(ratio*size[0])//2*2, int(ratio*size[1])//2*2)
 
-		else: size = ((size[0])//2*2, (size[1])//2*2) # set propainter size limit to 720 to reduce memory usage
+		# set propainter size limit to 720 to reduce memory usage
+		if max(size[0], size[1]) > 720:
+			scale = 720.0 / max(size[0], size[1])
+			# The ouput size should be divided by 2 so that it can encoded by libx264
+			size = (int(scale*size[0])//2*2, int(scale*size[1])//2*2)
 
 		frames_len = len(frames)
 		frames, size, out_size = resize_frames(frames, size)
@@ -246,7 +250,6 @@ class ProInpainter:
 					gt_flows_f_list.append(flows_f)
 					gt_flows_b_list.append(flows_b)
 					torch.cuda.empty_cache()
-					torch.cuda.ipc_collect()
 					
 				gt_flows_f = torch.cat(gt_flows_f_list, dim=1)
 				gt_flows_b = torch.cat(gt_flows_b_list, dim=1)
@@ -254,7 +257,6 @@ class ProInpainter:
 			else:
 				gt_flows_bi = self.fix_raft(frames, iters=raft_iter)
 				torch.cuda.empty_cache()
-				torch.cuda.ipc_collect()
 
 			if self.use_half:
 				frames, flow_masks, masks_dilated = frames.half(), flow_masks.half(), masks_dilated.half()
@@ -281,7 +283,7 @@ class ProInpainter:
 					pred_flows_f.append(pred_flows_bi_sub[0][:, pad_len_s:e_f-s_f-pad_len_e])
 					pred_flows_b.append(pred_flows_bi_sub[1][:, pad_len_s:e_f-s_f-pad_len_e])
 					torch.cuda.empty_cache()
-					torch.cuda.ipc_collect()
+					
 				pred_flows_f = torch.cat(pred_flows_f, dim=1)
 				pred_flows_b = torch.cat(pred_flows_b, dim=1)
 				pred_flows_bi = (pred_flows_f, pred_flows_b)
@@ -289,7 +291,6 @@ class ProInpainter:
 				pred_flows_bi, _ = self.fix_flow_complete.forward_bidirect_flow(gt_flows_bi, flow_masks)
 				pred_flows_bi = self.fix_flow_complete.combine_flow(gt_flows_bi, pred_flows_bi, flow_masks)
 				torch.cuda.empty_cache()
-				torch.cuda.ipc_collect()
 				
 			# ---- image propagation ----
 			masked_frames = frames * (1 - masks_dilated)
@@ -316,7 +317,7 @@ class ProInpainter:
 					updated_frames.append(updated_frames_sub[:, pad_len_s:e_f-s_f-pad_len_e])
 					updated_masks.append(updated_masks_sub[:, pad_len_s:e_f-s_f-pad_len_e])
 					torch.cuda.empty_cache()
-					torch.cuda.ipc_collect()
+					
 				updated_frames = torch.cat(updated_frames, dim=1)
 				updated_masks = torch.cat(updated_masks, dim=1)
 			else:
@@ -325,7 +326,6 @@ class ProInpainter:
 				updated_frames = frames * (1 - masks_dilated) + prop_imgs.view(b, t, 3, h, w) * masks_dilated
 				updated_masks = updated_local_masks.view(b, t, 1, h, w)
 				torch.cuda.empty_cache()
-				torch.cuda.ipc_collect()
 	
 		ori_frames = frames_inp
 		comp_frames = [None] * video_length
@@ -373,9 +373,8 @@ class ProInpainter:
 					comp_frames[idx] = comp_frames[idx].astype(np.uint8)
 			
 			torch.cuda.empty_cache()
-			torch.cuda.ipc_collect()
 
 		# need to return numpy array, T, H, W, 3
 		comp_frames = [cv2.resize(f, out_size) for f in comp_frames]
 
-		return comp_frames
+		return comp_frames