diff --git a/FOMM_PARTSWAP.json b/FOMM_PARTSWAP.json new file mode 100644 index 0000000..35b2897 --- /dev/null +++ b/FOMM_PARTSWAP.json @@ -0,0 +1,692 @@ +{ + "last_node_id": 28, + "last_link_id": 59, + "nodes": [ + { + "id": 1, + "type": "LoadImage", + "pos": [ + 111, + 131 + ], + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 47 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "Rick1.png", + "image" + ] + }, + { + "id": 24, + "type": "PreviewImage", + "pos": [ + 1423.509090909091, + 336.86363636363626 + ], + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 52 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 25, + "type": "PreviewImage", + "pos": [ + 1413.509090909091, + 625.8636363636363 + ], + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 53 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 16, + "type": "VHS_VideoInfoLoaded", + "pos": [ + 366, + 875 + ], + "size": { + "0": 235.1999969482422, + "1": 106 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "video_info", + "type": "VHS_VIDEOINFO", + "link": 12 + } + ], + "outputs": [ + { + "name": "fps🟦", + "type": "FLOAT", + "links": [ + 54 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "frame_count🟦", + "type": "INT", + "links": null, + "shape": 3 + }, + { + "name": "duration🟦", + "type": "FLOAT", + "links": null, + "shape": 3 + }, + { + "name": "width🟦", + "type": "INT", + "links": null, + "shape": 3 + }, + { + "name": "height🟦", + "type": "INT", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_VideoInfoLoaded" + }, + "widgets_values": {} + }, + { + "id": 13, + "type": "VHS_VideoCombine", + "pos": [ + 1048, + 349 + ], + "size": [ + 315, + 595 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 49 + }, + { + "name": "audio", + "type": "VHS_AUDIO", + "link": 50 + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null, + "slot_index": 2 + }, + { + "name": "frame_rate", + "type": "FLOAT", + "link": 51, + "widget": { + "name": "frame_rate" + } + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 8, + "loop_count": 0, + "filename_prefix": "Partswap", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 0, + "save_metadata": true, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "Partswap_00093-audio.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4" + } + } + } + }, + { + "id": 11, + "type": "VHS_LoadVideo", + "pos": [ + 117, + 506 + ], + "size": [ + 235.1999969482422, + 491.1999969482422 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 48 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "frame_count", + "type": "INT", + "links": null, + "shape": 3 + }, + { + "name": "audio", + "type": "VHS_AUDIO", + "links": [ + 45 + ], + "shape": 3, + "slot_index": 2 + }, + { + "name": "video_info", + "type": "VHS_VIDEOINFO", + "links": [ + 12 + ], + "shape": 3, + "slot_index": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_LoadVideo" + }, + "widgets_values": { + "video": "damedane.mp4", + "force_rate": 0, + "force_size": "Disabled", + "custom_width": 512, + "custom_height": 512, + "frame_load_cap": 10, + "skip_first_frames": 0, + "select_every_nth": 1, + "choose video to upload": "image", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "frame_load_cap": 10, + "skip_first_frames": 0, + "force_rate": 0, + "filename": "damedane.mp4", + "type": "input", + "format": "video/mp4", + "select_every_nth": 1 + } + } + } + }, + { + "id": 26, + "type": "FOMM_Seg5Chooser", + "pos": [ + 313, + 1078 + ], + "size": { + "0": 315, + "1": 178 + }, + "flags": {}, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "chosen_seg_indices", + "type": "STRING", + "links": [], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "FOMM_Seg5Chooser" + }, + "widgets_values": [ + false, + true, + true, + true, + true, + true + ] + }, + { + "id": 23, + "type": "FOMM_Seg10Chooser", + "pos": [ + 649, + 1080 + ], + "size": { + "0": 315, + "1": 298 + }, + "flags": {}, + "order": 3, + "mode": 0, + "outputs": [ + { + "name": "chosen_seg_indices", + "type": "STRING", + "links": [], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "FOMM_Seg10Chooser" + }, + "widgets_values": [ + false, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true + ] + }, + { + "id": 28, + "type": "FOMM_Seg15Chooser", + "pos": [ + 983, + 1081 + ], + "size": { + "0": 315, + "1": 418 + }, + "flags": {}, + "order": 4, + "mode": 0, + "outputs": [ + { + "name": "chosen_seg_indices", + "type": "STRING", + "links": [ + 59 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "FOMM_Seg15Chooser" + }, + "widgets_values": [ + false, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true + ] + }, + { + "id": 27, + "type": "Reroute", + "pos": [ + 709, + 725 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 59, + "widget": { + "name": "value" + } + } + ], + "outputs": [ + { + "name": "", + "type": "STRING", + "links": [ + 56 + ], + "slot_index": 0 + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 22, + "type": "FOMM_Partswap", + "pos": [ + 599, + 366 + ], + "size": { + "0": 393, + "1": 258 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "source_image", + "type": "IMAGE", + "link": 47 + }, + { + "name": "driving_video_input", + "type": "IMAGE", + "link": 48 + }, + { + "name": "audio", + "type": "VHS_AUDIO", + "link": 45 + }, + { + "name": "chosen_seg_indices", + "type": "STRING", + "link": 56, + "widget": { + "name": "chosen_seg_indices" + } + }, + { + "name": "frame_rate", + "type": "FLOAT", + "link": 54, + "widget": { + "name": "frame_rate" + } + } + ], + "outputs": [ + { + "name": "seg_src_viz", + "type": "IMAGE", + "links": [ + 52 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "seg_tgt_viz", + "type": "IMAGE", + "links": [ + 53 + ], + "shape": 3, + "slot_index": 1 + }, + { + "name": "images", + "type": "IMAGE", + "links": [ + 49 + ], + "shape": 3, + "slot_index": 2 + }, + { + "name": "audio", + "type": "VHS_AUDIO", + "links": [ + 50 + ], + "shape": 3, + "slot_index": 3 + }, + { + "name": "frame_rate", + "type": "FLOAT", + "links": [ + 51 + ], + "shape": 3, + "slot_index": 4 + } + ], + "properties": { + "Node name for S&R": "FOMM_Partswap" + }, + "widgets_values": [ + "vox-15segments", + 30, + 1, + false, + false, + false, + 1, + 0.6 + ] + } + ], + "links": [ + [ + 12, + 11, + 3, + 16, + 0, + "VHS_VIDEOINFO" + ], + [ + 45, + 11, + 2, + 22, + 2, + "VHS_AUDIO" + ], + [ + 47, + 1, + 0, + 22, + 0, + "IMAGE" + ], + [ + 48, + 11, + 0, + 22, + 1, + "IMAGE" + ], + [ + 49, + 22, + 2, + 13, + 0, + "IMAGE" + ], + [ + 50, + 22, + 3, + 13, + 1, + "VHS_AUDIO" + ], + [ + 51, + 22, + 4, + 13, + 3, + "FLOAT" + ], + [ + 52, + 22, + 0, + 24, + 0, + "IMAGE" + ], + [ + 53, + 22, + 1, + 25, + 0, + "IMAGE" + ], + [ + 54, + 16, + 0, + 22, + 4, + "FLOAT" + ], + [ + 56, + 27, + 0, + 22, + 3, + "STRING" + ], + [ + 59, + 28, + 0, + 27, + 0, + "*" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.8264462809917358, + "offset": [ + 101.0730714815582, + -279.5114901742915 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/README.md b/README.md index 7f44abd..d4c6b70 100644 --- a/README.md +++ b/README.md @@ -9,22 +9,37 @@ https://github.com/AliaksandrSiarohin/motion-cosegmentation ## Workflow: +### FOMM + [FOMM.json](FOMM.json) ![FOMM Workflow](workflow.png) - - https://github.com/FuouM/ComfyUI-FirstOrderMM/assets/57849916/36833a94-67f9-4dd9-b79b-181f86e83bfe +### Part Swap +[FOMM_PARTSWAP.json](FOMM_PARTSWAP.json) + +![Partswap Workflow](workflow_fomm_partswap.png) ## Arguments + +### FOMM + * `relative_movement`: Relative keypoint displacement (Inherit object proporions from the video) * `relative_jacobian`: No idea what this does, but usually == `relative_movement` * `adapt_movement_scale`: If disabled, will heavily distort the source face to match the driving face * `find_best_frame`: Find driving frame that best match the source. Split the batch into two halves, with the first half reversed. Gives mixed results. Needs to install `face-alignment` library. +### Part Swap + +* `blend_scale`: No idea, keeping at default = 1.0 seems to be fine +* `use_source_seg`: Whether to use the source's segmentation or the target's. May help if some of the target's segmentation regions are missing +* `hard_edges`: Whether to make the edges hard, instead of feathering +* `use_face_parser`: For Seg-based models, may help with cleaning up residual background (should only use `15seg` with this). TODO: Additional cleanup face_parser masks. Should definitely be used for FOMM models +* `viz_alpha`: Opacity of the segments in the visualization + ## Installation 1. Clone the repo to `ComfyUI/custom_nodes/` @@ -45,15 +60,39 @@ pip install face-alignment ## Model -Currently supporting `vox` and `vox-adv`. Models can and must be manually downloaded from: +FOMM currently supporting `vox` and `vox-adv`. Models can and must be manually downloaded from: * [AliaksandrSiarohin/first-order-model](https://github.com/AliaksandrSiarohin/first-order-model) * [graphemecluster/first-order-model-demo](https://github.com/graphemecluster/first-order-model-demo) +Part Swap currently supporting Seg-based models and FOMM (`vox` and `vox-adv`) models. +* `vox-5segments` +* `vox-10segments` +* `vox-15segments` +* `vox-first-order (partswap)` + +These models can be found in the original repository [Motion Supervised co-part Segmentation](https://github.com/AliaksandrSiarohin/motion-cosegmentation) + Place them in the `checkpoints` folder. It should look like this: ``` place_checkpoints_here.txt vox-adv-cpk.pth.tar vox-cpk.pth.tar + +vox-5segments.pth.tar +vox-10segments.pth.tar +vox-15segments.pth.tar +vox-first-order.pth.tar ``` -resnet18_url = "https://download.pytorch.org/models/resnet18-5c106cde.pth" +For Part Swap, Face-Parsing is also supported **(Optional)** (especially when using the FOMM or `vox-first-order` models) + +* resnet18 `resnet18-5c106cde`: https://download.pytorch.org/models/resnet18-5c106cde.pth +* face_parsing `79999_iter.pth`: https://github.com/zllrunning/face-makeup.PyTorch/tree/master/cp + +Place them in `face_parsing` folder: +``` +face_parsing_model.py +... +resnet18-5c106cde.pth +79999_iter.pth +``` \ No newline at end of file diff --git a/face_parsing/face_parsing_loader.py b/face_parsing/face_parsing_loader.py index fbb64c1..24ab6a8 100644 --- a/face_parsing/face_parsing_loader.py +++ b/face_parsing/face_parsing_loader.py @@ -1,13 +1,18 @@ import torch import numpy as np - +import os from .face_parsing_model import BiSeNet def load_face_parser_model(base_dir: str, face_parser_checkpoint_name: str, cpu=False): + face_parser_path = f"{base_dir}/face_parsing/{face_parser_checkpoint_name}" + if not os.path.isfile(face_parser_path): + print( + f"{face_parser_checkpoint_name} not found, revert to no face parsing, not fatal" + ) + return None print("Loading face_parser model") face_parser = BiSeNet(n_classes=19) - face_parser_path = f"{base_dir}/face_parsing/{face_parser_checkpoint_name}" mean = torch.Tensor(np.array([0.485, 0.456, 0.406], dtype=np.float32)).view( 1, 3, 1, 1 diff --git a/run.py b/run.py index 4fdde2d..c534baa 100644 --- a/run.py +++ b/run.py @@ -212,7 +212,7 @@ class FOMM_Partswap: "frame_rate": ("FLOAT", {"default": 30.0}), "blend_scale": ( "FLOAT", - {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1}, + {"default": 1.0, "min": 0.6, "max": 1.0, "step": 0.01}, ), "use_source_seg": ( "BOOLEAN", @@ -291,14 +291,15 @@ class FOMM_Partswap: ) face_parser_model = None - if model_name in partswap_face_parser_enforce_models: - assert use_face_parser, "Special supervised model, requires face_parser" + # if model_name in partswap_face_parser_enforce_models: + # assert use_face_parser, "Special supervised model, requires face_parser" if use_face_parser: face_parser_model = load_face_parser_model( base_dir, face_parser_checkpoint_name ) - face_parser_model.cuda() - face_parser_model.eval() + if face_parser_model is not None: + face_parser_model.cuda() + face_parser_model.eval() print(f"Chosen {chosen_seg_indices}") swap_indices = deserialize_integers(chosen_seg_indices) diff --git a/samples/sample_output.mp4 b/samples/sample_output.mp4 deleted file mode 100644 index a2031be..0000000 Binary files a/samples/sample_output.mp4 and /dev/null differ diff --git a/workflow_fomm_partswap.png b/workflow_fomm_partswap.png new file mode 100644 index 0000000..1696ae9 Binary files /dev/null and b/workflow_fomm_partswap.png differ