|
"""run bash scripts/download_models.sh first to prepare the weights file"""
|
|
import os
|
|
import shutil
|
|
from argparse import Namespace
|
|
from src.utils.preprocess import CropAndExtract
|
|
from src.test_audio2coeff import Audio2Coeff
|
|
from src.facerender.animate import AnimateFromCoeff
|
|
from src.generate_batch import get_data
|
|
from src.generate_facerender_batch import get_facerender_data
|
|
from cog import BasePredictor, Input, Path
|
|
|
|
checkpoints = "checkpoints"
|
|
|
|
|
|
class Predictor(BasePredictor):
|
|
def setup(self):
|
|
"""Load the model into memory to make running multiple predictions efficient"""
|
|
device = "cuda"
|
|
|
|
path_of_lm_croper = os.path.join(
|
|
checkpoints, "shape_predictor_68_face_landmarks.dat"
|
|
)
|
|
path_of_net_recon_model = os.path.join(checkpoints, "epoch_20.pth")
|
|
dir_of_BFM_fitting = os.path.join(checkpoints, "BFM_Fitting")
|
|
wav2lip_checkpoint = os.path.join(checkpoints, "wav2lip.pth")
|
|
|
|
audio2pose_checkpoint = os.path.join(checkpoints, "auido2pose_00140-model.pth")
|
|
audio2pose_yaml_path = os.path.join("src", "config", "auido2pose.yaml")
|
|
|
|
audio2exp_checkpoint = os.path.join(checkpoints, "auido2exp_00300-model.pth")
|
|
audio2exp_yaml_path = os.path.join("src", "config", "auido2exp.yaml")
|
|
|
|
free_view_checkpoint = os.path.join(
|
|
checkpoints, "facevid2vid_00189-model.pth.tar"
|
|
)
|
|
|
|
|
|
self.preprocess_model = CropAndExtract(
|
|
path_of_lm_croper, path_of_net_recon_model, dir_of_BFM_fitting, device
|
|
)
|
|
|
|
self.audio_to_coeff = Audio2Coeff(
|
|
audio2pose_checkpoint,
|
|
audio2pose_yaml_path,
|
|
audio2exp_checkpoint,
|
|
audio2exp_yaml_path,
|
|
wav2lip_checkpoint,
|
|
device,
|
|
)
|
|
|
|
self.animate_from_coeff = {
|
|
"full": AnimateFromCoeff(
|
|
free_view_checkpoint,
|
|
os.path.join(checkpoints, "mapping_00109-model.pth.tar"),
|
|
os.path.join("src", "config", "facerender_still.yaml"),
|
|
device,
|
|
),
|
|
"others": AnimateFromCoeff(
|
|
free_view_checkpoint,
|
|
os.path.join(checkpoints, "mapping_00229-model.pth.tar"),
|
|
os.path.join("src", "config", "facerender.yaml"),
|
|
device,
|
|
),
|
|
}
|
|
|
|
def predict(
|
|
self,
|
|
source_image: Path = Input(
|
|
description="Upload the source image, it can be video.mp4 or picture.png",
|
|
),
|
|
driven_audio: Path = Input(
|
|
description="Upload the driven audio, accepts .wav and .mp4 file",
|
|
),
|
|
enhancer: str = Input(
|
|
description="Choose a face enhancer",
|
|
choices=["gfpgan", "RestoreFormer"],
|
|
default="gfpgan",
|
|
),
|
|
preprocess: str = Input(
|
|
description="how to preprocess the images",
|
|
choices=["crop", "resize", "full"],
|
|
default="full",
|
|
),
|
|
ref_eyeblink: Path = Input(
|
|
description="path to reference video providing eye blinking",
|
|
default=None,
|
|
),
|
|
ref_pose: Path = Input(
|
|
description="path to reference video providing pose",
|
|
default=None,
|
|
),
|
|
still: bool = Input(
|
|
description="can crop back to the original videos for the full body aniamtion when preprocess is full",
|
|
default=True,
|
|
),
|
|
) -> Path:
|
|
"""Run a single prediction on the model"""
|
|
|
|
animate_from_coeff = (
|
|
self.animate_from_coeff["full"]
|
|
if preprocess == "full"
|
|
else self.animate_from_coeff["others"]
|
|
)
|
|
|
|
args = load_default()
|
|
args.pic_path = str(source_image)
|
|
args.audio_path = str(driven_audio)
|
|
device = "cuda"
|
|
args.still = still
|
|
args.ref_eyeblink = None if ref_eyeblink is None else str(ref_eyeblink)
|
|
args.ref_pose = None if ref_pose is None else str(ref_pose)
|
|
|
|
|
|
results_dir = "results"
|
|
if os.path.exists(results_dir):
|
|
shutil.rmtree(results_dir)
|
|
os.makedirs(results_dir)
|
|
first_frame_dir = os.path.join(results_dir, "first_frame_dir")
|
|
os.makedirs(first_frame_dir)
|
|
|
|
print("3DMM Extraction for source image")
|
|
first_coeff_path, crop_pic_path, crop_info = self.preprocess_model.generate(
|
|
args.pic_path, first_frame_dir, preprocess, source_image_flag=True
|
|
)
|
|
if first_coeff_path is None:
|
|
print("Can't get the coeffs of the input")
|
|
return
|
|
|
|
if ref_eyeblink is not None:
|
|
ref_eyeblink_videoname = os.path.splitext(os.path.split(ref_eyeblink)[-1])[
|
|
0
|
|
]
|
|
ref_eyeblink_frame_dir = os.path.join(results_dir, ref_eyeblink_videoname)
|
|
os.makedirs(ref_eyeblink_frame_dir, exist_ok=True)
|
|
print("3DMM Extraction for the reference video providing eye blinking")
|
|
ref_eyeblink_coeff_path, _, _ = self.preprocess_model.generate(
|
|
ref_eyeblink, ref_eyeblink_frame_dir
|
|
)
|
|
else:
|
|
ref_eyeblink_coeff_path = None
|
|
|
|
if ref_pose is not None:
|
|
if ref_pose == ref_eyeblink:
|
|
ref_pose_coeff_path = ref_eyeblink_coeff_path
|
|
else:
|
|
ref_pose_videoname = os.path.splitext(os.path.split(ref_pose)[-1])[0]
|
|
ref_pose_frame_dir = os.path.join(results_dir, ref_pose_videoname)
|
|
os.makedirs(ref_pose_frame_dir, exist_ok=True)
|
|
print("3DMM Extraction for the reference video providing pose")
|
|
ref_pose_coeff_path, _, _ = self.preprocess_model.generate(
|
|
ref_pose, ref_pose_frame_dir
|
|
)
|
|
else:
|
|
ref_pose_coeff_path = None
|
|
|
|
|
|
batch = get_data(
|
|
first_coeff_path,
|
|
args.audio_path,
|
|
device,
|
|
ref_eyeblink_coeff_path,
|
|
still=still,
|
|
)
|
|
coeff_path = self.audio_to_coeff.generate(
|
|
batch, results_dir, args.pose_style, ref_pose_coeff_path
|
|
)
|
|
|
|
print("coeff2video")
|
|
data = get_facerender_data(
|
|
coeff_path,
|
|
crop_pic_path,
|
|
first_coeff_path,
|
|
args.audio_path,
|
|
args.batch_size,
|
|
args.input_yaw,
|
|
args.input_pitch,
|
|
args.input_roll,
|
|
expression_scale=args.expression_scale,
|
|
still_mode=still,
|
|
preprocess=preprocess,
|
|
)
|
|
animate_from_coeff.generate(
|
|
data, results_dir, args.pic_path, crop_info,
|
|
enhancer=enhancer, background_enhancer=args.background_enhancer,
|
|
preprocess=preprocess)
|
|
|
|
output = "/tmp/out.mp4"
|
|
mp4_path = os.path.join(results_dir, [f for f in os.listdir(results_dir) if "enhanced.mp4" in f][0])
|
|
shutil.copy(mp4_path, output)
|
|
|
|
return Path(output)
|
|
|
|
|
|
def load_default():
|
|
return Namespace(
|
|
pose_style=0,
|
|
batch_size=2,
|
|
expression_scale=1.0,
|
|
input_yaw=None,
|
|
input_pitch=None,
|
|
input_roll=None,
|
|
background_enhancer=None,
|
|
face3dvis=False,
|
|
net_recon="resnet50",
|
|
init_path=None,
|
|
use_last_fc=False,
|
|
bfm_folder="./checkpoints/BFM_Fitting/",
|
|
bfm_model="BFM_model_front.mat",
|
|
focal=1015.0,
|
|
center=112.0,
|
|
camera_d=10.0,
|
|
z_near=5.0,
|
|
z_far=15.0,
|
|
)
|
|
|
