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")), )