app.py

from __future__ import annotations

import glob
import gc
import inspect
import os
import shutil
import tempfile
import time
import uuid
from functools import lru_cache
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import gradio as gr
import requests
import torch

APP_DIR = Path(__file__).resolve().parent
CHECKPOINTS_DIR = APP_DIR / "checkpoints"
DEFAULT_PACKED_PT = CHECKPOINTS_DIR / "PackedAvatar.pt"

# Hugging Face raw resolve URL for the PackedAvatar.pt file
HF_PACKED_PT_URL = "https://huggingface.co/HiMind/Packed-Avatar/resolve/main/PackedAvatar.pt"

# -------------------------------------------------------------------
# Avatar bank metadata for dropdown labels
# -------------------------------------------------------------------

AVATAR_META: Dict[str, Tuple[str, str]] = {
    # female / anime
    "Alison": ("female", "anime"),
    "Amber": ("female", "anime"),
    "Andrea": ("female", "anime"),
    "Angela": ("female", "anime"),
    "Christine": ("female", "anime"),
    "Cynthia": ("female", "anime"),
    "Heidi": ("female", "anime"),
    "Jennifer": ("female", "anime"),
    "Karla": ("female", "anime"),
    "Kristen": ("female", "anime"),
    "Laura": ("female", "anime"),
    "Nancy": ("female", "anime"),
    "Patricia": ("female", "anime"),
    "Rebecca": ("female", "anime"),
    "Sandra": ("female", "anime"),
    "Tara": ("female", "anime"),
    # female / cyber
    "Amanda": ("female", "cyber"),
    "Brenda": ("female", "cyber"),
    "Christina": ("female", "cyber"),
    "Janet": ("female", "cyber"),
    "Jill": ("female", "cyber"),
    "Julie": ("female", "cyber"),
    "Lisa": ("female", "cyber"),
    "Mallory": ("female", "cyber"),
    "Mandy": ("female", "cyber"),
    "Martha": ("female", "cyber"),
    "Melissa": ("female", "cyber"),
    "Michelle": ("female", "cyber"),
    "Regina": ("female", "cyber"),
    # female / drawn
    "Alyssa": ("female", "drawn"),
    "Danielle": ("female", "drawn"),
    "Joan": ("female", "drawn"),
    "Kaitlyn": ("female", "drawn"),
    "Kimberly": ("female", "drawn"),
    "Marie": ("female", "drawn"),
    "Samantha": ("female", "drawn"),
    "Veronica": ("female", "drawn"),
    # female / paint
    "Alejandra": ("female", "paint"),
    "Barbara": ("female", "paint"),
    "Briana": ("female", "paint"),
    "Brittany": ("female", "paint"),
    "Emily": ("female", "paint"),
    "Jacqueline": ("female", "paint"),
    "Jodi": ("female", "paint"),
    "Mary": ("female", "paint"),
    "Rhonda": ("female", "paint"),
    "Savannah": ("female", "paint"),
    "Tammy": ("female", "paint"),
    "Victoria": ("female", "paint"),
    "Yolanda": ("female", "paint"),
    # female / real
    "Amy": ("female", "real"),
    "Ann": ("female", "real"),
    "Ashley": ("female", "real"),
    "Colleen": ("female", "real"),
    "Heather": ("female", "real"),
    "Holly": ("female", "real"),
    "Jordan": ("female", "real"),
    "Kristin": ("female", "real"),
    "Kristine": ("female", "real"),
    "Mariah": ("female", "real"),
    "Pamela": ("female", "real"),
    "Sara": ("female", "real"),
    "Sharon": ("female", "real"),
    # male / anime
    "Brad": ("male", "anime"),
    "Brian": ("male", "anime"),
    "David": ("male", "anime"),
    "Gregory": ("male", "anime"),
    "John": ("male", "anime"),
    "Jose": ("male", "anime"),
    "Lawrence": ("male", "anime"),
    "Robert": ("male", "anime"),
    # male / cyber
    "Daniel": ("male", "cyber"),
    "Hayden": ("male", "cyber"),
    "James": ("male", "cyber"),
    "Jeremy": ("male", "cyber"),
    "Paul": ("male", "cyber"),
    "Ryan": ("male", "cyber"),
    "Sean": ("male", "cyber"),
    # male / drawn
    "Bobby": ("male", "drawn"),
    "George": ("male", "drawn"),
    "Gregg": ("male", "drawn"),
    "Kevin": ("male", "drawn"),
    "Matthew": ("male", "drawn"),
    "Ricky": ("male", "drawn"),
    "Thomas": ("male", "drawn"),
    # male / paint
    "Jacob": ("male", "paint"),
    "Justin": ("male", "paint"),
    "Michael": ("male", "paint"),
    "Nicholas": ("male", "paint"),
    "Steven": ("male", "paint"),
    "William": ("male", "paint"),
    "Zachary": ("male", "paint"),
    # male / real
    "Aaron": ("male", "real"),
    "Andrew": ("male", "real"),
    "Benjamin": ("male", "real"),
    "Christopher": ("male", "real"),
    "Derek": ("male", "real"),
    "Frank": ("male", "real"),
    "Jesse": ("male", "real"),
    "Joseph": ("male", "real"),
}

STYLE_ORDER = ["All", "anime", "cyber", "drawn", "paint", "real"]
GENDER_ORDER = ["All", "female", "male"]


def _to_path(value: Any) -> Optional[str]:
    if value is None:
        return None
    if isinstance(value, (str, os.PathLike)):
        return str(value)
    return str(value)


def _ensure_checkpoints_and_download():
    """
    Ensure the checkpoints directory exists and PackedAvatar.pt is present.
    If PackedAvatar.pt is missing, download it from the HF_PACKED_PT_URL.
    """
    CHECKPOINTS_DIR.mkdir(parents=True, exist_ok=True)
    if DEFAULT_PACKED_PT.exists():
        return str(DEFAULT_PACKED_PT)

    # Download the file into the checkpoints directory
    tmp_file = CHECKPOINTS_DIR / f"PackedAvatar.pt.part"
    final_file = DEFAULT_PACKED_PT
    try:
        with requests.get(HF_PACKED_PT_URL, stream=True, timeout=60) as r:
            r.raise_for_status()
            total = int(r.headers.get("content-length", 0))
            downloaded = 0
            with open(tmp_file, "wb") as f:
                for chunk in r.iter_content(chunk_size=8192):
                    if not chunk:
                        continue
                    f.write(chunk)
                    downloaded += len(chunk)
        tmp_file.replace(final_file)
    except Exception as e:
        # Clean up partial file if present
        try:
            tmp_file.unlink(missing_ok=True)
        except Exception:
            pass
        raise RuntimeError(f"Failed to download PackedAvatar.pt from {HF_PACKED_PT_URL}: {e}") from e

    return str(final_file)


@lru_cache(maxsize=1)
def get_model():
    # Ensure the bundle exists locally (download if necessary) before loading
    packed_path = _ensure_checkpoints_and_download()
    from PackedAvatar import PackedAvatar

    return PackedAvatar(packed_pt_path=packed_path)


def _call_with_supported_kwargs(fn, **kwargs):
    sig = inspect.signature(fn)
    filtered = {k: v for k, v in kwargs.items() if k in sig.parameters}
    return fn(**filtered)


def avatar_choices(gender: str = "All", style: str = "All") -> List[Tuple[str, str]]:
    items = []
    for name, (g, s) in sorted(AVATAR_META.items(), key=lambda kv: kv[0].lower()):
        if gender != "All" and g != gender:
            continue
        if style != "All" and s != style:
            continue
        items.append((f"{name}  ({g}/{s})", name))
    return [("Auto default", "")] + items


def summarize_bundle(bundle: Dict[str, Any]) -> str:
    keys = sorted(bundle.keys())
    lines = [f"- **Keys:** {', '.join(keys)}"]
    for key in ("coeff_path", "crop_pic_path", "crop_info", "crop_preview", "motion_3dmm", "full_3dmm", "full_3dmm_seq", "landmarks"):
        if key in bundle and bundle[key] is not None:
            value = bundle[key]
            if hasattr(value, "shape"):
                shape = tuple(value.shape)
                lines.append(f"- **{key}:** shape `{shape}`")
            elif isinstance(value, (str, os.PathLike)):
                lines.append(f"- **{key}:** `{value}`")
            else:
                lines.append(f"- **{key}:** present")
    return "\n".join(lines)


def extract_bundle(input_path: str, crop_or_resize: str, pic_size: int) -> Tuple[str, str]:
    model = get_model()
    input_path = _to_path(input_path)
    if not input_path:
        raise gr.Error("Please upload an image or video first.")

    extractor = getattr(model, "extract_embeddings", None) or getattr(model, "ExtractEmbeddings", None)
    if extractor is None:
        raise gr.Error("This build of PackedAvatar does not expose extract_embeddings().")

    bundle = _call_with_supported_kwargs(
        extractor,
        input_path=input_path,
        crop_or_resize=crop_or_resize,
        pic_size=pic_size,
    )

    out_dir = Path(tempfile.mkdtemp(prefix="packedavatar_bundle_"))
    out_path = out_dir / f"{Path(input_path).stem}_conditioning.pt"
    torch.save(bundle, out_path)

    summary = summarize_bundle(bundle)
    return str(out_path), summary


def _cleanup_old_tmp_files(root: Path, max_age_seconds: int = 60 * 60 * 6):
    now = time.time()
    for p in glob.glob(str(root / "packedavatar_*.mp4")):
        try:
            mtime = Path(p).stat().st_mtime
            if now - mtime > max_age_seconds:
                Path(p).unlink(missing_ok=True)
        except Exception:
            pass


def generate_video(
    source_image: Optional[str],
    driven_audio: Optional[str],
    avatar_id: str,
    avatar_condition: Optional[str],
    motion_condition: Optional[str],
    ref_video: Optional[str],
    use_ref_video: bool,
    ref_info: str,
    remove_background: bool,
    use_idle_mode: bool,
    length_of_audio: int,
    preprocess: str,
    size: int,
    pose_style: int,
    facerender: str,
    still_mode: bool,
    use_enhancer: bool,
    batch_size: int,
    exp_scale: float,
    use_blink: bool,
    result_dir: str,
) -> Tuple[Optional[str], str, Optional[str]]:
    """
    Generates a video and returns (preview_path, status, download_path).

    The generated file is copied into a temporary results folder (not the repo).
    """
    model = get_model()

    avatar_id = avatar_id or ""
    avatar_id = None if avatar_id == "" else avatar_id

    kwargs = {
        "source_image": _to_path(source_image),
        "driven_audio": _to_path(driven_audio),
        "preprocess": preprocess,
        "still_mode": still_mode,
        "use_enhancer": use_enhancer,
        "batch_size": batch_size,
        "size": size,
        "pose_style": pose_style,
        "facerender": facerender,
        "exp_scale": exp_scale,
        "use_ref_video": use_ref_video,
        "ref_video": _to_path(ref_video),
        "ref_info": ref_info if ref_info != "None" else None,
        "use_idle_mode": use_idle_mode,
        "length_of_audio": length_of_audio,
        "use_blink": use_blink,
        "result_dir": None,  # force runtime to return path but don't write into repo
        "avatar_id": avatar_id,
        "avatar_condition": _to_path(avatar_condition),
        "motion_condition": _to_path(motion_condition),
        "remove_background": remove_background,
        "use_wav2lip": False,
    }

    # Optional: update status in UI before long-running op (Gradio handles this automatically when using .click)
    try:
        result = _call_with_supported_kwargs(model.generate, **kwargs)
    except Exception as e:
        raise gr.Error(f"Generation failed: {e}") from e

    video_path = result[0] if isinstance(result, tuple) else result
    if video_path is None:
        raise gr.Error("Generation completed but no output path was returned by the runtime.")

    # Copy to a temp results directory so the Space repo is not modified
    tmp_root = Path(tempfile.gettempdir()) / "packedavatar_results"
    tmp_root.mkdir(parents=True, exist_ok=True)

    # cleanup old files opportunistically
    _cleanup_old_tmp_files(tmp_root)

    # create a unique filename to avoid collisions
    unique_name = f"packedavatar_{int(time.time())}_{uuid.uuid4().hex[:8]}.mp4"
    tmp_path = tmp_root / unique_name

    try:
        shutil.copy(video_path, tmp_path)
    except Exception:
        # If the runtime already produced an mp4 in a temp location, try to move/copy robustly
        if Path(video_path).exists():
            tmp_path = Path(video_path)
        else:
            raise gr.Error("Failed to copy generated video to temporary results folder.")

    status = f"Done. Preview below. Click Download to save the MP4. (Temp path: {tmp_path})"
    # Return preview path (for gr.Video), status text, and download path (for gr.File)
    return str(tmp_path), status, str(tmp_path)


def refresh_avatar_dropdown(gender: str, style: str):
    return gr.update(choices=avatar_choices(gender, style), value="")


INTRO = """
# PackedAvatar

A bundled SadTalker-based talking-head runtime for Hugging Face Spaces.

Upload a source image + audio, pick an AvatarBank identity, or feed in a reusable conditioning bundle.
The bundled Wav2Lip checkpoint is used automatically when enabled.
"""

AVATAR_TABLE = """
## Packed AvatarBank

**Female**
- **anime:** Alison, Amber, Andrea, Angela, Christine, Cynthia, Heidi, Jennifer, Karla, Kristen, Laura, Nancy, Patricia, Rebecca, Sandra, Tara
- **cyber:** Amanda, Brenda, Christina, Janet, Jill, Julie, Lisa, Mallory, Mandy, Martha, Melissa, Michelle, Regina
- **drawn:** Alyssa, Danielle, Joan, Kaitlyn, Kimberly, Marie, Samantha, Veronica
- **paint:** Alejandra, Barbara, Briana, Brittany, Emily, Jacqueline, Jodi, Mary, Rhonda, Savannah, Tammy, Victoria, Yolanda
- **real:** Amy, Ann, Ashley, Colleen, Heather, Holly, Jordan, Kristin, Kristine, Mariah, Pamela, Sara, Sharon

**Male**
- **anime:** Brad, Brian, David, Gregory, John, Jose, Lawrence, Robert
- **cyber:** Daniel, Hayden, James, Jeremy, Paul, Ryan, Sean
- **drawn:** Bobby, George, Gregg, Kevin, Matthew, Ricky, Thomas
- **paint:** Jacob, Justin, Michael, Nicholas, Steven, William, Zachary
- **real:** Aaron, Andrew, Benjamin, Christopher, Derek, Frank, Jesse, Joseph
"""

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(INTRO)

    with gr.Tabs():
        with gr.Tab("Generate"):
            with gr.Row():
                with gr.Column(scale=1):
                    source_image = gr.Image(
                        label="Source image",
                        type="filepath",
                        sources=["upload"],
                    )
                    driven_audio = gr.Audio(
                        label="Driven audio",
                        type="filepath",
                        sources=["upload"],
                    )

                    gr.Markdown("### AvatarBank selection")
                    with gr.Row():
                        gender_filter = gr.Dropdown(
                            label="Gender filter",
                            choices=GENDER_ORDER,
                            value="All",
                        )
                        style_filter = gr.Dropdown(
                            label="Style filter",
                            choices=STYLE_ORDER,
                            value="All",
                        )
                    avatar_id = gr.Dropdown(
                        label="Avatar ID",
                        choices=avatar_choices(),
                        value="",
                        allow_custom_value=False,
                    )

                    gr.Markdown("### Optional conditioning bundles")
                    avatar_condition = gr.File(
                        label="Avatar conditioning bundle (.pt / .pth / .mat)",
                        file_types=[".pt", ".pth", ".mat"],
                        type="filepath",
                    )
                    motion_condition = gr.File(
                        label="Motion conditioning bundle (.pt / .pth / .mat)",
                        file_types=[".pt", ".pth", ".mat"],
                        type="filepath",
                    )

                    gr.Markdown("### Reference video / idle mode")
                    use_ref_video = gr.Checkbox(label="Use reference video", value=False)
                    ref_video = gr.Video(label="Reference video", format="mp4")
                    ref_info = gr.Dropdown(
                        label="Reference mode",
                        choices=["pose", "blink", "pose+blink", "all", "None"],
                        value="pose",
                    )
                    use_idle_mode = gr.Checkbox(label="Use idle mode", value=False)
                    length_of_audio = gr.Slider(
                        label="Idle length (seconds)",
                        minimum=0,
                        maximum=30,
                        step=1,
                        value=0,
                    )

                    gr.Markdown("### Rendering options")
                    remove_background = gr.Checkbox(label="Remove background", value=False)
                    use_enhancer = gr.Checkbox(label="Use enhancer (GFPGAN)", value=False)
                    still_mode = gr.Checkbox(label="Still mode", value=False)
                    use_blink = gr.Checkbox(label="Use blink", value=True)

                    preprocess = gr.Dropdown(
                        label="Preprocess",
                        choices=["crop"],
                        value="crop",
                    )
                    facerender = gr.Dropdown(
                        label="Face renderer",
                        choices=["facevid2vid", "pirender"],
                        value="facevid2vid",
                    )
                    size = gr.Dropdown(
                        label="Size",
                        choices=[256, 512],
                        value=256,
                    )
                    pose_style = gr.Slider(
                        label="Pose style",
                        minimum=0,
                        maximum=45,
                        step=1,
                        value=0,
                    )
                    exp_scale = gr.Slider(
                        label="Expression scale",
                        minimum=0.5,
                        maximum=2.0,
                        step=0.1,
                        value=1.0,
                    )
                    batch_size = gr.Slider(
                        label="Batch size",
                        minimum=1,
                        maximum=8,
                        step=1,
                        value=1,
                    )
                    result_dir = gr.Textbox(label="Result directory (optional)", value="./results")

                    run_btn = gr.Button("Generate", variant="primary")

                with gr.Column(scale=1):
                    output_video = gr.Video(label="Output video", format="mp4")
                    status = gr.Textbox(label="Status", lines=4)
                    download_file = gr.File(label="Download generated video (click to save)", file_count="single")

            gender_filter.change(
                refresh_avatar_dropdown,
                inputs=[gender_filter, style_filter],
                outputs=[avatar_id],
            )
            style_filter.change(
                refresh_avatar_dropdown,
                inputs=[gender_filter, style_filter],
                outputs=[avatar_id],
            )

            run_btn.click(
                fn=generate_video,
                inputs=[
                    source_image,
                    driven_audio,
                    avatar_id,
                    avatar_condition,
                    motion_condition,
                    ref_video,
                    use_ref_video,
                    ref_info,
                    remove_background,
                    use_idle_mode,
                    length_of_audio,
                    preprocess,
                    size,
                    pose_style,
                    facerender,
                    still_mode,
                    use_enhancer,
                    batch_size,
                    exp_scale,
                    use_blink,
                    result_dir,
                ],
                outputs=[output_video, status, download_file],
            )

        with gr.Tab("Extract Conditioning"):
            gr.Markdown(
                """
                Upload a source image or reference video, extract a reusable conditioning bundle,
                then feed the saved `.pt` file back into `avatar_condition` or `motion_condition`.
                """
            )

            with gr.Row():
                with gr.Column(scale=1):
                    conditioning_input = gr.File(
                        label="Image or video",
                        file_types=[".png", ".jpg", ".jpeg", ".webp", ".mp4", ".mov", ".avi", ".mkv", ".webm"],
                        type="filepath",
                    )
                    crop_or_resize = gr.Dropdown(
                        label="Crop mode",
                        choices=["crop"],
                        value="crop",
                    )
                    pic_size = gr.Slider(
                        label="Picture size",
                        minimum=128,
                        maximum=512,
                        step=32,
                        value=256,
                    )
                    extract_btn = gr.Button("Extract bundle", variant="primary")

                with gr.Column(scale=1):
                    bundle_file = gr.File(label="Saved bundle")
                    bundle_summary = gr.Markdown()

            extract_btn.click(
                fn=extract_bundle,
                inputs=[conditioning_input, crop_or_resize, pic_size],
                outputs=[bundle_file, bundle_summary],
            )

        with gr.Tab("AvatarBank"):
            gr.Markdown(AVATAR_TABLE)



if __name__ == "__main__":
    # Attempt to ensure the bundle is present at startup (download if needed).
    try:
        _ensure_checkpoints_and_download()
    except Exception as e:
        # If download fails, print a warning but still launch the UI so the user can upload a local bundle.
        print(f"Warning: failed to ensure PackedAvatar.pt is present: {e}")

    # Optional cleanup at startup
    try:
        _cleanup_old_tmp_files(Path(tempfile.gettempdir()) / "packedavatar_results")
    except Exception:
        pass

    demo.queue(default_concurrency_limit=1)
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
    )