diff --git a/README.md b/README.md index a727e9a..bf79553 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,7 @@ Download the model [here](https://github.com/Yuehan717/RealViformer?tab=readme-o [workflow_realviformer.json](workflow/workflow_realviformer.json) -![realviformer_example](example_realviformer.png) +![realviformer_example](workflow_examples/example_realviformer.png) (Not a workflow-embedded image) @@ -77,7 +77,7 @@ To install, `pip install [].whl` [workflow_vfi_mamba.json](workflow/workflow_vfi_mamba.json) -![example_vfi_mamba](example_vfi_mamba.png) +![example_vfi_mamba](workflow_examples/example_vfi_mamba.png) (Not a workflow-embedded image) @@ -146,6 +146,13 @@ models/stabstitch/spatial_warp.pth models/stabstitch/smooth_warp.pth ``` +| StabStitch | StabStitch Stabilize | +|-|-| +| [stabstitch_stitch.json](workflow/stabstitch_stitch.json) (Example videos in `examples\stabstitch`) | [stabstich_stabilize.json](workflow/stabstitch_stabilize.json) | +| ![example_stabstitch_stitch](workflow_examples/example_stabstitch_stitch.png) | ![example_stabstitch_stabilize](workflow_examples/example_stabstitch_stabilize.png) | + +(Not workflow-embedded images) + ## Credits ``` diff --git a/__init__.py b/__init__.py index eb7cfff..8a3284d 100644 --- a/__init__.py +++ b/__init__.py @@ -1,17 +1,20 @@ from .run import ( WFEN, - RealViFormerSR, - ProPIH_Harmonizer, CoLIE_LowLight_Enhance, - StabStitch, + ProPIH_Harmonizer, + RealViFormerSR, ) - from .run_convir import ( + ConvIR_DefocusDeblur, ConvIR_DeHaze, ConvIR_DeRain, ConvIR_DeSnow, ConvIR_MotionDeBlur, - ConvIR_DefocusDeblur, +) +from .run_stabstitch import ( + StabStitch_Crop_Resize, + StabStitch_Stabilize, + StabStitch_Stitch, ) NODE_CLASS_MAPPINGS = { @@ -24,7 +27,9 @@ NODE_CLASS_MAPPINGS = { "ConvIR_DeSnow": ConvIR_DeSnow, "ConvIR_MotionDeBlur": ConvIR_MotionDeBlur, "ConvIR_DefocusDeblur": ConvIR_DefocusDeblur, - "StabStitch": StabStitch, + "StabStitch": StabStitch_Stitch, + "StabStitch_Stabilize": StabStitch_Stabilize, + "StabStitch_Crop_Resize": StabStitch_Crop_Resize, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -38,6 +43,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ConvIR_MotionDeBlur": "ConvIR Motion DeBlur", "ConvIR_DefocusDeblur": "ConvIR Defocus Deblur", "StabStitch": "StabStitch", + "StabStitch_Stabilize": "StabStitch Stabilize", + "StabStitch_Crop_Resize": "StabStitch Crop Resize", } try: diff --git a/examples/stabstitch/video1/0001.jpg b/examples/stabstitch/video1/0001.jpg new file mode 100644 index 0000000..27a5dd3 Binary files /dev/null and b/examples/stabstitch/video1/0001.jpg differ diff --git a/examples/stabstitch/video1/0002.jpg b/examples/stabstitch/video1/0002.jpg new file mode 100644 index 0000000..50dd30e Binary files /dev/null and b/examples/stabstitch/video1/0002.jpg differ diff --git a/examples/stabstitch/video1/0003.jpg b/examples/stabstitch/video1/0003.jpg new file mode 100644 index 0000000..5afc59f Binary files /dev/null and b/examples/stabstitch/video1/0003.jpg differ diff --git a/examples/stabstitch/video1/0004.jpg b/examples/stabstitch/video1/0004.jpg new file mode 100644 index 0000000..60d1416 Binary files /dev/null and b/examples/stabstitch/video1/0004.jpg differ diff --git a/examples/stabstitch/video1/0005.jpg b/examples/stabstitch/video1/0005.jpg new file mode 100644 index 0000000..187edcf Binary files /dev/null and b/examples/stabstitch/video1/0005.jpg differ diff --git a/examples/stabstitch/video1/0006.jpg b/examples/stabstitch/video1/0006.jpg new file mode 100644 index 0000000..a72b5ac Binary files /dev/null and b/examples/stabstitch/video1/0006.jpg differ diff --git a/examples/stabstitch/video1/0007.jpg b/examples/stabstitch/video1/0007.jpg new file mode 100644 index 0000000..906e064 Binary files /dev/null and b/examples/stabstitch/video1/0007.jpg differ diff --git a/examples/stabstitch/video1/0008.jpg b/examples/stabstitch/video1/0008.jpg new file mode 100644 index 0000000..b1ecbad Binary files /dev/null and b/examples/stabstitch/video1/0008.jpg differ diff --git a/examples/stabstitch/video1/0009.jpg b/examples/stabstitch/video1/0009.jpg new file mode 100644 index 0000000..57148e7 Binary files /dev/null and b/examples/stabstitch/video1/0009.jpg differ diff --git a/examples/stabstitch/video1/0010.jpg b/examples/stabstitch/video1/0010.jpg new file mode 100644 index 0000000..e196aee Binary files /dev/null and b/examples/stabstitch/video1/0010.jpg differ diff --git a/examples/stabstitch/video1/0011.jpg b/examples/stabstitch/video1/0011.jpg new file mode 100644 index 0000000..7cc3107 Binary files /dev/null and b/examples/stabstitch/video1/0011.jpg differ diff --git a/examples/stabstitch/video1/0012.jpg b/examples/stabstitch/video1/0012.jpg new file mode 100644 index 0000000..0142920 Binary files /dev/null and b/examples/stabstitch/video1/0012.jpg differ diff --git a/examples/stabstitch/video1/0013.jpg b/examples/stabstitch/video1/0013.jpg new file mode 100644 index 0000000..05cfea7 Binary files /dev/null and b/examples/stabstitch/video1/0013.jpg differ diff --git a/examples/stabstitch/video1/0014.jpg b/examples/stabstitch/video1/0014.jpg new file mode 100644 index 0000000..bbdfeb1 Binary files /dev/null and b/examples/stabstitch/video1/0014.jpg differ diff --git a/examples/stabstitch/video1/0015.jpg b/examples/stabstitch/video1/0015.jpg new file mode 100644 index 0000000..3226876 Binary files /dev/null and b/examples/stabstitch/video1/0015.jpg differ diff --git a/examples/stabstitch/video1/0016.jpg b/examples/stabstitch/video1/0016.jpg new file mode 100644 index 0000000..1a54631 Binary files /dev/null and b/examples/stabstitch/video1/0016.jpg differ diff --git a/examples/stabstitch/video1/0017.jpg b/examples/stabstitch/video1/0017.jpg new file mode 100644 index 0000000..7353ab1 Binary files /dev/null and b/examples/stabstitch/video1/0017.jpg differ diff --git a/examples/stabstitch/video1/0018.jpg b/examples/stabstitch/video1/0018.jpg new file mode 100644 index 0000000..37e10e4 Binary files /dev/null and b/examples/stabstitch/video1/0018.jpg differ diff --git a/examples/stabstitch/video1/0019.jpg b/examples/stabstitch/video1/0019.jpg new file mode 100644 index 0000000..c91af5f Binary files /dev/null and b/examples/stabstitch/video1/0019.jpg differ diff --git a/examples/stabstitch/video1/0020.jpg b/examples/stabstitch/video1/0020.jpg new file mode 100644 index 0000000..5b84d1d Binary files /dev/null and b/examples/stabstitch/video1/0020.jpg differ diff --git a/examples/stabstitch/video2/0001.jpg b/examples/stabstitch/video2/0001.jpg new file mode 100644 index 0000000..d1a4fcf Binary files /dev/null and b/examples/stabstitch/video2/0001.jpg differ diff --git a/examples/stabstitch/video2/0002.jpg b/examples/stabstitch/video2/0002.jpg new file mode 100644 index 0000000..c5fdde7 Binary files /dev/null and b/examples/stabstitch/video2/0002.jpg differ diff --git a/examples/stabstitch/video2/0003.jpg b/examples/stabstitch/video2/0003.jpg new file mode 100644 index 0000000..9e1b481 Binary files /dev/null and b/examples/stabstitch/video2/0003.jpg differ diff --git a/examples/stabstitch/video2/0004.jpg b/examples/stabstitch/video2/0004.jpg new file mode 100644 index 0000000..a993080 Binary files /dev/null and b/examples/stabstitch/video2/0004.jpg differ diff --git a/examples/stabstitch/video2/0005.jpg b/examples/stabstitch/video2/0005.jpg new file mode 100644 index 0000000..4cd1b32 Binary files /dev/null and b/examples/stabstitch/video2/0005.jpg differ diff --git a/examples/stabstitch/video2/0006.jpg b/examples/stabstitch/video2/0006.jpg new file mode 100644 index 0000000..a41ef28 Binary files /dev/null and b/examples/stabstitch/video2/0006.jpg differ diff --git a/examples/stabstitch/video2/0007.jpg b/examples/stabstitch/video2/0007.jpg new file mode 100644 index 0000000..0a7c10a Binary files /dev/null and b/examples/stabstitch/video2/0007.jpg differ diff --git a/examples/stabstitch/video2/0008.jpg b/examples/stabstitch/video2/0008.jpg new file mode 100644 index 0000000..efe7a83 Binary files /dev/null and b/examples/stabstitch/video2/0008.jpg differ diff --git a/examples/stabstitch/video2/0009.jpg b/examples/stabstitch/video2/0009.jpg new file mode 100644 index 0000000..7a979c1 Binary files /dev/null and b/examples/stabstitch/video2/0009.jpg differ diff --git a/examples/stabstitch/video2/0010.jpg b/examples/stabstitch/video2/0010.jpg new file mode 100644 index 0000000..70d2696 Binary files /dev/null and b/examples/stabstitch/video2/0010.jpg differ diff --git a/examples/stabstitch/video2/0011.jpg b/examples/stabstitch/video2/0011.jpg new file mode 100644 index 0000000..82aa21d Binary files /dev/null and b/examples/stabstitch/video2/0011.jpg differ diff --git a/examples/stabstitch/video2/0012.jpg b/examples/stabstitch/video2/0012.jpg new file mode 100644 index 0000000..69512c4 Binary files /dev/null and b/examples/stabstitch/video2/0012.jpg differ diff --git a/examples/stabstitch/video2/0013.jpg b/examples/stabstitch/video2/0013.jpg new file mode 100644 index 0000000..d67f6e5 Binary files /dev/null and b/examples/stabstitch/video2/0013.jpg differ diff --git a/examples/stabstitch/video2/0014.jpg b/examples/stabstitch/video2/0014.jpg new file mode 100644 index 0000000..fb9f985 Binary files /dev/null and b/examples/stabstitch/video2/0014.jpg differ diff --git a/examples/stabstitch/video2/0015.jpg b/examples/stabstitch/video2/0015.jpg new file mode 100644 index 0000000..8685204 Binary files /dev/null and b/examples/stabstitch/video2/0015.jpg differ diff --git a/examples/stabstitch/video2/0016.jpg b/examples/stabstitch/video2/0016.jpg new file mode 100644 index 0000000..64f29bd Binary files /dev/null and b/examples/stabstitch/video2/0016.jpg differ diff --git a/examples/stabstitch/video2/0017.jpg b/examples/stabstitch/video2/0017.jpg new file mode 100644 index 0000000..c940bd4 Binary files /dev/null and b/examples/stabstitch/video2/0017.jpg differ diff --git a/examples/stabstitch/video2/0018.jpg b/examples/stabstitch/video2/0018.jpg new file mode 100644 index 0000000..1abe614 Binary files /dev/null and b/examples/stabstitch/video2/0018.jpg differ diff --git a/examples/stabstitch/video2/0019.jpg b/examples/stabstitch/video2/0019.jpg new file mode 100644 index 0000000..0a7bdb0 Binary files /dev/null and b/examples/stabstitch/video2/0019.jpg differ diff --git a/examples/stabstitch/video2/0020.jpg b/examples/stabstitch/video2/0020.jpg new file mode 100644 index 0000000..9f54714 Binary files /dev/null and b/examples/stabstitch/video2/0020.jpg differ diff --git a/module_stabstitch/stabstitch_model.py b/module_stabstitch/stabstitch_model.py index 847c863..452edf4 100644 --- a/module_stabstitch/stabstitch_model.py +++ b/module_stabstitch/stabstitch_model.py @@ -1,12 +1,12 @@ import torch -from torchvision.transforms import GaussianBlur import torch.nn.functional as F import tqdm -from . import torch_tps_transform +from torchvision.transforms import GaussianBlur + +from . import torch_tps_transform, torch_tps_transform_point from .smooth_network import SmoothNet, build_SmoothNet from .spatial_network import SpatialNet, build_SpatialNet from .temporal_network import TemporalNet, build_TemporalNet -from . import torch_tps_transform_point grid_h = 6 grid_w = 8 @@ -15,6 +15,53 @@ img_h = 360 img_w = 480 +def crop_or_resize(tensor, target_h, target_w, crop_side="left"): + B, H, W, C = tensor.shape + aspect_ratio = W / H + target_aspect_ratio = target_w / target_h + + if aspect_ratio == target_aspect_ratio: + # Just resize + resized = F.interpolate( + tensor.permute(0, 3, 1, 2), + size=(target_h, target_w), + mode="bilinear", + align_corners=False, + ) + return resized.permute(0, 2, 3, 1) + + elif aspect_ratio > target_aspect_ratio: + # Crop width, then resize + new_w = int(H * target_aspect_ratio) + if crop_side == "left": + cropped = tensor[:, :, :new_w, :] + elif crop_side == "right": + cropped = tensor[:, :, -new_w:, :] + else: + raise ValueError("crop_side must be 'left' or 'right'") + + resized = F.interpolate( + cropped.permute(0, 3, 1, 2), + size=(target_h, target_w), + mode="bilinear", + align_corners=False, + ) + return resized.permute(0, 2, 3, 1) + + else: + # Crop height, then resize + new_h = int(W / target_aspect_ratio) + cropped = tensor[:, :new_h, :, :] + + resized = F.interpolate( + cropped.permute(0, 3, 1, 2), + size=(target_h, target_w), + mode="bilinear", + align_corners=False, + ) + return resized.permute(0, 2, 3, 1) + + def video_preprocess(src_video: torch.Tensor, h: int, w: int): # [B, H, W, C] [0, 1] tensor = src_video.permute(0, 3, 1, 2) # [B, C, H, W] @@ -36,39 +83,44 @@ def stabstitch_inference( temporal_path: str, smooth_path: str, device, + mode: str, + do_linear_blend: bool, ): - spatial_net = SpatialNet().to(device) - temporal_net = TemporalNet().to(device) - smooth_net = SmoothNet().to(device) - - spatial_net.load_state_dict(torch.load(spatial_path)["model"]) - temporal_net.load_state_dict(torch.load(temporal_path)["model"]) - smooth_net.load_state_dict(torch.load(smooth_path)["model"]) - - spatial_net.eval() - temporal_net.eval() - smooth_net.eval() + spatial_net, temporal_net, smooth_net = load_models( + spatial_path, temporal_path, smooth_path, device + ) video_1 = video_preprocess(src_video_1, img_h, img_w).to(device) video_2 = video_preprocess(src_video_2, img_h, img_w).to(device) + video_list_1 = video_to_list_ts(video_1) video_list_2 = video_to_list_ts(video_2) - num_frames = len(video_list_1) - smotion_tensor_list = [] - omask_tensor_list = [] + return inference( + mode, + do_linear_blend, + spatial_net, + temporal_net, + smooth_net, + video_list_1, + video_list_2, + ) + +def inference( + mode, + do_linear_blend, + spatial_net, + temporal_net, + smooth_net, + video_list_1, + video_list_2, +): + num_frames = len(video_list_2) # motion estimation - for k in tqdm.tqdm(range(num_frames), desc="Motion estimating"): - # step 1: spatial warp - with torch.no_grad(): - spatial_batch_out = build_SpatialNet( - spatial_net, video_list_1[k], video_list_2[k] - ) - smotion = spatial_batch_out["motion"] - omask = spatial_batch_out["overlap_mesh"] - smotion_tensor_list.append(smotion) - omask_tensor_list.append(omask) + smotion_tensor_list, omask_tensor_list = spatial_warping( + spatial_net, video_list_1, video_list_2, num_frames + ) # step 2: temporal warp with torch.no_grad(): @@ -78,37 +130,90 @@ def stabstitch_inference( ############################################## ############# data preparation ############ # converting tmotion (t-th frame) into tsmotion ( (t-1)-th frame ) - rigid_mesh = get_rigid_mesh(1, img_h, img_w) - norm_rigid_mesh = get_norm_mesh(rigid_mesh, img_h, img_w) - rigid_mesh = get_rigid_mesh(1, img_h, img_w) - norm_rigid_mesh = get_norm_mesh(rigid_mesh, img_h, img_w) - - smesh_list = [] - tsmotion_list = [] - - for k in tqdm.tqdm(range(len(tmotion_tensor_list)), desc="Temporal calculating"): - smotion = smotion_tensor_list[k] - smesh = rigid_mesh + smotion - if k == 0: - tsmotion = smotion.clone() * 0 - - else: - smotion_1 = smotion_tensor_list[k - 1] - smesh_1 = rigid_mesh + smotion_1 - tmotion = tmotion_tensor_list[k] - tmesh = rigid_mesh + tmotion - norm_smesh_1 = get_norm_mesh(smesh_1, img_h, img_w) - norm_tmesh = get_norm_mesh(tmesh, img_h, img_w) - tsmesh = torch_tps_transform_point.transformer( - norm_tmesh, norm_rigid_mesh, norm_smesh_1 - ) - tsmotion = recover_mesh(tsmesh, img_h, img_w) - smesh - # append - smesh_list.append(smesh) - tsmotion_list.append(tsmotion) + smesh_list, tsmotion_list = temporaling(smotion_tensor_list, tmotion_tensor_list) # step 3: smooth warp + target_mesh = smooth_meshing( + smooth_net, + smotion_tensor_list, + omask_tensor_list, + tmotion_tensor_list, + smesh_list, + tsmotion_list, + ) + + stable_list, mask_list_2, img_2_list, _, _ = get_stable_sqe( + video_list_1, video_list_2, target_mesh, mode, do_linear_blend + ) + + stable_list = [stable.unsqueeze(0)[:, [2, 1, 0], :, :] for stable in stable_list] + + return stable_list, img_2_list, mask_list_2 + + +def stabstitch_single_inference( + src_video_1: torch.Tensor, + spatial_path: str, + temporal_path: str, + smooth_path: str, + device, + mode: str, + do_linear_blend: bool, +): + spatial_net, temporal_net, smooth_net = load_models( + spatial_path, temporal_path, smooth_path, device + ) + + video_1 = video_preprocess(src_video_1[0].unsqueeze(0), img_h, img_w).to(device) + video_2 = video_preprocess(src_video_1, img_h, img_w).to(device) + + video_list_2 = video_to_list_ts(video_2) + + num_frames = len(video_list_2) + # motion estimation + smotion_tensor_list, omask_tensor_list = spatial_warping_single( + spatial_net, video_1, video_list_2, num_frames + ) + + # step 2: temporal warp + with torch.no_grad(): + temporal_batch_out = build_TemporalNet(temporal_net, video_list_2) + tmotion_tensor_list = temporal_batch_out["motion_list"] + + ############################################## + ############# data preparation ############ + # converting tmotion (t-th frame) into tsmotion ( (t-1)-th frame ) + + smesh_list, tsmotion_list = temporaling(smotion_tensor_list, tmotion_tensor_list) + + # step 3: smooth warp + target_mesh = smooth_meshing( + smooth_net, + smotion_tensor_list, + omask_tensor_list, + tmotion_tensor_list, + smesh_list, + tsmotion_list, + ) + + stable_list, mask_list_2, img_2_list, _, _ = get_stable_sqe( + video_1, video_list_2, target_mesh, mode, do_linear_blend, is_single=True + ) + + stable_list = [stable.unsqueeze(0)[:, [2, 1, 0], :, :] for stable in stable_list] + + return stable_list, img_2_list, mask_list_2 + + +def smooth_meshing( + smooth_net, + smotion_tensor_list, + omask_tensor_list, + tmotion_tensor_list, + smesh_list, + tsmotion_list, +): ori_mesh = 0 target_mesh = 0 @@ -138,18 +243,69 @@ def stabstitch_inference( (target_mesh, _target_mesh[:, -1, ...].unsqueeze(1)), 1 ) - stable_list, mesh_list_2, img_2_list, out_width, out_height = get_stable_sqe( - video_list_1, video_list_2, target_mesh - ) + return target_mesh - stable_list = [stable.unsqueeze(0)[:, [2, 1, 0], :, :] for stable in stable_list] - print(f"{stable_list[0].shape=}") # ([3, 380, 757]) - print(f"{img_2_list[0].shape=}") # ([3, 380, 757]) - if mesh_list_2: - print(f"{mesh_list_2[0].shape=}") +def temporaling(smotion_tensor_list, tmotion_tensor_list): + smesh_list = [] + tsmotion_list = [] - return stable_list, img_2_list, mesh_list_2 + rigid_mesh = get_rigid_mesh(1, img_h, img_w) + norm_rigid_mesh = get_norm_mesh(rigid_mesh, img_h, img_w) + + for k in tqdm.tqdm(range(len(tmotion_tensor_list)), desc="Temporal calculating"): + smotion = smotion_tensor_list[k] + smesh = rigid_mesh + smotion + if k == 0: + tsmotion = smotion.clone() * 0 + + else: + smotion_1 = smotion_tensor_list[k - 1] + smesh_1 = rigid_mesh + smotion_1 + tmotion = tmotion_tensor_list[k] + tmesh = rigid_mesh + tmotion + norm_smesh_1 = get_norm_mesh(smesh_1, img_h, img_w) + norm_tmesh = get_norm_mesh(tmesh, img_h, img_w) + tsmesh = torch_tps_transform_point.transformer( + norm_tmesh, norm_rigid_mesh, norm_smesh_1 + ) + tsmotion = recover_mesh(tsmesh, img_h, img_w) - smesh + # append + smesh_list.append(smesh) + tsmotion_list.append(tsmotion) + return smesh_list, tsmotion_list + + +def spatial_warping(spatial_net, video_list_1, video_list_2, num_frames): + smotion_tensor_list = [] + omask_tensor_list = [] + + for k in tqdm.tqdm(range(num_frames), desc="Motion estimating"): + # step 1: spatial warp + with torch.no_grad(): + spatial_batch_out = build_SpatialNet( + spatial_net, video_list_1[k], video_list_2[k] + ) + smotion = spatial_batch_out["motion"] + omask = spatial_batch_out["overlap_mesh"] + smotion_tensor_list.append(smotion) + omask_tensor_list.append(omask) + return smotion_tensor_list, omask_tensor_list + + +def spatial_warping_single(spatial_net, video_1, video_list_2, num_frames): + smotion_tensor_list = [] + omask_tensor_list = [] + + for k in tqdm.tqdm(range(num_frames), desc="Motion estimating"): + # step 1: spatial warp + with torch.no_grad(): + spatial_batch_out = build_SpatialNet(spatial_net, video_1, video_list_2[k]) + smotion = spatial_batch_out["motion"] + omask = spatial_batch_out["overlap_mesh"] + smotion_tensor_list.append(smotion) + omask_tensor_list.append(omask) + return smotion_tensor_list, omask_tensor_list def linear_blender(ref, tgt, ref_m, tgt_m, mask=False): @@ -224,30 +380,47 @@ def get_norm_mesh(mesh, height, width): # bs, T, h, w, 2 smooth_path -# def get_stable_sqe(img2_list, ori_mesh): -# batch_size, _, img_h, img_w = img2_list[0].shape +def get_stable_sqe( + img1_list, img2_list, ori_mesh, mode="FAST", do_linear_blend=False, is_single=False +): + norm_rigid_mesh, width_min, height_min, out_width, out_height = get_new_sizes( + img2_list, ori_mesh + ) -# rigid_mesh = get_rigid_mesh(batch_size, img_h, img_w) -# norm_rigid_mesh = get_norm_mesh(rigid_mesh, img_h, img_w) + img1_warp = torch.zeros((1, 3, out_height, out_width)).cuda() -# stable_list = [] -# for i in range(len(img2_list)): -# mesh = ori_mesh[:, i, :, :, :] -# norm_mesh = get_norm_mesh(mesh, img_h, img_w) -# img2 = (img2_list[i].cuda() + 1) / 2 -# print(img2) -# mask = torch.ones_like(img2).cuda() -# img2_warp = torch_tps_transform.transformer( -# torch.cat([img2, mask], 1), norm_mesh, norm_rigid_mesh, (img_h, img_w) -# ) -# stable_list.append(img2_warp) + stable_list = [] + mask_list_2 = [] + img_2_list = [] + for i in range(len(img2_list)): + if is_single: + img1 = (img1_list.cuda() + 1) / 2 + else: + img1 = (img1_list[i].cuda() + 1) / 2 + img2 = (img2_list[i].cuda() + 1) / 2 + mesh = ori_mesh[:, i, :, :, :] + img2_warp, mask2, ave_fusion = get_fused( + mode, + do_linear_blend, + norm_rigid_mesh, + width_min, + height_min, + out_width, + out_height, + img1_warp, + img1, + img2, + mesh, + ) -# return stable_list + img_2_list.append(img2_warp) + mask_list_2.append(mask2) + stable_list.append(ave_fusion) + + return stable_list, mask_list_2, img_2_list, out_width, out_height -# FAST -# NORMAL -def get_stable_sqe(img1_list, img2_list, ori_mesh, mode="FAST", do_linear_blend=False): +def get_new_sizes(img2_list, ori_mesh): batch_size, _, img_h, img_w = img2_list[0].shape rigid_mesh = get_rigid_mesh(batch_size, img_h, img_w) @@ -268,70 +441,85 @@ def get_stable_sqe(img1_list, img2_list, ori_mesh, mode="FAST", do_linear_blend= height_max = height_max.int() out_width = width_max - width_min + 1 out_height = height_max - height_min + 1 + return norm_rigid_mesh, width_min, height_min, out_width, out_height - print(out_width) - print(out_height) - img1_warp = torch.zeros((1, 3, out_height, out_width)).cuda() +def get_fused( + mode, + do_linear_blend, + norm_rigid_mesh, + width_min, + height_min, + out_width, + out_height, + img1_warp, + img1, + img2, + mesh, +): + mesh_trans = torch.stack([mesh[..., 0] - width_min, mesh[..., 1] - height_min], 3) + norm_mesh = get_norm_mesh(mesh_trans, out_height, out_width) - stable_list = [] - mask_list_2 = [] - img_2_list = [] - # mesh_tran_list = [] - for i in range(len(img2_list)): - mesh = ori_mesh[:, i, :, :, :] - mesh_trans = torch.stack( - [mesh[..., 0] - width_min, mesh[..., 1] - height_min], 3 - ) - norm_mesh = get_norm_mesh(mesh_trans, out_height, out_width) - img2 = (img2_list[i].cuda() + 1) / 2 + # mode = 'FAST': use F.grid_sample to interpolate. It's fast, but may produce thin black boundary. + # mode = 'NORMAL': use our implemented interpolation function. It's a bit slower, but avoid the black boundary. + img2_warp = torch_tps_transform.transformer( + img2, norm_mesh, norm_rigid_mesh, (out_height, out_width), mode=mode + ) + mask2 = torch_tps_transform.transformer( + torch.ones_like(img2).cuda(), + norm_mesh, + norm_rigid_mesh, + (out_height, out_width), + ) - # mode = 'FAST': use F.grid_sample to interpolate. It's fast, but may produce thin black boundary. - # mode = 'NORMAL': use our implemented interpolation function. It's a bit slower, but avoid the black boundary. - img2_warp = torch_tps_transform.transformer( - img2, norm_mesh, norm_rigid_mesh, (out_height, out_width), mode=mode - ) - img_2_list.append(img2_warp) + if not do_linear_blend: + # average blending + img1_warp[ + :, + :, + int(0 - height_min) : int(0 - height_min) + 360, + int(0 - width_min) : int(0 - width_min) + 480, + ] = img1 + + comp_1 = img1_warp[0] / (img1_warp[0] + img2_warp[0] + 1e-6) + comp_2 = img2_warp[0] / (img1_warp[0] + img2_warp[0] + 1e-6) + + ave_fusion = img1_warp[0] * comp_1 + img2_warp[0] * comp_2 + else: + img1_warp[ + :, + :, + int(0 - height_min) : int(0 - height_min) + 360, + int(0 - width_min) : int(0 - width_min) + 480, + ] = img1 + mask1 = img1_warp.clone() + mask1[ + :, + :, + int(0 - height_min) : int(0 - height_min) + 360, + int(0 - width_min) : int(0 - width_min) + 480, + ] = 1 mask2 = torch_tps_transform.transformer( torch.ones_like(img2).cuda(), norm_mesh, norm_rigid_mesh, (out_height, out_width), ) - mask_list_2.append(mask2) + ave_fusion = linear_blender(img1_warp, img2_warp, mask1, mask2) + ave_fusion = ave_fusion[0] + return img2_warp, mask2, ave_fusion - if not do_linear_blend: - # average blending - img1_warp[ - :, - :, - int(0 - height_min) : int(0 - height_min) + 360, - int(0 - width_min) : int(0 - width_min) + 480, - ] = (img1_list[i].cuda() + 1) / 2 - comp_1 = img1_warp[0] / (img1_warp[0] + img2_warp[0] + 1e-6) - comp_2 = img2_warp[0] / (img1_warp[0] + img2_warp[0] + 1e-6) +def load_models(spatial_path, temporal_path, smooth_path, device): + spatial_net = SpatialNet().to(device) + temporal_net = TemporalNet().to(device) + smooth_net = SmoothNet().to(device) - ave_fusion = img1_warp[0] * comp_1 + img2_warp[0] * comp_2 - else: - # linear blending - img1_warp[ - :, - int(0 - height_min) : int(0 - height_min) + 360, - int(0 - width_min) : int(0 - width_min) + 480, - ] = (img1_list[i][0] + 1) / 2 - mask1 = img1_warp.clone() - mask1[ - :, - int(0 - height_min) : int(0 - height_min) + 360, - int(0 - width_min) : int(0 - width_min) + 480, - ] = 1 + spatial_net.load_state_dict(torch.load(spatial_path)["model"]) + temporal_net.load_state_dict(torch.load(temporal_path)["model"]) + smooth_net.load_state_dict(torch.load(smooth_path)["model"]) - ave_fusion = linear_blender( - img1_warp.unsqueeze(0), img2_warp, mask1.unsqueeze(0), mask2 - ) - ave_fusion = ave_fusion[0] - - stable_list.append(ave_fusion) - - return stable_list, mask_list_2, img_2_list, out_width, out_height + spatial_net.eval() + temporal_net.eval() + smooth_net.eval() + return spatial_net, temporal_net, smooth_net diff --git a/pyproject.toml b/pyproject.toml index 715019a..d5f4f82 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "fm_nodes" -description = "An expanding collection of ComfyUI nodes focused on new papers. Including: WFEN, RealViFormer, ProPIH, CoLIE, VFIMamba, ConvIR" -version = "1.0.5" +description = "An expanding collection of ComfyUI nodes focused on new papers. Including: WFEN, RealViFormer, ProPIH, CoLIE, VFIMamba, ConvIR, StabStitch" +version = "1.0.6" license = {file = "LICENSE"} dependencies = ["torch", "einops"] diff --git a/run.py b/run.py index 86208bc..7c21dae 100644 --- a/run.py +++ b/run.py @@ -2,7 +2,7 @@ @author: Fuou Marinas @title: FM Nodes @nickname: FM_nodes -@description: A collection of nodes. WFEN Face Super Resolution. +@description: A collection of nodes. """ from pathlib import Path @@ -12,22 +12,17 @@ import torch.nn.functional as F import tqdm from comfy.utils import ProgressBar -from .module_stabstitch.stabstitch_model import stabstitch_inference - -from .module_colie.colie_model import CoLIE_Config, run_colie from .constants import ( PROPIH_G_MODEL_PATH, PROPIH_VGG_MODEL_PATH, REALVIFORMER_MODEL_PATH, - STAB_SPATIAL_PATH, - STAB_TEMPORAL_PATH, - STAB_SMOOTH_PATH, WFEN_MODEL_PATH, ) +from .module_colie.colie_model import CoLIE_Config, run_colie from .module_propih.propih_model import VGG19HRNetModel from .module_realviformer.realviformer_arch import RealViformer -from .utils import ensure_size, img_to_mask from .module_wfen.wfen_model import WFENModel +from .utils import ensure_size, img_to_mask base_dir = Path(__file__).resolve().parent @@ -288,54 +283,3 @@ class CoLIE_LowLight_Enhance: pbar.update_absolute(i, num_frames) return (torch.cat(result, dim=0),) - - -class StabStitch: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "video_1": ("IMAGE",), - "video_2": ("IMAGE",), - }, - } - - RETURN_TYPES = ( - "IMAGE", - "IMAGE", - "IMAGE", - ) - RETURN_NAMES = ( - "combined", - "stab", - "mask", - ) - FUNCTION = "todo" - CATEGORY = "FM_nodes" - - def todo( - self, - video_1: torch.Tensor, - video_2: torch.Tensor, - ): - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - stable_list, video_frames_2, mask_frames_2 = stabstitch_inference( - video_1, - video_2, - spatial_path=f"{base_dir}/{STAB_SPATIAL_PATH}", - temporal_path=f"{base_dir}/{STAB_TEMPORAL_PATH}", - smooth_path=f"{base_dir}/{STAB_SMOOTH_PATH}", - device=device, - ) - - out_video = torch.cat(stable_list, dim=0).permute(0, 2, 3, 1) - out_mask = torch.cat(mask_frames_2, dim=0).permute(0, 2, 3, 1) - out_img2 = torch.cat(video_frames_2, dim=0)[:, [2, 1, 0], :, :].permute( - 0, 2, 3, 1 - ) - - return ( - out_video, - out_img2, - out_mask, - ) diff --git a/run_stabstitch.py b/run_stabstitch.py new file mode 100644 index 0000000..ecc0a36 --- /dev/null +++ b/run_stabstitch.py @@ -0,0 +1,183 @@ +from pathlib import Path + +import torch + +from .constants import ( + STAB_CONSTRUCT_DEFAULT, + STAB_CONSTRUCT_MODES, + STAB_SMOOTH_PATH, + STAB_SPATIAL_PATH, + STAB_TEMPORAL_PATH, +) +from .module_stabstitch.stabstitch_model import ( + crop_or_resize, + stabstitch_inference, + stabstitch_single_inference, +) + +base_dir = Path(__file__).resolve().parent + + +class StabStitch_Stitch: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "video_1": ("IMAGE",), + "video_2": ("IMAGE",), + # mode = 'FAST': use F.grid_sample to interpolate. It's fast, but may produce thin black boundary. + # mode = 'NORMAL': use our implemented interpolation function. It's a bit slower, but avoid the black boundary. + "interpolate_mode": ( + STAB_CONSTRUCT_MODES, + {"default": STAB_CONSTRUCT_DEFAULT}, + ), + "do_linear_blend": ( + "BOOLEAN", + {"default": False}, + ), + }, + } + + RETURN_TYPES = ( + "IMAGE", + "IMAGE", + "IMAGE", + ) + RETURN_NAMES = ( + "combined", + "stab", + "mask", + ) + FUNCTION = "todo" + CATEGORY = "FM_nodes/StabStitch" + + def todo( + self, + video_1: torch.Tensor, + video_2: torch.Tensor, + interpolate_mode: str, + do_linear_blend: bool, + ): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + stable_list, video_frames_2, mask_frames_2 = stabstitch_inference( + video_1, + video_2, + spatial_path=f"{base_dir}/{STAB_SPATIAL_PATH}", + temporal_path=f"{base_dir}/{STAB_TEMPORAL_PATH}", + smooth_path=f"{base_dir}/{STAB_SMOOTH_PATH}", + device=device, + mode=interpolate_mode, + do_linear_blend=do_linear_blend, + ) + + out_video = torch.cat(stable_list, dim=0).permute(0, 2, 3, 1) + out_mask = torch.cat(mask_frames_2, dim=0).permute(0, 2, 3, 1) + out_img2 = torch.cat(video_frames_2, dim=0)[:, [2, 1, 0], :, :].permute( + 0, 2, 3, 1 + ) + + return ( + out_video, + out_img2, + out_mask, + ) + + +class StabStitch_Stabilize: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "video_1": ("IMAGE",), + # mode = 'FAST': use F.grid_sample to interpolate. It's fast, but may produce thin black boundary. + # mode = 'NORMAL': use our implemented interpolation function. It's a bit slower, but avoid the black boundary. + "interpolate_mode": ( + STAB_CONSTRUCT_MODES, + {"default": STAB_CONSTRUCT_DEFAULT}, + ), + "do_linear_blend": ( + "BOOLEAN", + {"default": False}, + ), + }, + } + + RETURN_TYPES = ( + "IMAGE", + "IMAGE", + "IMAGE", + ) + RETURN_NAMES = ( + "combined", + "stab", + "mask", + ) + FUNCTION = "todo" + CATEGORY = "FM_nodes/StabStitch" + + def todo( + self, + video_1: torch.Tensor, + interpolate_mode: str, + do_linear_blend: bool, + ): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + stable_list, video_frames_2, mask_frames_2 = stabstitch_single_inference( + video_1, + spatial_path=f"{base_dir}/{STAB_SPATIAL_PATH}", + temporal_path=f"{base_dir}/{STAB_TEMPORAL_PATH}", + smooth_path=f"{base_dir}/{STAB_SMOOTH_PATH}", + device=device, + mode=interpolate_mode, + do_linear_blend=do_linear_blend, + ) + + out_video = torch.cat(stable_list, dim=0).permute(0, 2, 3, 1) + out_mask = torch.cat(mask_frames_2, dim=0).permute(0, 2, 3, 1) + out_img2 = torch.cat(video_frames_2, dim=0)[:, [2, 1, 0], :, :].permute( + 0, 2, 3, 1 + ) + + return ( + out_video, + out_img2, + out_mask, + ) + + +class StabStitch_Crop_Resize: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "video_1": ("IMAGE",), + "crop_side": ( + ["left", "right"], + {"default": "left"}, + ), + "target_h": ( + "INT", + {"default": 360, "min": 1, "step": 1}, + ), + "target_w": ( + "INT", + {"default": 480, "min": 1, "step": 1}, + ), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("resized",) + FUNCTION = "todo" + CATEGORY = "FM_nodes/StabStitch" + + def todo( + self, + video_1: torch.Tensor, + crop_side: str, + target_h: int, + target_w: int, + ): + resized_video = crop_or_resize(video_1, target_h, target_w, crop_side) + + return (resized_video,) diff --git a/workflow/stabstitch_stabilize.json b/workflow/stabstitch_stabilize.json new file mode 100644 index 0000000..9c56384 --- /dev/null +++ b/workflow/stabstitch_stabilize.json @@ -0,0 +1,458 @@ +{ + "last_node_id": 19, + "last_link_id": 44, + "nodes": [ + { + "id": 6, + "type": "VHS_VideoCombine", + "pos": [ + 1348, + 27 + ], + "size": [ + 315, + 536.03125 + ], + "flags": { + "collapsed": false + }, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 44 + }, + { + "name": "audio", + "type": "AUDIO", + "link": null + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 20, + "loop_count": 0, + "filename_prefix": "stabmask", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 2, + "save_metadata": true, + "pingpong": false, + "save_output": false, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "stabmask_00001.mp4", + "subfolder": "", + "type": "temp", + "format": "video/h264-mp4", + "frame_rate": 20 + } + } + } + }, + { + "id": 19, + "type": "StabStitch_Stabilize", + "pos": [ + 296, + 55 + ], + "size": [ + 210, + 122 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "video_1", + "type": "IMAGE", + "link": 41 + } + ], + "outputs": [ + { + "name": "combined", + "type": "IMAGE", + "links": [ + 42 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "stab", + "type": "IMAGE", + "links": [ + 43 + ], + "shape": 3, + "slot_index": 1 + }, + { + "name": "mask", + "type": "IMAGE", + "links": [ + 44 + ], + "shape": 3, + "slot_index": 2 + } + ], + "properties": { + "Node name for S&R": "StabStitch_Stabilize" + }, + "widgets_values": [ + "FAST", + true + ] + }, + { + "id": 18, + "type": "StabStitch_Crop_Resize", + "pos": [ + 406, + 229 + ], + "size": [ + 210, + 106 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "video_1", + "type": "IMAGE", + "link": 40 + } + ], + "outputs": [ + { + "name": "resized", + "type": "IMAGE", + "links": [ + 41 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "StabStitch_Crop_Resize" + }, + "widgets_values": [ + "left", + 360, + 480 + ] + }, + { + "id": 17, + "type": "VHS_LoadVideo", + "pos": [ + 45, + 8 + ], + "size": [ + 235.1999969482422, + 398.1134866403979 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [ + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 40 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "frame_count", + "type": "INT", + "links": null, + "shape": 3 + }, + { + "name": "audio", + "type": "AUDIO", + "links": null, + "shape": 3 + }, + { + "name": "video_info", + "type": "VHS_VIDEOINFO", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_LoadVideo" + }, + "widgets_values": { + "video": "cat.mp4", + "force_rate": 0, + "force_size": "Disabled", + "custom_width": 512, + "custom_height": 512, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1, + "choose video to upload": "image", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "frame_load_cap": 0, + "skip_first_frames": 0, + "force_rate": 0, + "filename": "cat.mp4", + "type": "input", + "format": "video/mp4", + "select_every_nth": 1 + } + } + } + }, + { + "id": 15, + "type": "VHS_VideoCombine", + "pos": [ + 995, + 28 + ], + "size": [ + 315, + 536.03125 + ], + "flags": { + "collapsed": false + }, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 43 + }, + { + "name": "audio", + "type": "AUDIO", + "link": null + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 20, + "loop_count": 0, + "filename_prefix": "stab", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 2, + "save_metadata": true, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "stab_00001.mp4", + "subfolder": "", + "type": "temp", + "format": "video/h264-mp4", + "frame_rate": 20 + } + } + } + }, + { + "id": 3, + "type": "VHS_VideoCombine", + "pos": [ + 644, + 27 + ], + "size": [ + 315, + 536.03125 + ], + "flags": { + "collapsed": false + }, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 42 + }, + { + "name": "audio", + "type": "AUDIO", + "link": null + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 20, + "loop_count": 0, + "filename_prefix": "stabstitch", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 2, + "save_metadata": true, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "stabstitch_00007.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 20 + } + } + } + } + ], + "links": [ + [ + 40, + 17, + 0, + 18, + 0, + "IMAGE" + ], + [ + 41, + 18, + 0, + 19, + 0, + "IMAGE" + ], + [ + 42, + 19, + 0, + 3, + 0, + "IMAGE" + ], + [ + 43, + 19, + 1, + 15, + 0, + "IMAGE" + ], + [ + 44, + 19, + 2, + 6, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.9090909090909091, + "offset": [ + 86.31161530434264, + 196.97150602340395 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/workflow/stabstitch_stitch.json b/workflow/stabstitch_stitch.json new file mode 100644 index 0000000..f41e1f9 --- /dev/null +++ b/workflow/stabstitch_stitch.json @@ -0,0 +1,474 @@ +{ + "last_node_id": 18, + "last_link_id": 49, + "nodes": [ + { + "id": 15, + "type": "VHS_VideoCombine", + "pos": [ + 995, + 28 + ], + "size": [ + 315, + 473 + ], + "flags": { + "collapsed": false + }, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 48 + }, + { + "name": "audio", + "type": "AUDIO", + "link": null + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 20, + "loop_count": 0, + "filename_prefix": "stab", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 2, + "save_metadata": true, + "pingpong": false, + "save_output": false, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "stab_00002.mp4", + "subfolder": "", + "type": "temp", + "format": "video/h264-mp4", + "frame_rate": 20 + } + } + } + }, + { + "id": 6, + "type": "VHS_VideoCombine", + "pos": [ + 1348, + 27 + ], + "size": [ + 315, + 473 + ], + "flags": { + "collapsed": false + }, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 49 + }, + { + "name": "audio", + "type": "AUDIO", + "link": null + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 20, + "loop_count": 0, + "filename_prefix": "stabmask", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 2, + "save_metadata": true, + "pingpong": false, + "save_output": false, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "stabmask_00002.mp4", + "subfolder": "", + "type": "temp", + "format": "video/h264-mp4", + "frame_rate": 20 + } + } + } + }, + { + "id": 3, + "type": "VHS_VideoCombine", + "pos": [ + 649, + 27 + ], + "size": [ + 315, + 473 + ], + "flags": { + "collapsed": false + }, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 47 + }, + { + "name": "audio", + "type": "AUDIO", + "link": null + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 20, + "loop_count": 0, + "filename_prefix": "stabstitch", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 2, + "save_metadata": true, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "stabstitch_00006.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 20 + } + } + } + }, + { + "id": 9, + "type": "VHS_LoadImagesPath", + "pos": [ + 73, + 218 + ], + "size": [ + 226.8000030517578, + 363.2251838887286 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [ + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 45 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3, + "slot_index": 1 + }, + { + "name": "frame_count", + "type": "INT", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_LoadImagesPath" + }, + "widgets_values": { + "directory": "J:\\AI\\stabber\\video1", + "image_load_cap": 0, + "skip_first_images": 0, + "select_every_nth": 1, + "choose folder to upload": "image", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "frame_load_cap": 0, + "skip_first_images": 0, + "filename": "J:\\AI\\stabber\\video1", + "type": "path", + "format": "folder", + "select_every_nth": 1 + } + } + } + }, + { + "id": 10, + "type": "VHS_LoadImagesPath", + "pos": [ + 323, + 213 + ], + "size": [ + 226.8000030517578, + 363.2251838887286 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 46 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + }, + { + "name": "frame_count", + "type": "INT", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_LoadImagesPath" + }, + "widgets_values": { + "directory": "J:\\AI\\stabber\\video2", + "image_load_cap": 0, + "skip_first_images": 0, + "select_every_nth": 1, + "choose folder to upload": "image", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "frame_load_cap": 0, + "skip_first_images": 0, + "filename": "J:\\AI\\stabber\\video2", + "type": "path", + "format": "folder", + "select_every_nth": 1 + } + } + } + }, + { + "id": 18, + "type": "StabStitch", + "pos": [ + 330, + 37 + ], + "size": [ + 215.16154450854913, + 122 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "video_1", + "type": "IMAGE", + "link": 45 + }, + { + "name": "video_2", + "type": "IMAGE", + "link": 46 + } + ], + "outputs": [ + { + "name": "combined", + "type": "IMAGE", + "links": [ + 47 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "stab", + "type": "IMAGE", + "links": [ + 48 + ], + "shape": 3, + "slot_index": 1 + }, + { + "name": "mask", + "type": "IMAGE", + "links": [ + 49 + ], + "shape": 3, + "slot_index": 2 + } + ], + "properties": { + "Node name for S&R": "StabStitch" + }, + "widgets_values": [ + "FAST", + true + ] + } + ], + "links": [ + [ + 45, + 9, + 0, + 18, + 0, + "IMAGE" + ], + [ + 46, + 10, + 0, + 18, + 1, + "IMAGE" + ], + [ + 47, + 18, + 0, + 3, + 0, + "IMAGE" + ], + [ + 48, + 18, + 1, + 15, + 0, + "IMAGE" + ], + [ + 49, + 18, + 2, + 6, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1, + "offset": [ + 29.8384554914509, + 163.0904108785293 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/example_realviformer.png b/workflow_examples/example_realviformer.png similarity index 100% rename from example_realviformer.png rename to workflow_examples/example_realviformer.png diff --git a/workflow_examples/example_stabstitch_stabilize.png b/workflow_examples/example_stabstitch_stabilize.png new file mode 100644 index 0000000..dad276e Binary files /dev/null and b/workflow_examples/example_stabstitch_stabilize.png differ diff --git a/workflow_examples/example_stabstitch_stitch.png b/workflow_examples/example_stabstitch_stitch.png new file mode 100644 index 0000000..378c7a5 Binary files /dev/null and b/workflow_examples/example_stabstitch_stitch.png differ diff --git a/example_vfi_mamba.png b/workflow_examples/example_vfi_mamba.png similarity index 100% rename from example_vfi_mamba.png rename to workflow_examples/example_vfi_mamba.png