Spaces:
Sleeping
Sleeping
| """Talking head (SadTalker) on ZeroGPU — identity-locked performance. | |
| Why this exists (Tyler's quality bar, 2026-11-06): the performer's face must | |
| stay THE SAME — only the mouth (and head, subtly) may move. Generative i2v | |
| (LTX) re-dreams the face and morphs it; SadTalker instead WARPS the original | |
| photo using 3D facial coefficients driven by the song's audio, so identity | |
| is preserved by construction. Input match: one portrait + one audio slice. | |
| Proven locally (2026-11-06): on the zombie's prosthetic mask, landmark | |
| detection, 3DMM extraction and audio-to-expression ALL pass. The local CPU | |
| render died silently twice (sandbox limits), which is why this runs on the | |
| Space's GPU instead. | |
| Isolation design (the live app must never break because of this beta): | |
| SadTalker's era-sensitive stack (numpy 1.26, librosa, basicsr…) is installed | |
| into a private venv (--system-site-packages, so it borrows the app's CUDA | |
| torch) under FILES/sadtalker_vendor, and inference runs as a SUBPROCESS of | |
| that venv's python. requirements.txt and the app's own environment are | |
| untouched. Fallback if venv creation is unavailable: pip --target with the | |
| exact package set proven in the local spike, child = app python + | |
| PYTHONPATH. | |
| """ | |
| import json | |
| import os | |
| import subprocess | |
| import sys | |
| import urllib.request | |
| import zipfile | |
| from pathlib import Path | |
| import spaces | |
| from server import FILES | |
| VENDOR = FILES / "sadtalker_vendor" | |
| SRC_ZIP_URL = "https://codeload.github.com/OpenTalker/SadTalker/zip/refs/heads/main" | |
| CKPT_BASE = "https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc" | |
| CKPTS = ["SadTalker_V0.0.2_256.safetensors", "mapping_00109-model.pth.tar", | |
| "mapping_00229-model.pth.tar"] | |
| # Top-level pins for the venv install; pip resolves the transitive closure | |
| # (versions proven in the local spike venv where it matters). | |
| VENV_PINS = ["numpy==1.26.4", "librosa==0.10.2", "face-alignment==1.4.1", | |
| "kornia==0.7.2", "basicsr==1.4.2", "facexlib==0.3.0", "gfpgan==1.3.8", | |
| "resampy==0.4.3", "yacs==0.1.8", "pydub==0.25.1", "opencv-python", | |
| "safetensors", "imageio", "imageio-ffmpeg", "scikit-image", "scipy"] | |
| # Exact set from the local spike's pip freeze (minus torch/torchvision, | |
| # which the child borrows from the app env) — used only by the --target | |
| # fallback path. | |
| FREEZE_PINS = ("absl-py==2.5.1 addict==2.4.0 audioread==3.1.0 basicsr==1.4.2 " | |
| "certifi==2026.10.23 cffi==2.1.1 charset-normalizer==3.5.2 cloudpickle==3.1.2 " | |
| "contourpy==1.4.0 cycler==0.12.1 Cython==3.0.11 decorator==5.3.1 " | |
| "face-alignment==1.4.1 facexlib==0.3.0 filelock==4.2.4 filterpy==2.0.0 " | |
| "fonttools==4.68.0 fsspec==2026.9.0 future==1.0.0 gfpgan==1.3.8 grpcio==1.85.1 " | |
| "idna==3.20 ImageIO==2.38.1 imageio-ffmpeg==0.8.1 Jinja2==3.1.6 joblib==1.6.0 " | |
| "kiwisolver==1.5.1 kornia==0.7.2 kornia_rs==0.2.0 lazy-loader==0.6 " | |
| "librosa==0.10.2 llvmlite==0.50.0 lmdb==3.0.0 Markdown==3.11.3 MarkupSafe==3.0.3 " | |
| "matplotlib==3.11.2 mpmath==1.3.0 msgpack==1.2.4 narwhals==2.28.0 networkx==3.7 " | |
| "numba==0.68.1 numpy==1.26.4 opencv-python==4.11.0.86 packaging==26.3 " | |
| "pillow==13.0.0 platformdirs==4.14.5 pooch==1.9.0 protobuf==7.36.2 " | |
| "pycparser==3.12 pydub==0.25.1 pyparsing==3.3.3 python-dateutil==4.0.0 " | |
| "PyYAML==6.0.3 requests==2.34.2 resampy==0.4.3 safetensors==0.8.0 " | |
| "scikit-image==0.26.0 scikit-learn==1.9.1 scipy==1.17.1 setuptools==84.2.1 " | |
| "soundfile==0.14.0 soxr==1.1.0 sympy==1.14.0 tb-nightly==2.21.0a20251023 " | |
| "tensorboard-data-server==0.7.2 threadpoolctl==3.8.0 tifffile==2026.3.3 " | |
| "tqdm==4.70.1 typing_extensions==4.16.0 tzdata==2026.6 urllib3==2.8.0 " | |
| "Werkzeug==3.1.9 wheel==0.48.0 yacs==0.1.8 yapf==0.43.0").split() | |
| def _run(cmd, **kw): | |
| print("[talking] $", " ".join(str(c) for c in cmd)[:200], flush=True) | |
| return subprocess.run(cmd, capture_output=True, text=True, **kw) | |
| def _child(): | |
| """(python, env) for SadTalker subprocesses, per the READY marker.""" | |
| info = json.loads((VENDOR / "READY").read_text()) | |
| if info["mode"] == "venv": | |
| return str(VENDOR / "venv" / "bin" / "python"), None | |
| env = dict(os.environ) | |
| env["PYTHONPATH"] = str(VENDOR / "pkgs") + os.pathsep + env.get("PYTHONPATH", "") | |
| return sys.executable, env | |
| def ensure_talking(): | |
| """First use (CPU): download source + checkpoints (~1GB) and build the | |
| isolated SadTalker environment. Idempotent per container via READY.""" | |
| ready = VENDOR / "READY" | |
| if ready.exists(): | |
| print("[talking] vendor ready (cached)", flush=True) | |
| return json.loads(ready.read_text())["src"] | |
| VENDOR.mkdir(parents=True, exist_ok=True) | |
| src_root = VENDOR / "src" | |
| if not (src_root / "inference.py").exists(): | |
| print("[talking] downloading SadTalker source…", flush=True) | |
| zip_path = VENDOR / "sadtalker.zip" | |
| urllib.request.urlretrieve(SRC_ZIP_URL, zip_path) | |
| with zipfile.ZipFile(zip_path) as z: | |
| z.extractall(VENDOR / "extract") | |
| extracted = next((VENDOR / "extract").glob("SadTalker-*")) | |
| extracted.rename(src_root) | |
| zip_path.unlink(missing_ok=True) | |
| ck = src_root / "checkpoints" | |
| ck.mkdir(exist_ok=True) | |
| for name in CKPTS: | |
| dest = ck / name | |
| if not dest.exists() or dest.stat().st_size < 1_000_000: | |
| print(f"[talking] downloading checkpoint {name}…", flush=True) | |
| urllib.request.urlretrieve(f"{CKPT_BASE}/{name}", dest) | |
| mode = None | |
| print("[talking] building isolated environment (venv)…", flush=True) | |
| r = _run([sys.executable, "-m", "venv", "--system-site-packages", str(VENDOR / "venv")], | |
| timeout=300) | |
| if r.returncode == 0: | |
| vpy = str(VENDOR / "venv" / "bin" / "python") | |
| r = _run([vpy, "-m", "pip", "install", "-q", "--no-cache-dir"] + VENV_PINS, timeout=1500) | |
| if r.returncode == 0: | |
| mode = "venv" | |
| else: | |
| print("[talking] venv pip failed:", (r.stderr or "")[-600:], flush=True) | |
| if mode is None: | |
| print("[talking] falling back to pip --target…", flush=True) | |
| r = _run([sys.executable, "-m", "pip", "install", "-q", "--no-cache-dir", | |
| "--target", str(VENDOR / "pkgs"), "--no-deps"] + FREEZE_PINS, timeout=1500) | |
| if r.returncode != 0: | |
| raise RuntimeError(f"SadTalker env install failed: {(r.stderr or '')[-500:]}") | |
| mode = "target" | |
| ready.write_text(json.dumps({"mode": mode, "src": str(src_root)})) | |
| py, env = _child() | |
| probe = _run([py, "-c", "import cv2, librosa, face_alignment, safetensors; print('imports ok')"], | |
| env=env, timeout=180) | |
| print("[talking] import probe:", (probe.stdout or probe.stderr or "")[-300:], flush=True) | |
| if probe.returncode != 0: | |
| ready.unlink(missing_ok=True) | |
| raise RuntimeError(f"SadTalker imports failed: {(probe.stderr or '')[-500:]}") | |
| print(f"[talking] ready (mode={mode})", flush=True) | |
| return str(src_root) | |
| def sadtalker_gpu(image_path, audio_path, out_path): | |
| """Run SadTalker inference in the vendored child process. The child must | |
| see the GPU ZeroGPU attached for this call — probed first so a subprocess | |
| visibility failure is reported as itself, not as a timeout.""" | |
| src = json.loads((VENDOR / "READY").read_text())["src"] | |
| py, env = _child() | |
| probe = _run([py, "-c", "import torch; print('CUDA_OK' if torch.cuda.is_available() else 'CUDA_NO')"], | |
| env=env, timeout=60) | |
| verdict = (probe.stdout or "").strip() | |
| print(f"[talking] child CUDA probe: {verdict}", flush=True) | |
| if "CUDA_OK" not in verdict: | |
| raise RuntimeError("SadTalker child process cannot see the GPU " | |
| f"(probe: {verdict or probe.stderr[-200:]})") | |
| result_dir = Path(out_path).parent / f"sadtalker_out_{Path(out_path).stem}" | |
| cmd = [py, "inference.py", "--driven_audio", str(audio_path), | |
| "--source_image", str(image_path), "--checkpoint_dir", "checkpoints", | |
| "--result_dir", str(result_dir), "--size", "256", | |
| "--preprocess", "crop", "--batch_size", "4"] | |
| print("[talking] running inference…", flush=True) | |
| proc = _run(cmd, cwd=src, env=env, timeout=88) | |
| if proc.returncode != 0: | |
| tail = (proc.stderr or proc.stdout or "")[-1200:] | |
| raise RuntimeError(f"SadTalker inference failed: {tail}") | |
| videos = sorted(result_dir.rglob("*.mp4"), key=lambda p: p.stat().st_mtime, reverse=True) | |
| if not videos: | |
| raise RuntimeError("SadTalker produced no video") | |
| Path(out_path).write_bytes(videos[0].read_bytes()) | |
| print(f"[talking] wrote {out_path}", flush=True) | |
| return str(out_path) | |