|
|
|
|
|
|
|
from cog import BasePredictor, Input, Path
|
|
import os
|
|
import time
|
|
import subprocess
|
|
|
|
MODEL_CACHE = "checkpoints"
|
|
MODEL_URL = "https://weights.replicate.delivery/default/chunyu-li/LatentSync/model.tar"
|
|
|
|
def download_weights(url, dest):
|
|
start = time.time()
|
|
print("downloading url: ", url)
|
|
print("downloading to: ", dest)
|
|
subprocess.check_call(["pget", "-xf", url, dest], close_fds=False)
|
|
print("downloading took: ", time.time() - start)
|
|
|
|
class Predictor(BasePredictor):
|
|
def setup(self) -> None:
|
|
"""Load the model into memory to make running multiple predictions efficient"""
|
|
|
|
if not os.path.exists(MODEL_CACHE):
|
|
download_weights(MODEL_URL, MODEL_CACHE)
|
|
|
|
|
|
os.system("mkdir -p ~/.cache/torch/hub/checkpoints")
|
|
os.system("ln -s $(pwd)/checkpoints/auxiliary/2DFAN4-cd938726ad.zip ~/.cache/torch/hub/checkpoints/2DFAN4-cd938726ad.zip")
|
|
os.system("ln -s $(pwd)/checkpoints/auxiliary/s3fd-619a316812.pth ~/.cache/torch/hub/checkpoints/s3fd-619a316812.pth")
|
|
os.system("ln -s $(pwd)/checkpoints/auxiliary/vgg16-397923af.pth ~/.cache/torch/hub/checkpoints/vgg16-397923af.pth")
|
|
|
|
def predict(
|
|
self,
|
|
video: Path = Input(
|
|
description="Input video", default=None
|
|
),
|
|
audio: Path = Input(
|
|
description="Input audio to ", default=None
|
|
),
|
|
guidance_scale: float = Input(
|
|
description="Guidance scale", ge=0, le=10, default=1.0
|
|
),
|
|
seed: int = Input(
|
|
description="Set to 0 for Random seed", default=0
|
|
)
|
|
) -> Path:
|
|
"""Run a single prediction on the model"""
|
|
if seed <= 0:
|
|
seed = int.from_bytes(os.urandom(2), "big")
|
|
print(f"Using seed: {seed}")
|
|
|
|
video_path = str(video)
|
|
audio_path = str(audio)
|
|
config_path = "configs/unet/second_stage.yaml"
|
|
ckpt_path = "checkpoints/latentsync_unet.pt"
|
|
output_path = "/tmp/video_out.mp4"
|
|
|
|
|
|
os.system(f"python -m scripts.inference --unet_config_path {config_path} --inference_ckpt_path {ckpt_path} --guidance_scale {str(guidance_scale)} --video_path {video_path} --audio_path {audio_path} --video_out_path {output_path} --seed {seed}")
|
|
return Path(output_path)
|
|
|
