AnyImage/sync_models.py

"""Prepare AnyImage ONNX models and sync them to OmooLab R2."""

import argparse
import hashlib
import shutil
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path
from urllib.request import Request, urlopen


PROJECT_ROOT = Path(__file__).resolve().parent
MODELS_DIR = PROJECT_ROOT / "models"
SOURCE_ROOT = PROJECT_ROOT / "src" / "anyimage"
sys.path.insert(0, str(SOURCE_ROOT))
from server.model_catalog import DOWNLOADABLE_MODELS

R2_ENDPOINT = "https://models.omoolab.xyz"
CHUNK_SIZE = 1024 * 1024


@dataclass(frozen=True)
class ModelFile:
    folder: str
    size: int
    sha256: str
    sources: tuple[str, ...]
    filename: str = "model.onnx"

    @property
    def destination(self):
        return MODELS_DIR / self.folder / self.filename


def hf_url(repository, filename):
    return f"https://huggingface.co/{repository}/resolve/main/{filename}"


def r2_url(folder, filename="model.onnx"):
    return f"{R2_ENDPOINT}/{folder}/{filename}"


def catalog_model_files():
    files = []
    for model in DOWNLOADABLE_MODELS.values():
        for filename, size, sha256 in model.r2_files:
            relative_path = Path(filename)
            folder = str(Path(model.r2_directory) / relative_path.parent)
            sources = []
            if model.huggingface_repository:
                sources.append(
                    hf_url(model.huggingface_repository, filename)
                )
            sources.append(r2_url(model.r2_directory, filename))
            files.append(
                ModelFile(
                    folder,
                    size,
                    sha256,
                    tuple(sources),
                    relative_path.name,
                )
            )
    return tuple(files)


MODEL_FILES = catalog_model_files()


def file_sha256(path):
    digest = hashlib.sha256()
    with Path(path).open("rb") as stream:
        while chunk := stream.read(CHUNK_SIZE):
            digest.update(chunk)
    return digest.hexdigest()


def verify_existing(model):
    path = model.destination
    if not path.is_file():
        return False
    if path.stat().st_size != model.size:
        raise RuntimeError(f"Existing model has an unexpected size: {path}")
    if file_sha256(path) != model.sha256:
        raise RuntimeError(f"Existing model has an unexpected SHA-256: {path}")
    print(f"Skip existing: {path.relative_to(PROJECT_ROOT)}")
    return True


def download_once(model, url):
    destination = model.destination
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = destination.with_suffix(".onnx.part")
    temporary.unlink(missing_ok=True)
    request = Request(url, headers={"User-Agent": "AnyImage model sync"})
    print(f"Download: {url}")
    downloaded = 0
    digest = hashlib.sha256()
    try:
        with urlopen(request, timeout=60) as response, temporary.open("wb") as output:
            while chunk := response.read(CHUNK_SIZE):
                output.write(chunk)
                digest.update(chunk)
                downloaded += len(chunk)
                print(
                    f"\r  {downloaded / 1024**2:.1f} / "
                    f"{model.size / 1024**2:.1f} MiB",
                    end="",
                    flush=True,
                )
        print()
        if downloaded != model.size:
            raise RuntimeError(f"size mismatch: {downloaded} != {model.size}")
        if digest.hexdigest() != model.sha256:
            raise RuntimeError("SHA-256 mismatch")
        temporary.replace(destination)
    except Exception:
        temporary.unlink(missing_ok=True)
        raise


def prepare_model(model):
    if verify_existing(model):
        return
    failures = []
    for url in model.sources:
        try:
            download_once(model, url)
            return
        except Exception as error:
            failures.append(f"{url}: {error}")
            print(f"Source failed: {error}", file=sys.stderr)
    raise RuntimeError("All sources failed:\n" + "\n".join(failures))


def sync_models():
    executable = shutil.which("rclone")
    if executable is None:
        raise RuntimeError("rclone is not available on PATH")
    subprocess.run(
        [
            executable,
            "sync",
            "models",
            "omoolab-r2:ai-models",
            "--progress",
        ],
        cwd=PROJECT_ROOT,
        check=True,
    )


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--download-only",
        action="store_true",
        help="Prepare and verify models without running rclone",
    )
    return parser.parse_args()


def main():
    args = parse_args()
    for model in MODEL_FILES:
        prepare_model(model)
    if not args.download_only:
        sync_models()


if __name__ == "__main__":
    main()