diff --git a/.gitignore b/.gitignore index b233313..7e54636 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,8 @@ checkpoints/* !checkpoints/place_checkpoints_here.txt +face_parsing/*.pth +!face_parsing/put_model_safetensors_config_json_here.txt # Byte-compiled / optimized / DLL files __pycache__/ 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 1992cb2..e65f645 100644 --- a/README.md +++ b/README.md @@ -3,24 +3,43 @@ A ComfyUI-native node to run First Order Motion Model for Image Animation https://github.com/AliaksandrSiarohin/first-order-model +Now supports Face Swapping using Motion Supervised co-part Segmentation + +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` +* `relative_jacobian`: Only taken into effect when `relative_movement` is on, must also be on to avoid heavy deformation of the face (in a freaky way) * `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/` @@ -41,13 +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 ``` + +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/__init__.py b/__init__.py index 6c98dfa..ccb00b0 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,26 @@ -from .run import FOMM_Runner +from .run import ( + FOMM_Partswap, + FOMM_Runner, + FOMM_Seg5Chooser, + FOMM_Seg10Chooser, + FOMM_Seg15Chooser, +) NODE_CLASS_MAPPINGS = { "FOMM_Runner": FOMM_Runner, + "FOMM_Partswap": FOMM_Partswap, + "FOMM_Seg5Chooser": FOMM_Seg5Chooser, + "FOMM_Seg10Chooser": FOMM_Seg10Chooser, + "FOMM_Seg15Chooser": FOMM_Seg15Chooser, } NODE_DISPLAY_NAME_MAPPINGS = { "FOMM_Runner": "FOMM Runner", + "FOMM_Partswap": "FOMM Partswap", + "FOMM_Seg5Chooser": "FOMM Seg5 Chooser", + "FOMM_Seg10Chooser": "FOMM Seg10 Chooser", + "FOMM_Seg15Chooser": "FOMM Seg15 Chooser", } -__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/config/vox-256-sem-10segments.yaml b/config/vox-256-sem-10segments.yaml new file mode 100644 index 0000000..5871c97 --- /dev/null +++ b/config/vox-256-sem-10segments.yaml @@ -0,0 +1,92 @@ +# Dataset parameters +dataset_params: + # Path to data, data can be stored in several formats: .mp4 or .gif videos, stacked .png images or folders with frames. + root_dir: data/vox-png + # Image shape, needed for staked .png format. + image_shape: [256, 256, 3] + # In case of VoxCeleb or TaiChi single video can be splitted in many chunks, or the maybe several videos for single person. + # In this case epoch can be a pass over different videos (if id_sampling=True) or over different chunks (if id_sampling=False) + # If the name of the video '12335#adsbf.mp4' the id is assumed to be 12335 + id_sampling: True + # Augmentation parameters see augmentation.py for all posible augmentations + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + jitter_param: + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.1 + +# Defines model architecture +model_params: + common_params: + # Number of segments + num_segments: 10 + # Number of channels per image + num_channels: 3 + # Use only shift and no-affine part + estimate_affine_part: True + segmentation_module_params: + # Softmax temperature for shift heatmaps + temperature: 0.1 + # Number of features mutliplier + block_expansion: 32 + # Maximum allowed number of features + max_features: 1024 + # Number of block in Unet. Can be increased or decreased depending or resolution. + num_blocks: 5 + # Segmentations is predicted on smaller images for better performance, + # scale_factor=0.25 means that 256x256 image will be resized to 64x64 + scale_factor: 0.25 + reconstruction_module_params: + # Number of features mutliplier + block_expansion: 64 + # Maximum allowed number of features + max_features: 512 + # Number of downsampling blocks in Jonson architecture. + # Can be increased or decreased depending or resolution. + num_down_blocks: 2 + # Number of ResBlocks in Jonson architecture. + num_bottleneck_blocks: 6 + # Use visibility map or not + estimate_visibility: True + +# Parameters of training +train_params: + # Number of training epochs + num_epochs: 5 + # For better i/o performance when number of videos is small number of epochs can be multiplied by this number. + # Thus effectively with num_repeats=100 each epoch is 100 times larger. + num_repeats: 50 + # Learning rates + lr_segmentation: 2.0e-4 + lr_reconstruction: 2.0e-4 + batch_size: 20 + # Scales for perceptual pyramide loss. If scales = [1, 0.5, 0.25, 0.125] and image resolution is 256x256, + # than the loss will be computer on resolutions 256x256, 128x128, 64x64, 32x32. + scales: [1, 0.5, 0.25, 0.125] + # Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs. + checkpoint_freq: 50 + # Parameters of transform for equivariance loss + transform_params: + # Sigma for affine part + sigma_affine: 0.05 + # Sigma for deformation part + sigma_tps: 0.005 + # Number of point in the deformation grid + points_tps: 5 + loss_weights: + equivariance: 10 + perceptual: [10, 10, 10, 10, 10] + + +# Visualization parameters +visualizer_params: + # Draw keypoints (shifts in affine transformations) of this size, increase or decrease depending on resolution + kp_size: 5 + # Draw white border around images + draw_border: True + # Color map for keypoints + colormap: 'gist_rainbow' diff --git a/config/vox-256-sem-15segments.yaml b/config/vox-256-sem-15segments.yaml new file mode 100644 index 0000000..ade4f4d --- /dev/null +++ b/config/vox-256-sem-15segments.yaml @@ -0,0 +1,92 @@ +# Dataset parameters +dataset_params: + # Path to data, data can be stored in several formats: .mp4 or .gif videos, stacked .png images or folders with frames. + root_dir: data/vox-png + # Image shape, needed for staked .png format. + image_shape: [256, 256, 3] + # In case of VoxCeleb or TaiChi single video can be splitted in many chunks, or the maybe several videos for single person. + # In this case epoch can be a pass over different videos (if id_sampling=True) or over different chunks (if id_sampling=False) + # If the name of the video '12335#adsbf.mp4' the id is assumed to be 12335 + id_sampling: True + # Augmentation parameters see augmentation.py for all posible augmentations + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + jitter_param: + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.1 + +# Defines model architecture +model_params: + common_params: + # Number of segments + num_segments: 15 + # Number of channels per image + num_channels: 3 + # Use only shift and no-affine part + estimate_affine_part: True + segmentation_module_params: + # Softmax temperature for shift heatmaps + temperature: 0.1 + # Number of features mutliplier + block_expansion: 32 + # Maximum allowed number of features + max_features: 1024 + # Number of block in Unet. Can be increased or decreased depending or resolution. + num_blocks: 5 + # Segmentations is predicted on smaller images for better performance, + # scale_factor=0.25 means that 256x256 image will be resized to 64x64 + scale_factor: 0.25 + reconstruction_module_params: + # Number of features mutliplier + block_expansion: 64 + # Maximum allowed number of features + max_features: 512 + # Number of downsampling blocks in Jonson architecture. + # Can be increased or decreased depending or resolution. + num_down_blocks: 2 + # Number of ResBlocks in Jonson architecture. + num_bottleneck_blocks: 6 + # Use visibility map or not + estimate_visibility: True + +# Parameters of training +train_params: + # Number of training epochs + num_epochs: 5 + # For better i/o performance when number of videos is small number of epochs can be multiplied by this number. + # Thus effectively with num_repeats=100 each epoch is 100 times larger. + num_repeats: 50 + # Learning rates + lr_segmentation: 2.0e-4 + lr_reconstruction: 2.0e-4 + batch_size: 20 + # Scales for perceptual pyramide loss. If scales = [1, 0.5, 0.25, 0.125] and image resolution is 256x256, + # than the loss will be computer on resolutions 256x256, 128x128, 64x64, 32x32. + scales: [1, 0.5, 0.25, 0.125] + # Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs. + checkpoint_freq: 50 + # Parameters of transform for equivariance loss + transform_params: + # Sigma for affine part + sigma_affine: 0.05 + # Sigma for deformation part + sigma_tps: 0.005 + # Number of point in the deformation grid + points_tps: 5 + loss_weights: + equivariance: 10 + perceptual: [10, 10, 10, 10, 10] + + +# Visualization parameters +visualizer_params: + # Draw keypoints (shifts in affine transformations) of this size, increase or decrease depending on resolution + kp_size: 5 + # Draw white border around images + draw_border: True + # Color map for keypoints + colormap: 'gist_rainbow' diff --git a/config/vox-256-sem-5segments.yaml b/config/vox-256-sem-5segments.yaml new file mode 100644 index 0000000..974628f --- /dev/null +++ b/config/vox-256-sem-5segments.yaml @@ -0,0 +1,92 @@ +# Dataset parameters +dataset_params: + # Path to data, data can be stored in several formats: .mp4 or .gif videos, stacked .png images or folders with frames. + root_dir: data/vox-png + # Image shape, needed for staked .png format. + image_shape: [256, 256, 3] + # In case of VoxCeleb or TaiChi single video can be splitted in many chunks, or the maybe several videos for single person. + # In this case epoch can be a pass over different videos (if id_sampling=True) or over different chunks (if id_sampling=False) + # If the name of the video '12335#adsbf.mp4' the id is assumed to be 12335 + id_sampling: True + # Augmentation parameters see augmentation.py for all posible augmentations + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + jitter_param: + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.1 + +# Defines model architecture +model_params: + common_params: + # Number of segments + num_segments: 5 + # Number of channels per image + num_channels: 3 + # Use only shift and no-affine part + estimate_affine_part: True + segmentation_module_params: + # Softmax temperature for shift heatmaps + temperature: 0.1 + # Number of features mutliplier + block_expansion: 32 + # Maximum allowed number of features + max_features: 1024 + # Number of block in Unet. Can be increased or decreased depending or resolution. + num_blocks: 5 + # Segmentations is predicted on smaller images for better performance, + # scale_factor=0.25 means that 256x256 image will be resized to 64x64 + scale_factor: 0.25 + reconstruction_module_params: + # Number of features mutliplier + block_expansion: 64 + # Maximum allowed number of features + max_features: 512 + # Number of downsampling blocks in Jonson architecture. + # Can be increased or decreased depending or resolution. + num_down_blocks: 2 + # Number of ResBlocks in Jonson architecture. + num_bottleneck_blocks: 6 + # Use visibility map or not + estimate_visibility: True + +# Parameters of training +train_params: + # Number of training epochs + num_epochs: 5 + # For better i/o performance when number of videos is small number of epochs can be multiplied by this number. + # Thus effectively with num_repeats=100 each epoch is 100 times larger. + num_repeats: 50 + # Learning rates + lr_segmentation: 2.0e-4 + lr_reconstruction: 2.0e-4 + batch_size: 20 + # Scales for perceptual pyramide loss. If scales = [1, 0.5, 0.25, 0.125] and image resolution is 256x256, + # than the loss will be computer on resolutions 256x256, 128x128, 64x64, 32x32. + scales: [1, 0.5, 0.25, 0.125] + # Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs. + checkpoint_freq: 50 + # Parameters of transform for equivariance loss + transform_params: + # Sigma for affine part + sigma_affine: 0.05 + # Sigma for deformation part + sigma_tps: 0.005 + # Number of point in the deformation grid + points_tps: 5 + loss_weights: + equivariance: 10 + perceptual: [10, 10, 10, 10, 10] + + +# Visualization parameters +visualizer_params: + # Draw keypoints (shifts in affine transformations) of this size, increase or decrease depending on resolution + kp_size: 5 + # Draw white border around images + draw_border: True + # Color map for keypoints + colormap: 'gist_rainbow' diff --git a/constants.py b/constants.py index d69b199..2eade3a 100644 --- a/constants.py +++ b/constants.py @@ -9,15 +9,92 @@ supported_model_names = [ # "bair", ] -non_cpk_model_names = [ - "fashion" -] +non_cpk_model_names = ["fashion"] default_model_name = "vox-adv" checkpoint_folder_link = "https://github.com/graphemecluster/first-order-model-demo/releases/download/checkpoints" checkpoint_suffix = ".pth.tar" +config_suffix = "256.yaml" config_folder = "config" -checkpoint_folder = "checkpoints" \ No newline at end of file +checkpoint_folder = "checkpoints" + +face_parser_checkpoint_name = "79999_iter.pth" + +partswap_model_config_dict = { + "vox-5segments": "vox-256-sem-5segments.yaml", + "vox-10segments": "vox-256-sem-10segments.yaml", + "vox-15segments": "vox-256-sem-15segments.yaml", + "vox-cpk": "vox-256-sem-10segments.yaml", + "vox-adv-cpk": "vox-256-sem-10segments.yaml", + "vox-first-order": "vox-256-sem-10segments.yaml", +} + +partswap_face_parser_enforce_models = ["vox-first-order"] + +partswap_fomm_model_names = ["vox-cpk", "vox-adv-cpk", "vox-first-order"] + +partswap_model_length_dict = { + "vox-5segments": 6, + "vox-10segments": 11, + "vox-15segments": 16, + "vox-cpk": 11, + "vox-adv-cpk": 11, + "vox-first-order": 11, +} + +partswap_model_names = list(partswap_model_config_dict.keys()) + +default_partswap_model_name = "vox-10segments" + +seg10_segment_names = [ + "m_0_background", + "m_1_lf_ear", + "m_2_mouth", + "m_3_lf_fr_head", + "m_4_rt_ear", + "m_5_frnt_face", + "m_6_lt_cheek", + "m_7_eyes", + "m_8_rt_fr_head", + "m_9_shoulder", + "m_10_hair_top", +] + +seg5_segment_names = [ + "m_0_background", + "m_1_chin", + "m_2_frnt_face", + "m_3_shoulder", + "m_4_hair", + "m_5_fr_head", +] + +seg_15_segment_names = [ + "m_0_background", + "m_1_rt_fr_head", + "m_2_eyes", + "m_3_mouth", + "m_4_lwr_shoulder", + "m_5_shoulder", + "m_6_lt_fr_head", + "m_7", + "m_8_hair_tip", + "m_9_frnt_face", + "m_10_rt_ear", + "m_11_outer_hair", + "m_12_lt_ear", + "m_13_lt_hair", + "m_14_rt_hair", + "m_15_neck", +] + +seg_model_dict = { + "vox-5segments": seg5_segment_names, + "vox-10segments": seg10_segment_names, + "vox-15segments": seg_15_segment_names, + "vox": seg10_segment_names, + "vox-adv": seg10_segment_names, +} diff --git a/face_parsing/face_parsing_loader.py b/face_parsing/face_parsing_loader.py new file mode 100644 index 0000000..24ab6a8 --- /dev/null +++ b/face_parsing/face_parsing_loader.py @@ -0,0 +1,37 @@ +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) + + mean = torch.Tensor(np.array([0.485, 0.456, 0.406], dtype=np.float32)).view( + 1, 3, 1, 1 + ) + std = torch.Tensor(np.array([0.229, 0.224, 0.225], dtype=np.float32)).view( + 1, 3, 1, 1 + ) + + if not cpu: + face_parser.cuda() + face_parser.load_state_dict(torch.load(face_parser_path)) + face_parser.mean = mean.cuda() + face_parser.std = std.cuda() + else: + face_parser.load_state_dict( + torch.load(face_parser_path, map_location=torch.device("cpu")) + ) + face_parser.mean = mean + face_parser.std = std + print("face_parser model loaded!") + + return face_parser diff --git a/face_parsing/face_parsing_model.py b/face_parsing/face_parsing_model.py new file mode 100644 index 0000000..198f949 --- /dev/null +++ b/face_parsing/face_parsing_model.py @@ -0,0 +1,276 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .resnet import Resnet18 + + +class ConvBNReLU(nn.Module): + def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, *args, **kwargs): + super(ConvBNReLU, self).__init__() + self.conv = nn.Conv2d( + in_chan, + out_chan, + kernel_size=ks, + stride=stride, + padding=padding, + bias=False, + ) + self.bn = nn.BatchNorm2d(out_chan) + self.init_weight() + + def forward(self, x): + x = self.conv(x) + x = F.relu(self.bn(x)) + return x + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if ly.bias is not None: + nn.init.constant_(ly.bias, 0) + + +class BiSeNetOutput(nn.Module): + def __init__(self, in_chan, mid_chan, n_classes, *args, **kwargs): + super(BiSeNetOutput, self).__init__() + self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1) + self.conv_out = nn.Conv2d(mid_chan, n_classes, kernel_size=1, bias=False) + self.init_weight() + + def forward(self, x): + x = self.conv(x) + x = self.conv_out(x) + return x + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if ly.bias is not None: + nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if module.bias is not None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class AttentionRefinementModule(nn.Module): + def __init__(self, in_chan, out_chan, *args, **kwargs): + super(AttentionRefinementModule, self).__init__() + self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1) + self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size=1, bias=False) + self.bn_atten = nn.BatchNorm2d(out_chan) + self.sigmoid_atten = nn.Sigmoid() + self.init_weight() + + def forward(self, x): + feat = self.conv(x) + atten = F.avg_pool2d(feat, feat.size()[2:]) + atten = self.conv_atten(atten) + atten = self.bn_atten(atten) + atten = self.sigmoid_atten(atten) + out = torch.mul(feat, atten) + return out + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if ly.bias is not None: + nn.init.constant_(ly.bias, 0) + + +class ContextPath(nn.Module): + def __init__(self, *args, **kwargs): + super(ContextPath, self).__init__() + self.resnet = Resnet18() + self.arm16 = AttentionRefinementModule(256, 128) + self.arm32 = AttentionRefinementModule(512, 128) + self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1) + self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1) + self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0) + + self.init_weight() + + def forward(self, x): + H0, W0 = x.size()[2:] + feat8, feat16, feat32 = self.resnet(x) + H8, W8 = feat8.size()[2:] + H16, W16 = feat16.size()[2:] + H32, W32 = feat32.size()[2:] + + avg = F.avg_pool2d(feat32, feat32.size()[2:]) + avg = self.conv_avg(avg) + avg_up = F.interpolate(avg, (H32, W32), mode="nearest") + + feat32_arm = self.arm32(feat32) + feat32_sum = feat32_arm + avg_up + feat32_up = F.interpolate(feat32_sum, (H16, W16), mode="nearest") + feat32_up = self.conv_head32(feat32_up) + + feat16_arm = self.arm16(feat16) + feat16_sum = feat16_arm + feat32_up + feat16_up = F.interpolate(feat16_sum, (H8, W8), mode="nearest") + feat16_up = self.conv_head16(feat16_up) + + return feat8, feat16_up, feat32_up # x8, x8, x16 + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if ly.bias is not None: + nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, (nn.Linear, nn.Conv2d)): + wd_params.append(module.weight) + if module.bias is not None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +### This is not used, since I replace this with the resnet feature with the same size +class SpatialPath(nn.Module): + def __init__(self, *args, **kwargs): + super(SpatialPath, self).__init__() + self.conv1 = ConvBNReLU(3, 64, ks=7, stride=2, padding=3) + self.conv2 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1) + self.conv3 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1) + self.conv_out = ConvBNReLU(64, 128, ks=1, stride=1, padding=0) + self.init_weight() + + def forward(self, x): + feat = self.conv1(x) + feat = self.conv2(feat) + feat = self.conv3(feat) + feat = self.conv_out(feat) + return feat + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if ly.bias is not None: + nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if module.bias is not None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class FeatureFusionModule(nn.Module): + def __init__(self, in_chan, out_chan, *args, **kwargs): + super(FeatureFusionModule, self).__init__() + self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0) + self.conv1 = nn.Conv2d( + out_chan, out_chan // 4, kernel_size=1, stride=1, padding=0, bias=False + ) + self.conv2 = nn.Conv2d( + out_chan // 4, out_chan, kernel_size=1, stride=1, padding=0, bias=False + ) + self.relu = nn.ReLU(inplace=True) + self.sigmoid = nn.Sigmoid() + self.init_weight() + + def forward(self, fsp, fcp): + fcat = torch.cat([fsp, fcp], dim=1) + feat = self.convblk(fcat) + atten = F.avg_pool2d(feat, feat.size()[2:]) + atten = self.conv1(atten) + atten = self.relu(atten) + atten = self.conv2(atten) + atten = self.sigmoid(atten) + feat_atten = torch.mul(feat, atten) + feat_out = feat_atten + feat + return feat_out + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if ly.bias is not None: + nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if module.bias is not None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class BiSeNet(nn.Module): + def __init__(self, n_classes, *args, **kwargs): + super(BiSeNet, self).__init__() + self.cp = ContextPath() + ## here self.sp is deleted + self.ffm = FeatureFusionModule(256, 256) + self.conv_out = BiSeNetOutput(256, 256, n_classes) + self.conv_out16 = BiSeNetOutput(128, 64, n_classes) + self.conv_out32 = BiSeNetOutput(128, 64, n_classes) + self.init_weight() + + def forward(self, x): + H, W = x.size()[2:] + feat_res8, feat_cp8, feat_cp16 = self.cp(x) # here return res3b1 feature + feat_sp = feat_res8 # use res3b1 feature to replace spatial path feature + feat_fuse = self.ffm(feat_sp, feat_cp8) + + feat_out = self.conv_out(feat_fuse) + feat_out16 = self.conv_out16(feat_cp8) + feat_out32 = self.conv_out32(feat_cp16) + + feat_out = F.interpolate(feat_out, (H, W), mode="bilinear", align_corners=True) + feat_out16 = F.interpolate( + feat_out16, (H, W), mode="bilinear", align_corners=True + ) + feat_out32 = F.interpolate( + feat_out32, (H, W), mode="bilinear", align_corners=True + ) + return feat_out, feat_out16, feat_out32 + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if ly.bias is not None: + nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], [] + for name, child in self.named_children(): + child_wd_params, child_nowd_params = child.get_params() + if isinstance(child, FeatureFusionModule) or isinstance( + child, BiSeNetOutput + ): + lr_mul_wd_params += child_wd_params + lr_mul_nowd_params += child_nowd_params + else: + wd_params += child_wd_params + nowd_params += child_nowd_params + return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params diff --git a/face_parsing/resnet.py b/face_parsing/resnet.py new file mode 100644 index 0000000..ba9002d --- /dev/null +++ b/face_parsing/resnet.py @@ -0,0 +1,94 @@ +import torch.nn as nn +import torch.nn.functional as F +import torch +from pathlib import Path + + +def conv3x3(in_planes, out_planes, stride=1): + """3x3 convolution with padding""" + return nn.Conv2d( + in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False + ) + + +class BasicBlock(nn.Module): + def __init__(self, in_chan, out_chan, stride=1): + super(BasicBlock, self).__init__() + self.conv1 = conv3x3(in_chan, out_chan, stride) + self.bn1 = nn.BatchNorm2d(out_chan) + self.conv2 = conv3x3(out_chan, out_chan) + self.bn2 = nn.BatchNorm2d(out_chan) + self.relu = nn.ReLU(inplace=True) + self.downsample = None + if in_chan != out_chan or stride != 1: + self.downsample = nn.Sequential( + nn.Conv2d(in_chan, out_chan, kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(out_chan), + ) + + def forward(self, x): + residual = self.conv1(x) + residual = F.relu(self.bn1(residual)) + residual = self.conv2(residual) + residual = self.bn2(residual) + + shortcut = x + if self.downsample is not None: + shortcut = self.downsample(x) + + out = shortcut + residual + out = self.relu(out) + return out + + +def create_layer_basic(in_chan, out_chan, bnum, stride=1): + layers = [BasicBlock(in_chan, out_chan, stride=stride)] + for i in range(bnum - 1): + layers.append(BasicBlock(out_chan, out_chan, stride=1)) + return nn.Sequential(*layers) + + +class Resnet18(nn.Module): + def __init__(self): + super(Resnet18, self).__init__() + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False) + self.bn1 = nn.BatchNorm2d(64) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = create_layer_basic(64, 64, bnum=2, stride=1) + self.layer2 = create_layer_basic(64, 128, bnum=2, stride=2) + self.layer3 = create_layer_basic(128, 256, bnum=2, stride=2) + self.layer4 = create_layer_basic(256, 512, bnum=2, stride=2) + self.init_weight() + + def forward(self, x): + x = self.conv1(x) + x = F.relu(self.bn1(x)) + x = self.maxpool(x) + + x = self.layer1(x) + feat8 = self.layer2(x) # 1/8 + feat16 = self.layer3(feat8) # 1/16 + feat32 = self.layer4(feat16) # 1/32 + return feat8, feat16, feat32 + + def init_weight(self): + base_dir = Path(__file__).resolve().parent + print(f"resnet base dir {base_dir}") + state_dict = torch.load(f"{base_dir}/resnet18-5c106cde.pth") + self_state_dict = self.state_dict() + for k, v in state_dict.items(): + if "fc" in k: + continue + self_state_dict.update({k: v}) + self.load_state_dict(self_state_dict) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, (nn.Linear, nn.Conv2d)): + wd_params.append(module.weight) + if module.bias is not None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params diff --git a/fomm_inference.py b/fomm_inference.py new file mode 100644 index 0000000..2665374 --- /dev/null +++ b/fomm_inference.py @@ -0,0 +1,185 @@ +import torch +import numpy as np +from comfy.utils import ProgressBar +from tqdm import tqdm +from scipy.spatial import ConvexHull +from .modules.generator import OcclusionAwareGenerator +from .modules.keypoint_detector import KPDetector + +from .utils import tensor_to_numpy + +try: + import face_alignment + + face_alignment_available = True +except ImportError: + face_alignment_available = False + print("Warning: face_alignment module not found. fine_best_frame will not work. ") + + +def normalize_kp_best_frame(kp): + kp = kp - kp.mean(axis=0, keepdims=True) + area = ConvexHull(kp[:, :2]).volume + area = np.sqrt(area) + kp[:, :2] = kp[:, :2] / area + return kp + +def normalize_kp( + kp_source, + kp_driving, + kp_driving_initial, + adapt_movement_scale: bool, + use_relative_movement: bool, + use_relative_jacobian: bool, +): + if adapt_movement_scale: + source_area = ConvexHull(kp_source["value"][0].data.cpu().numpy()).volume + driving_area = ConvexHull( + kp_driving_initial["value"][0].data.cpu().numpy() + ).volume + adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area) + else: + adapt_movement_scale = 1 + kp_new = {k: v for k, v in kp_driving.items()} + + if use_relative_movement: + kp_value_diff = kp_driving["value"] - kp_driving_initial["value"] + kp_value_diff *= adapt_movement_scale + kp_new["value"] = kp_value_diff + kp_source["value"] + + if use_relative_jacobian: + jacobian_diff = torch.matmul( + kp_driving["jacobian"], torch.inverse(kp_driving_initial["jacobian"]) + ) + kp_new["jacobian"] = torch.matmul(jacobian_diff, kp_source["jacobian"]) + + return kp_new + +def find_best_frame(source_image, driving_video: list, cpu=False) -> int: + if not face_alignment_available: + print("Face Alignment module were not found. Best frame set to 0.") + return 0 + + print("Initializing Face Alignment...") + fa = face_alignment.FaceAlignment( + face_alignment.LandmarksType.TWO_D, + flip_input=True, + device="cpu" if cpu else "cuda", + ) + print("Face alignment model loaded.") + + source_np = tensor_to_numpy(source_image) + kp_source = fa.get_landmarks(255 * source_np)[0] + kp_source = normalize_kp_best_frame(kp_source) + + norm = float("inf") + frame_num = 0 + + num_frames = driving_video.shape[2] + + for i in tqdm(range(num_frames)): + # Extract single frame from driving video + frame = driving_video[0, :, i] # Shape: [3, 256, 256] + frame_np = frame.permute(1, 2, 0).cpu().numpy() + + kp_driving = fa.get_landmarks(255 * frame_np)[0] + kp_driving = normalize_kp_best_frame(kp_driving) + new_norm = (np.abs(kp_source - kp_driving) ** 2).sum() + if new_norm < norm: + norm = new_norm + frame_num = i + + print(f"Best frame: {frame_num}") + + return frame_num + +def inference( + source_image, + driving_video: list, + generator: OcclusionAwareGenerator, + kp_detector: KPDetector, + relative_movement: bool, + relative_jacobian: bool, + adapt_movement_scale: bool, + cpu=False, +) -> list: + with torch.no_grad(): + source = source_image + driving = driving_video + predictions = [] + + if not cpu: + source = source.cuda() + driving = driving.cuda() + print("Using GPU") + + kp_source = kp_detector(source) + kp_driving_initial = kp_detector(driving[:, :, 0]) + kp_norm = None + + num_frames = driving.shape[2] + pbar = ProgressBar(num_frames) + + for frame_idx in tqdm(range(num_frames)): + driving_frame = driving[:, :, frame_idx] + kp_driving = kp_detector(driving_frame) + kp_norm = normalize_kp( + kp_source=kp_source, + kp_driving=kp_driving, + kp_driving_initial=kp_driving_initial, + use_relative_movement=relative_movement, + use_relative_jacobian=relative_jacobian, + adapt_movement_scale=adapt_movement_scale, + ) + out = generator(source, kp_source=kp_source, kp_driving=kp_norm) + out_np = np.transpose(out["prediction"].data.cpu().numpy(), [0, 2, 3, 1])[0] + predictions.append(out_np) + pbar.update_absolute(frame_idx, num_frames) + + return predictions + + +def inference_best_frame( + source_image, + driving_video, + generator: OcclusionAwareGenerator, + kp_detector: KPDetector, + relative_movement: bool, + relative_jacobian: bool, + adapt_movement_scale: bool, + cpu=False, +) -> list: + best_frame_idx = find_best_frame(source_image, driving_video, cpu) + + first_half = driving_video[ + :, :, : best_frame_idx + 1 + ] # Include the best frame in the first half + first_half = torch.flip(first_half, dims=[2]) # Reverse the first half + second_half = driving_video[:, :, best_frame_idx + 1 :] + + predictions_first = inference( + source_image, + first_half, + generator, + kp_detector, + relative_movement, + relative_jacobian, + adapt_movement_scale, + cpu, + ) + predictions_second = inference( + source_image, + second_half, + generator, + kp_detector, + relative_movement, + relative_jacobian, + adapt_movement_scale, + cpu, + ) + + predictions = [] + predictions_first = predictions_first[::-1] # Reverse the first half back + predictions.extend(predictions_first) + predictions.extend(predictions_second) + return predictions \ No newline at end of file diff --git a/fomm_loader.py b/fomm_loader.py new file mode 100644 index 0000000..b100543 --- /dev/null +++ b/fomm_loader.py @@ -0,0 +1,35 @@ +import yaml +import torch + +from .modules.generator import OcclusionAwareGenerator +from .modules.keypoint_detector import KPDetector + +def load_checkpoint( + config_path: str, checkpoint_path: str, cpu=False +) -> tuple[OcclusionAwareGenerator, KPDetector]: + with open(config_path) as f: + config = yaml.full_load(f) + + generator = OcclusionAwareGenerator( + **config["model_params"]["generator_params"], + **config["model_params"]["common_params"], + ) + kp_detector = KPDetector( + **config["model_params"]["kp_detector_params"], + **config["model_params"]["common_params"], + ) + + if not cpu: + generator.cuda() + kp_detector.cuda() + checkpoint = torch.load(checkpoint_path) + else: + checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu")) + + generator.load_state_dict(checkpoint["generator"]) + kp_detector.load_state_dict(checkpoint["kp_detector"]) + + generator.eval() + kp_detector.eval() + + return generator, kp_detector \ No newline at end of file diff --git a/modules/dense_motion.py b/modules/dense_motion.py index 56d5ac5..1767bec 100644 --- a/modules/dense_motion.py +++ b/modules/dense_motion.py @@ -29,13 +29,13 @@ class DenseMotionNetwork(nn.Module): if self.scale_factor != 1: self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor) - def create_heatmap_representations(self, source_image, kp_driving, kp_source): + def create_heatmap_representations(self, source_image, kp_driving, kp_source, kp_key='value'): """ Eq 6. in the paper H_k(z) """ spatial_size = source_image.shape[2:] - gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=self.kp_variance) - gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=self.kp_variance) + gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=self.kp_variance, kp_key=kp_key) + gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=self.kp_variance, kp_key=kp_key) heatmap = gaussian_driving - gaussian_source #adding background feature @@ -44,22 +44,22 @@ class DenseMotionNetwork(nn.Module): heatmap = heatmap.unsqueeze(2) return heatmap - def create_sparse_motions(self, source_image, kp_driving, kp_source): + def create_sparse_motions(self, source_image, kp_driving, kp_source, kp_key='value', jacob_key='jacobian'): """ Eq 4. in the paper T_{s<-d}(z) """ bs, _, h, w = source_image.shape - identity_grid = make_coordinate_grid((h, w), type=kp_source['value'].type()) + identity_grid = make_coordinate_grid((h, w), type=kp_source[kp_key].type()) identity_grid = identity_grid.view(1, 1, h, w, 2) - coordinate_grid = identity_grid - kp_driving['value'].view(bs, self.num_kp, 1, 1, 2) - if 'jacobian' in kp_driving: - jacobian = torch.matmul(kp_source['jacobian'], torch.inverse(kp_driving['jacobian'])) + coordinate_grid = identity_grid - kp_driving[kp_key].view(bs, self.num_kp, 1, 1, 2) + if jacob_key in kp_driving: + jacobian = torch.matmul(kp_source[jacob_key], torch.inverse(kp_driving[jacob_key])) jacobian = jacobian.unsqueeze(-3).unsqueeze(-3) jacobian = jacobian.repeat(1, 1, h, w, 1, 1) coordinate_grid = torch.matmul(jacobian, coordinate_grid.unsqueeze(-1)) coordinate_grid = coordinate_grid.squeeze(-1) - driving_to_source = coordinate_grid + kp_source['value'].view(bs, self.num_kp, 1, 1, 2) + driving_to_source = coordinate_grid + kp_source[kp_key].view(bs, self.num_kp, 1, 1, 2) #adding background feature identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1) @@ -111,3 +111,36 @@ class DenseMotionNetwork(nn.Module): out_dict['occlusion_map'] = occlusion_map return out_dict + + def forward_partswap(self, source_image, seg_target, seg_source): + if self.scale_factor != 1: + source_image = self.down(source_image) + + bs, _, h, w = source_image.shape + + out_dict = dict() + heatmap_representation = self.create_heatmap_representations(source_image, seg_target, seg_source, kp_key='shift') + sparse_motion = self.create_sparse_motions(source_image, seg_target, seg_source, kp_key='shift', jacob_key='affine') + deformed_source = self.create_deformed_source_image(source_image, sparse_motion) + out_dict['sparse_deformed'] = deformed_source + + input = torch.cat([heatmap_representation, deformed_source], dim=2) + input = input.view(bs, -1, h, w) + + prediction = self.hourglass(input) + + mask = self.mask(prediction) + mask = F.softmax(mask, dim=1) + out_dict['mask'] = mask + mask = mask.unsqueeze(2) + sparse_motion = sparse_motion.permute(0, 1, 4, 2, 3) + deformation = (sparse_motion * mask).sum(dim=1) + deformation = deformation.permute(0, 2, 3, 1) + + out_dict['deformation'] = deformation + + if self.occlusion: + visibility = torch.sigmoid(self.occlusion(prediction)) + out_dict['visibility'] = visibility + + return out_dict diff --git a/modules/reconstruction_module.py b/modules/reconstruction_module.py new file mode 100644 index 0000000..bd7cf7d --- /dev/null +++ b/modules/reconstruction_module.py @@ -0,0 +1,109 @@ +import torch +from torch import nn +import torch.nn.functional as F +from .util import ResBlock2d, SameBlock2d, UpBlock2d, DownBlock2d +from .util import make_coordinate_grid + + +class ReconstructionModule(nn.Module): + """ + Reconstruct target from source and segmentation of target and part transformations + """ + + def __init__(self, num_channels, block_expansion, max_features, num_down_blocks, + num_bottleneck_blocks, num_segments, estimate_visibility=False, **kwargs): + super(ReconstructionModule, self).__init__() + + self.first = SameBlock2d(num_channels, block_expansion, kernel_size=(7, 7), padding=(3, 3)) + + down_blocks = [] + for i in range(num_down_blocks): + in_features = min(max_features, block_expansion * (2 ** i)) + out_features = min(max_features, block_expansion * (2 ** (i + 1))) + down_blocks.append(DownBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1))) + self.down_blocks = nn.ModuleList(down_blocks) + + up_blocks = [] + for i in range(num_down_blocks): + in_features = min(max_features, block_expansion * (2 ** (num_down_blocks - i))) + out_features = min(max_features, block_expansion * (2 ** (num_down_blocks - i - 1))) + up_blocks.append(UpBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1))) + self.up_blocks = nn.ModuleList(up_blocks) + + self.bottleneck = torch.nn.Sequential() + in_features = min(max_features, block_expansion * (2 ** num_down_blocks)) + for i in range(num_bottleneck_blocks): + self.bottleneck.add_module('r' + str(i), ResBlock2d(in_features, kernel_size=(3, 3), padding=(1, 1))) + + self.final = nn.Conv2d(block_expansion, num_channels, kernel_size=(7, 7), padding=(3, 3)) + self.estimate_visibility = estimate_visibility + self.num_channels = num_channels + self.num_segments = num_segments + + def segment_motion(self, seg_target, seg_source): + bs, _, h, w = seg_target['segmentation'].shape + identity_grid = make_coordinate_grid((h, w), type=seg_source['shift'].type()) + identity_grid = identity_grid.view(1, 1, h, w, 2) + coordinate_grid = identity_grid - seg_target['shift'].view(bs, self.num_segments, 1, 1, 2) + if 'affine' in seg_target: + affine = torch.matmul(seg_source['affine'], torch.inverse(seg_target['affine'])) + affine = affine.unsqueeze(-3).unsqueeze(-3) + affine = affine.repeat(1, 1, h, w, 1, 1) + coordinate_grid = torch.matmul(affine, coordinate_grid.unsqueeze(-1)) + coordinate_grid = coordinate_grid.squeeze(-1) + + target_to_source = coordinate_grid + seg_source['shift'].view(bs, self.num_segments, 1, 1, 2) + + # adding background feature + identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1) + return torch.cat([identity_grid, target_to_source], dim=1) + + def deform_input(self, inp, deformation): + _, h_old, w_old, _ = deformation.shape + _, _, h, w = inp.shape + if h_old != h or w_old != w: + deformation = deformation.permute(0, 3, 1, 2) + deformation = F.interpolate(deformation, size=(h, w), mode='bilinear') + deformation = deformation.permute(0, 2, 3, 1) + return F.grid_sample(inp, deformation) + + def forward(self, source_image, seg_target, seg_source): + # Encoding (downsampling) part + out = self.first(source_image) + for i in range(len(self.down_blocks)): + out = self.down_blocks[i](out) + + output_dict = {} + + # Computing segment motion + segment_motions = self.segment_motion(seg_target, seg_source) + segment_motions = segment_motions.permute(0, 1, 4, 2, 3) + mask = seg_target['segmentation'].unsqueeze(2) + deformation = (segment_motions * mask).sum(dim=1) + deformation = deformation.permute(0, 2, 3, 1) + + # Deform according to motion + out = self.deform_input(out, deformation) + output_dict["deformed"] = self.deform_input(source_image, deformation) + + if self.estimate_visibility: + visibility = seg_source['segmentation'][:, 1:].sum(dim=1, keepdim=True) * \ + (1 - seg_target['segmentation'][:, 1:].sum(dim=1, keepdim=True).detach()) + visibility = 1 - visibility + + if out.shape[2] != visibility.shape[2] or out.shape[3] != visibility.shape[3]: + visibility = F.interpolate(visibility, size=out.shape[2:], mode='bilinear') + out = out * visibility + output_dict['visibility_maps'] = visibility + + out = self.bottleneck(out) + + for i in range(len(self.up_blocks)): + out = self.up_blocks[i](out) + + out = self.final(out) + out = F.sigmoid(out) + + output_dict["prediction"] = out + + return output_dict diff --git a/modules/segmentation_module.py b/modules/segmentation_module.py new file mode 100644 index 0000000..e6d432b --- /dev/null +++ b/modules/segmentation_module.py @@ -0,0 +1,78 @@ +from torch import nn +import torch +import torch.nn.functional as F +from .util import Hourglass, make_coordinate_grid, AntiAliasInterpolation2d + + +class SegmentationModule(nn.Module): + """ + Computing a segmentation map and affine transformations. + """ + + def __init__(self, block_expansion, num_segments, num_channels, max_features, + num_blocks, temperature, estimate_affine_part=False, scale_factor=1): + super(SegmentationModule, self).__init__() + + self.predictor = Hourglass(block_expansion, in_features=num_channels, + max_features=max_features, num_blocks=num_blocks) + self.num_segments = num_segments + self.shift = nn.Conv2d(in_channels=self.predictor.out_filters, out_channels=num_segments, kernel_size=(7, 7), + padding=(3, 3)) + + if estimate_affine_part: + self.affine = nn.Conv2d(in_channels=self.predictor.out_filters, + out_channels=4 * num_segments, kernel_size=(7, 7), padding=(3, 3)) + self.affine.weight.data.zero_() + self.affine.bias.data.copy_(torch.tensor([1, 0, 0, 1] * num_segments, dtype=torch.float)) + else: + self.affine = None + + self.segmentation = nn.Conv2d(in_channels=self.predictor.out_filters, + out_channels=(1 + num_segments), kernel_size=(7, 7), padding=(3, 3)) + + self.temperature = temperature + self.scale_factor = scale_factor + if self.scale_factor != 1: + self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor) + + def gaussian2kp(self, heatmap): + """ + Extract the mean and the variance from a heatmap + """ + shape = heatmap.shape + heatmap = heatmap.unsqueeze(-1) + grid = make_coordinate_grid(shape[2:], heatmap.type()).unsqueeze_(0).unsqueeze_(0) + result = (heatmap * grid).sum(dim=(2, 3)) + + return result + + def forward(self, x): + if self.scale_factor != 1: + x = self.down(x) + + out = {} + feature_map = self.predictor(x) + out['segmentation'] = F.softmax(self.segmentation(feature_map), dim=1) + + prediction = self.shift(feature_map) + + final_shape = prediction.shape + heatmap = prediction.view(final_shape[0], final_shape[1], -1) + heatmap = F.softmax(heatmap / self.temperature, dim=2) + heatmap = heatmap.view(*final_shape) + + out['shift'] = self.gaussian2kp(heatmap) + + if self.affine is not None: + affine_map = self.affine(feature_map) + affine_map = affine_map.reshape(final_shape[0], self.num_segments, 4, final_shape[2], + final_shape[3]) + heatmap = heatmap.unsqueeze(2) + + affine = heatmap * affine_map + affine = affine.view(final_shape[0], final_shape[1], 4, -1) + affine = affine.sum(dim=-1) + affine = affine.view(affine.shape[0], affine.shape[1], 2, 2) + out['affine'] = affine + + return out diff --git a/modules/util.py b/modules/util.py index f9e5eaf..3b2e5d5 100644 --- a/modules/util.py +++ b/modules/util.py @@ -6,11 +6,11 @@ import torch from ..sync_batchnorm.batchnorm import SynchronizedBatchNorm2d as BatchNorm2d -def kp2gaussian(kp, spatial_size, kp_variance): +def kp2gaussian(kp, spatial_size, kp_variance, kp_key='value'): """ Transform a keypoint into gaussian like representation """ - mean = kp['value'] + mean = kp[kp_key] coordinate_grid = make_coordinate_grid(spatial_size, mean.type()) number_of_leading_dimensions = len(mean.shape) - 1 diff --git a/partswap_inference.py b/partswap_inference.py new file mode 100644 index 0000000..b062e1a --- /dev/null +++ b/partswap_inference.py @@ -0,0 +1,84 @@ +import numpy as np +import torch +import torch.nn.functional as F +from tqdm import tqdm + +from comfy.utils import ProgressBar +from .modules.segmentation_module import SegmentationModule +from .partswap_model import PartSwapGenerator + + +def face_parse_seg( + source_image, face_parser_model, in_size=(512, 512), out_size=(64, 64) +): + seg = F.interpolate(source_image, size=in_size) + seg = (seg - face_parser_model.mean) / face_parser_model.std + seg = torch.softmax(face_parser_model(seg)[0], dim=1) + seg = F.interpolate(seg, size=out_size) + return seg + + +def partswap_inference( + swap_indices: list[int], + source_image, + target_video, + reconstruction_module: PartSwapGenerator, + segmentation_module: SegmentationModule, + use_source_seg: bool, + face_parser_model=None, + hard_edges=False, + cpu=False, +) -> list: + with torch.no_grad(): + predictions = [] + seg_targets = [] + if not cpu: + source_image = source_image.cuda() + target_video = target_video.cuda() + print("Using GPU") + + seg_source = segmentation_module(source_image) + print(f"{seg_source['segmentation'].shape}") + + num_frames = target_video.shape[2] + pbar = ProgressBar(num_frames) + for frame_idx in tqdm(range(num_frames)): + target_frame = target_video[:, :, frame_idx] + seg_target = segmentation_module(target_frame) + seg_targets.append(seg_target) + + if face_parser_model is not None: + blend_mask = face_parse_seg( + source_image if use_source_seg else target_frame, + face_parser_model, + (512, 512), + (64, 64), + ) + + else: + blend_mask = ( + seg_source["segmentation"] + if use_source_seg + else seg_target["segmentation"] + ) + + blend_mask = blend_mask[:, swap_indices].sum(dim=1, keepdim=True) + if hard_edges: + blend_mask = (blend_mask > 0.5).type(blend_mask.type()) + print(f"{blend_mask.shape=}") + out = reconstruction_module( + source_image, + target_frame, + seg_source=seg_source, + seg_target=seg_target, + blend_mask=blend_mask, + use_source_segmentation=use_source_seg, + ) + + predictions.append( + np.transpose(out["prediction"].data.cpu().numpy(), [0, 2, 3, 1])[0] + ) + + pbar.update_absolute(frame_idx, num_frames) + + return seg_source, seg_targets, predictions diff --git a/partswap_loader.py b/partswap_loader.py new file mode 100644 index 0000000..55c599c --- /dev/null +++ b/partswap_loader.py @@ -0,0 +1,92 @@ +import torch +import yaml + +from .modules.segmentation_module import SegmentationModule +from .partswap_model import PartSwapGenerator +from .sync_batchnorm.replicate import DataParallelWithCallback + + +def partial_state_dict_load(module, state_dict): + own_state = module.state_dict() + for name, param in state_dict.items(): + if name not in own_state: + continue + + if isinstance(param, torch.nn.Parameter): + # backwards compatibility for serialized parameters + param = param.data + own_state[name].copy_(param) + + +def load_reconstruction_module(module, checkpoint): + if "generator" in checkpoint: + partial_state_dict_load(module, checkpoint["generator"]) + else: + module.load_state_dict(checkpoint["reconstruction_module"]) + + +def load_segmentation_module(module, checkpoint): + if "kp_detector" in checkpoint: + partial_state_dict_load(module, checkpoint["kp_detector"]) + module.state_dict()["affine.weight"].copy_( + checkpoint["kp_detector"]["jacobian.weight"] + ) + module.state_dict()["affine.bias"].copy_( + checkpoint["kp_detector"]["jacobian.bias"] + ) + module.state_dict()["shift.weight"].copy_( + checkpoint["kp_detector"]["kp.weight"] + ) + module.state_dict()["shift.bias"].copy_(checkpoint["kp_detector"]["kp.bias"]) + if "semantic_seg.weight" in checkpoint["kp_detector"]: + module.state_dict()["segmentation.weight"].copy_( + checkpoint["kp_detector"]["semantic_seg.weight"] + ) + module.state_dict()["segmentation.bias"].copy_( + checkpoint["kp_detector"]["semantic_seg.bias"] + ) + else: + print("Segmentation part initialized at random.") + else: + module.load_state_dict(checkpoint["segmentation_module"]) + + +def load_partswap_checkpoint( + config_path: str, + checkpoint_path: str, + blend_scale: float, + use_fomm: bool, + cpu=False, +): + with open(config_path) as f: + config = yaml.full_load(f) + + reconstruction_module = PartSwapGenerator( + blend_scale=blend_scale, + first_order_motion_model=use_fomm, + **config["model_params"]["reconstruction_module_params"], + **config["model_params"]["common_params"], + ) + + segmentation_module = SegmentationModule( + **config["model_params"]["segmentation_module_params"], + **config["model_params"]["common_params"], + ) + + if cpu: + checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu")) + else: + segmentation_module.cuda() + checkpoint = torch.load(checkpoint_path) + + load_reconstruction_module(reconstruction_module, checkpoint) + load_segmentation_module(segmentation_module, checkpoint) + + if not cpu: + reconstruction_module = DataParallelWithCallback(reconstruction_module) + segmentation_module = DataParallelWithCallback(segmentation_module) + + reconstruction_module.eval() + segmentation_module.eval() + + return reconstruction_module, segmentation_module diff --git a/partswap_model.py b/partswap_model.py new file mode 100644 index 0000000..ee55dea --- /dev/null +++ b/partswap_model.py @@ -0,0 +1,105 @@ +import torch.nn.functional as F + +from .modules.dense_motion import DenseMotionNetwork +from .modules.reconstruction_module import ReconstructionModule +from .modules.util import AntiAliasInterpolation2d + + +class PartSwapGenerator(ReconstructionModule): + def __init__(self, blend_scale=1, first_order_motion_model=False, **kwargs): + super(PartSwapGenerator, self).__init__(**kwargs) + if blend_scale == 1: + self.blend_downsample = lambda x: x + else: + self.blend_downsample = AntiAliasInterpolation2d(1, blend_scale) + + if first_order_motion_model: + self.dense_motion_network = DenseMotionNetwork( + block_expansion=64, + num_blocks=5, + max_features=1024, + num_kp=10, + num_channels=3, + estimate_occlusion_map=True, + scale_factor=0.25, + ) + else: + self.dense_motion_network = None + + def forward( + self, + source_image, + target_image, + seg_target, + seg_source, + blend_mask, + use_source_segmentation=False, + ): + # Encoding of source image + enc_source = self.first(source_image) + for i in range(len(self.down_blocks)): + enc_source = self.down_blocks[i](enc_source) + + # Encoding of target image + enc_target = self.first(target_image) + for i in range(len(self.down_blocks)): + enc_target = self.down_blocks[i](enc_target) + + output_dict = {} + # Compute flow field for source image + if self.dense_motion_network is None: + segment_motions = self.segment_motion(seg_target, seg_source) + segment_motions = segment_motions.permute(0, 1, 4, 2, 3) + mask = seg_target["segmentation"].unsqueeze(2) + deformation = (segment_motions * mask).sum(dim=1) + deformation = deformation.permute(0, 2, 3, 1) + else: + motion = self.dense_motion_network.forward_partswap( + source_image=source_image, seg_target=seg_target, seg_source=seg_source + ) + deformation = motion["deformation"] + + # Deform source encoding according to the motion + enc_source = self.deform_input(enc_source, deformation) + + if self.estimate_visibility: + if self.dense_motion_network is None: + visibility = seg_source["segmentation"][:, 1:].sum( + dim=1, keepdim=True + ) * ( + 1 + - seg_target["segmentation"][:, 1:] + .sum(dim=1, keepdim=True) + .detach() + ) + visibility = 1 - visibility + else: + visibility = motion["visibility"] + + if ( + enc_source.shape[2] != visibility.shape[2] + or enc_source.shape[3] != visibility.shape[3] + ): + visibility = F.interpolate( + visibility, size=enc_source.shape[2:], mode="bilinear" + ) + enc_source = enc_source * visibility + + blend_mask = self.blend_downsample(blend_mask) + # If source segmentation is provided use it should be deformed before blending + if use_source_segmentation: + blend_mask = self.deform_input(blend_mask, deformation) + + out = enc_target * (1 - blend_mask) + enc_source * blend_mask + + out = self.bottleneck(out) + + for i in range(len(self.up_blocks)): + out = self.up_blocks[i](out) + + out = self.final(out) + out = F.sigmoid(out) + + output_dict["prediction"] = out + + return output_dict diff --git a/requirements.txt b/requirements.txt index 054a554..fe18382 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,3 +2,4 @@ numpy torch scipy yaml +matplotlib \ No newline at end of file diff --git a/run.py b/run.py index b23fa25..16e10d9 100644 --- a/run.py +++ b/run.py @@ -9,13 +9,28 @@ from pathlib import Path import torch -from .constants import default_model_name, supported_model_names +from .constants import ( + config_folder, + default_model_name, + default_partswap_model_name, + face_parser_checkpoint_name, + partswap_fomm_model_names, + partswap_model_config_dict, + partswap_model_length_dict, + partswap_model_names, + supported_model_names, +) +from .face_parsing.face_parsing_loader import load_face_parser_model +from .fomm_inference import inference, inference_best_frame +from .fomm_loader import load_checkpoint +from .partswap_inference import partswap_inference +from .partswap_loader import load_partswap_checkpoint +from .seg_viz import visualize_frame from .utils import ( + build_seg_arguments, get_config_path, get_model_file_name, - inference, - inference_best_frame, - load_checkpoint, + get_partswap_model_file_name, out_video, reshape_image, ) @@ -67,8 +82,8 @@ class FOMM_Runner: def todo( self, - source_image: torch.Tensor, - driving_video_input: torch.Tensor, + source_image, + driving_video_input, model_name: str, frame_rate: float, relative_movement: bool, @@ -119,10 +134,213 @@ class FOMM_Runner: ) -NODE_CLASS_MAPPINGS = { - "FOMM_Runner": FOMM_Runner, -} +class FOMM_Seg5Chooser: + @classmethod + def INPUT_TYPES(s): + return build_seg_arguments("vox-5segments") -NODE_DISPLAY_NAME_MAPPINGS = { - "FOMM_Runner": "FOMM Runner", -} + RETURN_TYPES = ("STRING",) + + RETURN_NAMES = ("chosen_seg_indices",) + + FUNCTION = "todo" + CATEGORY = "FirstOrderMM" + + def todo(self, **args): + print(args) + assert len(args) == 6, f"Expected 6 arguments, got {len(args)}" + segments = list(args.values()) + seg_list = [i for i, seg in enumerate(segments) if seg] + seg_list = serialize_integers(seg_list) + return (seg_list,) + + +class FOMM_Seg10Chooser: + @classmethod + def INPUT_TYPES(s): + return build_seg_arguments("vox-10segments") + + RETURN_TYPES = ("STRING",) + + RETURN_NAMES = ("chosen_seg_indices",) + + FUNCTION = "todo" + CATEGORY = "FirstOrderMM" + + def todo(self, **args): + print(args) + assert len(args) == 11, f"Expected 11 arguments, got {len(args)}" + segments = list(args.values()) + seg_list = [i for i, seg in enumerate(segments) if seg] + seg_list = serialize_integers(seg_list) + return (seg_list,) + + +class FOMM_Seg15Chooser: + @classmethod + def INPUT_TYPES(s): + return build_seg_arguments("vox-15segments") + + RETURN_TYPES = ("STRING",) + + RETURN_NAMES = ("chosen_seg_indices",) + + FUNCTION = "todo" + CATEGORY = "FirstOrderMM" + + def todo(self, **args): + print(args) + assert len(args) == 16, f"Expected 16 arguments, got {len(args)}" + segments = list(args.values()) + seg_list = [i for i, seg in enumerate(segments) if seg] + seg_list = serialize_integers(seg_list) + return (seg_list,) + + +class FOMM_Partswap: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "source_image": ("IMAGE",), + "driving_video_input": ("IMAGE",), + "model_name": ( + partswap_model_names, + {"default": default_partswap_model_name}, + ), + "frame_rate": ("FLOAT", {"default": 30.0}), + "blend_scale": ( + "FLOAT", + {"default": 1.0, "min": 0.6, "max": 1.0, "step": 0.01}, + ), + "use_source_seg": ( + "BOOLEAN", + {"default": True}, + ), + "hard_edges": ( + "BOOLEAN", + {"default": False}, + ), + "use_face_parser": ( + "BOOLEAN", + {"default": False}, + ), + "chosen_seg_indices": ("STRING", {}), + "viz_alpha": ( + "FLOAT", + {"default": 0.6, "min": 0.0, "max": 1.0, "step": 0.1}, + ), + }, + "optional": {"audio": ("VHS_AUDIO",)}, + } + + RETURN_TYPES = ( + "IMAGE", + "IMAGE", + "IMAGE", + "VHS_AUDIO", + "FLOAT", + ) + RETURN_NAMES = ( + "seg_src_viz", + "seg_tgt_viz", + "images", + "audio", + "frame_rate", + ) + FUNCTION = "todo" + CATEGORY = "FirstOrderMM" + + def todo( + self, + source_image: torch.Tensor, + driving_video_input: torch.Tensor, + model_name: str, + frame_rate: float, + blend_scale: float, + use_source_seg: bool, + hard_edges: bool, + use_face_parser: bool, + chosen_seg_indices: str, + viz_alpha: float, + audio=None, + ): + print(f"{type(source_image)=}") + print(f"{type(driving_video_input)=}") + print(f"{source_image.shape=}") + print(f"{driving_video_input.shape=}") + print(f"{type(audio)=}") + print(base_dir) + + config_path = ( + f"{base_dir}/{config_folder}/{partswap_model_config_dict[model_name]}" + ) + checkpoint_path = f"{base_dir}/{get_partswap_model_file_name(model_name)}" + + source_image = reshape_image(source_image, (256, 256)) + driving_video_org = reshape_image(driving_video_input, (256, 256)) + driving_video = driving_video_org.unsqueeze(0).permute(0, 2, 1, 3, 4) + + use_fomm = False + if model_name in partswap_fomm_model_names: + use_fomm = True + + reconstruction_module, segmentation_module = load_partswap_checkpoint( + config_path, checkpoint_path, blend_scale, use_fomm=use_fomm + ) + + face_parser_model = None + # 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 + ) + 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) + swap_indices = [ + x for x in swap_indices if x < partswap_model_length_dict[model_name] + ] + print(f"Swapping {swap_indices}") + + params = { + "swap_indices": swap_indices, + "source_image": source_image, + "target_video": driving_video, + "reconstruction_module": reconstruction_module, + "segmentation_module": segmentation_module, + "use_source_seg": use_source_seg, + "face_parser_model": face_parser_model, + "hard_edges": hard_edges, + } + + seg_source, seg_targets, predictions = partswap_inference(**params) + + seg_src_viz = visualize_frame(source_image, seg_source, alpha=viz_alpha) + seg_tgt_viz = visualize_frame( + driving_video_org[0], seg_targets[0], alpha=viz_alpha + ) + + output_images = out_video(predictions) + + del face_parser_model + + return ( + seg_src_viz, + seg_tgt_viz, + output_images, + audio, + frame_rate, + ) + + +def serialize_integers(int_list): + return "_".join(map(str, int_list)) + + +def deserialize_integers(int_string): + return list(map(int, int_string.split("_"))) 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/seg_viz.py b/seg_viz.py new file mode 100644 index 0000000..953d998 --- /dev/null +++ b/seg_viz.py @@ -0,0 +1,88 @@ +import matplotlib.patches as mpatches +import matplotlib.pyplot as plt +from matplotlib.backends.backend_agg import FigureCanvasAgg +from matplotlib.figure import Figure +import numpy as np +import torch +import torch.nn.functional as F + + +def preprocess_image_and_mask(image, seg_source): + image = image.squeeze(0).cpu() # Remove batch dimension + mask = seg_source["segmentation"] + + if mask.shape[2:] != image.shape[1:]: + mask = F.interpolate(mask, size=image.shape[1:], mode="bilinear") + + return image, mask + + +def apply_hard_edges(mask): + return (torch.max(mask, dim=1, keepdim=True)[0] == mask).float() + + +def create_color_mask(mask, colormap="gist_rainbow"): + colormap = plt.get_cmap(colormap) + num_segments = mask.shape[1] + mask = mask.squeeze(0).permute(1, 2, 0).cpu().numpy() + color_mask = np.zeros((256, 256, 3), dtype=np.float32) + patches = [] + + for i in range(num_segments): + color = ( + np.array(colormap((i - 1) / (num_segments - 1)))[:3] + if i != 0 + else np.array((0, 0, 0)) + ) + patches.append(mpatches.Patch(color=color, label=str(i))) + color_mask += mask[..., i : (i + 1)] * color.reshape(1, 1, 3) + + return color_mask, patches + + +def overlay_mask(image, color_mask, alpha=0.6): + image_np = image.permute(1, 2, 0).numpy() + if image_np.max() > 1: + image_np = image_np / 255.0 + + overlaid_image = image_np * (1 - alpha) + color_mask * alpha + return np.clip(overlaid_image, 0, 1) + + +def create_figure(overlaid_image, patches): + fig = Figure(figsize=(10, 10)) + canvas = FigureCanvasAgg(fig) + ax = fig.add_subplot(111) + + ax.imshow(overlaid_image) + ax.axis("off") + ax.legend( + handles=patches, loc="upper right", bbox_to_anchor=(1, 1), fontsize="small" + ) + + fig.tight_layout(pad=0) + return fig, canvas + + +def render_figure_to_tensor(canvas): + canvas.draw() + result_image = np.frombuffer(canvas.tostring_rgb(), dtype="uint8") + result_image = result_image.reshape(canvas.get_width_height()[::-1] + (3,)) + result_tensor = torch.from_numpy(result_image).unsqueeze(0).float() / 255.0 + return result_tensor + + +def visualize_frame( + image, seg_source, hard_edges=True, colormap="gist_rainbow", alpha=0.6 +): + image, mask = preprocess_image_and_mask(image, seg_source) + + if hard_edges: + mask = apply_hard_edges(mask) + + color_mask, patches = create_color_mask(mask, colormap) + overlaid_image = overlay_mask(image, color_mask, alpha) + fig, canvas = create_figure(overlaid_image, patches) + result_tensor = render_figure_to_tensor(canvas) + + return result_tensor diff --git a/utils.py b/utils.py index acc77d9..213abb2 100644 --- a/utils.py +++ b/utils.py @@ -1,35 +1,28 @@ -try: - import face_alignment - - face_alignment_available = True -except ImportError: - face_alignment_available = False - print("Warning: face_alignment module not found. fine_best_frame will not work. ") - import numpy as np import torch import torch.nn.functional as F -import yaml -from comfy.utils import ProgressBar -from scipy.spatial import ConvexHull -from tqdm import tqdm from .constants import ( checkpoint_folder, checkpoint_folder_link, checkpoint_suffix, config_folder, + config_suffix, non_cpk_model_names, + partswap_model_names, + seg_model_dict, supported_model_names, ) -from .modules.generator import OcclusionAwareGenerator -from .modules.keypoint_detector import KPDetector + + +def check_supported(entity_name: str, entity_support_list: list[str]): + assert ( + entity_name in entity_support_list + ), f"Unknown {entity_name}, expected {entity_support_list}" def get_model_file_name(model_name: str) -> str: - assert ( - model_name in supported_model_names - ), f"Unknown model. Got {model_name}, expected {supported_model_names}" + check_supported(model_name, supported_model_names) if model_name in non_cpk_model_names: model_file_name = f"{checkpoint_folder}/{model_name}{checkpoint_suffix}" @@ -46,51 +39,23 @@ def get_model_link(model_name: str) -> str: def get_config_path(model_name: str) -> str: - config_path = f"{config_folder}/{model_name}-256.yaml" + config_path = f"{config_folder}/{model_name}-{config_suffix}" return config_path +def get_partswap_model_file_name(model_name: str): + return f"{checkpoint_folder}/{model_name}{checkpoint_suffix}" + + def out_video(predictions: list[np.ndarray]): out_tensor_list = [] for i in predictions: - # out_img = i.astype(np.float32) / 255.0 out_img = torch.from_numpy(i) out_tensor_list.append(out_img) images = torch.stack(out_tensor_list, dim=0) return images -def load_checkpoint( - config_path: str, checkpoint_path: str, cpu=False -) -> tuple[OcclusionAwareGenerator, KPDetector]: - with open(config_path) as f: - config = yaml.full_load(f) - - generator = OcclusionAwareGenerator( - **config["model_params"]["generator_params"], - **config["model_params"]["common_params"], - ) - kp_detector = KPDetector( - **config["model_params"]["kp_detector_params"], - **config["model_params"]["common_params"], - ) - - if not cpu: - generator.cuda() - kp_detector.cuda() - checkpoint = torch.load(checkpoint_path) - else: - checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu")) - - generator.load_state_dict(checkpoint["generator"]) - kp_detector.load_state_dict(checkpoint["kp_detector"]) - - generator.eval() - kp_detector.eval() - - return generator, kp_detector - - def reshape_image(image, target_size=(256, 256)): # Check if image is 4D (batch, height, width, channels) if image.dim() != 4: @@ -107,177 +72,21 @@ def reshape_image(image, target_size=(256, 256)): return resized_image -def normalize_kp_best_frame(kp): - kp = kp - kp.mean(axis=0, keepdims=True) - area = ConvexHull(kp[:, :2]).volume - area = np.sqrt(area) - kp[:, :2] = kp[:, :2] / area - return kp - - def tensor_to_numpy(tensor): # Convert [1, 3, H, W] tensor to [H, W, 3] numpy array return tensor.squeeze(0).permute(1, 2, 0).cpu().numpy() -def find_best_frame(source_image, driving_video, cpu=False) -> int: - if not face_alignment_available: - print("Face Alignment module were not found. Best frame set to 0.") - return 0 - print("Initializing Face Alignment...") - fa = face_alignment.FaceAlignment( - face_alignment.LandmarksType.TWO_D, - flip_input=True, - device="cpu" if cpu else "cuda", - ) - print("Face alignment model loaded.") +def build_seg_arguments(model_name: str): + check_supported(model_name, partswap_model_names) + segment_names = seg_model_dict[model_name] + required = dict() + for segment_name in segment_names: + default = True + if segment_name.endswith("background"): + default = False + required[segment_name] = ("BOOLEAN", {"default": default}) - source_np = tensor_to_numpy(source_image) - kp_source = fa.get_landmarks(255 * source_np)[0] - kp_source = normalize_kp_best_frame(kp_source) + arguments = {"required": required} - norm = float("inf") - frame_num = 0 - - num_frames = driving_video.shape[2] - - for i in tqdm(range(num_frames)): - # Extract single frame from driving video - frame = driving_video[0, :, i] # Shape: [3, 256, 256] - frame_np = frame.permute(1, 2, 0).cpu().numpy() - - kp_driving = fa.get_landmarks(255 * frame_np)[0] - kp_driving = normalize_kp_best_frame(kp_driving) - new_norm = (np.abs(kp_source - kp_driving) ** 2).sum() - if new_norm < norm: - norm = new_norm - frame_num = i - - print(f"Best frame: {frame_num}") - - return frame_num - - -def normalize_kp( - kp_source, - kp_driving, - kp_driving_initial, - adapt_movement_scale: bool, - use_relative_movement: bool, - use_relative_jacobian: bool, -): - if adapt_movement_scale: - source_area = ConvexHull(kp_source["value"][0].data.cpu().numpy()).volume - driving_area = ConvexHull( - kp_driving_initial["value"][0].data.cpu().numpy() - ).volume - adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area) - else: - adapt_movement_scale = 1 - # print(f"{adapt_movement_scale=}") - kp_new = {k: v for k, v in kp_driving.items()} - - if use_relative_movement: - kp_value_diff = kp_driving["value"] - kp_driving_initial["value"] - kp_value_diff *= adapt_movement_scale - kp_new["value"] = kp_value_diff + kp_source["value"] - - if use_relative_jacobian: - jacobian_diff = torch.matmul( - kp_driving["jacobian"], torch.inverse(kp_driving_initial["jacobian"]) - ) - kp_new["jacobian"] = torch.matmul(jacobian_diff, kp_source["jacobian"]) - - return kp_new - - -def inference( - source_image, - driving_video: list, - generator: OcclusionAwareGenerator, - kp_detector: KPDetector, - relative_movement: bool, - relative_jacobian: bool, - adapt_movement_scale: bool, - cpu=False, -) -> list: - with torch.no_grad(): - source = source_image - driving = driving_video - predictions = [] - - if not cpu: - source = source.cuda() - driving = driving.cuda() - print("Using GPU") - - kp_source = kp_detector(source) - kp_driving_initial = kp_detector(driving[:, :, 0]) - kp_norm = None - - num_frames = driving.shape[2] - pbar = ProgressBar(num_frames) - - for frame_idx in tqdm(range(num_frames)): - driving_frame = driving[:, :, frame_idx] - kp_driving = kp_detector(driving_frame) - kp_norm = normalize_kp( - kp_source=kp_source, - kp_driving=kp_driving, - kp_driving_initial=kp_driving_initial, - use_relative_movement=relative_movement, - use_relative_jacobian=relative_jacobian, - adapt_movement_scale=adapt_movement_scale, - ) - out = generator(source, kp_source=kp_source, kp_driving=kp_norm) - out_np = np.transpose(out["prediction"].data.cpu().numpy(), [0, 2, 3, 1])[0] - predictions.append(out_np) - pbar.update_absolute(frame_idx, num_frames) - - return predictions - - -def inference_best_frame( - source_image, - driving_video, - generator: OcclusionAwareGenerator, - kp_detector: KPDetector, - relative_movement: bool, - relative_jacobian: bool, - adapt_movement_scale: bool, - cpu=False, -) -> list: - best_frame_idx = find_best_frame(source_image, driving_video, cpu) - - first_half = driving_video[ - :, :, : best_frame_idx + 1 - ] # Include the best frame in the first half - first_half = torch.flip(first_half, dims=[2]) # Reverse the first half - second_half = driving_video[:, :, best_frame_idx + 1 :] - - predictions_first = inference( - source_image, - first_half, - generator, - kp_detector, - relative_movement, - relative_jacobian, - adapt_movement_scale, - cpu, - ) - predictions_second = inference( - source_image, - second_half, - generator, - kp_detector, - relative_movement, - relative_jacobian, - adapt_movement_scale, - cpu, - ) - - predictions = [] - predictions_first = predictions_first[::-1] # Reverse the first half back - predictions.extend(predictions_first) - predictions.extend(predictions_second) - return predictions + return arguments 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