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

web-demos/hugging_face/inpainter/base_inpainter.py

@@ -246,6 +246,7 @@ 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)
@@ -253,6 +254,7 @@ 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()
@@ -279,7 +281,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)
@@ -287,6 +289,7 @@ 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)
@@ -313,7 +316,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:
@@ -322,6 +325,7 @@ 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
@@ -369,6 +373,7 @@ 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]