AnyImage/tools/models/cli.py

"""Prepare AnyImage ONNX artifacts and upload them to OmooLab R2."""

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


PROJECT_ROOT = Path(__file__).resolve().parents[2]
MODELS_DIR = PROJECT_ROOT / "models"
MODEL_CACHE_DIR = PROJECT_ROOT / ".model-cache"
SOURCE_ROOT = PROJECT_ROOT / "src" / "anyimage"
sys.path.insert(0, str(SOURCE_ROOT))
from server.model_catalog import DOWNLOADABLE_MODELS

R2_REMOTE = "omoolab-r2:ai-models"
CHUNK_SIZE = 1024 * 1024


@dataclass(frozen=True)
class BuildSpec:
    source_url: str
    source_filename: str
    source_size: int
    source_sha256: str
    exporter: str


@dataclass(frozen=True)
class ModelFile:
    key: str
    folder: str
    filename: str
    size: int
    sha256: str
    source_url: str = ""
    build: BuildSpec | None = None

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

    @property
    def remote(self):
        return f"{R2_REMOTE}/{self.folder}/{self.filename}"


MODEL_BUILDS = {
    "REALESRGAN_GENERAL_WDN_X4V3": BuildSpec(
        source_url=(
            "https://github.com/xinntao/Real-ESRGAN/releases/download/"
            "v0.2.5.0/realesr-general-wdn-x4v3.pth"
        ),
        source_filename="realesr-general-wdn-x4v3.pth",
        source_size=4885111,
        source_sha256=(
            "1641f8c4464b9f097c9fdda5589273713f67cf59f3d909e0bd688f0cee269dca"
        ),
        exporter="realesrgan_wdn_x4",
    ),
}

ISOLATED_EXPORTERS = {
    "MOGE3_VITL": ("moge3_export.py", ()),
    "HAT_GAN_X4_SHARPER": ("hat_export.py", ()),
    "BIREFNET_LITE": ("birefnet_export.py", ("--variant", "lite")),
    "BIREFNET_HR_MATTING": ("birefnet_export.py", ("--variant", "hr-matting")),
}


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


def catalog_model_files():
    files = []
    for key, model in DOWNLOADABLE_MODELS.items():
        for filename, size, sha256 in model.r2_files:
            source_url = ""
            if model.huggingface_repository:
                source_url = hf_url(model.huggingface_repository, filename)
            files.append(
                ModelFile(
                    key=key,
                    folder=model.r2_directory,
                    filename=filename,
                    size=size,
                    sha256=sha256,
                    source_url=source_url,
                    build=MODEL_BUILDS.get(key),
                )
            )
    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_file(path, size, sha256):
    path = Path(path)
    if not path.is_file():
        return False
    if path.stat().st_size != size:
        raise RuntimeError(f"File has an unexpected size: {path}")
    if file_sha256(path) != sha256:
        raise RuntimeError(f"File has an unexpected SHA-256: {path}")
    return True


def download_file(url, destination, size, sha256):
    destination = Path(destination)
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = destination.with_name(f"{destination.name}.part")
    temporary.unlink(missing_ok=True)
    request = Request(url, headers={"User-Agent": "AnyImage model build"})
    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} / {size / 1024**2:.1f} MiB",
                    end="",
                    flush=True,
                )
        print()
        if downloaded != size:
            raise RuntimeError(f"size mismatch: {downloaded} != {size}")
        if digest.hexdigest() != sha256:
            raise RuntimeError("SHA-256 mismatch")
        temporary.replace(destination)
    except Exception:
        temporary.unlink(missing_ok=True)
        raise


def prepare_source(build):
    source = MODEL_CACHE_DIR / build.source_filename
    if verify_file(source, build.source_size, build.source_sha256):
        return source
    download_file(
        build.source_url,
        source,
        build.source_size,
        build.source_sha256,
    )
    return source


def build_model(model):
    from tools.models.exports import export_model

    source = prepare_source(model.build)
    destination = model.destination
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = destination.with_name(f"{destination.name}.part")
    temporary.unlink(missing_ok=True)
    print(f"Build: {destination.relative_to(PROJECT_ROOT)}")
    try:
        export_model(model.build.exporter, source, temporary)
        verify_file(temporary, model.size, model.sha256)
        temporary.replace(destination)
    except Exception:
        temporary.unlink(missing_ok=True)
        raise


def prepare_model(model):
    if verify_file(model.destination, model.size, model.sha256):
        print(f"Skip existing: {model.destination.relative_to(PROJECT_ROOT)}")
        return
    if model.build:
        build_model(model)
        return
    if model.key in ISOLATED_EXPORTERS:
        model.destination.parent.mkdir(parents=True, exist_ok=True)
        exporter, arguments = ISOLATED_EXPORTERS[model.key]
        with TemporaryDirectory(prefix=f".{model.folder}-", dir=MODELS_DIR) as directory:
            subprocess.run(
                ["uv", "run", "--script", str(Path(__file__).with_name(exporter)),
                 *arguments, "--destination", directory],
                cwd=PROJECT_ROOT, check=True,
            )
            models = [item for item in MODEL_FILES if item.key == model.key]
            for item in models:
                if not verify_file(Path(directory) / item.filename, item.size, item.sha256):
                    raise RuntimeError(f"{model.key} export differs from the published catalog")
            for item in models:
                path = Path(directory) / item.filename
                destination = MODELS_DIR / item.folder / item.filename
                destination.parent.mkdir(parents=True, exist_ok=True)
                path.replace(destination)
        return
    if not model.source_url:
        raise RuntimeError(f"No upstream source or builder for {model.key}")
    download_file(
        model.source_url,
        model.destination,
        model.size,
        model.sha256,
    )


def upload_models(models=MODEL_FILES):
    executable = shutil.which("rclone")
    if executable is None:
        raise RuntimeError("rclone is not available on PATH")
    for model in models:
        print(f"Upload: {model.remote}")
        subprocess.run(
            [
                executable,
                "copyto",
                str(model.destination),
                model.remote,
                "--progress",
            ],
            cwd=PROJECT_ROOT,
            check=True,
        )


def parse_args():
    parser = argparse.ArgumentParser(description="Prepare models locally or sync them to R2.")
    commands = parser.add_subparsers(dest="command")
    commands.add_parser("prepare", help="Download or export models and verify them locally")
    commands.add_parser("sync", help="Prepare and verify all models, then upload them to R2")
    args = parser.parse_args()
    if args.command is None:
        parser.print_help()
    return args


def main():
    args = parse_args()
    if args.command is None:
        return
    for model in MODEL_FILES:
        prepare_model(model)
    if args.command == "sync":
        upload_models()


if __name__ == "__main__":
    main()