From 8d5c84c0eccd5fae1f7b539fcd8b95fdce6a102b Mon Sep 17 00:00:00 2001 From: FuouM Date: Sat, 10 Aug 2024 19:36:06 +0700 Subject: [PATCH] refactor --- README.md | 6 + fomm_loader.py | 35 --- ...te_inference.py => inference_articulate.py | 0 fomm_inference.py => inference_fomm.py | 45 ++- inference_partswap.py | 274 ++++++++++++++++++ spline_inference.py => inference_spline.py | 2 +- partswap_inference.py | 84 ------ partswap_loader.py | 92 ------ partswap_model.py | 105 ------- pyproject.toml | 2 +- requirements.txt | 2 +- run.py | 12 +- 12 files changed, 328 insertions(+), 331 deletions(-) delete mode 100644 fomm_loader.py rename articulate_inference.py => inference_articulate.py (100%) rename fomm_inference.py => inference_fomm.py (85%) create mode 100644 inference_partswap.py rename spline_inference.py => inference_spline.py (99%) delete mode 100644 partswap_inference.py delete mode 100644 partswap_loader.py delete mode 100644 partswap_model.py diff --git a/README.md b/README.md index a62b4d1..2595756 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,7 @@ Now supports: 1. Face Swapping using [Motion Supervised co-part Segmentation](https://github.com/AliaksandrSiarohin/motion-cosegmentation): 2. [Motion Representations for Articulated Animation](https://github.com/snap-research/articulated-animation) 3. [Thin-Plate Spline Motion Model for Image Animation](https://github.com/yoyo-nb/thin-plate-spline-motion-model) +4. [Learning Motion Refinement for Unsupervised Face Animation](https://github.com/JialeTao/MRFA/) More will come soon @@ -69,6 +70,9 @@ Doesn't need any * `avd`: similar to `relative`, may yield better but more "jittery/jumpy" result * `find_best_frame`: Same as FOMM + +### MRFA + ## Installation 1. Clone the repo to `ComfyUI/custom_nodes/` @@ -129,3 +133,5 @@ resnet18-5c106cde.pth For **Articulate**, download the model from [Pre-trained checkpoints](https://github.com/snap-research/articulated-animation?tab=readme-ov-file#pre-trained-checkpoints) section and place it here: `articulate_module/models/vox256.pth` For **Spline**, download the model from [Pre-trained models](https://github.com/yoyo-nb/thin-plate-spline-motion-model?tab=readme-ov-file#pre-trained-models) section and place it here: `spline_module/models/vox.pth.tar`. To use `find_best_frame`, install `face-alignment`. + +For **MRFA**, download the model from [Pretrained models](https://github.com/JialeTao/MRFA/?tab=readme-ov-file#pretrained-models) section and place it here: `mrfa_module/models/` \ No newline at end of file diff --git a/fomm_loader.py b/fomm_loader.py deleted file mode 100644 index b100543..0000000 --- a/fomm_loader.py +++ /dev/null @@ -1,35 +0,0 @@ -import yaml -import torch - -from .modules.generator import OcclusionAwareGenerator -from .modules.keypoint_detector import KPDetector - -def load_checkpoint( - config_path: str, checkpoint_path: str, cpu=False -) -> tuple[OcclusionAwareGenerator, KPDetector]: - with open(config_path) as f: - config = yaml.full_load(f) - - generator = OcclusionAwareGenerator( - **config["model_params"]["generator_params"], - **config["model_params"]["common_params"], - ) - kp_detector = KPDetector( - **config["model_params"]["kp_detector_params"], - **config["model_params"]["common_params"], - ) - - if not cpu: - generator.cuda() - kp_detector.cuda() - checkpoint = torch.load(checkpoint_path) - else: - checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu")) - - generator.load_state_dict(checkpoint["generator"]) - kp_detector.load_state_dict(checkpoint["kp_detector"]) - - generator.eval() - kp_detector.eval() - - return generator, kp_detector \ No newline at end of file diff --git a/articulate_inference.py b/inference_articulate.py similarity index 100% rename from articulate_inference.py rename to inference_articulate.py diff --git a/fomm_inference.py b/inference_fomm.py similarity index 85% rename from fomm_inference.py rename to inference_fomm.py index 2665374..150d6c3 100644 --- a/fomm_inference.py +++ b/inference_fomm.py @@ -1,11 +1,12 @@ -import torch import numpy as np +import torch +import yaml from comfy.utils import ProgressBar -from tqdm import tqdm from scipy.spatial import ConvexHull +from tqdm import tqdm + from .modules.generator import OcclusionAwareGenerator from .modules.keypoint_detector import KPDetector - from .utils import tensor_to_numpy try: @@ -24,6 +25,7 @@ def normalize_kp_best_frame(kp): kp[:, :2] = kp[:, :2] / area return kp + def normalize_kp( kp_source, kp_driving, @@ -55,11 +57,12 @@ def normalize_kp( 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, @@ -93,6 +96,7 @@ def find_best_frame(source_image, driving_video: list, cpu=False) -> int: return frame_num + def inference( source_image, driving_video: list, @@ -182,4 +186,35 @@ def inference_best_frame( predictions_first = predictions_first[::-1] # Reverse the first half back predictions.extend(predictions_first) predictions.extend(predictions_second) - return predictions \ No newline at end of file + return predictions + + +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 diff --git a/inference_partswap.py b/inference_partswap.py new file mode 100644 index 0000000..08f7150 --- /dev/null +++ b/inference_partswap.py @@ -0,0 +1,274 @@ +import numpy as np +import torch +import torch.nn.functional as F +import yaml +from comfy.utils import ProgressBar +from tqdm import tqdm + +from .modules.dense_motion import DenseMotionNetwork +from .modules.reconstruction_module import ReconstructionModule +from .modules.segmentation_module import SegmentationModule +from .modules.util import AntiAliasInterpolation2d +from .sync_batchnorm.replicate import DataParallelWithCallback + + +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 + + +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 + + +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 + + +def partial_state_dict_load(module, state_dict): + own_state = module.state_dict() + for name, param in state_dict.items(): + if name not in own_state: + continue + + if isinstance(param, torch.nn.Parameter): + # backwards compatibility for serialized parameters + param = param.data + own_state[name].copy_(param) + + +def load_reconstruction_module(module, checkpoint): + if "generator" in checkpoint: + partial_state_dict_load(module, checkpoint["generator"]) + else: + module.load_state_dict(checkpoint["reconstruction_module"]) + + +def load_segmentation_module(module, checkpoint): + if "kp_detector" in checkpoint: + partial_state_dict_load(module, checkpoint["kp_detector"]) + module.state_dict()["affine.weight"].copy_( + checkpoint["kp_detector"]["jacobian.weight"] + ) + module.state_dict()["affine.bias"].copy_( + checkpoint["kp_detector"]["jacobian.bias"] + ) + module.state_dict()["shift.weight"].copy_( + checkpoint["kp_detector"]["kp.weight"] + ) + module.state_dict()["shift.bias"].copy_(checkpoint["kp_detector"]["kp.bias"]) + if "semantic_seg.weight" in checkpoint["kp_detector"]: + module.state_dict()["segmentation.weight"].copy_( + checkpoint["kp_detector"]["semantic_seg.weight"] + ) + module.state_dict()["segmentation.bias"].copy_( + checkpoint["kp_detector"]["semantic_seg.bias"] + ) + else: + print("Segmentation part initialized at random.") + else: + module.load_state_dict(checkpoint["segmentation_module"]) + + +def load_partswap_checkpoint( + config_path: str, + checkpoint_path: str, + blend_scale: float, + use_fomm: bool, + cpu=False, +): + with open(config_path) as f: + config = yaml.full_load(f) + + reconstruction_module = PartSwapGenerator( + blend_scale=blend_scale, + first_order_motion_model=use_fomm, + **config["model_params"]["reconstruction_module_params"], + **config["model_params"]["common_params"], + ) + + segmentation_module = SegmentationModule( + **config["model_params"]["segmentation_module_params"], + **config["model_params"]["common_params"], + ) + + if cpu: + checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu")) + else: + segmentation_module.cuda() + checkpoint = torch.load(checkpoint_path) + + load_reconstruction_module(reconstruction_module, checkpoint) + load_segmentation_module(segmentation_module, checkpoint) + + if not cpu: + reconstruction_module = DataParallelWithCallback(reconstruction_module) + segmentation_module = DataParallelWithCallback(segmentation_module) + + reconstruction_module.eval() + segmentation_module.eval() + + return reconstruction_module, segmentation_module diff --git a/spline_inference.py b/inference_spline.py similarity index 99% rename from spline_inference.py rename to inference_spline.py index c432318..80bec5f 100644 --- a/spline_inference.py +++ b/inference_spline.py @@ -5,7 +5,7 @@ import yaml import tqdm from comfy.utils import ProgressBar -from .fomm_inference import find_best_frame +from .inference_fomm import find_best_frame from .spline_module.avd_network import AVDNetwork from .spline_module.dense_motion import DenseMotionNetwork diff --git a/partswap_inference.py b/partswap_inference.py deleted file mode 100644 index 012575f..0000000 --- a/partswap_inference.py +++ /dev/null @@ -1,84 +0,0 @@ -import numpy as np -import torch -import torch.nn.functional as F -from tqdm import tqdm - -from comfy.utils import ProgressBar -from .modules.segmentation_module import SegmentationModule -from .partswap_model import PartSwapGenerator - - -def face_parse_seg( - source_image, face_parser_model, in_size=(512, 512), out_size=(64, 64) -): - seg = F.interpolate(source_image, size=in_size) - seg = (seg - face_parser_model.mean) / face_parser_model.std - seg = torch.softmax(face_parser_model(seg)[0], dim=1) - seg = F.interpolate(seg, size=out_size) - return seg - - -def partswap_inference( - swap_indices: list[int], - source_image, - target_video, - reconstruction_module: PartSwapGenerator, - segmentation_module: SegmentationModule, - use_source_seg: bool, - face_parser_model=None, - hard_edges=False, - cpu=False, -) -> list: - with torch.no_grad(): - predictions = [] - seg_targets = [] - if not cpu: - source_image = source_image.cuda() - target_video = target_video.cuda() - print("Using GPU") - - seg_source = segmentation_module(source_image) - # print(f"{seg_source['segmentation'].shape}") - - num_frames = target_video.shape[2] - pbar = ProgressBar(num_frames) - for frame_idx in tqdm(range(num_frames)): - target_frame = target_video[:, :, frame_idx] - seg_target = segmentation_module(target_frame) - seg_targets.append(seg_target) - - if face_parser_model is not None: - blend_mask = face_parse_seg( - source_image if use_source_seg else target_frame, - face_parser_model, - (512, 512), - (64, 64), - ) - - else: - blend_mask = ( - seg_source["segmentation"] - if use_source_seg - else seg_target["segmentation"] - ) - - blend_mask = blend_mask[:, swap_indices].sum(dim=1, keepdim=True) - if hard_edges: - blend_mask = (blend_mask > 0.5).type(blend_mask.type()) - print(f"{blend_mask.shape=}") - out = reconstruction_module( - source_image, - target_frame, - seg_source=seg_source, - seg_target=seg_target, - blend_mask=blend_mask, - use_source_segmentation=use_source_seg, - ) - - predictions.append( - np.transpose(out["prediction"].data.cpu().numpy(), [0, 2, 3, 1])[0] - ) - - pbar.update_absolute(frame_idx, num_frames) - - return seg_source, seg_targets, predictions diff --git a/partswap_loader.py b/partswap_loader.py deleted file mode 100644 index 55c599c..0000000 --- a/partswap_loader.py +++ /dev/null @@ -1,92 +0,0 @@ -import torch -import yaml - -from .modules.segmentation_module import SegmentationModule -from .partswap_model import PartSwapGenerator -from .sync_batchnorm.replicate import DataParallelWithCallback - - -def partial_state_dict_load(module, state_dict): - own_state = module.state_dict() - for name, param in state_dict.items(): - if name not in own_state: - continue - - if isinstance(param, torch.nn.Parameter): - # backwards compatibility for serialized parameters - param = param.data - own_state[name].copy_(param) - - -def load_reconstruction_module(module, checkpoint): - if "generator" in checkpoint: - partial_state_dict_load(module, checkpoint["generator"]) - else: - module.load_state_dict(checkpoint["reconstruction_module"]) - - -def load_segmentation_module(module, checkpoint): - if "kp_detector" in checkpoint: - partial_state_dict_load(module, checkpoint["kp_detector"]) - module.state_dict()["affine.weight"].copy_( - checkpoint["kp_detector"]["jacobian.weight"] - ) - module.state_dict()["affine.bias"].copy_( - checkpoint["kp_detector"]["jacobian.bias"] - ) - module.state_dict()["shift.weight"].copy_( - checkpoint["kp_detector"]["kp.weight"] - ) - module.state_dict()["shift.bias"].copy_(checkpoint["kp_detector"]["kp.bias"]) - if "semantic_seg.weight" in checkpoint["kp_detector"]: - module.state_dict()["segmentation.weight"].copy_( - checkpoint["kp_detector"]["semantic_seg.weight"] - ) - module.state_dict()["segmentation.bias"].copy_( - checkpoint["kp_detector"]["semantic_seg.bias"] - ) - else: - print("Segmentation part initialized at random.") - else: - module.load_state_dict(checkpoint["segmentation_module"]) - - -def load_partswap_checkpoint( - config_path: str, - checkpoint_path: str, - blend_scale: float, - use_fomm: bool, - cpu=False, -): - with open(config_path) as f: - config = yaml.full_load(f) - - reconstruction_module = PartSwapGenerator( - blend_scale=blend_scale, - first_order_motion_model=use_fomm, - **config["model_params"]["reconstruction_module_params"], - **config["model_params"]["common_params"], - ) - - segmentation_module = SegmentationModule( - **config["model_params"]["segmentation_module_params"], - **config["model_params"]["common_params"], - ) - - if cpu: - checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu")) - else: - segmentation_module.cuda() - checkpoint = torch.load(checkpoint_path) - - load_reconstruction_module(reconstruction_module, checkpoint) - load_segmentation_module(segmentation_module, checkpoint) - - if not cpu: - reconstruction_module = DataParallelWithCallback(reconstruction_module) - segmentation_module = DataParallelWithCallback(segmentation_module) - - reconstruction_module.eval() - segmentation_module.eval() - - return reconstruction_module, segmentation_module diff --git a/partswap_model.py b/partswap_model.py deleted file mode 100644 index ee55dea..0000000 --- a/partswap_model.py +++ /dev/null @@ -1,105 +0,0 @@ -import torch.nn.functional as F - -from .modules.dense_motion import DenseMotionNetwork -from .modules.reconstruction_module import ReconstructionModule -from .modules.util import AntiAliasInterpolation2d - - -class PartSwapGenerator(ReconstructionModule): - def __init__(self, blend_scale=1, first_order_motion_model=False, **kwargs): - super(PartSwapGenerator, self).__init__(**kwargs) - if blend_scale == 1: - self.blend_downsample = lambda x: x - else: - self.blend_downsample = AntiAliasInterpolation2d(1, blend_scale) - - if first_order_motion_model: - self.dense_motion_network = DenseMotionNetwork( - block_expansion=64, - num_blocks=5, - max_features=1024, - num_kp=10, - num_channels=3, - estimate_occlusion_map=True, - scale_factor=0.25, - ) - else: - self.dense_motion_network = None - - def forward( - self, - source_image, - target_image, - seg_target, - seg_source, - blend_mask, - use_source_segmentation=False, - ): - # Encoding of source image - enc_source = self.first(source_image) - for i in range(len(self.down_blocks)): - enc_source = self.down_blocks[i](enc_source) - - # Encoding of target image - enc_target = self.first(target_image) - for i in range(len(self.down_blocks)): - enc_target = self.down_blocks[i](enc_target) - - output_dict = {} - # Compute flow field for source image - if self.dense_motion_network is None: - segment_motions = self.segment_motion(seg_target, seg_source) - segment_motions = segment_motions.permute(0, 1, 4, 2, 3) - mask = seg_target["segmentation"].unsqueeze(2) - deformation = (segment_motions * mask).sum(dim=1) - deformation = deformation.permute(0, 2, 3, 1) - else: - motion = self.dense_motion_network.forward_partswap( - source_image=source_image, seg_target=seg_target, seg_source=seg_source - ) - deformation = motion["deformation"] - - # Deform source encoding according to the motion - enc_source = self.deform_input(enc_source, deformation) - - if self.estimate_visibility: - if self.dense_motion_network is None: - visibility = seg_source["segmentation"][:, 1:].sum( - dim=1, keepdim=True - ) * ( - 1 - - seg_target["segmentation"][:, 1:] - .sum(dim=1, keepdim=True) - .detach() - ) - visibility = 1 - visibility - else: - visibility = motion["visibility"] - - if ( - enc_source.shape[2] != visibility.shape[2] - or enc_source.shape[3] != visibility.shape[3] - ): - visibility = F.interpolate( - visibility, size=enc_source.shape[2:], mode="bilinear" - ) - enc_source = enc_source * visibility - - blend_mask = self.blend_downsample(blend_mask) - # If source segmentation is provided use it should be deformed before blending - if use_source_segmentation: - blend_mask = self.deform_input(blend_mask, deformation) - - out = enc_target * (1 - blend_mask) + enc_source * blend_mask - - out = self.bottleneck(out) - - for i in range(len(self.up_blocks)): - out = self.up_blocks[i](out) - - out = self.final(out) - out = F.sigmoid(out) - - output_dict["prediction"] = out - - return output_dict diff --git a/pyproject.toml b/pyproject.toml index c90ad79..cfa910a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "comfyui-firstordermm" -description = "A ComfyUI-native node to run First Order Motion Model for Image Animation. [a/https://github.com/AliaksandrSiarohin/first-order-model](https://github.com/AliaksandrSiarohin/first-order-model)" +description = "ComfyUI-native nodes to run First Order Motion Model for Image Animation and its non-diffusion-based successors. [a/https://github.com/AliaksandrSiarohin/first-order-model](https://github.com/AliaksandrSiarohin/first-order-model)" version = "1.0.3" license = { file = "LICENSE" } dependencies = ["numpy", "torch", "scipy", "yaml", "matplotlib"] diff --git a/requirements.txt b/requirements.txt index fe18382..5eecb23 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ numpy torch scipy -yaml +pyyaml matplotlib \ No newline at end of file diff --git a/run.py b/run.py index 42c7645..340dae2 100644 --- a/run.py +++ b/run.py @@ -2,14 +2,13 @@ @author: Fuou Marinas @title: ComfyUI-FirstOrderMM @nickname: FOMM -@description: Run First Order Motion Model for Image Animation in ComfyUI. +@description: ComfyUI-native nodes to run First Order Motion Model for Image Animation and its non-diffusion-based successors. """ from pathlib import Path import torch -from .articulate_inference import articulate_inference from .constants import ( ARTICULATE_CFG_PATH, ARTICULATE_MODEL_PATH, @@ -28,12 +27,11 @@ from .constants import ( 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 .inference_articulate import articulate_inference +from .inference_fomm import inference, inference_best_frame, load_checkpoint +from .inference_partswap import load_partswap_checkpoint, partswap_inference +from .inference_spline import spline_inference from .seg_viz import visualize_frame -from .spline_inference import spline_inference from .utils import ( build_seg_arguments, get_config_path,