Update app.py

app.py

@@ -1,6 +1,5 @@
-import os
-import shutil
 import subprocess
+import threading
 from pathlib import Path
 import gradio as gr
 
@@ -9,40 +8,98 @@ SETUP_FLAG = Path("setup_done.txt")
 OUTPUT_DIR = Path("outputs")
 OUTPUT_DIR.mkdir(exist_ok=True)
 
+# Hugging Face mirror containing the checkpoint files (fast + reliable)
+HF_BASE = "https://huggingface.co/camenduru/video-retalking/resolve/main"
+
+# Files needed by Video-Retalking pipeline (common minimal set)
+CHECKPOINT_FILES = [
+    "30_net_gen.pth",
+    "BFM.zip",
+    "DNet.pt",
+    "ENet.pth",
+    "LNet.pth",
+    "ParseNet-latest.pth",
+    "RetinaFace-R50.pth",
+    "expression.mat",
+    "face3d_pretrain_epoch_20.pth",
+    "GFPGANv1.3.pth",
+    "GPEN-BFR-512.pth",
+    "shape_predictor_68_face_landmarks.dat",
+]
+
+_setup_lock = threading.Lock()
+
+
 def run(cmd, cwd=None):
-    """
-    Run a shell command safely and raise a helpful error if it fails.
-    """
     p = subprocess.run(
         cmd,
         cwd=str(cwd) if cwd else None,
         stdout=subprocess.PIPE,
         stderr=subprocess.STDOUT,
         text=True,
-        check=False
+        check=False,
     )
     if p.returncode != 0:
-        raise RuntimeError(f"Command failed ({p.returncode}): {' '.join(cmd)}\n\n{p.stdout}")
+        raise RuntimeError(
+            f"Command failed ({p.returncode}): {' '.join(cmd)}\n\n{p.stdout}"
+        )
     return p.stdout
 
-def setup():
+
+def download_file(url: str, dest: Path):
     """
-    One-time setup: clone repo + download models.
-    Uses a flag file so it doesn't redo work every request.
+    Download using curl with resume support (-C -) and follow redirects (-L).
     """
-    if SETUP_FLAG.exists() and REPO_DIR.exists():
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    # Skip if already present
+    if dest.exists() and dest.stat().st_size > 0:
         return
+    run(["bash", "-lc", f'curl -L -C - --retry 5 --retry-delay 2 -o "{dest}" "{url}"'])
+
 
-    print("Setting up Video-Retalking...")
+def ensure_checkpoints():
+    ckpt_dir = REPO_DIR / "checkpoints"
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
 
-    if not REPO_DIR.exists():
-        run(["git", "clone", "https://github.com/OpenTalker/video-retalking.git", str(REPO_DIR)])
+    # Download required checkpoint files
+    for fname in CHECKPOINT_FILES:
+        url = f"{HF_BASE}/{fname}"
+        download_file(url, ckpt_dir / fname)
 
-    # Download models (idempotent in most cases)
-    run(["bash", "scripts/download_models.sh"], cwd=REPO_DIR)
+    # Unzip BFM.zip -> checkpoints/BFM/
+    bfm_zip = ckpt_dir / "BFM.zip"
+    bfm_dir = ckpt_dir / "BFM"
+    if not bfm_dir.exists():
+        bfm_dir.mkdir(parents=True, exist_ok=True)
+        run(["unzip", "-q", str(bfm_zip), "-d", str(bfm_dir)])
+
+
+def setup():
+    """
+    One-time setup: clone repo, pull LFS (best effort), download checkpoints.
+    Guarded by a lock to prevent multiple concurrent setup runs.
+    """
+    with _setup_lock:
+        if SETUP_FLAG.exists() and REPO_DIR.exists():
+            return
+
+        print("Setting up Video-Retalking...")
+
+        if not REPO_DIR.exists():
+            run(["git", "clone", "https://github.com/OpenTalker/video-retalking.git", str(REPO_DIR)])
+
+        # Best-effort: fetch any git-lfs files if repo uses them
+        try:
+            run(["git", "lfs", "pull"], cwd=REPO_DIR)
+        except Exception:
+            # Not fatal; we'll fetch checkpoints from HF mirror anyway
+            pass
+
+        ensure_checkpoints()
+
+        SETUP_FLAG.touch()
+        print("✅ Setup complete!")
 
-    SETUP_FLAG.touch()
-    print("✅ Setup complete!")
 
 def generate(image_path, audio_path):
     if not image_path or not audio_path:
@@ -59,14 +116,10 @@ def generate(image_path, audio_path):
         if not audio_path.exists():
             return None, f"❌ Audio not found: {audio_path}"
 
-        # Unique output name per run
         out_path = (OUTPUT_DIR / "result.mp4").resolve()
-
-        # Clean old output if any
         if out_path.exists():
             out_path.unlink()
 
-        # Run inference
         cmd = [
             "python",
             "inference.py",
@@ -78,11 +131,12 @@ def generate(image_path, audio_path):
 
         if out_path.exists():
             return str(out_path), "✅ Video generated successfully!"
-        return None, "❌ Failed to generate video (no output file created)."
+        return None, "❌ Failed (no output file created)."
 
     except Exception as e:
         return None, f"❌ Error: {e}"
 
+
 demo = gr.Interface(
     fn=generate,
     inputs=[
@@ -91,7 +145,7 @@ demo = gr.Interface(
     ],
     outputs=[
         gr.Video(label="📹 Generated Video"),
-        gr.Textbox(label="Status", lines=3),
+        gr.Textbox(label="Status", lines=4),
     ],
     title="🎬 Video-Retalking Lip Sync",
     description="Upload a face image and audio to generate a lip-synced video.",