Wan2GP/shared/ffmpeg_setup.py

import os
import sys
import shutil
import tarfile
import tempfile
import typing
import importlib.util
from pathlib import Path

import requests
from tqdm import tqdm
import zipfile


__all__ = ["download_ffmpeg"]


_DLL_DIRECTORY_HANDLES = {}


def download_ffmpeg(bin_directory: typing.Optional[typing.Union[str, Path]] = None):
    """Ensure ffmpeg/ffprobe binaries (and their shared libs) are downloaded and on PATH."""
    required_binaries = ["ffmpeg", "ffprobe"]
    if os.name == "nt":
        required_binaries.append("ffplay")

    if bin_directory is None:
        repo_root = Path(__file__).resolve().parents[1]
        bin_dir = repo_root / "ffmpeg_bins"
    else:
        bin_dir = Path(bin_directory)

    bin_dir.mkdir(parents=True, exist_ok=True)
    repo_root = bin_dir.parent

    def _candidate_name(name: str) -> str:
        if os.name == "nt" and not name.endswith(".exe"):
            return f"{name}.exe"
        return name

    def _quarantine_root_ffmpeg():
        root_ffmpeg = repo_root / _candidate_name("ffmpeg")
        if not root_ffmpeg.is_file():
            return
        quarantine_dir = repo_root / "ffmpeg_quarantine"
        quarantine_dir.mkdir(parents=True, exist_ok=True)
        target_path = quarantine_dir / root_ffmpeg.name
        if target_path.exists():
            stem = target_path.stem
            suffix = target_path.suffix
            idx = 1
            while True:
                candidate = quarantine_dir / f"{stem}_{idx}{suffix}"
                if not candidate.exists():
                    target_path = candidate
                    break
                idx += 1
        shutil.move(str(root_ffmpeg), str(target_path))
        print(
            f"[FFmpeg] Quarantined root binary: {root_ffmpeg} -> {target_path}. "
            "Reason: ffmpeg.exe in the project root can be picked from CWD and break TorchCodec DLL loading on Windows. Quarantined file can be deleted if unused."
        )

    def _ensure_bin_dir_on_path():
        current_path = os.environ.get("PATH", "")
        path_parts = current_path.split(os.pathsep) if current_path else []

        def _normalize(p: str) -> str:
            p = os.path.normpath(p)
            return os.path.normcase(p) if os.name == "nt" else p

        prioritized = []
        seen = set()
        for d in [bin_dir, repo_root]:
            key = _normalize(str(d))
            if key not in seen:
                prioritized.append(str(d))
                seen.add(key)

        filtered = [p for p in path_parts if _normalize(p) not in seen]
        os.environ["PATH"] = os.pathsep.join(prioritized + filtered)

    def _add_windows_dll_directory(path: Path):
        if os.name != "nt" or not hasattr(os, "add_dll_directory") or not path.is_dir():
            return
        normalized = os.path.normcase(os.path.normpath(str(path)))
        if normalized in _DLL_DIRECTORY_HANDLES:
            return
        _DLL_DIRECTORY_HANDLES[normalized] = os.add_dll_directory(str(path))

    def _ensure_decord_dll_directory():
        if os.name != "nt":
            return
        decord_spec = importlib.util.find_spec("decord")
        if not decord_spec or not decord_spec.submodule_search_locations:
            return
        for location in decord_spec.submodule_search_locations:
            decord_dir = Path(location)
            if any(decord_dir.glob("avcodec-*.dll")):
                _add_windows_dll_directory(decord_dir)
                break

    def _ensure_library_path():
        if os.name == "nt":
            return
        current_ld = os.environ.get("LD_LIBRARY_PATH", "")
        ld_parts = [p for p in current_ld.split(os.pathsep) if p]
        if str(bin_dir) not in ld_parts:
            os.environ["LD_LIBRARY_PATH"] = os.pathsep.join([str(bin_dir)] + ld_parts) if current_ld else str(bin_dir)

    _quarantine_root_ffmpeg()
    _ensure_bin_dir_on_path()
    _ensure_decord_dll_directory()
    _ensure_library_path()

    def _resolve_path(name: str) -> typing.Optional[Path]:
        # Check ffmpeg_bins folder first
        candidate = bin_dir / _candidate_name(name)
        if candidate.exists():
            return candidate
        # Check repo root folder (some users put ffmpeg there)
        repo_root = bin_dir.parent
        candidate_root = repo_root / _candidate_name(name)
        if candidate_root.exists():
            return candidate_root
        # Fall back to system PATH
        resolved = shutil.which(name)
        return Path(resolved) if resolved else None

    def _binary_exists(name: str) -> bool:
        return _resolve_path(name) is not None

    def _local_binary_exists(name: str) -> bool:
        candidate = bin_dir / _candidate_name(name)
        return candidate.exists()

    def _libs_present() -> bool:
        if os.name == "nt":
            return True
        ffmpeg_path = _resolve_path("ffmpeg")
        if ffmpeg_path and ffmpeg_path.parent != bin_dir:
            # System-installed FFmpeg should already have its shared libs available.
            return True
        return any(bin_dir.glob("libavdevice.so*"))

    def _set_env_vars():
        ffmpeg_path = _resolve_path("ffmpeg")
        ffprobe_path = _resolve_path("ffprobe")
        ffplay_path = _resolve_path("ffplay") if "ffplay" in required_binaries else None
        if ffmpeg_path:
            os.environ["FFMPEG_BINARY"] = str(ffmpeg_path)
        if ffprobe_path:
            os.environ["FFPROBE_BINARY"] = str(ffprobe_path)
        if ffplay_path:
            os.environ["FFPLAY_BINARY"] = str(ffplay_path)

    if os.name == "nt":
        missing = [binary for binary in required_binaries if not _local_binary_exists(binary)]
        libs_ok = True
    else:
        missing = [binary for binary in required_binaries if not _binary_exists(binary)]
        libs_ok = _libs_present()
    if not missing and libs_ok:
        _set_env_vars()
        return

    def _download_file(url: str, destination: Path):
        with requests.get(url, stream=True, timeout=120) as response:
            response.raise_for_status()
            total = int(response.headers.get("Content-Length", 0))
            with open(destination, "wb") as file_handle, tqdm(
                total=total if total else None,
                unit="B",
                unit_scale=True,
                desc=f"Downloading {destination.name}"
            ) as progress:
                for chunk in response.iter_content(chunk_size=8192):
                    if not chunk:
                        continue
                    file_handle.write(chunk)
                    progress.update(len(chunk))

    def _download_windows_build():
        exes = [_candidate_name(name) for name in required_binaries]
        api_url = "https://api.github.com/repos/GyanD/codexffmpeg/releases/latest"
        response = requests.get(api_url, headers={"Accept": "application/vnd.github+json"}, timeout=30)
        response.raise_for_status()
        assets = response.json().get("assets", [])
        zip_asset = next((asset for asset in assets if asset.get("name", "").endswith("essentials_build.zip")), None)
        if not zip_asset:
            raise RuntimeError("Unable to locate FFmpeg essentials build for Windows.")
        zip_path = bin_dir / zip_asset["name"]
        _download_file(zip_asset["browser_download_url"], zip_path)

        try:
            with zipfile.ZipFile(zip_path) as archive:
                for member in archive.namelist():
                    normalized = member.replace("\\", "/")
                    if "/bin/" not in normalized:
                        continue
                    base_name = os.path.basename(normalized)
                    if base_name not in exes and not base_name.lower().endswith(".dll"):
                        continue
                    destination = bin_dir / base_name
                    with archive.open(member) as source, open(destination, "wb") as target:
                        shutil.copyfileobj(source, target)
                    destination.chmod(0o755)
        finally:
            try:
                zip_path.unlink(missing_ok=True)
            except TypeError:
                if zip_path.exists():
                    zip_path.unlink()

    def _download_posix_build():
        api_url = "https://api.github.com/repos/BtbN/FFmpeg-Builds/releases/latest"
        response = requests.get(api_url, headers={"Accept": "application/vnd.github+json"}, timeout=30)
        response.raise_for_status()
        assets = response.json().get("assets", [])
        if sys.platform.startswith("linux"):
            keywords = ["linux64", "gpl"]
        elif sys.platform == "darwin":
            keywords = ["macos64", "gpl"]
        else:
            raise RuntimeError("Unsupported platform for automatic FFmpeg download.")

        tar_asset = next(
            (
                asset for asset in assets
                if asset.get("name", "").endswith(".tar.xz") and all(k in asset.get("name", "") for k in keywords)
            ),
            None
        )
        if not tar_asset:
            raise RuntimeError("Unable to locate a suitable FFmpeg build for this platform.")

        tar_path = bin_dir / tar_asset["name"]
        _download_file(tar_asset["browser_download_url"], tar_path)
        try:
            with tempfile.TemporaryDirectory() as tmp_dir:
                tmp_path = Path(tmp_dir)
                with tarfile.open(tar_path, "r:xz") as archive:
                    archive.extractall(tmp_path)

                build_root = None
                for candidate in tmp_path.iterdir():
                    if (candidate / "bin").exists():
                        build_root = candidate
                        break

                if build_root is None:
                    raise RuntimeError("Unable to locate FFmpeg bin directory in downloaded archive.")

                bin_source = build_root / "bin"
                lib_source = build_root / "lib"

                for binary in required_binaries:
                    source_file = bin_source / binary
                    if not source_file.exists():
                        continue
                    destination = bin_dir / binary
                    shutil.copy2(source_file, destination)
                    destination.chmod(0o755)

                if lib_source.exists():
                    for lib_file in lib_source.rglob("*.so*"):
                        destination = bin_dir / lib_file.name
                        shutil.copy2(lib_file, destination)
                        destination.chmod(0o755)
        finally:
            try:
                tar_path.unlink(missing_ok=True)
            except TypeError:
                if tar_path.exists():
                    tar_path.unlink()

    try:
        if os.name == "nt":
            _download_windows_build()
        else:
            _download_posix_build()
    except Exception as exc:
        print(f"Failed to download FFmpeg binaries automatically: {exc}")
        return

    if os.name == "nt":
        if not all(_local_binary_exists(binary) for binary in required_binaries):
            print("FFmpeg binaries are still missing after download; please install them manually.")
            return
    else:
        if not all(_binary_exists(binary) for binary in required_binaries):
            print("FFmpeg binaries are still missing after download; please install them manually.")
            return

    _ensure_bin_dir_on_path()
    _ensure_library_path()
    _set_env_vars()