"""lipsync_models.py — download any MISSING LatentSync models into the node's folder.

Run with the ComfyUI python:
    C:\\ComfyUI\\.venv\\Scripts\\python.exe C:\\malin\\lipsync_models.py

Safe: it SKIPS any file that already exists, so it will not overwrite models you already have.
Source = the public chunyu-li/LatentSync HuggingFace repo (no login needed).
"""
import os, urllib.request

BASE = r"C:\ComfyUI\models\lipsync\latentsync"
FILES = {
    "whisper/tiny.pt":       "https://huggingface.co/chunyu-li/LatentSync/resolve/main/whisper/tiny.pt",
    "latentsync_unet.pt":    "https://huggingface.co/chunyu-li/LatentSync/resolve/main/latentsync_unet.pt",
    "latentsync_syncnet.pt": "https://huggingface.co/chunyu-li/LatentSync/resolve/main/latentsync_syncnet.pt",
    # The node loads the VAE as a local diffusers folder named exactly "sd-vae-ft-mse"
    "sd-vae-ft-mse/config.json":
        "https://huggingface.co/stabilityai/sd-vae-ft-mse/resolve/main/config.json",
    "sd-vae-ft-mse/diffusion_pytorch_model.safetensors":
        "https://huggingface.co/stabilityai/sd-vae-ft-mse/resolve/main/diffusion_pytorch_model.safetensors",
}


def fetch(url, dst, rel):
    os.makedirs(os.path.dirname(dst), exist_ok=True)
    if os.path.exists(dst) and os.path.getsize(dst) > 100:
        print(f"[skip - already there] {rel}")
        return
    print(f"[downloading] {rel} ...")
    req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
    with urllib.request.urlopen(req) as r:
        total = int(r.headers.get("content-length", 0))
        done = 0
        with open(dst, "wb") as f:
            while True:
                chunk = r.read(1 << 20)
                if not chunk:
                    break
                f.write(chunk)
                done += len(chunk)
                if total:
                    print(f"\r   {done * 100 // total}%  ({done // (1 << 20)}/{total // (1 << 20)} MB)",
                          end="", flush=True)
    print(f"\n[done] {rel}")


for rel, url in FILES.items():
    fetch(url, os.path.join(BASE, rel.replace("/", os.sep)), rel)

print("\nAll LatentSync models in place. Re-run your workflow.")