nodes/FuouM-ComfyUI-FirstOrderMM
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #3 from FuouM/partswap

Browse Source 
Partswap integration
This commit is contained in:
Fuou Marinas
2024-07-01 15:41:06 +07:00
committed by GitHub
parent dfcb5203dc 314560aab5
commit b3b4f26cc1
26 changed files with 2601 additions and 250 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2
View File
@@ -3,6 +3,8 @@
checkpoints/*
!checkpoints/place_checkpoints_here.txt
face_parsing/*.pth
!face_parsing/put_model_safetensors_config_json_here.txt
# Byte-compiled / optimized / DLL files
__pycache__/
FOMM_PARTSWAP.json
+692
View File
@@ -0,0 +1,692 @@
{
"last_node_id": 28,
"last_link_id": 59,
"nodes": [
{
"id": 1,
"type": "LoadImage",
"pos": [
111,
131
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
47
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"Rick1.png",
"image"
]
},
{
"id": 24,
"type": "PreviewImage",
"pos": [
1423.509090909091,
336.86363636363626
],
"size": {
"0": 210,
"1": 246
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 52
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 25,
"type": "PreviewImage",
"pos": [
1413.509090909091,
625.8636363636363
],
"size": {
"0": 210,
"1": 246
},
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 53
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 16,
"type": "VHS_VideoInfoLoaded",
"pos": [
366,
875
],
"size": {
"0": 235.1999969482422,
"1": 106
},
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "video_info",
"type": "VHS_VIDEOINFO",
"link": 12
}
],
"outputs": [
{
"name": "fps🟦",
"type": "FLOAT",
"links": [
54
],
"shape": 3,
"slot_index": 0
},
{
"name": "frame_count🟦",
"type": "INT",
"links": null,
"shape": 3
},
{
"name": "duration🟦",
"type": "FLOAT",
"links": null,
"shape": 3
},
{
"name": "width🟦",
"type": "INT",
"links": null,
"shape": 3
},
{
"name": "height🟦",
"type": "INT",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "VHS_VideoInfoLoaded"
},
"widgets_values": {}
},
{
"id": 13,
"type": "VHS_VideoCombine",
"pos": [
1048,
349
],
"size": [
315,
595
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 49
},
{
"name": "audio",
"type": "VHS_AUDIO",
"link": 50
},
{
"name": "meta_batch",
"type": "VHS_BatchManager",
"link": null,
"slot_index": 2
},
{
"name": "frame_rate",
"type": "FLOAT",
"link": 51,
"widget": {
"name": "frame_rate"
}
}
],
"outputs": [
{
"name": "Filenames",
"type": "VHS_FILENAMES",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "VHS_VideoCombine"
},
"widgets_values": {
"frame_rate": 8,
"loop_count": 0,
"filename_prefix": "Partswap",
"format": "video/h264-mp4",
"pix_fmt": "yuv420p",
"crf": 0,
"save_metadata": true,
"pingpong": false,
"save_output": true,
"videopreview": {
"hidden": false,
"paused": false,
"params": {
"filename": "Partswap_00093-audio.mp4",
"subfolder": "",
"type": "output",
"format": "video/h264-mp4"
}
}
}
},
{
"id": 11,
"type": "VHS_LoadVideo",
"pos": [
117,
506
],
"size": [
235.1999969482422,
491.1999969482422
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"name": "meta_batch",
"type": "VHS_BatchManager",
"link": null
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
48
],
"shape": 3,
"slot_index": 0
},
{
"name": "frame_count",
"type": "INT",
"links": null,
"shape": 3
},
{
"name": "audio",
"type": "VHS_AUDIO",
"links": [
45
],
"shape": 3,
"slot_index": 2
},
{
"name": "video_info",
"type": "VHS_VIDEOINFO",
"links": [
12
],
"shape": 3,
"slot_index": 3
}
],
"properties": {
"Node name for S&R": "VHS_LoadVideo"
},
"widgets_values": {
"video": "damedane.mp4",
"force_rate": 0,
"force_size": "Disabled",
"custom_width": 512,
"custom_height": 512,
"frame_load_cap": 10,
"skip_first_frames": 0,
"select_every_nth": 1,
"choose video to upload": "image",
"videopreview": {
"hidden": false,
"paused": false,
"params": {
"frame_load_cap": 10,
"skip_first_frames": 0,
"force_rate": 0,
"filename": "damedane.mp4",
"type": "input",
"format": "video/mp4",
"select_every_nth": 1
}
}
}
},
{
"id": 26,
"type": "FOMM_Seg5Chooser",
"pos": [
313,
1078
],
"size": {
"0": 315,
"1": 178
},
"flags": {},
"order": 2,
"mode": 0,
"outputs": [
{
"name": "chosen_seg_indices",
"type": "STRING",
"links": [],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "FOMM_Seg5Chooser"
},
"widgets_values": [
false,
true,
true,
true,
true,
true
]
},
{
"id": 23,
"type": "FOMM_Seg10Chooser",
"pos": [
649,
1080
],
"size": {
"0": 315,
"1": 298
},
"flags": {},
"order": 3,
"mode": 0,
"outputs": [
{
"name": "chosen_seg_indices",
"type": "STRING",
"links": [],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "FOMM_Seg10Chooser"
},
"widgets_values": [
false,
true,
true,
true,
true,
true,
true,
true,
true,
true,
true
]
},
{
"id": 28,
"type": "FOMM_Seg15Chooser",
"pos": [
983,
1081
],
"size": {
"0": 315,
"1": 418
},
"flags": {},
"order": 4,
"mode": 0,
"outputs": [
{
"name": "chosen_seg_indices",
"type": "STRING",
"links": [
59
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "FOMM_Seg15Chooser"
},
"widgets_values": [
false,
true,
true,
true,
true,
true,
true,
true,
true,
true,
true,
true,
true,
true,
true,
true
]
},
{
"id": 27,
"type": "Reroute",
"pos": [
709,
725
],
"size": [
75,
26
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "",
"type": "*",
"link": 59,
"widget": {
"name": "value"
}
}
],
"outputs": [
{
"name": "",
"type": "STRING",
"links": [
56
],
"slot_index": 0
}
],
"properties": {
"showOutputText": false,
"horizontal": false
}
},
{
"id": 22,
"type": "FOMM_Partswap",
"pos": [
599,
366
],
"size": {
"0": 393,
"1": 258
},
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "source_image",
"type": "IMAGE",
"link": 47
},
{
"name": "driving_video_input",
"type": "IMAGE",
"link": 48
},
{
"name": "audio",
"type": "VHS_AUDIO",
"link": 45
},
{
"name": "chosen_seg_indices",
"type": "STRING",
"link": 56,
"widget": {
"name": "chosen_seg_indices"
}
},
{
"name": "frame_rate",
"type": "FLOAT",
"link": 54,
"widget": {
"name": "frame_rate"
}
}
],
"outputs": [
{
"name": "seg_src_viz",
"type": "IMAGE",
"links": [
52
],
"shape": 3,
"slot_index": 0
},
{
"name": "seg_tgt_viz",
"type": "IMAGE",
"links": [
53
],
"shape": 3,
"slot_index": 1
},
{
"name": "images",
"type": "IMAGE",
"links": [
49
],
"shape": 3,
"slot_index": 2
},
{
"name": "audio",
"type": "VHS_AUDIO",
"links": [
50
],
"shape": 3,
"slot_index": 3
},
{
"name": "frame_rate",
"type": "FLOAT",
"links": [
51
],
"shape": 3,
"slot_index": 4
}
],
"properties": {
"Node name for S&R": "FOMM_Partswap"
},
"widgets_values": [
"vox-15segments",
30,
1,
false,
false,
false,
1,
0.6
]
}
],
"links": [
[
12,
11,
3,
16,
0,
"VHS_VIDEOINFO"
],
[
45,
11,
2,
22,
2,
"VHS_AUDIO"
],
[
47,
1,
0,
22,
0,
"IMAGE"
],
[
48,
11,
0,
22,
1,
"IMAGE"
],
[
49,
22,
2,
13,
0,
"IMAGE"
],
[
50,
22,
3,
13,
1,
"VHS_AUDIO"
],
[
51,
22,
4,
13,
3,
"FLOAT"
],
[
52,
22,
0,
24,
0,
"IMAGE"
],
[
53,
22,
1,
25,
0,
"IMAGE"
],
[
54,
16,
0,
22,
4,
"FLOAT"
],
[
56,
27,
0,
22,
3,
"STRING"
],
[
59,
28,
0,
27,
0,
"*"
]
],
"groups": [],
"config": {},
"extra": {
"ds": {
"scale": 0.8264462809917358,
"offset": [
101.0730714815582,
-279.5114901742915
]
}
},
"version": 0.4
}
README.md
+49 -4
View File
@@ -3,24 +3,43 @@ A ComfyUI-native node to run First Order Motion Model for Image Animation
https://github.com/AliaksandrSiarohin/first-order-model
Now supports Face Swapping using Motion Supervised co-part Segmentation
https://github.com/AliaksandrSiarohin/motion-cosegmentation
## Workflow:
### FOMM
[FOMM.json](FOMM.json)
![FOMM Workflow](workflow.png)
https://github.com/FuouM/ComfyUI-FirstOrderMM/assets/57849916/36833a94-67f9-4dd9-b79b-181f86e83bfe
### Part Swap
[FOMM_PARTSWAP.json](FOMM_PARTSWAP.json)
![Partswap Workflow](workflow_fomm_partswap.png)
## Arguments
### FOMM
* `relative_movement`: Relative keypoint displacement (Inherit object proporions from the video)
* `relative_jacobian`: No idea what this does, but usually == `relative_movement`
* `relative_jacobian`: Only taken into effect when `relative_movement` is on, must also be on to avoid heavy deformation of the face (in a freaky way)
* `adapt_movement_scale`: If disabled, will heavily distort the source face to match the driving face
* `find_best_frame`: Find driving frame that best match the source. Split the batch into two halves, with the first half reversed. Gives mixed results. Needs to install `face-alignment` library.
### Part Swap
* `blend_scale`: No idea, keeping at default = 1.0 seems to be fine
* `use_source_seg`: Whether to use the source's segmentation or the target's. May help if some of the target's segmentation regions are missing
* `hard_edges`: Whether to make the edges hard, instead of feathering
* `use_face_parser`: For Seg-based models, may help with cleaning up residual background (should only use `15seg` with this). TODO: Additional cleanup face_parser masks. Should definitely be used for FOMM models
* `viz_alpha`: Opacity of the segments in the visualization
## Installation
1. Clone the repo to `ComfyUI/custom_nodes/`
@@ -41,13 +60,39 @@ pip install face-alignment
## Model
Currently supporting `vox` and `vox-adv`. Models can and must be manually downloaded from:
FOMM currently supporting `vox` and `vox-adv`. Models can and must be manually downloaded from:
* [AliaksandrSiarohin/first-order-model](https://github.com/AliaksandrSiarohin/first-order-model)
* [graphemecluster/first-order-model-demo](https://github.com/graphemecluster/first-order-model-demo)
Part Swap currently supporting Seg-based models and FOMM (`vox` and `vox-adv`) models.
* `vox-5segments`
* `vox-10segments`
* `vox-15segments`
* `vox-first-order (partswap)`
These models can be found in the original repository [Motion Supervised co-part Segmentation](https://github.com/AliaksandrSiarohin/motion-cosegmentation)
Place them in the `checkpoints` folder. It should look like this:
```
place_checkpoints_here.txt
vox-adv-cpk.pth.tar
vox-cpk.pth.tar
vox-5segments.pth.tar
vox-10segments.pth.tar
vox-15segments.pth.tar
vox-first-order.pth.tar
```
For Part Swap, Face-Parsing is also supported **(Optional)** (especially when using the FOMM or `vox-first-order` models)
* resnet18 `resnet18-5c106cde`: https://download.pytorch.org/models/resnet18-5c106cde.pth
* face_parsing `79999_iter.pth`: https://github.com/zllrunning/face-makeup.PyTorch/tree/master/cp
Place them in `face_parsing` folder:
```
face_parsing_model.py
...
resnet18-5c106cde.pth
79999_iter.pth
```
__init__.py
+17 -2
View File
@@ -1,11 +1,26 @@
from .run import FOMM_Runner
from .run import (
FOMM_Partswap,
FOMM_Runner,
FOMM_Seg5Chooser,
FOMM_Seg10Chooser,
FOMM_Seg15Chooser,
)
NODE_CLASS_MAPPINGS = {
"FOMM_Runner": FOMM_Runner,
"FOMM_Partswap": FOMM_Partswap,
"FOMM_Seg5Chooser": FOMM_Seg5Chooser,
"FOMM_Seg10Chooser": FOMM_Seg10Chooser,
"FOMM_Seg15Chooser": FOMM_Seg15Chooser,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"FOMM_Runner": "FOMM Runner",
"FOMM_Partswap": "FOMM Partswap",
"FOMM_Seg5Chooser": "FOMM Seg5 Chooser",
"FOMM_Seg10Chooser": "FOMM Seg10 Chooser",
"FOMM_Seg15Chooser": "FOMM Seg15 Chooser",
}
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
config/vox-256-sem-10segments.yaml
+92
View File
@@ -0,0 +1,92 @@
# Dataset parameters
dataset_params:
# Path to data, data can be stored in several formats: .mp4 or .gif videos, stacked .png images or folders with frames.
root_dir: data/vox-png
# Image shape, needed for staked .png format.
image_shape: [256, 256, 3]
# In case of VoxCeleb or TaiChi single video can be splitted in many chunks, or the maybe several videos for single person.
# In this case epoch can be a pass over different videos (if id_sampling=True) or over different chunks (if id_sampling=False)
# If the name of the video '12335#adsbf.mp4' the id is assumed to be 12335
id_sampling: True
# Augmentation parameters see augmentation.py for all posible augmentations
augmentation_params:
flip_param:
horizontal_flip: True
time_flip: True
jitter_param:
brightness: 0.1
contrast: 0.1
saturation: 0.1
hue: 0.1
# Defines model architecture
model_params:
common_params:
# Number of segments
num_segments: 10
# Number of channels per image
num_channels: 3
# Use only shift and no-affine part
estimate_affine_part: True
segmentation_module_params:
# Softmax temperature for shift heatmaps
temperature: 0.1
# Number of features mutliplier
block_expansion: 32
# Maximum allowed number of features
max_features: 1024
# Number of block in Unet. Can be increased or decreased depending or resolution.
num_blocks: 5
# Segmentations is predicted on smaller images for better performance,
# scale_factor=0.25 means that 256x256 image will be resized to 64x64
scale_factor: 0.25
reconstruction_module_params:
# Number of features mutliplier
block_expansion: 64
# Maximum allowed number of features
max_features: 512
# Number of downsampling blocks in Jonson architecture.
# Can be increased or decreased depending or resolution.
num_down_blocks: 2
# Number of ResBlocks in Jonson architecture.
num_bottleneck_blocks: 6
# Use visibility map or not
estimate_visibility: True
# Parameters of training
train_params:
# Number of training epochs
num_epochs: 5
# For better i/o performance when number of videos is small number of epochs can be multiplied by this number.
# Thus effectively with num_repeats=100 each epoch is 100 times larger.
num_repeats: 50
# Learning rates
lr_segmentation: 2.0e-4
lr_reconstruction: 2.0e-4
batch_size: 20
# Scales for perceptual pyramide loss. If scales = [1, 0.5, 0.25, 0.125] and image resolution is 256x256,
# than the loss will be computer on resolutions 256x256, 128x128, 64x64, 32x32.
scales: [1, 0.5, 0.25, 0.125]
# Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs.
checkpoint_freq: 50
# Parameters of transform for equivariance loss
transform_params:
# Sigma for affine part
sigma_affine: 0.05
# Sigma for deformation part
sigma_tps: 0.005
# Number of point in the deformation grid
points_tps: 5
loss_weights:
equivariance: 10
perceptual: [10, 10, 10, 10, 10]
# Visualization parameters
visualizer_params:
# Draw keypoints (shifts in affine transformations) of this size, increase or decrease depending on resolution
kp_size: 5
# Draw white border around images
draw_border: True
# Color map for keypoints
colormap: 'gist_rainbow'
config/vox-256-sem-15segments.yaml
+92
View File
@@ -0,0 +1,92 @@
# Dataset parameters
dataset_params:
# Path to data, data can be stored in several formats: .mp4 or .gif videos, stacked .png images or folders with frames.
root_dir: data/vox-png
# Image shape, needed for staked .png format.
image_shape: [256, 256, 3]
# In case of VoxCeleb or TaiChi single video can be splitted in many chunks, or the maybe several videos for single person.
# In this case epoch can be a pass over different videos (if id_sampling=True) or over different chunks (if id_sampling=False)
# If the name of the video '12335#adsbf.mp4' the id is assumed to be 12335
id_sampling: True
# Augmentation parameters see augmentation.py for all posible augmentations
augmentation_params:
flip_param:
horizontal_flip: True
time_flip: True
jitter_param:
brightness: 0.1
contrast: 0.1
saturation: 0.1
hue: 0.1
# Defines model architecture
model_params:
common_params:
# Number of segments
num_segments: 15
# Number of channels per image
num_channels: 3
# Use only shift and no-affine part
estimate_affine_part: True
segmentation_module_params:
# Softmax temperature for shift heatmaps
temperature: 0.1
# Number of features mutliplier
block_expansion: 32
# Maximum allowed number of features
max_features: 1024
# Number of block in Unet. Can be increased or decreased depending or resolution.
num_blocks: 5
# Segmentations is predicted on smaller images for better performance,
# scale_factor=0.25 means that 256x256 image will be resized to 64x64
scale_factor: 0.25
reconstruction_module_params:
# Number of features mutliplier
block_expansion: 64
# Maximum allowed number of features
max_features: 512
# Number of downsampling blocks in Jonson architecture.
# Can be increased or decreased depending or resolution.
num_down_blocks: 2
# Number of ResBlocks in Jonson architecture.
num_bottleneck_blocks: 6
# Use visibility map or not
estimate_visibility: True
# Parameters of training
train_params:
# Number of training epochs
num_epochs: 5
# For better i/o performance when number of videos is small number of epochs can be multiplied by this number.
# Thus effectively with num_repeats=100 each epoch is 100 times larger.
num_repeats: 50
# Learning rates
lr_segmentation: 2.0e-4
lr_reconstruction: 2.0e-4
batch_size: 20
# Scales for perceptual pyramide loss. If scales = [1, 0.5, 0.25, 0.125] and image resolution is 256x256,
# than the loss will be computer on resolutions 256x256, 128x128, 64x64, 32x32.
scales: [1, 0.5, 0.25, 0.125]
# Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs.
checkpoint_freq: 50
# Parameters of transform for equivariance loss
transform_params:
# Sigma for affine part
sigma_affine: 0.05
# Sigma for deformation part
sigma_tps: 0.005
# Number of point in the deformation grid
points_tps: 5
loss_weights:
equivariance: 10
perceptual: [10, 10, 10, 10, 10]
# Visualization parameters
visualizer_params:
# Draw keypoints (shifts in affine transformations) of this size, increase or decrease depending on resolution
kp_size: 5
# Draw white border around images
draw_border: True
# Color map for keypoints
colormap: 'gist_rainbow'
config/vox-256-sem-5segments.yaml
+92
View File
@@ -0,0 +1,92 @@
# Dataset parameters
dataset_params:
# Path to data, data can be stored in several formats: .mp4 or .gif videos, stacked .png images or folders with frames.
root_dir: data/vox-png
# Image shape, needed for staked .png format.
image_shape: [256, 256, 3]
# In case of VoxCeleb or TaiChi single video can be splitted in many chunks, or the maybe several videos for single person.
# In this case epoch can be a pass over different videos (if id_sampling=True) or over different chunks (if id_sampling=False)
# If the name of the video '12335#adsbf.mp4' the id is assumed to be 12335
id_sampling: True
# Augmentation parameters see augmentation.py for all posible augmentations
augmentation_params:
flip_param:
horizontal_flip: True
time_flip: True
jitter_param:
brightness: 0.1
contrast: 0.1
saturation: 0.1
hue: 0.1
# Defines model architecture
model_params:
common_params:
# Number of segments
num_segments: 5
# Number of channels per image
num_channels: 3
# Use only shift and no-affine part
estimate_affine_part: True
segmentation_module_params:
# Softmax temperature for shift heatmaps
temperature: 0.1
# Number of features mutliplier
block_expansion: 32
# Maximum allowed number of features
max_features: 1024
# Number of block in Unet. Can be increased or decreased depending or resolution.
num_blocks: 5
# Segmentations is predicted on smaller images for better performance,
# scale_factor=0.25 means that 256x256 image will be resized to 64x64
scale_factor: 0.25
reconstruction_module_params:
# Number of features mutliplier
block_expansion: 64
# Maximum allowed number of features
max_features: 512
# Number of downsampling blocks in Jonson architecture.
# Can be increased or decreased depending or resolution.
num_down_blocks: 2
# Number of ResBlocks in Jonson architecture.
num_bottleneck_blocks: 6
# Use visibility map or not
estimate_visibility: True
# Parameters of training
train_params:
# Number of training epochs
num_epochs: 5
# For better i/o performance when number of videos is small number of epochs can be multiplied by this number.
# Thus effectively with num_repeats=100 each epoch is 100 times larger.
num_repeats: 50
# Learning rates
lr_segmentation: 2.0e-4
lr_reconstruction: 2.0e-4
batch_size: 20
# Scales for perceptual pyramide loss. If scales = [1, 0.5, 0.25, 0.125] and image resolution is 256x256,
# than the loss will be computer on resolutions 256x256, 128x128, 64x64, 32x32.
scales: [1, 0.5, 0.25, 0.125]
# Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs.
checkpoint_freq: 50
# Parameters of transform for equivariance loss
transform_params:
# Sigma for affine part
sigma_affine: 0.05
# Sigma for deformation part
sigma_tps: 0.005
# Number of point in the deformation grid
points_tps: 5
loss_weights:
equivariance: 10
perceptual: [10, 10, 10, 10, 10]
# Visualization parameters
visualizer_params:
# Draw keypoints (shifts in affine transformations) of this size, increase or decrease depending on resolution
kp_size: 5
# Draw white border around images
draw_border: True
# Color map for keypoints
colormap: 'gist_rainbow'
constants.py
+81 -4
View File
@@ -9,15 +9,92 @@ supported_model_names = [
# "bair",
]
non_cpk_model_names = [
"fashion"
]
non_cpk_model_names = ["fashion"]
default_model_name = "vox-adv"
checkpoint_folder_link = "https://github.com/graphemecluster/first-order-model-demo/releases/download/checkpoints"
checkpoint_suffix = ".pth.tar"
config_suffix = "256.yaml"
config_folder = "config"
checkpoint_folder = "checkpoints"
checkpoint_folder = "checkpoints"
face_parser_checkpoint_name = "79999_iter.pth"
partswap_model_config_dict = {
"vox-5segments": "vox-256-sem-5segments.yaml",
"vox-10segments": "vox-256-sem-10segments.yaml",
"vox-15segments": "vox-256-sem-15segments.yaml",
"vox-cpk": "vox-256-sem-10segments.yaml",
"vox-adv-cpk": "vox-256-sem-10segments.yaml",
"vox-first-order": "vox-256-sem-10segments.yaml",
}
partswap_face_parser_enforce_models = ["vox-first-order"]
partswap_fomm_model_names = ["vox-cpk", "vox-adv-cpk", "vox-first-order"]
partswap_model_length_dict = {
"vox-5segments": 6,
"vox-10segments": 11,
"vox-15segments": 16,
"vox-cpk": 11,
"vox-adv-cpk": 11,
"vox-first-order": 11,
}
partswap_model_names = list(partswap_model_config_dict.keys())
default_partswap_model_name = "vox-10segments"
seg10_segment_names = [
"m_0_background",
"m_1_lf_ear",
"m_2_mouth",
"m_3_lf_fr_head",
"m_4_rt_ear",
"m_5_frnt_face",
"m_6_lt_cheek",
"m_7_eyes",
"m_8_rt_fr_head",
"m_9_shoulder",
"m_10_hair_top",
]
seg5_segment_names = [
"m_0_background",
"m_1_chin",
"m_2_frnt_face",
"m_3_shoulder",
"m_4_hair",
"m_5_fr_head",
]
seg_15_segment_names = [
"m_0_background",
"m_1_rt_fr_head",
"m_2_eyes",
"m_3_mouth",
"m_4_lwr_shoulder",
"m_5_shoulder",
"m_6_lt_fr_head",
"m_7",
"m_8_hair_tip",
"m_9_frnt_face",
"m_10_rt_ear",
"m_11_outer_hair",
"m_12_lt_ear",
"m_13_lt_hair",
"m_14_rt_hair",
"m_15_neck",
]
seg_model_dict = {
"vox-5segments": seg5_segment_names,
"vox-10segments": seg10_segment_names,
"vox-15segments": seg_15_segment_names,
"vox": seg10_segment_names,
"vox-adv": seg10_segment_names,
}
face_parsing/face_parsing_loader.py
+37
View File
@@ -0,0 +1,37 @@
import torch
import numpy as np
import os
from .face_parsing_model import BiSeNet
def load_face_parser_model(base_dir: str, face_parser_checkpoint_name: str, cpu=False):
face_parser_path = f"{base_dir}/face_parsing/{face_parser_checkpoint_name}"
if not os.path.isfile(face_parser_path):
print(
f"{face_parser_checkpoint_name} not found, revert to no face parsing, not fatal"
)
return None
print("Loading face_parser model")
face_parser = BiSeNet(n_classes=19)
mean = torch.Tensor(np.array([0.485, 0.456, 0.406], dtype=np.float32)).view(
1, 3, 1, 1
)
std = torch.Tensor(np.array([0.229, 0.224, 0.225], dtype=np.float32)).view(
1, 3, 1, 1
)
if not cpu:
face_parser.cuda()
face_parser.load_state_dict(torch.load(face_parser_path))
face_parser.mean = mean.cuda()
face_parser.std = std.cuda()
else:
face_parser.load_state_dict(
torch.load(face_parser_path, map_location=torch.device("cpu"))
)
face_parser.mean = mean
face_parser.std = std
print("face_parser model loaded!")
return face_parser
face_parsing/face_parsing_model.py
+276
View File
@@ -0,0 +1,276 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
from .resnet import Resnet18
class ConvBNReLU(nn.Module):
def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, *args, **kwargs):
super(ConvBNReLU, self).__init__()
self.conv = nn.Conv2d(
in_chan,
out_chan,
kernel_size=ks,
stride=stride,
padding=padding,
bias=False,
)
self.bn = nn.BatchNorm2d(out_chan)
self.init_weight()
def forward(self, x):
x = self.conv(x)
x = F.relu(self.bn(x))
return x
def init_weight(self):
for ly in self.children():
if isinstance(ly, nn.Conv2d):
nn.init.kaiming_normal_(ly.weight, a=1)
if ly.bias is not None:
nn.init.constant_(ly.bias, 0)
class BiSeNetOutput(nn.Module):
def __init__(self, in_chan, mid_chan, n_classes, *args, **kwargs):
super(BiSeNetOutput, self).__init__()
self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1)
self.conv_out = nn.Conv2d(mid_chan, n_classes, kernel_size=1, bias=False)
self.init_weight()
def forward(self, x):
x = self.conv(x)
x = self.conv_out(x)
return x
def init_weight(self):
for ly in self.children():
if isinstance(ly, nn.Conv2d):
nn.init.kaiming_normal_(ly.weight, a=1)
if ly.bias is not None:
nn.init.constant_(ly.bias, 0)
def get_params(self):
wd_params, nowd_params = [], []
for name, module in self.named_modules():
if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
wd_params.append(module.weight)
if module.bias is not None:
nowd_params.append(module.bias)
elif isinstance(module, nn.BatchNorm2d):
nowd_params += list(module.parameters())
return wd_params, nowd_params
class AttentionRefinementModule(nn.Module):
def __init__(self, in_chan, out_chan, *args, **kwargs):
super(AttentionRefinementModule, self).__init__()
self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1)
self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size=1, bias=False)
self.bn_atten = nn.BatchNorm2d(out_chan)
self.sigmoid_atten = nn.Sigmoid()
self.init_weight()
def forward(self, x):
feat = self.conv(x)
atten = F.avg_pool2d(feat, feat.size()[2:])
atten = self.conv_atten(atten)
atten = self.bn_atten(atten)
atten = self.sigmoid_atten(atten)
out = torch.mul(feat, atten)
return out
def init_weight(self):
for ly in self.children():
if isinstance(ly, nn.Conv2d):
nn.init.kaiming_normal_(ly.weight, a=1)
if ly.bias is not None:
nn.init.constant_(ly.bias, 0)
class ContextPath(nn.Module):
def __init__(self, *args, **kwargs):
super(ContextPath, self).__init__()
self.resnet = Resnet18()
self.arm16 = AttentionRefinementModule(256, 128)
self.arm32 = AttentionRefinementModule(512, 128)
self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0)
self.init_weight()
def forward(self, x):
H0, W0 = x.size()[2:]
feat8, feat16, feat32 = self.resnet(x)
H8, W8 = feat8.size()[2:]
H16, W16 = feat16.size()[2:]
H32, W32 = feat32.size()[2:]
avg = F.avg_pool2d(feat32, feat32.size()[2:])
avg = self.conv_avg(avg)
avg_up = F.interpolate(avg, (H32, W32), mode="nearest")
feat32_arm = self.arm32(feat32)
feat32_sum = feat32_arm + avg_up
feat32_up = F.interpolate(feat32_sum, (H16, W16), mode="nearest")
feat32_up = self.conv_head32(feat32_up)
feat16_arm = self.arm16(feat16)
feat16_sum = feat16_arm + feat32_up
feat16_up = F.interpolate(feat16_sum, (H8, W8), mode="nearest")
feat16_up = self.conv_head16(feat16_up)
return feat8, feat16_up, feat32_up # x8, x8, x16
def init_weight(self):
for ly in self.children():
if isinstance(ly, nn.Conv2d):
nn.init.kaiming_normal_(ly.weight, a=1)
if ly.bias is not None:
nn.init.constant_(ly.bias, 0)
def get_params(self):
wd_params, nowd_params = [], []
for name, module in self.named_modules():
if isinstance(module, (nn.Linear, nn.Conv2d)):
wd_params.append(module.weight)
if module.bias is not None:
nowd_params.append(module.bias)
elif isinstance(module, nn.BatchNorm2d):
nowd_params += list(module.parameters())
return wd_params, nowd_params
### This is not used, since I replace this with the resnet feature with the same size
class SpatialPath(nn.Module):
def __init__(self, *args, **kwargs):
super(SpatialPath, self).__init__()
self.conv1 = ConvBNReLU(3, 64, ks=7, stride=2, padding=3)
self.conv2 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)
self.conv3 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)
self.conv_out = ConvBNReLU(64, 128, ks=1, stride=1, padding=0)
self.init_weight()
def forward(self, x):
feat = self.conv1(x)
feat = self.conv2(feat)
feat = self.conv3(feat)
feat = self.conv_out(feat)
return feat
def init_weight(self):
for ly in self.children():
if isinstance(ly, nn.Conv2d):
nn.init.kaiming_normal_(ly.weight, a=1)
if ly.bias is not None:
nn.init.constant_(ly.bias, 0)
def get_params(self):
wd_params, nowd_params = [], []
for name, module in self.named_modules():
if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
wd_params.append(module.weight)
if module.bias is not None:
nowd_params.append(module.bias)
elif isinstance(module, nn.BatchNorm2d):
nowd_params += list(module.parameters())
return wd_params, nowd_params
class FeatureFusionModule(nn.Module):
def __init__(self, in_chan, out_chan, *args, **kwargs):
super(FeatureFusionModule, self).__init__()
self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0)
self.conv1 = nn.Conv2d(
out_chan, out_chan // 4, kernel_size=1, stride=1, padding=0, bias=False
)
self.conv2 = nn.Conv2d(
out_chan // 4, out_chan, kernel_size=1, stride=1, padding=0, bias=False
)
self.relu = nn.ReLU(inplace=True)
self.sigmoid = nn.Sigmoid()
self.init_weight()
def forward(self, fsp, fcp):
fcat = torch.cat([fsp, fcp], dim=1)
feat = self.convblk(fcat)
atten = F.avg_pool2d(feat, feat.size()[2:])
atten = self.conv1(atten)
atten = self.relu(atten)
atten = self.conv2(atten)
atten = self.sigmoid(atten)
feat_atten = torch.mul(feat, atten)
feat_out = feat_atten + feat
return feat_out
def init_weight(self):
for ly in self.children():
if isinstance(ly, nn.Conv2d):
nn.init.kaiming_normal_(ly.weight, a=1)
if ly.bias is not None:
nn.init.constant_(ly.bias, 0)
def get_params(self):
wd_params, nowd_params = [], []
for name, module in self.named_modules():
if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
wd_params.append(module.weight)
if module.bias is not None:
nowd_params.append(module.bias)
elif isinstance(module, nn.BatchNorm2d):
nowd_params += list(module.parameters())
return wd_params, nowd_params
class BiSeNet(nn.Module):
def __init__(self, n_classes, *args, **kwargs):
super(BiSeNet, self).__init__()
self.cp = ContextPath()
## here self.sp is deleted
self.ffm = FeatureFusionModule(256, 256)
self.conv_out = BiSeNetOutput(256, 256, n_classes)
self.conv_out16 = BiSeNetOutput(128, 64, n_classes)
self.conv_out32 = BiSeNetOutput(128, 64, n_classes)
self.init_weight()
def forward(self, x):
H, W = x.size()[2:]
feat_res8, feat_cp8, feat_cp16 = self.cp(x) # here return res3b1 feature
feat_sp = feat_res8 # use res3b1 feature to replace spatial path feature
feat_fuse = self.ffm(feat_sp, feat_cp8)
feat_out = self.conv_out(feat_fuse)
feat_out16 = self.conv_out16(feat_cp8)
feat_out32 = self.conv_out32(feat_cp16)
feat_out = F.interpolate(feat_out, (H, W), mode="bilinear", align_corners=True)
feat_out16 = F.interpolate(
feat_out16, (H, W), mode="bilinear", align_corners=True
)
feat_out32 = F.interpolate(
feat_out32, (H, W), mode="bilinear", align_corners=True
)
return feat_out, feat_out16, feat_out32
def init_weight(self):
for ly in self.children():
if isinstance(ly, nn.Conv2d):
nn.init.kaiming_normal_(ly.weight, a=1)
if ly.bias is not None:
nn.init.constant_(ly.bias, 0)
def get_params(self):
wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], []
for name, child in self.named_children():
child_wd_params, child_nowd_params = child.get_params()
if isinstance(child, FeatureFusionModule) or isinstance(
child, BiSeNetOutput
):
lr_mul_wd_params += child_wd_params
lr_mul_nowd_params += child_nowd_params
else:
wd_params += child_wd_params
nowd_params += child_nowd_params
return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params
face_parsing/resnet.py
+94
View File
@@ -0,0 +1,94 @@
import torch.nn as nn
import torch.nn.functional as F
import torch
from pathlib import Path
def conv3x3(in_planes, out_planes, stride=1):
"""3x3 convolution with padding"""
return nn.Conv2d(
in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False
)
class BasicBlock(nn.Module):
def __init__(self, in_chan, out_chan, stride=1):
super(BasicBlock, self).__init__()
self.conv1 = conv3x3(in_chan, out_chan, stride)
self.bn1 = nn.BatchNorm2d(out_chan)
self.conv2 = conv3x3(out_chan, out_chan)
self.bn2 = nn.BatchNorm2d(out_chan)
self.relu = nn.ReLU(inplace=True)
self.downsample = None
if in_chan != out_chan or stride != 1:
self.downsample = nn.Sequential(
nn.Conv2d(in_chan, out_chan, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(out_chan),
)
def forward(self, x):
residual = self.conv1(x)
residual = F.relu(self.bn1(residual))
residual = self.conv2(residual)
residual = self.bn2(residual)
shortcut = x
if self.downsample is not None:
shortcut = self.downsample(x)
out = shortcut + residual
out = self.relu(out)
return out
def create_layer_basic(in_chan, out_chan, bnum, stride=1):
layers = [BasicBlock(in_chan, out_chan, stride=stride)]
for i in range(bnum - 1):
layers.append(BasicBlock(out_chan, out_chan, stride=1))
return nn.Sequential(*layers)
class Resnet18(nn.Module):
def __init__(self):
super(Resnet18, self).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = create_layer_basic(64, 64, bnum=2, stride=1)
self.layer2 = create_layer_basic(64, 128, bnum=2, stride=2)
self.layer3 = create_layer_basic(128, 256, bnum=2, stride=2)
self.layer4 = create_layer_basic(256, 512, bnum=2, stride=2)
self.init_weight()
def forward(self, x):
x = self.conv1(x)
x = F.relu(self.bn1(x))
x = self.maxpool(x)
x = self.layer1(x)
feat8 = self.layer2(x) # 1/8
feat16 = self.layer3(feat8) # 1/16
feat32 = self.layer4(feat16) # 1/32
return feat8, feat16, feat32
def init_weight(self):
base_dir = Path(__file__).resolve().parent
print(f"resnet base dir {base_dir}")
state_dict = torch.load(f"{base_dir}/resnet18-5c106cde.pth")
self_state_dict = self.state_dict()
for k, v in state_dict.items():
if "fc" in k:
continue
self_state_dict.update({k: v})
self.load_state_dict(self_state_dict)
def get_params(self):
wd_params, nowd_params = [], []
for name, module in self.named_modules():
if isinstance(module, (nn.Linear, nn.Conv2d)):
wd_params.append(module.weight)
if module.bias is not None:
nowd_params.append(module.bias)
elif isinstance(module, nn.BatchNorm2d):
nowd_params += list(module.parameters())
return wd_params, nowd_params
fomm_inference.py
+185
View File
@@ -0,0 +1,185 @@
import torch
import numpy as np
from comfy.utils import ProgressBar
from tqdm import tqdm
from scipy.spatial import ConvexHull
from .modules.generator import OcclusionAwareGenerator
from .modules.keypoint_detector import KPDetector
from .utils import tensor_to_numpy
try:
import face_alignment
face_alignment_available = True
except ImportError:
face_alignment_available = False
print("Warning: face_alignment module not found. fine_best_frame will not work. ")
def normalize_kp_best_frame(kp):
kp = kp - kp.mean(axis=0, keepdims=True)
area = ConvexHull(kp[:, :2]).volume
area = np.sqrt(area)
kp[:, :2] = kp[:, :2] / area
return kp
def normalize_kp(
kp_source,
kp_driving,
kp_driving_initial,
adapt_movement_scale: bool,
use_relative_movement: bool,
use_relative_jacobian: bool,
):
if adapt_movement_scale:
source_area = ConvexHull(kp_source["value"][0].data.cpu().numpy()).volume
driving_area = ConvexHull(
kp_driving_initial["value"][0].data.cpu().numpy()
).volume
adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area)
else:
adapt_movement_scale = 1
kp_new = {k: v for k, v in kp_driving.items()}
if use_relative_movement:
kp_value_diff = kp_driving["value"] - kp_driving_initial["value"]
kp_value_diff *= adapt_movement_scale
kp_new["value"] = kp_value_diff + kp_source["value"]
if use_relative_jacobian:
jacobian_diff = torch.matmul(
kp_driving["jacobian"], torch.inverse(kp_driving_initial["jacobian"])
)
kp_new["jacobian"] = torch.matmul(jacobian_diff, kp_source["jacobian"])
return kp_new
def find_best_frame(source_image, driving_video: list, cpu=False) -> int:
if not face_alignment_available:
print("Face Alignment module were not found. Best frame set to 0.")
return 0
print("Initializing Face Alignment...")
fa = face_alignment.FaceAlignment(
face_alignment.LandmarksType.TWO_D,
flip_input=True,
device="cpu" if cpu else "cuda",
)
print("Face alignment model loaded.")
source_np = tensor_to_numpy(source_image)
kp_source = fa.get_landmarks(255 * source_np)[0]
kp_source = normalize_kp_best_frame(kp_source)
norm = float("inf")
frame_num = 0
num_frames = driving_video.shape[2]
for i in tqdm(range(num_frames)):
# Extract single frame from driving video
frame = driving_video[0, :, i] # Shape: [3, 256, 256]
frame_np = frame.permute(1, 2, 0).cpu().numpy()
kp_driving = fa.get_landmarks(255 * frame_np)[0]
kp_driving = normalize_kp_best_frame(kp_driving)
new_norm = (np.abs(kp_source - kp_driving) ** 2).sum()
if new_norm < norm:
norm = new_norm
frame_num = i
print(f"Best frame: {frame_num}")
return frame_num
def inference(
source_image,
driving_video: list,
generator: OcclusionAwareGenerator,
kp_detector: KPDetector,
relative_movement: bool,
relative_jacobian: bool,
adapt_movement_scale: bool,
cpu=False,
) -> list:
with torch.no_grad():
source = source_image
driving = driving_video
predictions = []
if not cpu:
source = source.cuda()
driving = driving.cuda()
print("Using GPU")
kp_source = kp_detector(source)
kp_driving_initial = kp_detector(driving[:, :, 0])
kp_norm = None
num_frames = driving.shape[2]
pbar = ProgressBar(num_frames)
for frame_idx in tqdm(range(num_frames)):
driving_frame = driving[:, :, frame_idx]
kp_driving = kp_detector(driving_frame)
kp_norm = normalize_kp(
kp_source=kp_source,
kp_driving=kp_driving,
kp_driving_initial=kp_driving_initial,
use_relative_movement=relative_movement,
use_relative_jacobian=relative_jacobian,
adapt_movement_scale=adapt_movement_scale,
)
out = generator(source, kp_source=kp_source, kp_driving=kp_norm)
out_np = np.transpose(out["prediction"].data.cpu().numpy(), [0, 2, 3, 1])[0]
predictions.append(out_np)
pbar.update_absolute(frame_idx, num_frames)
return predictions
def inference_best_frame(
source_image,
driving_video,
generator: OcclusionAwareGenerator,
kp_detector: KPDetector,
relative_movement: bool,
relative_jacobian: bool,
adapt_movement_scale: bool,
cpu=False,
) -> list:
best_frame_idx = find_best_frame(source_image, driving_video, cpu)
first_half = driving_video[
:, :, : best_frame_idx + 1
] # Include the best frame in the first half
first_half = torch.flip(first_half, dims=[2]) # Reverse the first half
second_half = driving_video[:, :, best_frame_idx + 1 :]
predictions_first = inference(
source_image,
first_half,
generator,
kp_detector,
relative_movement,
relative_jacobian,
adapt_movement_scale,
cpu,
)
predictions_second = inference(
source_image,
second_half,
generator,
kp_detector,
relative_movement,
relative_jacobian,
adapt_movement_scale,
cpu,
)
predictions = []
predictions_first = predictions_first[::-1] # Reverse the first half back
predictions.extend(predictions_first)
predictions.extend(predictions_second)
return predictions
fomm_loader.py
+35
View File
@@ -0,0 +1,35 @@
import yaml
import torch
from .modules.generator import OcclusionAwareGenerator
from .modules.keypoint_detector import KPDetector
def load_checkpoint(
config_path: str, checkpoint_path: str, cpu=False
) -> tuple[OcclusionAwareGenerator, KPDetector]:
with open(config_path) as f:
config = yaml.full_load(f)
generator = OcclusionAwareGenerator(
**config["model_params"]["generator_params"],
**config["model_params"]["common_params"],
)
kp_detector = KPDetector(
**config["model_params"]["kp_detector_params"],
**config["model_params"]["common_params"],
)
if not cpu:
generator.cuda()
kp_detector.cuda()
checkpoint = torch.load(checkpoint_path)
else:
checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
generator.load_state_dict(checkpoint["generator"])
kp_detector.load_state_dict(checkpoint["kp_detector"])
generator.eval()
kp_detector.eval()
return generator, kp_detector
modules/dense_motion.py
+42 -9
View File
@@ -29,13 +29,13 @@ class DenseMotionNetwork(nn.Module):
if self.scale_factor != 1:
self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor)
def create_heatmap_representations(self, source_image, kp_driving, kp_source):
def create_heatmap_representations(self, source_image, kp_driving, kp_source, kp_key='value'):
"""
Eq 6. in the paper H_k(z)
"""
spatial_size = source_image.shape[2:]
gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=self.kp_variance)
gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=self.kp_variance)
gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=self.kp_variance, kp_key=kp_key)
gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=self.kp_variance, kp_key=kp_key)
heatmap = gaussian_driving - gaussian_source
#adding background feature
@@ -44,22 +44,22 @@ class DenseMotionNetwork(nn.Module):
heatmap = heatmap.unsqueeze(2)
return heatmap
def create_sparse_motions(self, source_image, kp_driving, kp_source):
def create_sparse_motions(self, source_image, kp_driving, kp_source, kp_key='value', jacob_key='jacobian'):
"""
Eq 4. in the paper T_{s<-d}(z)
"""
bs, _, h, w = source_image.shape
identity_grid = make_coordinate_grid((h, w), type=kp_source['value'].type())
identity_grid = make_coordinate_grid((h, w), type=kp_source[kp_key].type())
identity_grid = identity_grid.view(1, 1, h, w, 2)
coordinate_grid = identity_grid - kp_driving['value'].view(bs, self.num_kp, 1, 1, 2)
if 'jacobian' in kp_driving:
jacobian = torch.matmul(kp_source['jacobian'], torch.inverse(kp_driving['jacobian']))
coordinate_grid = identity_grid - kp_driving[kp_key].view(bs, self.num_kp, 1, 1, 2)
if jacob_key in kp_driving:
jacobian = torch.matmul(kp_source[jacob_key], torch.inverse(kp_driving[jacob_key]))
jacobian = jacobian.unsqueeze(-3).unsqueeze(-3)
jacobian = jacobian.repeat(1, 1, h, w, 1, 1)
coordinate_grid = torch.matmul(jacobian, coordinate_grid.unsqueeze(-1))
coordinate_grid = coordinate_grid.squeeze(-1)
driving_to_source = coordinate_grid + kp_source['value'].view(bs, self.num_kp, 1, 1, 2)
driving_to_source = coordinate_grid + kp_source[kp_key].view(bs, self.num_kp, 1, 1, 2)
#adding background feature
identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1)
@@ -111,3 +111,36 @@ class DenseMotionNetwork(nn.Module):
out_dict['occlusion_map'] = occlusion_map
return out_dict
def forward_partswap(self, source_image, seg_target, seg_source):
if self.scale_factor != 1:
source_image = self.down(source_image)
bs, _, h, w = source_image.shape
out_dict = dict()
heatmap_representation = self.create_heatmap_representations(source_image, seg_target, seg_source, kp_key='shift')
sparse_motion = self.create_sparse_motions(source_image, seg_target, seg_source, kp_key='shift', jacob_key='affine')
deformed_source = self.create_deformed_source_image(source_image, sparse_motion)
out_dict['sparse_deformed'] = deformed_source
input = torch.cat([heatmap_representation, deformed_source], dim=2)
input = input.view(bs, -1, h, w)
prediction = self.hourglass(input)
mask = self.mask(prediction)
mask = F.softmax(mask, dim=1)
out_dict['mask'] = mask
mask = mask.unsqueeze(2)
sparse_motion = sparse_motion.permute(0, 1, 4, 2, 3)
deformation = (sparse_motion * mask).sum(dim=1)
deformation = deformation.permute(0, 2, 3, 1)
out_dict['deformation'] = deformation
if self.occlusion:
visibility = torch.sigmoid(self.occlusion(prediction))
out_dict['visibility'] = visibility
return out_dict
modules/reconstruction_module.py
+109
View File
@@ -0,0 +1,109 @@
import torch
from torch import nn
import torch.nn.functional as F
from .util import ResBlock2d, SameBlock2d, UpBlock2d, DownBlock2d
from .util import make_coordinate_grid
class ReconstructionModule(nn.Module):
"""
Reconstruct target from source and segmentation of target and part transformations
"""
def __init__(self, num_channels, block_expansion, max_features, num_down_blocks,
num_bottleneck_blocks, num_segments, estimate_visibility=False, **kwargs):
super(ReconstructionModule, self).__init__()
self.first = SameBlock2d(num_channels, block_expansion, kernel_size=(7, 7), padding=(3, 3))
down_blocks = []
for i in range(num_down_blocks):
in_features = min(max_features, block_expansion * (2 ** i))
out_features = min(max_features, block_expansion * (2 ** (i + 1)))
down_blocks.append(DownBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1)))
self.down_blocks = nn.ModuleList(down_blocks)
up_blocks = []
for i in range(num_down_blocks):
in_features = min(max_features, block_expansion * (2 ** (num_down_blocks - i)))
out_features = min(max_features, block_expansion * (2 ** (num_down_blocks - i - 1)))
up_blocks.append(UpBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1)))
self.up_blocks = nn.ModuleList(up_blocks)
self.bottleneck = torch.nn.Sequential()
in_features = min(max_features, block_expansion * (2 ** num_down_blocks))
for i in range(num_bottleneck_blocks):
self.bottleneck.add_module('r' + str(i), ResBlock2d(in_features, kernel_size=(3, 3), padding=(1, 1)))
self.final = nn.Conv2d(block_expansion, num_channels, kernel_size=(7, 7), padding=(3, 3))
self.estimate_visibility = estimate_visibility
self.num_channels = num_channels
self.num_segments = num_segments
def segment_motion(self, seg_target, seg_source):
bs, _, h, w = seg_target['segmentation'].shape
identity_grid = make_coordinate_grid((h, w), type=seg_source['shift'].type())
identity_grid = identity_grid.view(1, 1, h, w, 2)
coordinate_grid = identity_grid - seg_target['shift'].view(bs, self.num_segments, 1, 1, 2)
if 'affine' in seg_target:
affine = torch.matmul(seg_source['affine'], torch.inverse(seg_target['affine']))
affine = affine.unsqueeze(-3).unsqueeze(-3)
affine = affine.repeat(1, 1, h, w, 1, 1)
coordinate_grid = torch.matmul(affine, coordinate_grid.unsqueeze(-1))
coordinate_grid = coordinate_grid.squeeze(-1)
target_to_source = coordinate_grid + seg_source['shift'].view(bs, self.num_segments, 1, 1, 2)
# adding background feature
identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1)
return torch.cat([identity_grid, target_to_source], dim=1)
def deform_input(self, inp, deformation):
_, h_old, w_old, _ = deformation.shape
_, _, h, w = inp.shape
if h_old != h or w_old != w:
deformation = deformation.permute(0, 3, 1, 2)
deformation = F.interpolate(deformation, size=(h, w), mode='bilinear')
deformation = deformation.permute(0, 2, 3, 1)
return F.grid_sample(inp, deformation)
def forward(self, source_image, seg_target, seg_source):
# Encoding (downsampling) part
out = self.first(source_image)
for i in range(len(self.down_blocks)):
out = self.down_blocks[i](out)
output_dict = {}
# Computing segment motion
segment_motions = self.segment_motion(seg_target, seg_source)
segment_motions = segment_motions.permute(0, 1, 4, 2, 3)
mask = seg_target['segmentation'].unsqueeze(2)
deformation = (segment_motions * mask).sum(dim=1)
deformation = deformation.permute(0, 2, 3, 1)
# Deform according to motion
out = self.deform_input(out, deformation)
output_dict["deformed"] = self.deform_input(source_image, deformation)
if self.estimate_visibility:
visibility = seg_source['segmentation'][:, 1:].sum(dim=1, keepdim=True) * \
(1 - seg_target['segmentation'][:, 1:].sum(dim=1, keepdim=True).detach())
visibility = 1 - visibility
if out.shape[2] != visibility.shape[2] or out.shape[3] != visibility.shape[3]:
visibility = F.interpolate(visibility, size=out.shape[2:], mode='bilinear')
out = out * visibility
output_dict['visibility_maps'] = visibility
out = self.bottleneck(out)
for i in range(len(self.up_blocks)):
out = self.up_blocks[i](out)
out = self.final(out)
out = F.sigmoid(out)
output_dict["prediction"] = out
return output_dict
modules/segmentation_module.py
+78
View File
@@ -0,0 +1,78 @@
from torch import nn
import torch
import torch.nn.functional as F
from .util import Hourglass, make_coordinate_grid, AntiAliasInterpolation2d
class SegmentationModule(nn.Module):
"""
Computing a segmentation map and affine transformations.
"""
def __init__(self, block_expansion, num_segments, num_channels, max_features,
num_blocks, temperature, estimate_affine_part=False, scale_factor=1):
super(SegmentationModule, self).__init__()
self.predictor = Hourglass(block_expansion, in_features=num_channels,
max_features=max_features, num_blocks=num_blocks)
self.num_segments = num_segments
self.shift = nn.Conv2d(in_channels=self.predictor.out_filters, out_channels=num_segments, kernel_size=(7, 7),
padding=(3, 3))
if estimate_affine_part:
self.affine = nn.Conv2d(in_channels=self.predictor.out_filters,
out_channels=4 * num_segments, kernel_size=(7, 7), padding=(3, 3))
self.affine.weight.data.zero_()
self.affine.bias.data.copy_(torch.tensor([1, 0, 0, 1] * num_segments, dtype=torch.float))
else:
self.affine = None
self.segmentation = nn.Conv2d(in_channels=self.predictor.out_filters,
out_channels=(1 + num_segments), kernel_size=(7, 7), padding=(3, 3))
self.temperature = temperature
self.scale_factor = scale_factor
if self.scale_factor != 1:
self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor)
def gaussian2kp(self, heatmap):
"""
Extract the mean and the variance from a heatmap
"""
shape = heatmap.shape
heatmap = heatmap.unsqueeze(-1)
grid = make_coordinate_grid(shape[2:], heatmap.type()).unsqueeze_(0).unsqueeze_(0)
result = (heatmap * grid).sum(dim=(2, 3))
return result
def forward(self, x):
if self.scale_factor != 1:
x = self.down(x)
out = {}
feature_map = self.predictor(x)
out['segmentation'] = F.softmax(self.segmentation(feature_map), dim=1)
prediction = self.shift(feature_map)
final_shape = prediction.shape
heatmap = prediction.view(final_shape[0], final_shape[1], -1)
heatmap = F.softmax(heatmap / self.temperature, dim=2)
heatmap = heatmap.view(*final_shape)
out['shift'] = self.gaussian2kp(heatmap)
if self.affine is not None:
affine_map = self.affine(feature_map)
affine_map = affine_map.reshape(final_shape[0], self.num_segments, 4, final_shape[2],
final_shape[3])
heatmap = heatmap.unsqueeze(2)
affine = heatmap * affine_map
affine = affine.view(final_shape[0], final_shape[1], 4, -1)
affine = affine.sum(dim=-1)
affine = affine.view(affine.shape[0], affine.shape[1], 2, 2)
out['affine'] = affine
return out
modules/util.py
+2 -2
View File
@@ -6,11 +6,11 @@ import torch
from ..sync_batchnorm.batchnorm import SynchronizedBatchNorm2d as BatchNorm2d
def kp2gaussian(kp, spatial_size, kp_variance):
def kp2gaussian(kp, spatial_size, kp_variance, kp_key='value'):
"""
Transform a keypoint into gaussian like representation
"""
mean = kp['value']
mean = kp[kp_key]
coordinate_grid = make_coordinate_grid(spatial_size, mean.type())
number_of_leading_dimensions = len(mean.shape) - 1
partswap_inference.py
+84
View File
@@ -0,0 +1,84 @@
import numpy as np
import torch
import torch.nn.functional as F
from tqdm import tqdm
from comfy.utils import ProgressBar
from .modules.segmentation_module import SegmentationModule
from .partswap_model import PartSwapGenerator
def face_parse_seg(
source_image, face_parser_model, in_size=(512, 512), out_size=(64, 64)
):
seg = F.interpolate(source_image, size=in_size)
seg = (seg - face_parser_model.mean) / face_parser_model.std
seg = torch.softmax(face_parser_model(seg)[0], dim=1)
seg = F.interpolate(seg, size=out_size)
return seg
def partswap_inference(
swap_indices: list[int],
source_image,
target_video,
reconstruction_module: PartSwapGenerator,
segmentation_module: SegmentationModule,
use_source_seg: bool,
face_parser_model=None,
hard_edges=False,
cpu=False,
) -> list:
with torch.no_grad():
predictions = []
seg_targets = []
if not cpu:
source_image = source_image.cuda()
target_video = target_video.cuda()
print("Using GPU")
seg_source = segmentation_module(source_image)
print(f"{seg_source['segmentation'].shape}")
num_frames = target_video.shape[2]
pbar = ProgressBar(num_frames)
for frame_idx in tqdm(range(num_frames)):
target_frame = target_video[:, :, frame_idx]
seg_target = segmentation_module(target_frame)
seg_targets.append(seg_target)
if face_parser_model is not None:
blend_mask = face_parse_seg(
source_image if use_source_seg else target_frame,
face_parser_model,
(512, 512),
(64, 64),
)
else:
blend_mask = (
seg_source["segmentation"]
if use_source_seg
else seg_target["segmentation"]
)
blend_mask = blend_mask[:, swap_indices].sum(dim=1, keepdim=True)
if hard_edges:
blend_mask = (blend_mask > 0.5).type(blend_mask.type())
print(f"{blend_mask.shape=}")
out = reconstruction_module(
source_image,
target_frame,
seg_source=seg_source,
seg_target=seg_target,
blend_mask=blend_mask,
use_source_segmentation=use_source_seg,
)
predictions.append(
np.transpose(out["prediction"].data.cpu().numpy(), [0, 2, 3, 1])[0]
)
pbar.update_absolute(frame_idx, num_frames)
return seg_source, seg_targets, predictions
partswap_loader.py
+92
View File
@@ -0,0 +1,92 @@
import torch
import yaml
from .modules.segmentation_module import SegmentationModule
from .partswap_model import PartSwapGenerator
from .sync_batchnorm.replicate import DataParallelWithCallback
def partial_state_dict_load(module, state_dict):
own_state = module.state_dict()
for name, param in state_dict.items():
if name not in own_state:
continue
if isinstance(param, torch.nn.Parameter):
# backwards compatibility for serialized parameters
param = param.data
own_state[name].copy_(param)
def load_reconstruction_module(module, checkpoint):
if "generator" in checkpoint:
partial_state_dict_load(module, checkpoint["generator"])
else:
module.load_state_dict(checkpoint["reconstruction_module"])
def load_segmentation_module(module, checkpoint):
if "kp_detector" in checkpoint:
partial_state_dict_load(module, checkpoint["kp_detector"])
module.state_dict()["affine.weight"].copy_(
checkpoint["kp_detector"]["jacobian.weight"]
)
module.state_dict()["affine.bias"].copy_(
checkpoint["kp_detector"]["jacobian.bias"]
)
module.state_dict()["shift.weight"].copy_(
checkpoint["kp_detector"]["kp.weight"]
)
module.state_dict()["shift.bias"].copy_(checkpoint["kp_detector"]["kp.bias"])
if "semantic_seg.weight" in checkpoint["kp_detector"]:
module.state_dict()["segmentation.weight"].copy_(
checkpoint["kp_detector"]["semantic_seg.weight"]
)
module.state_dict()["segmentation.bias"].copy_(
checkpoint["kp_detector"]["semantic_seg.bias"]
)
else:
print("Segmentation part initialized at random.")
else:
module.load_state_dict(checkpoint["segmentation_module"])
def load_partswap_checkpoint(
config_path: str,
checkpoint_path: str,
blend_scale: float,
use_fomm: bool,
cpu=False,
):
with open(config_path) as f:
config = yaml.full_load(f)
reconstruction_module = PartSwapGenerator(
blend_scale=blend_scale,
first_order_motion_model=use_fomm,
**config["model_params"]["reconstruction_module_params"],
**config["model_params"]["common_params"],
)
segmentation_module = SegmentationModule(
**config["model_params"]["segmentation_module_params"],
**config["model_params"]["common_params"],
)
if cpu:
checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
else:
segmentation_module.cuda()
checkpoint = torch.load(checkpoint_path)
load_reconstruction_module(reconstruction_module, checkpoint)
load_segmentation_module(segmentation_module, checkpoint)
if not cpu:
reconstruction_module = DataParallelWithCallback(reconstruction_module)
segmentation_module = DataParallelWithCallback(segmentation_module)
reconstruction_module.eval()
segmentation_module.eval()
return reconstruction_module, segmentation_module
partswap_model.py
+105
View File
@@ -0,0 +1,105 @@
import torch.nn.functional as F
from .modules.dense_motion import DenseMotionNetwork
from .modules.reconstruction_module import ReconstructionModule
from .modules.util import AntiAliasInterpolation2d
class PartSwapGenerator(ReconstructionModule):
def __init__(self, blend_scale=1, first_order_motion_model=False, **kwargs):
super(PartSwapGenerator, self).__init__(**kwargs)
if blend_scale == 1:
self.blend_downsample = lambda x: x
else:
self.blend_downsample = AntiAliasInterpolation2d(1, blend_scale)
if first_order_motion_model:
self.dense_motion_network = DenseMotionNetwork(
block_expansion=64,
num_blocks=5,
max_features=1024,
num_kp=10,
num_channels=3,
estimate_occlusion_map=True,
scale_factor=0.25,
)
else:
self.dense_motion_network = None
def forward(
self,
source_image,
target_image,
seg_target,
seg_source,
blend_mask,
use_source_segmentation=False,
):
# Encoding of source image
enc_source = self.first(source_image)
for i in range(len(self.down_blocks)):
enc_source = self.down_blocks[i](enc_source)
# Encoding of target image
enc_target = self.first(target_image)
for i in range(len(self.down_blocks)):
enc_target = self.down_blocks[i](enc_target)
output_dict = {}
# Compute flow field for source image
if self.dense_motion_network is None:
segment_motions = self.segment_motion(seg_target, seg_source)
segment_motions = segment_motions.permute(0, 1, 4, 2, 3)
mask = seg_target["segmentation"].unsqueeze(2)
deformation = (segment_motions * mask).sum(dim=1)
deformation = deformation.permute(0, 2, 3, 1)
else:
motion = self.dense_motion_network.forward_partswap(
source_image=source_image, seg_target=seg_target, seg_source=seg_source
)
deformation = motion["deformation"]
# Deform source encoding according to the motion
enc_source = self.deform_input(enc_source, deformation)
if self.estimate_visibility:
if self.dense_motion_network is None:
visibility = seg_source["segmentation"][:, 1:].sum(
dim=1, keepdim=True
) * (
1
- seg_target["segmentation"][:, 1:]
.sum(dim=1, keepdim=True)
.detach()
)
visibility = 1 - visibility
else:
visibility = motion["visibility"]
if (
enc_source.shape[2] != visibility.shape[2]
or enc_source.shape[3] != visibility.shape[3]
):
visibility = F.interpolate(
visibility, size=enc_source.shape[2:], mode="bilinear"
)
enc_source = enc_source * visibility
blend_mask = self.blend_downsample(blend_mask)
# If source segmentation is provided use it should be deformed before blending
if use_source_segmentation:
blend_mask = self.deform_input(blend_mask, deformation)
out = enc_target * (1 - blend_mask) + enc_source * blend_mask
out = self.bottleneck(out)
for i in range(len(self.up_blocks)):
out = self.up_blocks[i](out)
out = self.final(out)
out = F.sigmoid(out)
output_dict["prediction"] = out
return output_dict
requirements.txt
+1
View File
@@ -2,3 +2,4 @@ numpy
torch
scipy
yaml
matplotlib
run.py
+230 -12
View File
@@ -9,13 +9,28 @@ from pathlib import Path
import torch
from .constants import default_model_name, supported_model_names
from .constants import (
config_folder,
default_model_name,
default_partswap_model_name,
face_parser_checkpoint_name,
partswap_fomm_model_names,
partswap_model_config_dict,
partswap_model_length_dict,
partswap_model_names,
supported_model_names,
)
from .face_parsing.face_parsing_loader import load_face_parser_model
from .fomm_inference import inference, inference_best_frame
from .fomm_loader import load_checkpoint
from .partswap_inference import partswap_inference
from .partswap_loader import load_partswap_checkpoint
from .seg_viz import visualize_frame
from .utils import (
build_seg_arguments,
get_config_path,
get_model_file_name,
inference,
inference_best_frame,
load_checkpoint,
get_partswap_model_file_name,
out_video,
reshape_image,
)
@@ -67,8 +82,8 @@ class FOMM_Runner:
def todo(
self,
source_image: torch.Tensor,
driving_video_input: torch.Tensor,
source_image,
driving_video_input,
model_name: str,
frame_rate: float,
relative_movement: bool,
@@ -119,10 +134,213 @@ class FOMM_Runner:
)
NODE_CLASS_MAPPINGS = {
"FOMM_Runner": FOMM_Runner,
}
class FOMM_Seg5Chooser:
@classmethod
def INPUT_TYPES(s):
return build_seg_arguments("vox-5segments")
NODE_DISPLAY_NAME_MAPPINGS = {
"FOMM_Runner": "FOMM Runner",
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("chosen_seg_indices",)
FUNCTION = "todo"
CATEGORY = "FirstOrderMM"
def todo(self, **args):
print(args)
assert len(args) == 6, f"Expected 6 arguments, got {len(args)}"
segments = list(args.values())
seg_list = [i for i, seg in enumerate(segments) if seg]
seg_list = serialize_integers(seg_list)
return (seg_list,)
class FOMM_Seg10Chooser:
@classmethod
def INPUT_TYPES(s):
return build_seg_arguments("vox-10segments")
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("chosen_seg_indices",)
FUNCTION = "todo"
CATEGORY = "FirstOrderMM"
def todo(self, **args):
print(args)
assert len(args) == 11, f"Expected 11 arguments, got {len(args)}"
segments = list(args.values())
seg_list = [i for i, seg in enumerate(segments) if seg]
seg_list = serialize_integers(seg_list)
return (seg_list,)
class FOMM_Seg15Chooser:
@classmethod
def INPUT_TYPES(s):
return build_seg_arguments("vox-15segments")
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("chosen_seg_indices",)
FUNCTION = "todo"
CATEGORY = "FirstOrderMM"
def todo(self, **args):
print(args)
assert len(args) == 16, f"Expected 16 arguments, got {len(args)}"
segments = list(args.values())
seg_list = [i for i, seg in enumerate(segments) if seg]
seg_list = serialize_integers(seg_list)
return (seg_list,)
class FOMM_Partswap:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"source_image": ("IMAGE",),
"driving_video_input": ("IMAGE",),
"model_name": (
partswap_model_names,
{"default": default_partswap_model_name},
),
"frame_rate": ("FLOAT", {"default": 30.0}),
"blend_scale": (
"FLOAT",
{"default": 1.0, "min": 0.6, "max": 1.0, "step": 0.01},
),
"use_source_seg": (
"BOOLEAN",
{"default": True},
),
"hard_edges": (
"BOOLEAN",
{"default": False},
),
"use_face_parser": (
"BOOLEAN",
{"default": False},
),
"chosen_seg_indices": ("STRING", {}),
"viz_alpha": (
"FLOAT",
{"default": 0.6, "min": 0.0, "max": 1.0, "step": 0.1},
),
},
"optional": {"audio": ("VHS_AUDIO",)},
}
RETURN_TYPES = (
"IMAGE",
"IMAGE",
"IMAGE",
"VHS_AUDIO",
"FLOAT",
)
RETURN_NAMES = (
"seg_src_viz",
"seg_tgt_viz",
"images",
"audio",
"frame_rate",
)
FUNCTION = "todo"
CATEGORY = "FirstOrderMM"
def todo(
self,
source_image: torch.Tensor,
driving_video_input: torch.Tensor,
model_name: str,
frame_rate: float,
blend_scale: float,
use_source_seg: bool,
hard_edges: bool,
use_face_parser: bool,
chosen_seg_indices: str,
viz_alpha: float,
audio=None,
):
print(f"{type(source_image)=}")
print(f"{type(driving_video_input)=}")
print(f"{source_image.shape=}")
print(f"{driving_video_input.shape=}")
print(f"{type(audio)=}")
print(base_dir)
config_path = (
f"{base_dir}/{config_folder}/{partswap_model_config_dict[model_name]}"
)
checkpoint_path = f"{base_dir}/{get_partswap_model_file_name(model_name)}"
source_image = reshape_image(source_image, (256, 256))
driving_video_org = reshape_image(driving_video_input, (256, 256))
driving_video = driving_video_org.unsqueeze(0).permute(0, 2, 1, 3, 4)
use_fomm = False
if model_name in partswap_fomm_model_names:
use_fomm = True
reconstruction_module, segmentation_module = load_partswap_checkpoint(
config_path, checkpoint_path, blend_scale, use_fomm=use_fomm
)
face_parser_model = None
# if model_name in partswap_face_parser_enforce_models:
# assert use_face_parser, "Special supervised model, requires face_parser"
if use_face_parser:
face_parser_model = load_face_parser_model(
base_dir, face_parser_checkpoint_name
)
if face_parser_model is not None:
face_parser_model.cuda()
face_parser_model.eval()
print(f"Chosen {chosen_seg_indices}")
swap_indices = deserialize_integers(chosen_seg_indices)
swap_indices = [
x for x in swap_indices if x < partswap_model_length_dict[model_name]
]
print(f"Swapping {swap_indices}")
params = {
"swap_indices": swap_indices,
"source_image": source_image,
"target_video": driving_video,
"reconstruction_module": reconstruction_module,
"segmentation_module": segmentation_module,
"use_source_seg": use_source_seg,
"face_parser_model": face_parser_model,
"hard_edges": hard_edges,
}
seg_source, seg_targets, predictions = partswap_inference(**params)
seg_src_viz = visualize_frame(source_image, seg_source, alpha=viz_alpha)
seg_tgt_viz = visualize_frame(
driving_video_org[0], seg_targets[0], alpha=viz_alpha
)
output_images = out_video(predictions)
del face_parser_model
return (
seg_src_viz,
seg_tgt_viz,
output_images,
audio,
frame_rate,
)
def serialize_integers(int_list):
return "_".join(map(str, int_list))
def deserialize_integers(int_string):
return list(map(int, int_string.split("_")))
samples/sample_output.mp4
BIN
View File
Binary file not shown.
seg_viz.py
+88
View File
@@ -0,0 +1,88 @@
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg
from matplotlib.figure import Figure
import numpy as np
import torch
import torch.nn.functional as F
def preprocess_image_and_mask(image, seg_source):
image = image.squeeze(0).cpu() # Remove batch dimension
mask = seg_source["segmentation"]
if mask.shape[2:] != image.shape[1:]:
mask = F.interpolate(mask, size=image.shape[1:], mode="bilinear")
return image, mask
def apply_hard_edges(mask):
return (torch.max(mask, dim=1, keepdim=True)[0] == mask).float()
def create_color_mask(mask, colormap="gist_rainbow"):
colormap = plt.get_cmap(colormap)
num_segments = mask.shape[1]
mask = mask.squeeze(0).permute(1, 2, 0).cpu().numpy()
color_mask = np.zeros((256, 256, 3), dtype=np.float32)
patches = []
for i in range(num_segments):
color = (
np.array(colormap((i - 1) / (num_segments - 1)))[:3]
if i != 0
else np.array((0, 0, 0))
)
patches.append(mpatches.Patch(color=color, label=str(i)))
color_mask += mask[..., i : (i + 1)] * color.reshape(1, 1, 3)
return color_mask, patches
def overlay_mask(image, color_mask, alpha=0.6):
image_np = image.permute(1, 2, 0).numpy()
if image_np.max() > 1:
image_np = image_np / 255.0
overlaid_image = image_np * (1 - alpha) + color_mask * alpha
return np.clip(overlaid_image, 0, 1)
def create_figure(overlaid_image, patches):
fig = Figure(figsize=(10, 10))
canvas = FigureCanvasAgg(fig)
ax = fig.add_subplot(111)
ax.imshow(overlaid_image)
ax.axis("off")
ax.legend(
handles=patches, loc="upper right", bbox_to_anchor=(1, 1), fontsize="small"
)
fig.tight_layout(pad=0)
return fig, canvas
def render_figure_to_tensor(canvas):
canvas.draw()
result_image = np.frombuffer(canvas.tostring_rgb(), dtype="uint8")
result_image = result_image.reshape(canvas.get_width_height()[::-1] + (3,))
result_tensor = torch.from_numpy(result_image).unsqueeze(0).float() / 255.0
return result_tensor
def visualize_frame(
image, seg_source, hard_edges=True, colormap="gist_rainbow", alpha=0.6
):
image, mask = preprocess_image_and_mask(image, seg_source)
if hard_edges:
mask = apply_hard_edges(mask)
color_mask, patches = create_color_mask(mask, colormap)
overlaid_image = overlay_mask(image, color_mask, alpha)
fig, canvas = create_figure(overlaid_image, patches)
result_tensor = render_figure_to_tensor(canvas)
return result_tensor
utils.py
+26 -217
View File
@@ -1,35 +1,28 @@
try:
import face_alignment
face_alignment_available = True
except ImportError:
face_alignment_available = False
print("Warning: face_alignment module not found. fine_best_frame will not work. ")
import numpy as np
import torch
import torch.nn.functional as F
import yaml
from comfy.utils import ProgressBar
from scipy.spatial import ConvexHull
from tqdm import tqdm
from .constants import (
checkpoint_folder,
checkpoint_folder_link,
checkpoint_suffix,
config_folder,
config_suffix,
non_cpk_model_names,
partswap_model_names,
seg_model_dict,
supported_model_names,
)
from .modules.generator import OcclusionAwareGenerator
from .modules.keypoint_detector import KPDetector
def check_supported(entity_name: str, entity_support_list: list[str]):
assert (
entity_name in entity_support_list
), f"Unknown {entity_name}, expected {entity_support_list}"
def get_model_file_name(model_name: str) -> str:
assert (
model_name in supported_model_names
), f"Unknown model. Got {model_name}, expected {supported_model_names}"
check_supported(model_name, supported_model_names)
if model_name in non_cpk_model_names:
model_file_name = f"{checkpoint_folder}/{model_name}{checkpoint_suffix}"
@@ -46,51 +39,23 @@ def get_model_link(model_name: str) -> str:
def get_config_path(model_name: str) -> str:
config_path = f"{config_folder}/{model_name}-256.yaml"
config_path = f"{config_folder}/{model_name}-{config_suffix}"
return config_path
def get_partswap_model_file_name(model_name: str):
return f"{checkpoint_folder}/{model_name}{checkpoint_suffix}"
def out_video(predictions: list[np.ndarray]):
out_tensor_list = []
for i in predictions:
# out_img = i.astype(np.float32) / 255.0
out_img = torch.from_numpy(i)
out_tensor_list.append(out_img)
images = torch.stack(out_tensor_list, dim=0)
return images
def load_checkpoint(
config_path: str, checkpoint_path: str, cpu=False
) -> tuple[OcclusionAwareGenerator, KPDetector]:
with open(config_path) as f:
config = yaml.full_load(f)
generator = OcclusionAwareGenerator(
**config["model_params"]["generator_params"],
**config["model_params"]["common_params"],
)
kp_detector = KPDetector(
**config["model_params"]["kp_detector_params"],
**config["model_params"]["common_params"],
)
if not cpu:
generator.cuda()
kp_detector.cuda()
checkpoint = torch.load(checkpoint_path)
else:
checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
generator.load_state_dict(checkpoint["generator"])
kp_detector.load_state_dict(checkpoint["kp_detector"])
generator.eval()
kp_detector.eval()
return generator, kp_detector
def reshape_image(image, target_size=(256, 256)):
# Check if image is 4D (batch, height, width, channels)
if image.dim() != 4:
@@ -107,177 +72,21 @@ def reshape_image(image, target_size=(256, 256)):
return resized_image
def normalize_kp_best_frame(kp):
kp = kp - kp.mean(axis=0, keepdims=True)
area = ConvexHull(kp[:, :2]).volume
area = np.sqrt(area)
kp[:, :2] = kp[:, :2] / area
return kp
def tensor_to_numpy(tensor):
# Convert [1, 3, H, W] tensor to [H, W, 3] numpy array
return tensor.squeeze(0).permute(1, 2, 0).cpu().numpy()
def find_best_frame(source_image, driving_video, cpu=False) -> int:
if not face_alignment_available:
print("Face Alignment module were not found. Best frame set to 0.")
return 0
print("Initializing Face Alignment...")
fa = face_alignment.FaceAlignment(
face_alignment.LandmarksType.TWO_D,
flip_input=True,
device="cpu" if cpu else "cuda",
)
print("Face alignment model loaded.")
def build_seg_arguments(model_name: str):
check_supported(model_name, partswap_model_names)
segment_names = seg_model_dict[model_name]
required = dict()
for segment_name in segment_names:
default = True
if segment_name.endswith("background"):
default = False
required[segment_name] = ("BOOLEAN", {"default": default})
source_np = tensor_to_numpy(source_image)
kp_source = fa.get_landmarks(255 * source_np)[0]
kp_source = normalize_kp_best_frame(kp_source)
arguments = {"required": required}
norm = float("inf")
frame_num = 0
num_frames = driving_video.shape[2]
for i in tqdm(range(num_frames)):
# Extract single frame from driving video
frame = driving_video[0, :, i] # Shape: [3, 256, 256]
frame_np = frame.permute(1, 2, 0).cpu().numpy()
kp_driving = fa.get_landmarks(255 * frame_np)[0]
kp_driving = normalize_kp_best_frame(kp_driving)
new_norm = (np.abs(kp_source - kp_driving) ** 2).sum()
if new_norm < norm:
norm = new_norm
frame_num = i
print(f"Best frame: {frame_num}")
return frame_num
def normalize_kp(
kp_source,
kp_driving,
kp_driving_initial,
adapt_movement_scale: bool,
use_relative_movement: bool,
use_relative_jacobian: bool,
):
if adapt_movement_scale:
source_area = ConvexHull(kp_source["value"][0].data.cpu().numpy()).volume
driving_area = ConvexHull(
kp_driving_initial["value"][0].data.cpu().numpy()
).volume
adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area)
else:
adapt_movement_scale = 1
# print(f"{adapt_movement_scale=}")
kp_new = {k: v for k, v in kp_driving.items()}
if use_relative_movement:
kp_value_diff = kp_driving["value"] - kp_driving_initial["value"]
kp_value_diff *= adapt_movement_scale
kp_new["value"] = kp_value_diff + kp_source["value"]
if use_relative_jacobian:
jacobian_diff = torch.matmul(
kp_driving["jacobian"], torch.inverse(kp_driving_initial["jacobian"])
)
kp_new["jacobian"] = torch.matmul(jacobian_diff, kp_source["jacobian"])
return kp_new
def inference(
source_image,
driving_video: list,
generator: OcclusionAwareGenerator,
kp_detector: KPDetector,
relative_movement: bool,
relative_jacobian: bool,
adapt_movement_scale: bool,
cpu=False,
) -> list:
with torch.no_grad():
source = source_image
driving = driving_video
predictions = []
if not cpu:
source = source.cuda()
driving = driving.cuda()
print("Using GPU")
kp_source = kp_detector(source)
kp_driving_initial = kp_detector(driving[:, :, 0])
kp_norm = None
num_frames = driving.shape[2]
pbar = ProgressBar(num_frames)
for frame_idx in tqdm(range(num_frames)):
driving_frame = driving[:, :, frame_idx]
kp_driving = kp_detector(driving_frame)
kp_norm = normalize_kp(
kp_source=kp_source,
kp_driving=kp_driving,
kp_driving_initial=kp_driving_initial,
use_relative_movement=relative_movement,
use_relative_jacobian=relative_jacobian,
adapt_movement_scale=adapt_movement_scale,
)
out = generator(source, kp_source=kp_source, kp_driving=kp_norm)
out_np = np.transpose(out["prediction"].data.cpu().numpy(), [0, 2, 3, 1])[0]
predictions.append(out_np)
pbar.update_absolute(frame_idx, num_frames)
return predictions
def inference_best_frame(
source_image,
driving_video,
generator: OcclusionAwareGenerator,
kp_detector: KPDetector,
relative_movement: bool,
relative_jacobian: bool,
adapt_movement_scale: bool,
cpu=False,
) -> list:
best_frame_idx = find_best_frame(source_image, driving_video, cpu)
first_half = driving_video[
:, :, : best_frame_idx + 1
] # Include the best frame in the first half
first_half = torch.flip(first_half, dims=[2]) # Reverse the first half
second_half = driving_video[:, :, best_frame_idx + 1 :]
predictions_first = inference(
source_image,
first_half,
generator,
kp_detector,
relative_movement,
relative_jacobian,
adapt_movement_scale,
cpu,
)
predictions_second = inference(
source_image,
second_half,
generator,
kp_detector,
relative_movement,
relative_jacobian,
adapt_movement_scale,
cpu,
)
predictions = []
predictions_first = predictions_first[::-1] # Reverse the first half back
predictions.extend(predictions_first)
predictions.extend(predictions_second)
return predictions
return arguments
workflow_fomm_partswap.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 546 KiB