263 lines
7.9 KiB
Python
263 lines
7.9 KiB
Python
"""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()
|