starred/ProPainter
1
0 Fork 0
mirror of https://github.com/sczhou/ProPainter.git synced 2026-10-18 20:59:51 -05:00
Code Issues Packages Projects Releases 1 Wiki Activity

support --subvideo_length for memory-efficient inference.

Browse Source
This commit is contained in:
sczhou committed 2023-09-22 02:02:49 +08:00
1 parent 2dc07818f7
commit 8bdb4cbbc3
2 files changed
+85 -23
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

README.md
+2 -2
View File
@@ -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
</tr>
<tr>
<td>Masks</td>
<td colspan="2"> [<a href="https://drive.google.com/file/d/1dFTneS_zaJAHjglxU10gYzr1-xALgHa4/view?usp=sharing">Google Drive</a>] [<a href="https://pan.baidu.com/s/1JC-UKmlQfjhVtD81196cxA?pwd=87e3">Baidu Disk</a>] (For reproducing paper results; provided in <a href="https://github.com/MCG-NKU/E2FGVI">E<sup>2</sup>FGVI</a> paper) </td>
<td colspan="2"> [<a href="https://drive.google.com/file/d/1dFTneS_zaJAHjglxU10gYzr1-xALgHa4/view?usp=sharing">Google Drive</a>] [<a href="https://pan.baidu.com/s/1JC-UKmlQfjhVtD81196cxA?pwd=87e3">Baidu Disk</a>] (For reproducing paper results; provided in <a href="https://arxiv.org/abs/2309.03897">ProPainter</a> paper) </td>
</tr>
</tbody>
</table>
inference_propainter.py
+83 -21
View File
@@ -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