diff --git a/README.md b/README.md index 603f25a..766bd7b 100644 --- a/README.md +++ b/README.md @@ -163,7 +163,7 @@ Video inpainting typically requires a significant amount of GPU memory. Here, we - Set the `--resize_ratio` (default 1.0) to resize the processing video. - Set a smaller video size via specifying the `--width` and `--height`. - Set `--fp16` to use fp16 (half precision) during inference. - - Reduce the frames of sub-videos `--subvideo_lentgh` (default 80), which effectively decouples GPU memory costs and video length. (coming soon!) + - Reduce the frames of sub-videos `--subvideo_length` (default 80), which effectively decouples GPU memory costs and video length. Blow shows the estimated GPU memory requirements for different sub-video lengths with fp32/fp16 precision: @@ -196,7 +196,7 @@ Blow shows the estimated GPU memory requirements for different sub-video lengths Masks - [Google Drive] [Baidu Disk] (For reproducing paper results; provided in E2FGVI paper) + [Google Drive] [Baidu Disk] (For reproducing paper results; provided in ProPainter paper) diff --git a/inference_propainter.py b/inference_propainter.py index f201c7e..698333d 100644 --- a/inference_propainter.py +++ b/inference_propainter.py @@ -156,15 +156,15 @@ def extrapolation(video_ori, scale): return frames, flow_masks, masks_dilated, (imgW_extr, imgH_extr) -def get_ref_index(neighbor_ids, length, ref_stride=10, ref_num=-1): +def get_ref_index(mid_neighbor_id, neighbor_ids, length, ref_stride=10, ref_num=-1): ref_index = [] if ref_num == -1: for i in range(0, length, ref_stride): if i not in neighbor_ids: ref_index.append(i) else: - start_idx = max(0, f - ref_stride * (ref_num // 2)) - end_idx = min(length, f + ref_stride * (ref_num // 2)) + start_idx = max(0, mid_neighbor_id - ref_stride * (ref_num // 2)) + end_idx = min(length, mid_neighbor_id + ref_stride * (ref_num // 2)) for i in range(start_idx, end_idx + 1, ref_stride): if i not in neighbor_ids: if len(ref_index) > ref_num: @@ -173,6 +173,7 @@ def get_ref_index(neighbor_ids, length, ref_stride=10, ref_num=-1): return ref_index + if __name__ == '__main__': # device = torch.device("cuda" if torch.cuda.is_available() else "cpu") device = get_device() @@ -194,12 +195,10 @@ if __name__ == '__main__': '--mask_dilation', type=int, default=4, help='Mask dilation for video and flow masking.') parser.add_argument( "--ref_stride", type=int, default=10, help='Stride of global reference frames.') - parser.add_argument( - "--ref_num", type=int, default=-1, help='Number of reference frames.') parser.add_argument( "--neighbor_length", type=int, default=10, help='Length of local neighboring frames.') parser.add_argument( - "--subvideo_lentgh", type=int, default=80, help='Length of sub-video for long video inference.') + "--subvideo_length", type=int, default=80, help='Length of sub-video for long video inference.') parser.add_argument( "--raft_iter", type=int, default=20, help='Iterations for RAFT inference.') parser.add_argument( @@ -306,7 +305,8 @@ if __name__ == '__main__': short_clip_len = 4 else: short_clip_len = 2 - + + # use fp32 for RAFT if frames.size(1) > short_clip_len: gt_flows_f_list, gt_flows_b_list = [], [] for f in range(0, video_length, short_clip_len): @@ -318,12 +318,14 @@ if __name__ == '__main__': gt_flows_f_list.append(flows_f) gt_flows_b_list.append(flows_b) + torch.cuda.empty_cache() gt_flows_f = torch.cat(gt_flows_f_list, dim=1) gt_flows_b = torch.cat(gt_flows_b_list, dim=1) gt_flows_bi = (gt_flows_f, gt_flows_b) else: gt_flows_bi = fix_raft(frames, iters=args.raft_iter) + torch.cuda.empty_cache() if use_half: @@ -332,32 +334,90 @@ if __name__ == '__main__': fix_flow_complete = fix_flow_complete.half() model = model.half() - # ---- complete flow ---- - pred_flows_bi, _ = fix_flow_complete.forward_bidirect_flow(gt_flows_bi, flow_masks) - pred_flows_bi = fix_flow_complete.combine_flow(gt_flows_bi, pred_flows_bi, flow_masks) - - # ---- temporal propagation ---- - masked_frames = frames * (1 - masks_dilated) - prop_imgs, updated_local_masks = model.img_propagation(masked_frames, pred_flows_bi, masks_dilated, 'nearest') - - b, t, _, _, _ = masks_dilated.size() - updated_masks = updated_local_masks.view(b, t, 1, h, w) - updated_frames = frames * (1-masks_dilated) + prop_imgs.view(b, t, 3, h, w) * masks_dilated # merge - del gt_flows_bi, frames, prop_imgs, updated_local_masks - torch.cuda.empty_cache() + # ---- complete flow ---- + flow_length = gt_flows_bi[0].size(1) + if flow_length > args.subvideo_length: + pred_flows_f, pred_flows_b = [], [] + pad_len = 5 + for f in range(0, flow_length, args.subvideo_length): + s_f = max(0, f - pad_len) + e_f = min(flow_length, f + args.subvideo_length + pad_len) + pad_len_s = max(0, f) - s_f + pad_len_e = e_f - min(flow_length, f + args.subvideo_length) + pred_flows_bi_sub, _ = fix_flow_complete.forward_bidirect_flow( + (gt_flows_bi[0][:, s_f:e_f], gt_flows_bi[1][:, s_f:e_f]), + flow_masks[:, s_f:e_f+1]) + pred_flows_bi_sub = fix_flow_complete.combine_flow( + (gt_flows_bi[0][:, s_f:e_f], gt_flows_bi[1][:, s_f:e_f]), + pred_flows_bi_sub, + flow_masks[:, s_f:e_f+1]) + 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() + + 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) + else: + pred_flows_bi, _ = fix_flow_complete.forward_bidirect_flow(gt_flows_bi, flow_masks) + pred_flows_bi = fix_flow_complete.combine_flow(gt_flows_bi, pred_flows_bi, flow_masks) + torch.cuda.empty_cache() + + # ---- image propagation ---- + masked_frames = frames * (1 - masks_dilated) + subvideo_length_img_prop = min(100, args.subvideo_length) # ensure a minimum of 100 frames for image propagation + if video_length > subvideo_length_img_prop: + updated_frames, updated_masks = [], [] + pad_len = 10 + for f in range(0, video_length, subvideo_length_img_prop): + s_f = max(0, f - pad_len) + e_f = min(video_length, f + subvideo_length_img_prop + pad_len) + pad_len_s = max(0, f) - s_f + pad_len_e = e_f - min(video_length, f + subvideo_length_img_prop) + + b, t, _, _, _ = masks_dilated[:, s_f:e_f].size() + pred_flows_bi_sub = (pred_flows_bi[0][:, s_f:e_f-1], pred_flows_bi[1][:, s_f:e_f-1]) + prop_imgs_sub, updated_local_masks_sub = model.img_propagation(masked_frames[:, s_f:e_f], + pred_flows_bi_sub, + masks_dilated[:, s_f:e_f], + 'nearest') + updated_frames_sub = frames[:, s_f:e_f] * (1 - masks_dilated[:, s_f:e_f]) + \ + prop_imgs_sub.view(b, t, 3, h, w) * masks_dilated[:, s_f:e_f] + updated_masks_sub = updated_local_masks_sub.view(b, t, 1, h, w) + + 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() + + updated_frames = torch.cat(updated_frames, dim=1) + updated_masks = torch.cat(updated_masks, dim=1) + else: + b, t, _, _, _ = masks_dilated.size() + prop_imgs, updated_local_masks = model.img_propagation(masked_frames, pred_flows_bi, masks_dilated, 'nearest') + 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() + + ori_frames = frames_inp comp_frames = [None] * video_length neighbor_stride = args.neighbor_length // 2 + if video_length > args.subvideo_length: + ref_num = args.subvideo_length // args.ref_stride + else: + ref_num = -1 + + # ---- feature propagation + transformer ---- for f in tqdm(range(0, video_length, neighbor_stride)): neighbor_ids = [ i for i in range(max(0, f - neighbor_stride), min(video_length, f + neighbor_stride + 1)) ] - ref_ids = get_ref_index(neighbor_ids, video_length, args.ref_stride, args.ref_num) + ref_ids = get_ref_index(f, neighbor_ids, video_length, args.ref_stride, ref_num) selected_imgs = updated_frames[:, neighbor_ids + ref_ids, :, :, :] selected_masks = masks_dilated[:, neighbor_ids + ref_ids, :, :, :] selected_update_masks = updated_masks[:, neighbor_ids + ref_ids, :, :, :] @@ -367,7 +427,9 @@ if __name__ == '__main__': # 1.0 indicates mask l_t = len(neighbor_ids) + # pred_img = selected_imgs # results of image propagation pred_img = model(selected_imgs, selected_pred_flows_bi, selected_masks, selected_update_masks, l_t) + pred_img = pred_img.view(-1, 3, h, w) pred_img = (pred_img + 1) / 2