File size: 2,729 Bytes
148f3e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | import os
from typing import Dict, Any
from PIL import Image
import tempfile
import base64
import shutil
from src.gradio_demo import SadTalker
from huggingface_hub import snapshot_download
# Download model checkpoints at init
REPO_ID = "vinthony/SadTalker-V002rc"
snapshot_download(
repo_id=REPO_ID, local_dir="./checkpoints", local_dir_use_symlinks=True
)
# Load SadTalker
sad_talker = SadTalker(lazy_load=True)
def save_base64_to_file(base64_str, suffix):
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
data = base64.b64decode(base64_str)
with open(temp_file.name, "wb") as f:
f.write(data)
return temp_file.name
def predict(inputs: Dict[str, Any]) -> Dict[str, Any]:
"""
Expected input JSON:
{
"image": base64_image_string,
"audio": base64_audio_string,
"ref_video": base64_video_string or None,
"use_ref_video": true/false,
"use_idle_mode": true/false,
"length": 5,
"pose_style": 0,
"expression_scale": 1.0,
"preprocess": "crop",
...
}
"""
try:
image_path = save_base64_to_file(inputs["image"], ".png")
audio_path = (
save_base64_to_file(inputs["audio"], ".wav")
if not inputs.get("use_idle_mode", False)
else None
)
ref_video_path = (
save_base64_to_file(inputs["ref_video"], ".mp4")
if inputs.get("use_ref_video")
else None
)
output_path = sad_talker.test(
source_image=image_path,
driven_audio=audio_path,
preprocess=inputs.get("preprocess", "crop"),
still_mode=inputs.get("still_mode", False),
enhancer=inputs.get("enhancer", False),
batch_size=inputs.get("batch_size", 1),
size_of_image=inputs.get("size", 256),
pose_style=inputs.get("pose_style", 0),
facerender=inputs.get("facerender", "facevid2vid"),
expression_scale=inputs.get("expression_scale", 1.0),
use_ref_video=inputs.get("use_ref_video", False),
ref_video=ref_video_path,
ref_info=inputs.get("ref_info", "pose"),
use_idle_mode=inputs.get("use_idle_mode", False),
length_of_audio=inputs.get("length", 5),
blink_every=inputs.get("blink", True),
)
with open(output_path, "rb") as f:
video_bytes = f.read()
video_b64 = base64.b64encode(video_bytes).decode("utf-8")
return {"video": video_b64}
except Exception as e:
return {"error": str(e)}
|