"""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 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", ), "SAFMN_REAL_X4": BuildSpec( source_url=( "https://github.com/sunny2109/SAFMN/releases/download/" "v0.1.0/SAFMN_L_Real_LSDIR_x4-v2.pth" ), source_filename="SAFMN_L_Real_LSDIR_x4-v2.pth", source_size=45002097, source_sha256=( "f1ac0ee3ee143fbbc49aff6584cb7b48b71ba8d1961dfb8ad076b98ca4799ae2" ), exporter="safmn_real_x4", ), } 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 == "PISA_SR": subprocess.run( ["uv", "run", "--script", str(Path(__file__).with_name("pisa_sr_export.py"))], cwd=PROJECT_ROOT, check=True, ) if not verify_file(model.destination, model.size, model.sha256): raise RuntimeError("PiSA-SR export differs from the published catalog") 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() parser.add_argument( "--prepare-only", action="store_true", help="Prepare and verify models without uploading them", ) return parser.parse_args() def main(): args = parse_args() for model in MODEL_FILES: prepare_model(model) if not args.prepare_only: upload_models() if __name__ == "__main__": main()