| 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
|
|
|
|
|
|
|
| REPO_ID = "vinthony/SadTalker-V002rc"
|
| snapshot_download(
|
| repo_id=REPO_ID, local_dir="./checkpoints", local_dir_use_symlinks=True
|
| )
|
|
|
|
|
| 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)}
|
|
|
