Works on my machine
This commit is contained in:
FuouM
@@ -1,3 +1,9 @@
|
||||
# Stuff
|
||||
|
||||
checkpoints/*
|
||||
!checkpoints/place_checkpoints_here.txt
|
||||
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
|
||||
@@ -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}
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
## 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
|
||||
```
|
||||
+11
@@ -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']
|
||||
@@ -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'
|
||||
@@ -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'
|
||||
@@ -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'
|
||||
@@ -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'
|
||||
@@ -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'
|
||||
@@ -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'
|
||||
@@ -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'
|
||||
@@ -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'
|
||||
@@ -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"
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
+245
@@ -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
|
||||
@@ -0,0 +1,4 @@
|
||||
numpy
|
||||
torch
|
||||
scipy
|
||||
yaml
|
||||
@@ -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",
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 82 KiB |
Binary file not shown.
Binary file not shown.
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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())
|
||||
)
|
||||
@@ -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
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 280 KiB |
Reference in New Issue
Block a user