diff --git a/.gitignore b/.gitignore index 82f9275..b233313 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,9 @@ +# Stuff + +checkpoints/* +!checkpoints/place_checkpoints_here.txt + + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/FOMM.json b/FOMM.json new file mode 100644 index 0000000..967d9a3 --- /dev/null +++ b/FOMM.json @@ -0,0 +1 @@ +{"last_node_id": 21, "last_link_id": 42, "nodes": [{"id": 13, "type": "VHS_VideoCombine", "pos": [1046, 347], "size": [315, 595], "flags": {}, "order": 4, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 39}, {"name": "audio", "type": "VHS_AUDIO", "link": 40}, {"name": "meta_batch", "type": "VHS_BatchManager", "link": null, "slot_index": 2}, {"name": "frame_rate", "type": "FLOAT", "link": 41, "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": "FOMM", "format": "video/h264-mp4", "pix_fmt": "yuv420p", "crf": 0, "save_metadata": true, "pingpong": false, "save_output": true, "videopreview": {"hidden": false, "paused": false, "params": {"filename": "FOMM_01042-audio.mp4", "subfolder": "", "type": "output", "format": "video/h264-mp4"}}}}, {"id": 1, "type": "LoadImage", "pos": [111, 131], "size": {"0": 315, "1": 314}, "flags": {}, "order": 0, "mode": 0, "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [36], "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": 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": [37], "shape": 3, "slot_index": 0}, {"name": "frame_count", "type": "INT", "links": null, "shape": 3}, {"name": "audio", "type": "VHS_AUDIO", "links": [38], "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": 0, "skip_first_frames": 0, "select_every_nth": 1, "choose video to upload": "image", "videopreview": {"hidden": false, "paused": false, "params": {"frame_load_cap": 0, "skip_first_frames": 0, "force_rate": 0, "filename": "damedane.mp4", "type": "input", "format": "video/mp4", "select_every_nth": 1}}}}, {"id": 16, "type": "VHS_VideoInfoLoaded", "pos": [366, 875], "size": {"0": 235.1999969482422, "1": 106}, "flags": {}, "order": 2, "mode": 0, "inputs": [{"name": "video_info", "type": "VHS_VIDEOINFO", "link": 12}], "outputs": [{"name": "fps\ud83d\udfe6", "type": "FLOAT", "links": [42], "shape": 3, "slot_index": 0}, {"name": "frame_count\ud83d\udfe6", "type": "INT", "links": null, "shape": 3}, {"name": "duration\ud83d\udfe6", "type": "FLOAT", "links": null, "shape": 3}, {"name": "width\ud83d\udfe6", "type": "INT", "links": null, "shape": 3}, {"name": "height\ud83d\udfe6", "type": "INT", "links": null, "shape": 3}], "properties": {"Node name for S&R": "VHS_VideoInfoLoaded"}, "widgets_values": {}}, {"id": 21, "type": "FOMM_Runner", "pos": [638, 354], "size": {"0": 380.4000244140625, "1": 218}, "flags": {}, "order": 3, "mode": 0, "inputs": [{"name": "source_image", "type": "IMAGE", "link": 36}, {"name": "driving_video_input", "type": "IMAGE", "link": 37}, {"name": "audio", "type": "VHS_AUDIO", "link": 38}, {"name": "frame_rate", "type": "FLOAT", "link": 42, "widget": {"name": "frame_rate"}}], "outputs": [{"name": "images", "type": "IMAGE", "links": [39], "shape": 3, "slot_index": 0}, {"name": "audio", "type": "VHS_AUDIO", "links": [40], "shape": 3, "slot_index": 1}, {"name": "frame_rate", "type": "FLOAT", "links": [41], "shape": 3, "slot_index": 2}], "properties": {"Node name for S&R": "FOMM_Runner"}, "widgets_values": ["vox-adv", 30, true, true, true, true]}], "links": [[12, 11, 3, 16, 0, "VHS_VIDEOINFO"], [36, 1, 0, 21, 0, "IMAGE"], [37, 11, 0, 21, 1, "IMAGE"], [38, 11, 2, 21, 2, "VHS_AUDIO"], [39, 21, 0, 13, 0, "IMAGE"], [40, 21, 1, 13, 1, "VHS_AUDIO"], [41, 21, 2, 13, 3, "FLOAT"], [42, 16, 0, 21, 3, "FLOAT"]], "groups": [], "config": {}, "extra": {"ds": {"scale": 0.9090909090909091, "offset": [196.27130578512396, -60.60085123966945]}}, "version": 0.4} \ No newline at end of file diff --git a/README.md b/README.md index cf3f31b..7e1c2d3 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,47 @@ # ComfyUI-FirstOrderMM -A ComfyUI-native nodes to run First Order Motion Model for Image Animation +A ComfyUI-native node to run First Order Motion Model for Image Animation + +https://github.com/AliaksandrSiarohin/first-order-model + +## Workflow: + +[FOMM.json](FOMM.json) + +![FOMM Workflow](workflow.png) + +## Arguments +* `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. + +## Installation + +1. Clone the repo to `ComfyUI/custom_nodes/` +``` +git clone https://github.com/FuouM/ComfyUI-FirstOrderMM.git +``` + +2. Install required dependencies +``` +pip install -r requirements.txt +``` + +**Optional**: Install [face-alignment](https://github.com/1adrianb/face-alignment) to use the `find_best_frame` feature: + +``` +pip install face-alignment +``` + +## Model + +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) + +Place them in the `checkpoints` folder. It should look like this: +``` +place_checkpoints_here.txt +vox-adv-cpk.pth.tar +vox-cpk.pth.tar +``` \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..6c98dfa --- /dev/null +++ b/__init__.py @@ -0,0 +1,11 @@ +from .run import FOMM_Runner + +NODE_CLASS_MAPPINGS = { + "FOMM_Runner": FOMM_Runner, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "FOMM_Runner": "FOMM Runner", +} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/checkpoints/place_checkpoints_here.txt b/checkpoints/place_checkpoints_here.txt new file mode 100644 index 0000000..e69de29 diff --git a/config/bair-256.yaml b/config/bair-256.yaml new file mode 100644 index 0000000..daffb16 --- /dev/null +++ b/config/bair-256.yaml @@ -0,0 +1,82 @@ +dataset_params: + root_dir: data/bair + frame_shape: [256, 256, 3] + id_sampling: False + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + jitter_param: + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.1 + + +model_params: + common_params: + num_kp: 10 + num_channels: 3 + estimate_jacobian: True + kp_detector_params: + temperature: 0.1 + block_expansion: 32 + max_features: 1024 + scale_factor: 0.25 + num_blocks: 5 + generator_params: + block_expansion: 64 + max_features: 512 + num_down_blocks: 2 + num_bottleneck_blocks: 6 + estimate_occlusion_map: True + dense_motion_params: + block_expansion: 64 + max_features: 1024 + num_blocks: 5 + scale_factor: 0.25 + discriminator_params: + scales: [1] + block_expansion: 32 + max_features: 512 + num_blocks: 4 + sn: True + +train_params: + num_epochs: 20 + num_repeats: 1 + epoch_milestones: [12, 18] + lr_generator: 2.0e-4 + lr_discriminator: 2.0e-4 + lr_kp_detector: 2.0e-4 + batch_size: 36 + scales: [1, 0.5, 0.25, 0.125] + checkpoint_freq: 10 + transform_params: + sigma_affine: 0.05 + sigma_tps: 0.005 + points_tps: 5 + loss_weights: + generator_gan: 1 + discriminator_gan: 1 + feature_matching: [10, 10, 10, 10] + perceptual: [10, 10, 10, 10, 10] + equivariance_value: 10 + equivariance_jacobian: 10 + +reconstruction_params: + num_videos: 1000 + format: '.mp4' + +animate_params: + num_pairs: 50 + format: '.mp4' + normalization_params: + adapt_movement_scale: False + use_relative_movement: True + use_relative_jacobian: True + +visualizer_params: + kp_size: 5 + draw_border: True + colormap: 'gist_rainbow' diff --git a/config/fashion-256.yaml b/config/fashion-256.yaml new file mode 100644 index 0000000..760034d --- /dev/null +++ b/config/fashion-256.yaml @@ -0,0 +1,77 @@ +dataset_params: + root_dir: data/fashion-png + frame_shape: [256, 256, 3] + id_sampling: False + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + jitter_param: + hue: 0.1 + +model_params: + common_params: + num_kp: 10 + num_channels: 3 + estimate_jacobian: True + kp_detector_params: + temperature: 0.1 + block_expansion: 32 + max_features: 1024 + scale_factor: 0.25 + num_blocks: 5 + generator_params: + block_expansion: 64 + max_features: 512 + num_down_blocks: 2 + num_bottleneck_blocks: 6 + estimate_occlusion_map: True + dense_motion_params: + block_expansion: 64 + max_features: 1024 + num_blocks: 5 + scale_factor: 0.25 + discriminator_params: + scales: [1] + block_expansion: 32 + max_features: 512 + num_blocks: 4 + +train_params: + num_epochs: 100 + num_repeats: 50 + epoch_milestones: [60, 90] + lr_generator: 2.0e-4 + lr_discriminator: 2.0e-4 + lr_kp_detector: 2.0e-4 + batch_size: 27 + scales: [1, 0.5, 0.25, 0.125] + checkpoint_freq: 50 + transform_params: + sigma_affine: 0.05 + sigma_tps: 0.005 + points_tps: 5 + loss_weights: + generator_gan: 1 + discriminator_gan: 1 + feature_matching: [10, 10, 10, 10] + perceptual: [10, 10, 10, 10, 10] + equivariance_value: 10 + equivariance_jacobian: 10 + +reconstruction_params: + num_videos: 1000 + format: '.mp4' + +animate_params: + num_pairs: 50 + format: '.mp4' + normalization_params: + adapt_movement_scale: False + use_relative_movement: True + use_relative_jacobian: True + +visualizer_params: + kp_size: 5 + draw_border: True + colormap: 'gist_rainbow' diff --git a/config/mgif-256.yaml b/config/mgif-256.yaml new file mode 100644 index 0000000..598c60b --- /dev/null +++ b/config/mgif-256.yaml @@ -0,0 +1,84 @@ +dataset_params: + root_dir: data/moving-gif + frame_shape: [256, 256, 3] + id_sampling: False + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + crop_param: + size: [256, 256] + resize_param: + ratio: [0.9, 1.1] + jitter_param: + hue: 0.5 + +model_params: + common_params: + num_kp: 10 + num_channels: 3 + estimate_jacobian: True + kp_detector_params: + temperature: 0.1 + block_expansion: 32 + max_features: 1024 + scale_factor: 0.25 + num_blocks: 5 + single_jacobian_map: True + generator_params: + block_expansion: 64 + max_features: 512 + num_down_blocks: 2 + num_bottleneck_blocks: 6 + estimate_occlusion_map: True + dense_motion_params: + block_expansion: 64 + max_features: 1024 + num_blocks: 5 + scale_factor: 0.25 + discriminator_params: + scales: [1] + block_expansion: 32 + max_features: 512 + num_blocks: 4 + sn: True + +train_params: + num_epochs: 100 + num_repeats: 25 + epoch_milestones: [60, 90] + lr_generator: 2.0e-4 + lr_discriminator: 2.0e-4 + lr_kp_detector: 2.0e-4 + + batch_size: 36 + scales: [1, 0.5, 0.25, 0.125] + checkpoint_freq: 100 + transform_params: + sigma_affine: 0.05 + sigma_tps: 0.005 + points_tps: 5 + loss_weights: + generator_gan: 1 + discriminator_gan: 1 + feature_matching: [10, 10, 10, 10] + perceptual: [10, 10, 10, 10, 10] + equivariance_value: 10 + equivariance_jacobian: 10 + +reconstruction_params: + num_videos: 1000 + format: '.mp4' + +animate_params: + num_pairs: 50 + format: '.mp4' + normalization_params: + adapt_movement_scale: False + use_relative_movement: True + use_relative_jacobian: True + +visualizer_params: + kp_size: 5 + draw_border: True + colormap: 'gist_rainbow' diff --git a/config/nemo-256.yaml b/config/nemo-256.yaml new file mode 100644 index 0000000..1bd44b9 --- /dev/null +++ b/config/nemo-256.yaml @@ -0,0 +1,76 @@ +dataset_params: + root_dir: data/nemo-png + frame_shape: [256, 256, 3] + id_sampling: False + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + +model_params: + common_params: + num_kp: 10 + num_channels: 3 + estimate_jacobian: True + kp_detector_params: + temperature: 0.1 + block_expansion: 32 + max_features: 1024 + scale_factor: 0.25 + num_blocks: 5 + generator_params: + block_expansion: 64 + max_features: 512 + num_down_blocks: 2 + num_bottleneck_blocks: 6 + estimate_occlusion_map: True + dense_motion_params: + block_expansion: 64 + max_features: 1024 + num_blocks: 5 + scale_factor: 0.25 + discriminator_params: + scales: [1] + block_expansion: 32 + max_features: 512 + num_blocks: 4 + sn: True + +train_params: + num_epochs: 100 + num_repeats: 8 + epoch_milestones: [60, 90] + lr_generator: 2.0e-4 + lr_discriminator: 2.0e-4 + lr_kp_detector: 2.0e-4 + batch_size: 36 + scales: [1, 0.5, 0.25, 0.125] + checkpoint_freq: 50 + transform_params: + sigma_affine: 0.05 + sigma_tps: 0.005 + points_tps: 5 + loss_weights: + generator_gan: 1 + discriminator_gan: 1 + feature_matching: [10, 10, 10, 10] + perceptual: [10, 10, 10, 10, 10] + equivariance_value: 10 + equivariance_jacobian: 10 + +reconstruction_params: + num_videos: 1000 + format: '.mp4' + +animate_params: + num_pairs: 50 + format: '.mp4' + normalization_params: + adapt_movement_scale: False + use_relative_movement: True + use_relative_jacobian: True + +visualizer_params: + kp_size: 5 + draw_border: True + colormap: 'gist_rainbow' diff --git a/config/taichi-256.yaml b/config/taichi-256.yaml new file mode 100644 index 0000000..44b7839 --- /dev/null +++ b/config/taichi-256.yaml @@ -0,0 +1,157 @@ +# Dataset parameters +# Each dataset should contain 2 folders train and test +# Each video can be represented as: +# - an image of concatenated frames +# - '.mp4' or '.gif' +# - folder with all frames from a specific video +# In case of Taichi. Same (youtube) video can be splitted in many parts (chunks). Each part has a following +# format (id)#other#info.mp4. For example '12335#adsbf.mp4' has an id 12335. In case of TaiChi id stands for youtube +# video id. +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/taichi-png + # Image shape, needed for staked .png format. + frame_shape: [256, 256, 3] + # In case of 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 + # List with pairs for animation, None for random pairs + pairs_list: data/taichi256.csv + # 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 keypoint + num_kp: 10 + # Number of channels per image + num_channels: 3 + # Using first or zero order model + estimate_jacobian: True + kp_detector_params: + # Softmax temperature for keypoint 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 + # Keypioint 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 + generator_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 occlusion map or not + estimate_occlusion_map: True + + dense_motion_params: + # Number of features mutliplier + block_expansion: 64 + # Maximum allowed number of features + max_features: 1024 + # Number of block in Unet. Can be increased or decreased depending or resolution. + num_blocks: 5 + # Dense motion 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 + discriminator_params: + # Discriminator can be multiscale, if you want 2 discriminator on original + # resolution and half of the original, specify scales: [1, 0.5] + scales: [1] + # Number of features mutliplier + block_expansion: 32 + # Maximum allowed number of features + max_features: 512 + # Number of blocks. Can be increased or decreased depending or resolution. + num_blocks: 4 + +# Parameters of training +train_params: + # Number of training epochs + num_epochs: 100 + # For better i/o performance when number of videos is small number of epochs can be multiplied by this number. + # Thus effectivlly with num_repeats=100 each epoch is 100 times larger. + num_repeats: 150 + # Drop learning rate by 10 times after this epochs + epoch_milestones: [60, 90] + # Initial learing rate for all modules + lr_generator: 2.0e-4 + lr_discriminator: 2.0e-4 + lr_kp_detector: 2.0e-4 + batch_size: 30 + # 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: + # Weight for LSGAN loss in generator, 0 for no adversarial loss. + generator_gan: 0 + # Weight for LSGAN loss in discriminator + discriminator_gan: 1 + # Weights for feature matching loss, the number should be the same as number of blocks in discriminator. + feature_matching: [10, 10, 10, 10] + # Weights for perceptual loss. + perceptual: [10, 10, 10, 10, 10] + # Weights for value equivariance. + equivariance_value: 10 + # Weights for jacobian equivariance. + equivariance_jacobian: 10 + +# Parameters of reconstruction +reconstruction_params: + # Maximum number of videos for reconstruction + num_videos: 1000 + # Format for visualization, note that results will be also stored in staked .png. + format: '.mp4' + +# Parameters of animation +animate_params: + # Maximum number of pairs for animation, the pairs will be either taken from pairs_list or random. + num_pairs: 50 + # Format for visualization, note that results will be also stored in staked .png. + format: '.mp4' + # Normalization of diriving keypoints + normalization_params: + # Increase or decrease relative movement scale depending on the size of the object + adapt_movement_scale: False + # Apply only relative displacement of the keypoint + use_relative_movement: True + # Apply only relative change in jacobian + use_relative_jacobian: True + +# Visualization parameters +visualizer_params: + # Draw keypoints 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/taichi-adv-256.yaml b/config/taichi-adv-256.yaml new file mode 100644 index 0000000..cd1941e --- /dev/null +++ b/config/taichi-adv-256.yaml @@ -0,0 +1,150 @@ +# 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/taichi-png + # Image shape, needed for staked .png format. + frame_shape: [256, 256, 3] + # In case of 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 + # List with pairs for animation, None for random pairs + pairs_list: data/taichi256.csv + # 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 keypoint + num_kp: 10 + # Number of channels per image + num_channels: 3 + # Using first or zero order model + estimate_jacobian: True + kp_detector_params: + # Softmax temperature for keypoint 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 + # Keypioint 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 + generator_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 occlusion map or not + estimate_occlusion_map: True + + dense_motion_params: + # Number of features mutliplier + block_expansion: 64 + # Maximum allowed number of features + max_features: 1024 + # Number of block in Unet. Can be increased or decreased depending or resolution. + num_blocks: 5 + # Dense motion 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 + discriminator_params: + # Discriminator can be multiscale, if you want 2 discriminator on original + # resolution and half of the original, specify scales: [1, 0.5] + scales: [1] + # Number of features mutliplier + block_expansion: 32 + # Maximum allowed number of features + max_features: 512 + # Number of blocks. Can be increased or decreased depending or resolution. + num_blocks: 4 + use_kp: True + +# Parameters of training +train_params: + # Number of training epochs + num_epochs: 150 + # For better i/o performance when number of videos is small number of epochs can be multiplied by this number. + # Thus effectivlly with num_repeats=100 each epoch is 100 times larger. + num_repeats: 150 + # Drop learning rate by 10 times after this epochs + epoch_milestones: [] + # Initial learing rate for all modules + lr_generator: 2.0e-4 + lr_discriminator: 2.0e-4 + lr_kp_detector: 0 + batch_size: 27 + # 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: + # Weight for LSGAN loss in generator + generator_gan: 1 + # Weight for LSGAN loss in discriminator + discriminator_gan: 1 + # Weights for feature matching loss, the number should be the same as number of blocks in discriminator. + feature_matching: [10, 10, 10, 10] + # Weights for perceptual loss. + perceptual: [10, 10, 10, 10, 10] + # Weights for value equivariance. + equivariance_value: 10 + # Weights for jacobian equivariance. + equivariance_jacobian: 10 + +# Parameters of reconstruction +reconstruction_params: + # Maximum number of videos for reconstruction + num_videos: 1000 + # Format for visualization, note that results will be also stored in staked .png. + format: '.mp4' + +# Parameters of animation +animate_params: + # Maximum number of pairs for animation, the pairs will be either taken from pairs_list or random. + num_pairs: 50 + # Format for visualization, note that results will be also stored in staked .png. + format: '.mp4' + # Normalization of diriving keypoints + normalization_params: + # Increase or decrease relative movement scale depending on the size of the object + adapt_movement_scale: False + # Apply only relative displacement of the keypoint + use_relative_movement: True + # Apply only relative change in jacobian + use_relative_jacobian: True + +# Visualization parameters +visualizer_params: + # Draw keypoints 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.yaml b/config/vox-256.yaml new file mode 100644 index 0000000..abfe9a2 --- /dev/null +++ b/config/vox-256.yaml @@ -0,0 +1,83 @@ +dataset_params: + root_dir: data/vox-png + frame_shape: [256, 256, 3] + id_sampling: True + pairs_list: data/vox256.csv + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + jitter_param: + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.1 + + +model_params: + common_params: + num_kp: 10 + num_channels: 3 + estimate_jacobian: True + kp_detector_params: + temperature: 0.1 + block_expansion: 32 + max_features: 1024 + scale_factor: 0.25 + num_blocks: 5 + generator_params: + block_expansion: 64 + max_features: 512 + num_down_blocks: 2 + num_bottleneck_blocks: 6 + estimate_occlusion_map: True + dense_motion_params: + block_expansion: 64 + max_features: 1024 + num_blocks: 5 + scale_factor: 0.25 + discriminator_params: + scales: [1] + block_expansion: 32 + max_features: 512 + num_blocks: 4 + sn: True + +train_params: + num_epochs: 100 + num_repeats: 75 + epoch_milestones: [60, 90] + lr_generator: 2.0e-4 + lr_discriminator: 2.0e-4 + lr_kp_detector: 2.0e-4 + batch_size: 40 + scales: [1, 0.5, 0.25, 0.125] + checkpoint_freq: 50 + transform_params: + sigma_affine: 0.05 + sigma_tps: 0.005 + points_tps: 5 + loss_weights: + generator_gan: 0 + discriminator_gan: 1 + feature_matching: [10, 10, 10, 10] + perceptual: [10, 10, 10, 10, 10] + equivariance_value: 10 + equivariance_jacobian: 10 + +reconstruction_params: + num_videos: 1000 + format: '.mp4' + +animate_params: + num_pairs: 50 + format: '.mp4' + normalization_params: + adapt_movement_scale: False + use_relative_movement: True + use_relative_jacobian: True + +visualizer_params: + kp_size: 5 + draw_border: True + colormap: 'gist_rainbow' diff --git a/config/vox-adv-256.yaml b/config/vox-adv-256.yaml new file mode 100644 index 0000000..ed89890 --- /dev/null +++ b/config/vox-adv-256.yaml @@ -0,0 +1,84 @@ +dataset_params: + root_dir: data/vox-png + frame_shape: [256, 256, 3] + id_sampling: True + pairs_list: data/vox256.csv + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + jitter_param: + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.1 + + +model_params: + common_params: + num_kp: 10 + num_channels: 3 + estimate_jacobian: True + kp_detector_params: + temperature: 0.1 + block_expansion: 32 + max_features: 1024 + scale_factor: 0.25 + num_blocks: 5 + generator_params: + block_expansion: 64 + max_features: 512 + num_down_blocks: 2 + num_bottleneck_blocks: 6 + estimate_occlusion_map: True + dense_motion_params: + block_expansion: 64 + max_features: 1024 + num_blocks: 5 + scale_factor: 0.25 + discriminator_params: + scales: [1] + block_expansion: 32 + max_features: 512 + num_blocks: 4 + use_kp: True + + +train_params: + num_epochs: 150 + num_repeats: 75 + epoch_milestones: [] + lr_generator: 2.0e-4 + lr_discriminator: 2.0e-4 + lr_kp_detector: 2.0e-4 + batch_size: 36 + scales: [1, 0.5, 0.25, 0.125] + checkpoint_freq: 50 + transform_params: + sigma_affine: 0.05 + sigma_tps: 0.005 + points_tps: 5 + loss_weights: + generator_gan: 1 + discriminator_gan: 1 + feature_matching: [10, 10, 10, 10] + perceptual: [10, 10, 10, 10, 10] + equivariance_value: 10 + equivariance_jacobian: 10 + +reconstruction_params: + num_videos: 1000 + format: '.mp4' + +animate_params: + num_pairs: 50 + format: '.mp4' + normalization_params: + adapt_movement_scale: False + use_relative_movement: True + use_relative_jacobian: True + +visualizer_params: + kp_size: 5 + draw_border: True + colormap: 'gist_rainbow' diff --git a/constants.py b/constants.py new file mode 100644 index 0000000..d69b199 --- /dev/null +++ b/constants.py @@ -0,0 +1,23 @@ +supported_model_names = [ + "vox-adv", + "vox", + # "taichi", + # "taichi-adv", + # "nemo", + # "mgif", + # "fashion", + # "bair", +] + +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_folder = "config" +checkpoint_folder = "checkpoints" \ No newline at end of file diff --git a/modules/__init__.py b/modules/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/modules/dense_motion.py b/modules/dense_motion.py new file mode 100644 index 0000000..56d5ac5 --- /dev/null +++ b/modules/dense_motion.py @@ -0,0 +1,113 @@ +from torch import nn +import torch.nn.functional as F +import torch +from .util import Hourglass, AntiAliasInterpolation2d, make_coordinate_grid, kp2gaussian + + +class DenseMotionNetwork(nn.Module): + """ + Module that predicting a dense motion from sparse motion representation given by kp_source and kp_driving + """ + + def __init__(self, block_expansion, num_blocks, max_features, num_kp, num_channels, estimate_occlusion_map=False, + scale_factor=1, kp_variance=0.01): + super(DenseMotionNetwork, self).__init__() + self.hourglass = Hourglass(block_expansion=block_expansion, in_features=(num_kp + 1) * (num_channels + 1), + max_features=max_features, num_blocks=num_blocks) + + self.mask = nn.Conv2d(self.hourglass.out_filters, num_kp + 1, kernel_size=(7, 7), padding=(3, 3)) + + if estimate_occlusion_map: + self.occlusion = nn.Conv2d(self.hourglass.out_filters, 1, kernel_size=(7, 7), padding=(3, 3)) + else: + self.occlusion = None + + self.num_kp = num_kp + self.scale_factor = scale_factor + self.kp_variance = kp_variance + + if self.scale_factor != 1: + self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor) + + def create_heatmap_representations(self, source_image, kp_driving, kp_source): + """ + 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) + heatmap = gaussian_driving - gaussian_source + + #adding background feature + zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1]).type(heatmap.type()) + heatmap = torch.cat([zeros, heatmap], dim=1) + heatmap = heatmap.unsqueeze(2) + return heatmap + + def create_sparse_motions(self, source_image, kp_driving, kp_source): + """ + 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 = 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'])) + 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) + + #adding background feature + identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1) + sparse_motions = torch.cat([identity_grid, driving_to_source], dim=1) + return sparse_motions + + def create_deformed_source_image(self, source_image, sparse_motions): + """ + Eq 7. in the paper \hat{T}_{s<-d}(z) + """ + bs, _, h, w = source_image.shape + source_repeat = source_image.unsqueeze(1).unsqueeze(1).repeat(1, self.num_kp + 1, 1, 1, 1, 1) + source_repeat = source_repeat.view(bs * (self.num_kp + 1), -1, h, w) + sparse_motions = sparse_motions.view((bs * (self.num_kp + 1), h, w, -1)) + sparse_deformed = F.grid_sample(source_repeat, sparse_motions) + sparse_deformed = sparse_deformed.view((bs, self.num_kp + 1, -1, h, w)) + return sparse_deformed + + def forward(self, source_image, kp_driving, kp_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, kp_driving, kp_source) + sparse_motion = self.create_sparse_motions(source_image, kp_driving, kp_source) + 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 + + # Sec. 3.2 in the paper + if self.occlusion: + occlusion_map = torch.sigmoid(self.occlusion(prediction)) + out_dict['occlusion_map'] = occlusion_map + + return out_dict diff --git a/modules/discriminator.py b/modules/discriminator.py new file mode 100644 index 0000000..07c1692 --- /dev/null +++ b/modules/discriminator.py @@ -0,0 +1,95 @@ +from torch import nn +import torch.nn.functional as F +from .util import kp2gaussian +import torch + + +class DownBlock2d(nn.Module): + """ + Simple block for processing video (encoder). + """ + + def __init__(self, in_features, out_features, norm=False, kernel_size=4, pool=False, sn=False): + super(DownBlock2d, self).__init__() + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size) + + if sn: + self.conv = nn.utils.spectral_norm(self.conv) + + if norm: + self.norm = nn.InstanceNorm2d(out_features, affine=True) + else: + self.norm = None + self.pool = pool + + def forward(self, x): + out = x + out = self.conv(out) + if self.norm: + out = self.norm(out) + out = F.leaky_relu(out, 0.2) + if self.pool: + out = F.avg_pool2d(out, (2, 2)) + return out + + +class Discriminator(nn.Module): + """ + Discriminator similar to Pix2Pix + """ + + def __init__(self, num_channels=3, block_expansion=64, num_blocks=4, max_features=512, + sn=False, use_kp=False, num_kp=10, kp_variance=0.01, **kwargs): + super(Discriminator, self).__init__() + + down_blocks = [] + for i in range(num_blocks): + down_blocks.append( + DownBlock2d(num_channels + num_kp * use_kp if i == 0 else min(max_features, block_expansion * (2 ** i)), + min(max_features, block_expansion * (2 ** (i + 1))), + norm=(i != 0), kernel_size=4, pool=(i != num_blocks - 1), sn=sn)) + + self.down_blocks = nn.ModuleList(down_blocks) + self.conv = nn.Conv2d(self.down_blocks[-1].conv.out_channels, out_channels=1, kernel_size=1) + if sn: + self.conv = nn.utils.spectral_norm(self.conv) + self.use_kp = use_kp + self.kp_variance = kp_variance + + def forward(self, x, kp=None): + feature_maps = [] + out = x + if self.use_kp: + heatmap = kp2gaussian(kp, x.shape[2:], self.kp_variance) + out = torch.cat([out, heatmap], dim=1) + + for down_block in self.down_blocks: + feature_maps.append(down_block(out)) + out = feature_maps[-1] + prediction_map = self.conv(out) + + return feature_maps, prediction_map + + +class MultiScaleDiscriminator(nn.Module): + """ + Multi-scale (scale) discriminator + """ + + def __init__(self, scales=(), **kwargs): + super(MultiScaleDiscriminator, self).__init__() + self.scales = scales + discs = {} + for scale in scales: + discs[str(scale).replace('.', '-')] = Discriminator(**kwargs) + self.discs = nn.ModuleDict(discs) + + def forward(self, x, kp=None): + out_dict = {} + for scale, disc in self.discs.items(): + scale = str(scale).replace('-', '.') + key = 'prediction_' + scale + feature_maps, prediction_map = disc(x[key], kp) + out_dict['feature_maps_' + scale] = feature_maps + out_dict['prediction_map_' + scale] = prediction_map + return out_dict diff --git a/modules/generator.py b/modules/generator.py new file mode 100644 index 0000000..691ef5a --- /dev/null +++ b/modules/generator.py @@ -0,0 +1,97 @@ +import torch +from torch import nn +import torch.nn.functional as F +from .util import ResBlock2d, SameBlock2d, UpBlock2d, DownBlock2d +from .dense_motion import DenseMotionNetwork + + +class OcclusionAwareGenerator(nn.Module): + """ + Generator that given source image and and keypoints try to transform image according to movement trajectories + induced by keypoints. Generator follows Johnson architecture. + """ + + def __init__(self, num_channels, num_kp, block_expansion, max_features, num_down_blocks, + num_bottleneck_blocks, estimate_occlusion_map=False, dense_motion_params=None, estimate_jacobian=False): + super(OcclusionAwareGenerator, self).__init__() + + if dense_motion_params is not None: + self.dense_motion_network = DenseMotionNetwork(num_kp=num_kp, num_channels=num_channels, + estimate_occlusion_map=estimate_occlusion_map, + **dense_motion_params) + else: + self.dense_motion_network = None + + 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_occlusion_map = estimate_occlusion_map + self.num_channels = num_channels + + 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, kp_driving, kp_source): + # Encoding (downsampling) part + out = self.first(source_image) + for i in range(len(self.down_blocks)): + out = self.down_blocks[i](out) + + # Transforming feature representation according to deformation and occlusion + output_dict = {} + if self.dense_motion_network is not None: + dense_motion = self.dense_motion_network(source_image=source_image, kp_driving=kp_driving, + kp_source=kp_source) + output_dict['mask'] = dense_motion['mask'] + output_dict['sparse_deformed'] = dense_motion['sparse_deformed'] + + if 'occlusion_map' in dense_motion: + occlusion_map = dense_motion['occlusion_map'] + output_dict['occlusion_map'] = occlusion_map + else: + occlusion_map = None + deformation = dense_motion['deformation'] + out = self.deform_input(out, deformation) + + if occlusion_map is not None: + if out.shape[2] != occlusion_map.shape[2] or out.shape[3] != occlusion_map.shape[3]: + occlusion_map = F.interpolate(occlusion_map, size=out.shape[2:], mode='bilinear') + out = out * occlusion_map + + output_dict["deformed"] = self.deform_input(source_image, deformation) + + # Decoding part + 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/keypoint_detector.py b/modules/keypoint_detector.py new file mode 100644 index 0000000..b938c20 --- /dev/null +++ b/modules/keypoint_detector.py @@ -0,0 +1,75 @@ +from torch import nn +import torch +import torch.nn.functional as F +from .util import Hourglass, make_coordinate_grid, AntiAliasInterpolation2d + + +class KPDetector(nn.Module): + """ + Detecting a keypoints. Return keypoint position and jacobian near each keypoint. + """ + + def __init__(self, block_expansion, num_kp, num_channels, max_features, + num_blocks, temperature, estimate_jacobian=False, scale_factor=1, + single_jacobian_map=False, pad=0): + super(KPDetector, self).__init__() + + self.predictor = Hourglass(block_expansion, in_features=num_channels, + max_features=max_features, num_blocks=num_blocks) + + self.kp = nn.Conv2d(in_channels=self.predictor.out_filters, out_channels=num_kp, kernel_size=(7, 7), + padding=pad) + + if estimate_jacobian: + self.num_jacobian_maps = 1 if single_jacobian_map else num_kp + self.jacobian = nn.Conv2d(in_channels=self.predictor.out_filters, + out_channels=4 * self.num_jacobian_maps, kernel_size=(7, 7), padding=pad) + self.jacobian.weight.data.zero_() + self.jacobian.bias.data.copy_(torch.tensor([1, 0, 0, 1] * self.num_jacobian_maps, dtype=torch.float)) + else: + self.jacobian = None + + 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 from a heatmap + """ + shape = heatmap.shape + heatmap = heatmap.unsqueeze(-1) + grid = make_coordinate_grid(shape[2:], heatmap.type()).unsqueeze_(0).unsqueeze_(0) + value = (heatmap * grid).sum(dim=(2, 3)) + kp = {'value': value} + + return kp + + def forward(self, x): + if self.scale_factor != 1: + x = self.down(x) + + feature_map = self.predictor(x) + prediction = self.kp(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 = self.gaussian2kp(heatmap) + + if self.jacobian is not None: + jacobian_map = self.jacobian(feature_map) + jacobian_map = jacobian_map.reshape(final_shape[0], self.num_jacobian_maps, 4, final_shape[2], + final_shape[3]) + heatmap = heatmap.unsqueeze(2) + + jacobian = heatmap * jacobian_map + jacobian = jacobian.view(final_shape[0], final_shape[1], 4, -1) + jacobian = jacobian.sum(dim=-1) + jacobian = jacobian.view(jacobian.shape[0], jacobian.shape[1], 2, 2) + out['jacobian'] = jacobian + + return out diff --git a/modules/model.py b/modules/model.py new file mode 100644 index 0000000..29a3a44 --- /dev/null +++ b/modules/model.py @@ -0,0 +1,259 @@ +from torch import nn +import torch +import torch.nn.functional as F +from .util import AntiAliasInterpolation2d, make_coordinate_grid +from torchvision import models +import numpy as np +from torch.autograd import grad + + +class Vgg19(torch.nn.Module): + """ + Vgg19 network for perceptual loss. See Sec 3.3. + """ + def __init__(self, requires_grad=False): + super(Vgg19, self).__init__() + vgg_pretrained_features = models.vgg19(pretrained=True).features + self.slice1 = torch.nn.Sequential() + self.slice2 = torch.nn.Sequential() + self.slice3 = torch.nn.Sequential() + self.slice4 = torch.nn.Sequential() + self.slice5 = torch.nn.Sequential() + for x in range(2): + self.slice1.add_module(str(x), vgg_pretrained_features[x]) + for x in range(2, 7): + self.slice2.add_module(str(x), vgg_pretrained_features[x]) + for x in range(7, 12): + self.slice3.add_module(str(x), vgg_pretrained_features[x]) + for x in range(12, 21): + self.slice4.add_module(str(x), vgg_pretrained_features[x]) + for x in range(21, 30): + self.slice5.add_module(str(x), vgg_pretrained_features[x]) + + self.mean = torch.nn.Parameter(data=torch.Tensor(np.array([0.485, 0.456, 0.406]).reshape((1, 3, 1, 1))), + requires_grad=False) + self.std = torch.nn.Parameter(data=torch.Tensor(np.array([0.229, 0.224, 0.225]).reshape((1, 3, 1, 1))), + requires_grad=False) + + if not requires_grad: + for param in self.parameters(): + param.requires_grad = False + + def forward(self, X): + X = (X - self.mean) / self.std + h_relu1 = self.slice1(X) + h_relu2 = self.slice2(h_relu1) + h_relu3 = self.slice3(h_relu2) + h_relu4 = self.slice4(h_relu3) + h_relu5 = self.slice5(h_relu4) + out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5] + return out + + +class ImagePyramide(torch.nn.Module): + """ + Create image pyramide for computing pyramide perceptual loss. See Sec 3.3 + """ + def __init__(self, scales, num_channels): + super(ImagePyramide, self).__init__() + downs = {} + for scale in scales: + downs[str(scale).replace('.', '-')] = AntiAliasInterpolation2d(num_channels, scale) + self.downs = nn.ModuleDict(downs) + + def forward(self, x): + out_dict = {} + for scale, down_module in self.downs.items(): + out_dict['prediction_' + str(scale).replace('-', '.')] = down_module(x) + return out_dict + + +class Transform: + """ + Random tps transformation for equivariance constraints. See Sec 3.3 + """ + def __init__(self, bs, **kwargs): + noise = torch.normal(mean=0, std=kwargs['sigma_affine'] * torch.ones([bs, 2, 3])) + self.theta = noise + torch.eye(2, 3).view(1, 2, 3) + self.bs = bs + + if ('sigma_tps' in kwargs) and ('points_tps' in kwargs): + self.tps = True + self.control_points = make_coordinate_grid((kwargs['points_tps'], kwargs['points_tps']), type=noise.type()) + self.control_points = self.control_points.unsqueeze(0) + self.control_params = torch.normal(mean=0, + std=kwargs['sigma_tps'] * torch.ones([bs, 1, kwargs['points_tps'] ** 2])) + else: + self.tps = False + + def transform_frame(self, frame): + grid = make_coordinate_grid(frame.shape[2:], type=frame.type()).unsqueeze(0) + grid = grid.view(1, frame.shape[2] * frame.shape[3], 2) + grid = self.warp_coordinates(grid).view(self.bs, frame.shape[2], frame.shape[3], 2) + return F.grid_sample(frame, grid, padding_mode="reflection") + + def warp_coordinates(self, coordinates): + theta = self.theta.type(coordinates.type()) + theta = theta.unsqueeze(1) + transformed = torch.matmul(theta[:, :, :, :2], coordinates.unsqueeze(-1)) + theta[:, :, :, 2:] + transformed = transformed.squeeze(-1) + + if self.tps: + control_points = self.control_points.type(coordinates.type()) + control_params = self.control_params.type(coordinates.type()) + distances = coordinates.view(coordinates.shape[0], -1, 1, 2) - control_points.view(1, 1, -1, 2) + distances = torch.abs(distances).sum(-1) + + result = distances ** 2 + result = result * torch.log(distances + 1e-6) + result = result * control_params + result = result.sum(dim=2).view(self.bs, coordinates.shape[1], 1) + transformed = transformed + result + + return transformed + + def jacobian(self, coordinates): + new_coordinates = self.warp_coordinates(coordinates) + grad_x = grad(new_coordinates[..., 0].sum(), coordinates, create_graph=True) + grad_y = grad(new_coordinates[..., 1].sum(), coordinates, create_graph=True) + jacobian = torch.cat([grad_x[0].unsqueeze(-2), grad_y[0].unsqueeze(-2)], dim=-2) + return jacobian + + +def detach_kp(kp): + return {key: value.detach() for key, value in kp.items()} + + +class GeneratorFullModel(torch.nn.Module): + """ + Merge all generator related updates into single model for better multi-gpu usage + """ + + def __init__(self, kp_extractor, generator, discriminator, train_params): + super(GeneratorFullModel, self).__init__() + self.kp_extractor = kp_extractor + self.generator = generator + self.discriminator = discriminator + self.train_params = train_params + self.scales = train_params['scales'] + self.disc_scales = self.discriminator.scales + self.pyramid = ImagePyramide(self.scales, generator.num_channels) + if torch.cuda.is_available(): + self.pyramid = self.pyramid.cuda() + + self.loss_weights = train_params['loss_weights'] + + if sum(self.loss_weights['perceptual']) != 0: + self.vgg = Vgg19() + if torch.cuda.is_available(): + self.vgg = self.vgg.cuda() + + def forward(self, x): + kp_source = self.kp_extractor(x['source']) + kp_driving = self.kp_extractor(x['driving']) + + generated = self.generator(x['source'], kp_source=kp_source, kp_driving=kp_driving) + generated.update({'kp_source': kp_source, 'kp_driving': kp_driving}) + + loss_values = {} + + pyramide_real = self.pyramid(x['driving']) + pyramide_generated = self.pyramid(generated['prediction']) + + if sum(self.loss_weights['perceptual']) != 0: + value_total = 0 + for scale in self.scales: + x_vgg = self.vgg(pyramide_generated['prediction_' + str(scale)]) + y_vgg = self.vgg(pyramide_real['prediction_' + str(scale)]) + + for i, weight in enumerate(self.loss_weights['perceptual']): + value = torch.abs(x_vgg[i] - y_vgg[i].detach()).mean() + value_total += self.loss_weights['perceptual'][i] * value + loss_values['perceptual'] = value_total + + if self.loss_weights['generator_gan'] != 0: + discriminator_maps_generated = self.discriminator(pyramide_generated, kp=detach_kp(kp_driving)) + discriminator_maps_real = self.discriminator(pyramide_real, kp=detach_kp(kp_driving)) + value_total = 0 + for scale in self.disc_scales: + key = 'prediction_map_%s' % scale + value = ((1 - discriminator_maps_generated[key]) ** 2).mean() + value_total += self.loss_weights['generator_gan'] * value + loss_values['gen_gan'] = value_total + + if sum(self.loss_weights['feature_matching']) != 0: + value_total = 0 + for scale in self.disc_scales: + key = 'feature_maps_%s' % scale + for i, (a, b) in enumerate(zip(discriminator_maps_real[key], discriminator_maps_generated[key])): + if self.loss_weights['feature_matching'][i] == 0: + continue + value = torch.abs(a - b).mean() + value_total += self.loss_weights['feature_matching'][i] * value + loss_values['feature_matching'] = value_total + + if (self.loss_weights['equivariance_value'] + self.loss_weights['equivariance_jacobian']) != 0: + transform = Transform(x['driving'].shape[0], **self.train_params['transform_params']) + transformed_frame = transform.transform_frame(x['driving']) + transformed_kp = self.kp_extractor(transformed_frame) + + generated['transformed_frame'] = transformed_frame + generated['transformed_kp'] = transformed_kp + + ## Value loss part + if self.loss_weights['equivariance_value'] != 0: + value = torch.abs(kp_driving['value'] - transform.warp_coordinates(transformed_kp['value'])).mean() + loss_values['equivariance_value'] = self.loss_weights['equivariance_value'] * value + + ## jacobian loss part + if self.loss_weights['equivariance_jacobian'] != 0: + jacobian_transformed = torch.matmul(transform.jacobian(transformed_kp['value']), + transformed_kp['jacobian']) + + normed_driving = torch.inverse(kp_driving['jacobian']) + normed_transformed = jacobian_transformed + value = torch.matmul(normed_driving, normed_transformed) + + eye = torch.eye(2).view(1, 1, 2, 2).type(value.type()) + + value = torch.abs(eye - value).mean() + loss_values['equivariance_jacobian'] = self.loss_weights['equivariance_jacobian'] * value + + return loss_values, generated + + +class DiscriminatorFullModel(torch.nn.Module): + """ + Merge all discriminator related updates into single model for better multi-gpu usage + """ + + def __init__(self, kp_extractor, generator, discriminator, train_params): + super(DiscriminatorFullModel, self).__init__() + self.kp_extractor = kp_extractor + self.generator = generator + self.discriminator = discriminator + self.train_params = train_params + self.scales = self.discriminator.scales + self.pyramid = ImagePyramide(self.scales, generator.num_channels) + if torch.cuda.is_available(): + self.pyramid = self.pyramid.cuda() + + self.loss_weights = train_params['loss_weights'] + + def forward(self, x, generated): + pyramide_real = self.pyramid(x['driving']) + pyramide_generated = self.pyramid(generated['prediction'].detach()) + + kp_driving = generated['kp_driving'] + discriminator_maps_generated = self.discriminator(pyramide_generated, kp=detach_kp(kp_driving)) + discriminator_maps_real = self.discriminator(pyramide_real, kp=detach_kp(kp_driving)) + + loss_values = {} + value_total = 0 + for scale in self.scales: + key = 'prediction_map_%s' % scale + value = (1 - discriminator_maps_real[key]) ** 2 + discriminator_maps_generated[key] ** 2 + value_total += self.loss_weights['discriminator_gan'] * value.mean() + loss_values['disc_gan'] = value_total + + return loss_values diff --git a/modules/util.py b/modules/util.py new file mode 100644 index 0000000..f9e5eaf --- /dev/null +++ b/modules/util.py @@ -0,0 +1,245 @@ +from torch import nn + +import torch.nn.functional as F +import torch + +from ..sync_batchnorm.batchnorm import SynchronizedBatchNorm2d as BatchNorm2d + + +def kp2gaussian(kp, spatial_size, kp_variance): + """ + Transform a keypoint into gaussian like representation + """ + mean = kp['value'] + + coordinate_grid = make_coordinate_grid(spatial_size, mean.type()) + number_of_leading_dimensions = len(mean.shape) - 1 + shape = (1,) * number_of_leading_dimensions + coordinate_grid.shape + coordinate_grid = coordinate_grid.view(*shape) + repeats = mean.shape[:number_of_leading_dimensions] + (1, 1, 1) + coordinate_grid = coordinate_grid.repeat(*repeats) + + # Preprocess kp shape + shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 2) + mean = mean.view(*shape) + + mean_sub = (coordinate_grid - mean) + + out = torch.exp(-0.5 * (mean_sub ** 2).sum(-1) / kp_variance) + + return out + + +def make_coordinate_grid(spatial_size, type): + """ + Create a meshgrid [-1,1] x [-1,1] of given spatial_size. + """ + h, w = spatial_size + x = torch.arange(w).type(type) + y = torch.arange(h).type(type) + + x = (2 * (x / (w - 1)) - 1) + y = (2 * (y / (h - 1)) - 1) + + yy = y.view(-1, 1).repeat(1, w) + xx = x.view(1, -1).repeat(h, 1) + + meshed = torch.cat([xx.unsqueeze_(2), yy.unsqueeze_(2)], 2) + + return meshed + + +class ResBlock2d(nn.Module): + """ + Res block, preserve spatial resolution. + """ + + def __init__(self, in_features, kernel_size, padding): + super(ResBlock2d, self).__init__() + self.conv1 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, + padding=padding) + self.conv2 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, + padding=padding) + self.norm1 = BatchNorm2d(in_features, affine=True) + self.norm2 = BatchNorm2d(in_features, affine=True) + + def forward(self, x): + out = self.norm1(x) + out = F.relu(out) + out = self.conv1(out) + out = self.norm2(out) + out = F.relu(out) + out = self.conv2(out) + out += x + return out + + +class UpBlock2d(nn.Module): + """ + Upsampling block for use in decoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(UpBlock2d, self).__init__() + + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups) + self.norm = BatchNorm2d(out_features, affine=True) + + def forward(self, x): + out = F.interpolate(x, scale_factor=2) + out = self.conv(out) + out = self.norm(out) + out = F.relu(out) + return out + + +class DownBlock2d(nn.Module): + """ + Downsampling block for use in encoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(DownBlock2d, self).__init__() + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups) + self.norm = BatchNorm2d(out_features, affine=True) + self.pool = nn.AvgPool2d(kernel_size=(2, 2)) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + out = self.pool(out) + return out + + +class SameBlock2d(nn.Module): + """ + Simple block, preserve spatial resolution. + """ + + def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1): + super(SameBlock2d, self).__init__() + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, + kernel_size=kernel_size, padding=padding, groups=groups) + self.norm = BatchNorm2d(out_features, affine=True) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + return out + + +class Encoder(nn.Module): + """ + Hourglass Encoder + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Encoder, self).__init__() + + down_blocks = [] + for i in range(num_blocks): + down_blocks.append(DownBlock2d(in_features if i == 0 else min(max_features, block_expansion * (2 ** i)), + min(max_features, block_expansion * (2 ** (i + 1))), + kernel_size=3, padding=1)) + self.down_blocks = nn.ModuleList(down_blocks) + + def forward(self, x): + outs = [x] + for down_block in self.down_blocks: + outs.append(down_block(outs[-1])) + return outs + + +class Decoder(nn.Module): + """ + Hourglass Decoder + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Decoder, self).__init__() + + up_blocks = [] + + for i in range(num_blocks)[::-1]: + in_filters = (1 if i == num_blocks - 1 else 2) * min(max_features, block_expansion * (2 ** (i + 1))) + out_filters = min(max_features, block_expansion * (2 ** i)) + up_blocks.append(UpBlock2d(in_filters, out_filters, kernel_size=3, padding=1)) + + self.up_blocks = nn.ModuleList(up_blocks) + self.out_filters = block_expansion + in_features + + def forward(self, x): + out = x.pop() + for up_block in self.up_blocks: + out = up_block(out) + skip = x.pop() + out = torch.cat([out, skip], dim=1) + return out + + +class Hourglass(nn.Module): + """ + Hourglass architecture. + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Hourglass, self).__init__() + self.encoder = Encoder(block_expansion, in_features, num_blocks, max_features) + self.decoder = Decoder(block_expansion, in_features, num_blocks, max_features) + self.out_filters = self.decoder.out_filters + + def forward(self, x): + return self.decoder(self.encoder(x)) + + +class AntiAliasInterpolation2d(nn.Module): + """ + Band-limited downsampling, for better preservation of the input signal. + """ + def __init__(self, channels, scale): + super(AntiAliasInterpolation2d, self).__init__() + sigma = (1 / scale - 1) / 2 + kernel_size = 2 * round(sigma * 4) + 1 + self.ka = kernel_size // 2 + self.kb = self.ka - 1 if kernel_size % 2 == 0 else self.ka + + kernel_size = [kernel_size, kernel_size] + sigma = [sigma, sigma] + # The gaussian kernel is the product of the + # gaussian function of each dimension. + kernel = 1 + meshgrids = torch.meshgrid( + [ + torch.arange(size, dtype=torch.float32) + for size in kernel_size + ] + ) + for size, std, mgrid in zip(kernel_size, sigma, meshgrids): + mean = (size - 1) / 2 + kernel *= torch.exp(-(mgrid - mean) ** 2 / (2 * std ** 2)) + + # Make sure sum of values in gaussian kernel equals 1. + kernel = kernel / torch.sum(kernel) + # Reshape to depthwise convolutional weight + kernel = kernel.view(1, 1, *kernel.size()) + kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1)) + + self.register_buffer('weight', kernel) + self.groups = channels + self.scale = scale + inv_scale = 1 / scale + self.int_inv_scale = int(inv_scale) + + def forward(self, input): + if self.scale == 1.0: + return input + + out = F.pad(input, (self.ka, self.kb, self.ka, self.kb)) + out = F.conv2d(out, weight=self.weight, groups=self.groups) + out = out[:, :, ::self.int_inv_scale, ::self.int_inv_scale] + + return out diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..054a554 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +numpy +torch +scipy +yaml diff --git a/run.py b/run.py new file mode 100644 index 0000000..63aeba9 --- /dev/null +++ b/run.py @@ -0,0 +1,121 @@ +from pathlib import Path + +import torch + +from .constants import default_model_name, supported_model_names +from .utils import ( + get_config_path, + get_model_file_name, + inference, + inference_best_frame, + load_checkpoint, + out_video, + reshape_image, +) + +base_dir = Path(__file__).resolve().parent + + +class FOMM_Runner: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "source_image": ("IMAGE",), + "driving_video_input": ("IMAGE",), + "model_name": (supported_model_names, {"default": default_model_name}), + "frame_rate": ("FLOAT", {"default": 30.0}), + "relative_movement": ( + "BOOLEAN", + {"default": True}, + ), # Relative keypoint displacement (Inherit object proporions from the video) + "relative_jacobian": ( + "BOOLEAN", + {"default": True}, + ), # No idea + "adapt_movement_scale": ( + "BOOLEAN", + {"default": True}, + ), # Adapt movement scale + "find_best_frame": ( + "BOOLEAN", + {"default": False}, + ), # Generate from the frame that is the most alligned with source + }, + "optional": {"audio": ("VHS_AUDIO",)}, + } + + RETURN_TYPES = ( + "IMAGE", + "VHS_AUDIO", + "FLOAT", + ) + RETURN_NAMES = ( + "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, + relative_movement: bool, + relative_jacobian: bool, + adapt_movement_scale: bool, + find_best_frame: bool, + 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}/{get_config_path(model_name)}" + checkpoint_path = f"{base_dir}/{get_model_file_name(model_name)}" + + generator, kp_detector = load_checkpoint(config_path, checkpoint_path) + source_image = reshape_image(source_image, (256, 256)) + driving_video = reshape_image(driving_video_input, (256, 256)).unsqueeze(0) + driving_video = driving_video.permute(0, 2, 1, 3, 4) + + print("After reshaping") + print(f"{source_image.shape=}") + print(f"{driving_video.shape=}") + params = { + "source_image": source_image, + "driving_video": driving_video, + "generator": generator, + "kp_detector": kp_detector, + "relative_movement": relative_movement, + "relative_jacobian": relative_jacobian, + "adapt_movement_scale": adapt_movement_scale, + } + + if find_best_frame: + predictions = inference_best_frame(**params) + else: + predictions = inference(**params) + + output_images = out_video(predictions) + + return ( + output_images, + audio, + frame_rate, + ) + + +NODE_CLASS_MAPPINGS = { + "FOMM_Runner": FOMM_Runner, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "FOMM_Runner": "FOMM Runner", +} diff --git a/samples/Rick1.png b/samples/Rick1.png new file mode 100644 index 0000000..e6224fb Binary files /dev/null and b/samples/Rick1.png differ diff --git a/samples/damedane.mp4 b/samples/damedane.mp4 new file mode 100644 index 0000000..9d28f78 Binary files /dev/null and b/samples/damedane.mp4 differ diff --git a/samples/sample_output.webm b/samples/sample_output.webm new file mode 100644 index 0000000..1a5434a Binary files /dev/null and b/samples/sample_output.webm differ diff --git a/sync_batchnorm/__init__.py b/sync_batchnorm/__init__.py new file mode 100644 index 0000000..41cff07 --- /dev/null +++ b/sync_batchnorm/__init__.py @@ -0,0 +1,12 @@ +# -*- coding: utf-8 -*- +# File : __init__.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +# from .batchnorm import SynchronizedBatchNorm1d, SynchronizedBatchNorm2d, SynchronizedBatchNorm3d +# from .replicate import DataParallelWithCallback, patch_replication_callback diff --git a/sync_batchnorm/batchnorm.py b/sync_batchnorm/batchnorm.py new file mode 100644 index 0000000..a5f8ff0 --- /dev/null +++ b/sync_batchnorm/batchnorm.py @@ -0,0 +1,315 @@ +# -*- coding: utf-8 -*- +# File : batchnorm.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import collections + +# import torch +import torch.nn.functional as F + +from torch.nn.modules.batchnorm import _BatchNorm +from torch.nn.parallel._functions import ReduceAddCoalesced, Broadcast + +from .comm import SyncMaster + +__all__ = ['SynchronizedBatchNorm1d', 'SynchronizedBatchNorm2d', 'SynchronizedBatchNorm3d'] + + +def _sum_ft(tensor): + """sum over the first and last dimention""" + return tensor.sum(dim=0).sum(dim=-1) + + +def _unsqueeze_ft(tensor): + """add new dementions at the front and the tail""" + return tensor.unsqueeze(0).unsqueeze(-1) + + +_ChildMessage = collections.namedtuple('_ChildMessage', ['sum', 'ssum', 'sum_size']) +_MasterMessage = collections.namedtuple('_MasterMessage', ['sum', 'inv_std']) + + +class _SynchronizedBatchNorm(_BatchNorm): + def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True): + super(_SynchronizedBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine) + + self._sync_master = SyncMaster(self._data_parallel_master) + + self._is_parallel = False + self._parallel_id = None + self._slave_pipe = None + + def forward(self, input): + # If it is not parallel computation or is in evaluation mode, use PyTorch's implementation. + if not (self._is_parallel and self.training): + return F.batch_norm( + input, self.running_mean, self.running_var, self.weight, self.bias, + self.training, self.momentum, self.eps) + + # Resize the input to (B, C, -1). + input_shape = input.size() + input = input.view(input.size(0), self.num_features, -1) + + # Compute the sum and square-sum. + sum_size = input.size(0) * input.size(2) + input_sum = _sum_ft(input) + input_ssum = _sum_ft(input ** 2) + + # Reduce-and-broadcast the statistics. + if self._parallel_id == 0: + mean, inv_std = self._sync_master.run_master(_ChildMessage(input_sum, input_ssum, sum_size)) + else: + mean, inv_std = self._slave_pipe.run_slave(_ChildMessage(input_sum, input_ssum, sum_size)) + + # Compute the output. + if self.affine: + # MJY:: Fuse the multiplication for speed. + output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std * self.weight) + _unsqueeze_ft(self.bias) + else: + output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std) + + # Reshape it. + return output.view(input_shape) + + def __data_parallel_replicate__(self, ctx, copy_id): + self._is_parallel = True + self._parallel_id = copy_id + + # parallel_id == 0 means master device. + if self._parallel_id == 0: + ctx.sync_master = self._sync_master + else: + self._slave_pipe = ctx.sync_master.register_slave(copy_id) + + def _data_parallel_master(self, intermediates): + """Reduce the sum and square-sum, compute the statistics, and broadcast it.""" + + # Always using same "device order" makes the ReduceAdd operation faster. + # Thanks to:: Tete Xiao (http://tetexiao.com/) + intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device()) + + to_reduce = [i[1][:2] for i in intermediates] + to_reduce = [j for i in to_reduce for j in i] # flatten + target_gpus = [i[1].sum.get_device() for i in intermediates] + + sum_size = sum([i[1].sum_size for i in intermediates]) + sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce) + mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size) + + broadcasted = Broadcast.apply(target_gpus, mean, inv_std) + + outputs = [] + for i, rec in enumerate(intermediates): + outputs.append((rec[0], _MasterMessage(*broadcasted[i*2:i*2+2]))) + + return outputs + + def _compute_mean_std(self, sum_, ssum, size): + """Compute the mean and standard-deviation with sum and square-sum. This method + also maintains the moving average on the master device.""" + assert size > 1, 'BatchNorm computes unbiased standard-deviation, which requires size > 1.' + mean = sum_ / size + sumvar = ssum - sum_ * mean + unbias_var = sumvar / (size - 1) + bias_var = sumvar / size + + self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data + self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data + + return mean, bias_var.clamp(self.eps) ** -0.5 + + +class SynchronizedBatchNorm1d(_SynchronizedBatchNorm): + r"""Applies Synchronized Batch Normalization over a 2d or 3d input that is seen as a + mini-batch. + + .. math:: + + y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta + + This module differs from the built-in PyTorch BatchNorm1d as the mean and + standard-deviation are reduced across all devices during training. + + For example, when one uses `nn.DataParallel` to wrap the network during + training, PyTorch's implementation normalize the tensor on each device using + the statistics only on that device, which accelerated the computation and + is also easy to implement, but the statistics might be inaccurate. + Instead, in this synchronized version, the statistics will be computed + over all training samples distributed on multiple devices. + + Note that, for one-GPU or CPU-only case, this module behaves exactly same + as the built-in PyTorch implementation. + + The mean and standard-deviation are calculated per-dimension over + the mini-batches and gamma and beta are learnable parameter vectors + of size C (where C is the input size). + + During training, this layer keeps a running estimate of its computed mean + and variance. The running sum is kept with a default momentum of 0.1. + + During evaluation, this running mean/variance is used for normalization. + + Because the BatchNorm is done over the `C` dimension, computing statistics + on `(N, L)` slices, it's common terminology to call this Temporal BatchNorm + + Args: + num_features: num_features from an expected input of size + `batch_size x num_features [x width]` + eps: a value added to the denominator for numerical stability. + Default: 1e-5 + momentum: the value used for the running_mean and running_var + computation. Default: 0.1 + affine: a boolean value that when set to ``True``, gives the layer learnable + affine parameters. Default: ``True`` + + Shape: + - Input: :math:`(N, C)` or :math:`(N, C, L)` + - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input) + + Examples: + >>> # With Learnable Parameters + >>> m = SynchronizedBatchNorm1d(100) + >>> # Without Learnable Parameters + >>> m = SynchronizedBatchNorm1d(100, affine=False) + >>> input = torch.autograd.Variable(torch.randn(20, 100)) + >>> output = m(input) + """ + + def _check_input_dim(self, input): + if input.dim() != 2 and input.dim() != 3: + raise ValueError('expected 2D or 3D input (got {}D input)' + .format(input.dim())) + super(SynchronizedBatchNorm1d, self)._check_input_dim(input) + + +class SynchronizedBatchNorm2d(_SynchronizedBatchNorm): + r"""Applies Batch Normalization over a 4d input that is seen as a mini-batch + of 3d inputs + + .. math:: + + y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta + + This module differs from the built-in PyTorch BatchNorm2d as the mean and + standard-deviation are reduced across all devices during training. + + For example, when one uses `nn.DataParallel` to wrap the network during + training, PyTorch's implementation normalize the tensor on each device using + the statistics only on that device, which accelerated the computation and + is also easy to implement, but the statistics might be inaccurate. + Instead, in this synchronized version, the statistics will be computed + over all training samples distributed on multiple devices. + + Note that, for one-GPU or CPU-only case, this module behaves exactly same + as the built-in PyTorch implementation. + + The mean and standard-deviation are calculated per-dimension over + the mini-batches and gamma and beta are learnable parameter vectors + of size C (where C is the input size). + + During training, this layer keeps a running estimate of its computed mean + and variance. The running sum is kept with a default momentum of 0.1. + + During evaluation, this running mean/variance is used for normalization. + + Because the BatchNorm is done over the `C` dimension, computing statistics + on `(N, H, W)` slices, it's common terminology to call this Spatial BatchNorm + + Args: + num_features: num_features from an expected input of + size batch_size x num_features x height x width + eps: a value added to the denominator for numerical stability. + Default: 1e-5 + momentum: the value used for the running_mean and running_var + computation. Default: 0.1 + affine: a boolean value that when set to ``True``, gives the layer learnable + affine parameters. Default: ``True`` + + Shape: + - Input: :math:`(N, C, H, W)` + - Output: :math:`(N, C, H, W)` (same shape as input) + + Examples: + >>> # With Learnable Parameters + >>> m = SynchronizedBatchNorm2d(100) + >>> # Without Learnable Parameters + >>> m = SynchronizedBatchNorm2d(100, affine=False) + >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45)) + >>> output = m(input) + """ + + def _check_input_dim(self, input): + if input.dim() != 4: + raise ValueError('expected 4D input (got {}D input)' + .format(input.dim())) + super(SynchronizedBatchNorm2d, self)._check_input_dim(input) + + +class SynchronizedBatchNorm3d(_SynchronizedBatchNorm): + r"""Applies Batch Normalization over a 5d input that is seen as a mini-batch + of 4d inputs + + .. math:: + + y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta + + This module differs from the built-in PyTorch BatchNorm3d as the mean and + standard-deviation are reduced across all devices during training. + + For example, when one uses `nn.DataParallel` to wrap the network during + training, PyTorch's implementation normalize the tensor on each device using + the statistics only on that device, which accelerated the computation and + is also easy to implement, but the statistics might be inaccurate. + Instead, in this synchronized version, the statistics will be computed + over all training samples distributed on multiple devices. + + Note that, for one-GPU or CPU-only case, this module behaves exactly same + as the built-in PyTorch implementation. + + The mean and standard-deviation are calculated per-dimension over + the mini-batches and gamma and beta are learnable parameter vectors + of size C (where C is the input size). + + During training, this layer keeps a running estimate of its computed mean + and variance. The running sum is kept with a default momentum of 0.1. + + During evaluation, this running mean/variance is used for normalization. + + Because the BatchNorm is done over the `C` dimension, computing statistics + on `(N, D, H, W)` slices, it's common terminology to call this Volumetric BatchNorm + or Spatio-temporal BatchNorm + + Args: + num_features: num_features from an expected input of + size batch_size x num_features x depth x height x width + eps: a value added to the denominator for numerical stability. + Default: 1e-5 + momentum: the value used for the running_mean and running_var + computation. Default: 0.1 + affine: a boolean value that when set to ``True``, gives the layer learnable + affine parameters. Default: ``True`` + + Shape: + - Input: :math:`(N, C, D, H, W)` + - Output: :math:`(N, C, D, H, W)` (same shape as input) + + Examples: + >>> # With Learnable Parameters + >>> m = SynchronizedBatchNorm3d(100) + >>> # Without Learnable Parameters + >>> m = SynchronizedBatchNorm3d(100, affine=False) + >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45, 10)) + >>> output = m(input) + """ + + def _check_input_dim(self, input): + if input.dim() != 5: + raise ValueError('expected 5D input (got {}D input)' + .format(input.dim())) + super(SynchronizedBatchNorm3d, self)._check_input_dim(input) diff --git a/sync_batchnorm/comm.py b/sync_batchnorm/comm.py new file mode 100644 index 0000000..922f8c4 --- /dev/null +++ b/sync_batchnorm/comm.py @@ -0,0 +1,137 @@ +# -*- coding: utf-8 -*- +# File : comm.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import queue +import collections +import threading + +__all__ = ['FutureResult', 'SlavePipe', 'SyncMaster'] + + +class FutureResult(object): + """A thread-safe future implementation. Used only as one-to-one pipe.""" + + def __init__(self): + self._result = None + self._lock = threading.Lock() + self._cond = threading.Condition(self._lock) + + def put(self, result): + with self._lock: + assert self._result is None, 'Previous result has\'t been fetched.' + self._result = result + self._cond.notify() + + def get(self): + with self._lock: + if self._result is None: + self._cond.wait() + + res = self._result + self._result = None + return res + + +_MasterRegistry = collections.namedtuple('MasterRegistry', ['result']) +_SlavePipeBase = collections.namedtuple('_SlavePipeBase', ['identifier', 'queue', 'result']) + + +class SlavePipe(_SlavePipeBase): + """Pipe for master-slave communication.""" + + def run_slave(self, msg): + self.queue.put((self.identifier, msg)) + ret = self.result.get() + self.queue.put(True) + return ret + + +class SyncMaster(object): + """An abstract `SyncMaster` object. + + - During the replication, as the data parallel will trigger an callback of each module, all slave devices should + call `register(id)` and obtain an `SlavePipe` to communicate with the master. + - During the forward pass, master device invokes `run_master`, all messages from slave devices will be collected, + and passed to a registered callback. + - After receiving the messages, the master device should gather the information and determine to message passed + back to each slave devices. + """ + + def __init__(self, master_callback): + """ + + Args: + master_callback: a callback to be invoked after having collected messages from slave devices. + """ + self._master_callback = master_callback + self._queue = queue.Queue() + self._registry = collections.OrderedDict() + self._activated = False + + def __getstate__(self): + return {'master_callback': self._master_callback} + + def __setstate__(self, state): + self.__init__(state['master_callback']) + + def register_slave(self, identifier): + """ + Register an slave device. + + Args: + identifier: an identifier, usually is the device id. + + Returns: a `SlavePipe` object which can be used to communicate with the master device. + + """ + if self._activated: + assert self._queue.empty(), 'Queue is not clean before next initialization.' + self._activated = False + self._registry.clear() + future = FutureResult() + self._registry[identifier] = _MasterRegistry(future) + return SlavePipe(identifier, self._queue, future) + + def run_master(self, master_msg): + """ + Main entry for the master device in each forward pass. + The messages were first collected from each devices (including the master device), and then + an callback will be invoked to compute the message to be sent back to each devices + (including the master device). + + Args: + master_msg: the message that the master want to send to itself. This will be placed as the first + message when calling `master_callback`. For detailed usage, see `_SynchronizedBatchNorm` for an example. + + Returns: the message to be sent back to the master device. + + """ + self._activated = True + + intermediates = [(0, master_msg)] + for i in range(self.nr_slaves): + intermediates.append(self._queue.get()) + + results = self._master_callback(intermediates) + assert results[0][0] == 0, 'The first result should belongs to the master.' + + for i, res in results: + if i == 0: + continue + self._registry[i].result.put(res) + + for i in range(self.nr_slaves): + assert self._queue.get() is True + + return results[0][1] + + @property + def nr_slaves(self): + return len(self._registry) diff --git a/sync_batchnorm/replicate.py b/sync_batchnorm/replicate.py new file mode 100644 index 0000000..b71c7b8 --- /dev/null +++ b/sync_batchnorm/replicate.py @@ -0,0 +1,94 @@ +# -*- coding: utf-8 -*- +# File : replicate.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import functools + +from torch.nn.parallel.data_parallel import DataParallel + +__all__ = [ + 'CallbackContext', + 'execute_replication_callbacks', + 'DataParallelWithCallback', + 'patch_replication_callback' +] + + +class CallbackContext(object): + pass + + +def execute_replication_callbacks(modules): + """ + Execute an replication callback `__data_parallel_replicate__` on each module created by original replication. + + The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)` + + Note that, as all modules are isomorphism, we assign each sub-module with a context + (shared among multiple copies of this module on different devices). + Through this context, different copies can share some information. + + We guarantee that the callback on the master copy (the first copy) will be called ahead of calling the callback + of any slave copies. + """ + master_copy = modules[0] + nr_modules = len(list(master_copy.modules())) + ctxs = [CallbackContext() for _ in range(nr_modules)] + + for i, module in enumerate(modules): + for j, m in enumerate(module.modules()): + if hasattr(m, '__data_parallel_replicate__'): + m.__data_parallel_replicate__(ctxs[j], i) + + +class DataParallelWithCallback(DataParallel): + """ + Data Parallel with a replication callback. + + An replication callback `__data_parallel_replicate__` of each module will be invoked after being created by + original `replicate` function. + The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)` + + Examples: + > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) + > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) + # sync_bn.__data_parallel_replicate__ will be invoked. + """ + + def replicate(self, module, device_ids): + modules = super(DataParallelWithCallback, self).replicate(module, device_ids) + execute_replication_callbacks(modules) + return modules + + +def patch_replication_callback(data_parallel): + """ + Monkey-patch an existing `DataParallel` object. Add the replication callback. + Useful when you have customized `DataParallel` implementation. + + Examples: + > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) + > sync_bn = DataParallel(sync_bn, device_ids=[0, 1]) + > patch_replication_callback(sync_bn) + # this is equivalent to + > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) + > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) + """ + + assert isinstance(data_parallel, DataParallel) + + old_replicate = data_parallel.replicate + + @functools.wraps(old_replicate) + def new_replicate(module, device_ids): + modules = old_replicate(module, device_ids) + execute_replication_callbacks(modules) + return modules + + data_parallel.replicate = new_replicate diff --git a/sync_batchnorm/unittest.py b/sync_batchnorm/unittest.py new file mode 100644 index 0000000..0675c02 --- /dev/null +++ b/sync_batchnorm/unittest.py @@ -0,0 +1,29 @@ +# -*- coding: utf-8 -*- +# File : unittest.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import unittest + +import numpy as np +from torch.autograd import Variable + + +def as_numpy(v): + if isinstance(v, Variable): + v = v.data + return v.cpu().numpy() + + +class TorchTestCase(unittest.TestCase): + def assertTensorClose(self, a, b, atol=1e-3, rtol=1e-3): + npa, npb = as_numpy(a), as_numpy(b) + self.assertTrue( + np.allclose(npa, npb, atol=atol), + 'Tensor close check failed\n{}\n{}\nadiff={}, rdiff={}'.format(a, b, np.abs(npa - npb).max(), np.abs((npa - npb) / np.fmax(npa, 1e-5)).max()) + ) diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..12894cd --- /dev/null +++ b/utils.py @@ -0,0 +1,283 @@ +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, + non_cpk_model_names, + supported_model_names, +) +from .modules.generator import OcclusionAwareGenerator +from .modules.keypoint_detector import KPDetector + + +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}" + + if model_name in non_cpk_model_names: + model_file_name = f"{checkpoint_folder}/{model_name}{checkpoint_suffix}" + else: + model_file_name = f"{checkpoint_folder}/{model_name}-cpk{checkpoint_suffix}" + + return model_file_name + + +def get_model_link(model_name: str) -> str: + model_file_name = get_model_file_name(model_name) + model_dl_link = f"{checkpoint_folder_link}/{model_file_name}" + return model_dl_link + + +def get_config_path(model_name: str) -> str: + config_path = f"{config_folder}/{model_name}-256.yaml" + return config_path + + +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: + raise ValueError("Expected 4D input (batch, height, width, channels)") + + # Permute dimensions to [B, C, H, W] for resizing + image = image.permute(0, 3, 1, 2) + + # Resize to target size + resized_image = F.interpolate( + image, size=target_size, mode="bilinear", align_corners=False + ) + + 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: 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[1] + + 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 diff --git a/workflow.png b/workflow.png new file mode 100644 index 0000000..da13848 Binary files /dev/null and b/workflow.png differ