import sys, os, tempfile

os.environ["CUDA_VISIBLE_DEVICES"] = ""

from huggingface_hub import snapshot_download

weights_dir = "/app/LivePortrait/pretrained_weights"
if not os.path.exists(os.path.join(weights_dir, "appearance_feature_extractor.safetensors")):
    print("Downloading weights (~2GB)...")
    snapshot_download(
        repo_id="KlingTeam/LivePortrait",
        local_dir=weights_dir,
        ignore_patterns=["*.git*", "README.md", "docs/*"],
    )
    print("Weights ready ✓")

sys.path.insert(0, "/app/LivePortrait")
os.chdir("/app/LivePortrait")

from src.config.argument_config import ArgumentConfig
from src.config.inference_config import InferenceConfig
from src.config.crop_config import CropConfig
from src.live_portrait_pipeline import LivePortraitPipeline

infer_cfg = InferenceConfig()
infer_cfg.flag_force_cpu = True
pipeline = LivePortraitPipeline(inference_cfg=infer_cfg, crop_cfg=CropConfig())

def animate(source_image, driving_video):
    if source_image is None or driving_video is None:
        return None
    out_dir = tempfile.mkdtemp()
    args = ArgumentConfig(source=source_image, driving=driving_video, output_dir=out_dir)
    wfp, _ = pipeline.execute(args)
    return wfp

import gradio as gr

# Patch get_api_info to never fail
import gradio.blocks as _gb
_orig_get_api_info = _gb.Blocks.get_api_info
def _safe_get_api_info(self):
    try:
        return _orig_get_api_info(self)
    except Exception:
        return {}
_gb.Blocks.get_api_info = _safe_get_api_info

demo = gr.Interface(
    fn=animate,
    inputs=[
        gr.Image(label="Source Portrait", type="filepath"),
        gr.Video(label="Driving Video"),
    ],
    outputs=gr.Video(label="Animated Result"),
    title="🎭 LivePortrait",
    description="Animate any portrait with a driving video.\n⚠️ CPU mode — takes a few minutes.",
)

demo.launch(server_name="0.0.0.0", server_port=7860, show_api=False)