201 lines
6.7 KiB
Python
201 lines
6.7 KiB
Python
import numpy as np
|
|
import torch
|
|
import tqdm
|
|
import yaml
|
|
from comfy.utils import ProgressBar
|
|
from scipy.spatial import ConvexHull
|
|
|
|
from .module_articulate.avd_network import AVDNetwork
|
|
from .module_articulate.bg_motion_predictor import BGMotionPredictor
|
|
from .module_articulate.generator import Generator
|
|
from .module_articulate.region_predictor import RegionPredictor
|
|
from .sync_batchnorm.replicate import DataParallelWithCallback
|
|
|
|
|
|
def articulate_inference(
|
|
source_image,
|
|
driving_video: list,
|
|
config_path: str,
|
|
checkpoint_path: str,
|
|
estimate_affine=True,
|
|
pca_based=True,
|
|
):
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
with open(config_path) as f:
|
|
config = yaml.full_load(f)
|
|
|
|
generator, region_predictor, bg_predictor, avd_network = init_models(
|
|
config, estimate_affine, pca_based
|
|
)
|
|
generator = generator.to(device)
|
|
region_predictor = region_predictor.to(device)
|
|
bg_predictor = bg_predictor.to(device)
|
|
avd_network = avd_network.to(device)
|
|
|
|
animate_params = config["animate_params"]
|
|
|
|
load_cpk(checkpoint_path, generator, region_predictor, bg_predictor, avd_network)
|
|
|
|
if torch.cuda.is_available():
|
|
generator = DataParallelWithCallback(generator)
|
|
region_predictor = DataParallelWithCallback(region_predictor)
|
|
avd_network = DataParallelWithCallback(avd_network)
|
|
|
|
generator.eval()
|
|
region_predictor.eval()
|
|
avd_network.eval()
|
|
|
|
source_frame = source_image
|
|
driving = driving_video
|
|
predictions = []
|
|
|
|
num_frames = driving.shape[2]
|
|
pbar = ProgressBar(num_frames)
|
|
|
|
with torch.no_grad():
|
|
source_region_params = region_predictor(source_frame)
|
|
driving_region_params_initial = region_predictor(driving_video[:, :, 0])
|
|
|
|
for frame_idx in tqdm.tqdm(range(num_frames)):
|
|
driving_frame = driving[:, :, frame_idx]
|
|
driving_region_params = region_predictor(driving_frame)
|
|
new_region_params = get_animation_region_params(
|
|
source_region_params,
|
|
driving_region_params,
|
|
driving_region_params_initial,
|
|
mode=animate_params["mode"],
|
|
avd_network=avd_network,
|
|
)
|
|
out = generator(
|
|
source_frame,
|
|
source_region_params=source_region_params,
|
|
driving_region_params=new_region_params,
|
|
)
|
|
|
|
out["driving_region_params"] = driving_region_params
|
|
out["source_region_params"] = source_region_params
|
|
out["new_region_params"] = new_region_params
|
|
|
|
# visualization = Visualizer(**config["visualizer_params"]).visualize(
|
|
# source=source_frame, driving=driving_frame, out=out
|
|
# ) / 255.0
|
|
|
|
prediction = out["prediction"].data.cpu().numpy()
|
|
prediction = np.transpose(prediction, [0, 2, 3, 1]).squeeze(0)
|
|
|
|
# visualizations.append(visualization)
|
|
predictions.append(prediction)
|
|
|
|
pbar.update_absolute(frame_idx, num_frames)
|
|
|
|
# print(f"{predictions[0].shape=}")
|
|
# print(f"{visualizations[0].shape=}")
|
|
|
|
# return predictions, visualizations
|
|
return predictions
|
|
|
|
|
|
def init_models(config, estimate_affine, pca_based):
|
|
generator = Generator(
|
|
num_regions=config["model_params"]["num_regions"],
|
|
num_channels=config["model_params"]["num_channels"],
|
|
revert_axis_swap=config["model_params"]["revert_axis_swap"],
|
|
**config["model_params"]["generator_params"],
|
|
)
|
|
|
|
config["model_params"]["region_predictor_params"]["pca_based"] = pca_based
|
|
|
|
region_predictor = RegionPredictor(
|
|
num_regions=config["model_params"]["num_regions"],
|
|
num_channels=config["model_params"]["num_channels"],
|
|
estimate_affine=estimate_affine,
|
|
**config["model_params"]["region_predictor_params"],
|
|
)
|
|
|
|
bg_predictor = BGMotionPredictor(
|
|
num_channels=config["model_params"]["num_channels"],
|
|
**config["model_params"]["bg_predictor_params"],
|
|
)
|
|
|
|
avd_network = AVDNetwork(
|
|
num_regions=config["model_params"]["num_regions"],
|
|
**config["model_params"]["avd_network_params"],
|
|
)
|
|
|
|
return generator, region_predictor, bg_predictor, avd_network
|
|
|
|
|
|
def load_cpk(
|
|
checkpoint_path,
|
|
generator=None,
|
|
region_predictor=None,
|
|
bg_predictor=None,
|
|
avd_network=None,
|
|
optimizer_reconstruction=None,
|
|
optimizer_avd=None,
|
|
):
|
|
checkpoint = torch.load(checkpoint_path)
|
|
# print(checkpoint.keys())
|
|
if generator is not None:
|
|
generator.load_state_dict(checkpoint["generator"])
|
|
if region_predictor is not None:
|
|
region_predictor.load_state_dict(checkpoint["region_predictor"])
|
|
if bg_predictor is not None:
|
|
bg_predictor.load_state_dict(checkpoint["bg_predictor"])
|
|
if avd_network is not None:
|
|
if "avd_network" in checkpoint:
|
|
avd_network.load_state_dict(checkpoint["avd_network"])
|
|
|
|
if optimizer_reconstruction is not None:
|
|
optimizer_reconstruction.load_state_dict(checkpoint["optimizer_reconstruction"])
|
|
return checkpoint["epoch_reconstruction"]
|
|
|
|
if optimizer_avd is not None:
|
|
if "optimizer_avd" in checkpoint:
|
|
optimizer_avd.load_state_dict(checkpoint["optimizer_avd"])
|
|
return checkpoint["epoch_avd"]
|
|
return 0
|
|
|
|
return 0
|
|
|
|
|
|
def get_animation_region_params(
|
|
source_region_params,
|
|
driving_region_params,
|
|
driving_region_params_initial,
|
|
mode="standard",
|
|
avd_network=None,
|
|
adapt_movement_scale=True,
|
|
):
|
|
assert mode in ["standard", "relative", "avd"]
|
|
new_region_params = {k: v for k, v in driving_region_params.items()}
|
|
if mode == "standard":
|
|
return new_region_params
|
|
elif mode == "relative":
|
|
source_area = ConvexHull(
|
|
source_region_params["shift"][0].data.cpu().numpy()
|
|
).volume
|
|
driving_area = ConvexHull(
|
|
driving_region_params_initial["shift"][0].data.cpu().numpy()
|
|
).volume
|
|
movement_scale = np.sqrt(source_area) / np.sqrt(driving_area)
|
|
|
|
shift_diff = (
|
|
driving_region_params["shift"] - driving_region_params_initial["shift"]
|
|
)
|
|
shift_diff *= movement_scale
|
|
new_region_params["shift"] = shift_diff + source_region_params["shift"]
|
|
|
|
affine_diff = torch.matmul(
|
|
driving_region_params["affine"],
|
|
torch.inverse(driving_region_params_initial["affine"]),
|
|
)
|
|
new_region_params["affine"] = torch.matmul(
|
|
affine_diff, source_region_params["affine"]
|
|
)
|
|
return new_region_params
|
|
elif mode == "avd":
|
|
new_region_params = avd_network(source_region_params, driving_region_params)
|
|
return new_region_params
|