nan
89087fc2e2
- Implemented the "O Image Relief Plane" Geometry Nodes group to create a fixed-base image relief using depth images and various parameters. - Created the "O Shadeless" shader node group to display color without scene lighting while preserving shadows and reflections. - Introduced a preview feature for node groups, allowing for interactive visualization of node arrangements in an HTML format. - Added a template HTML file for rendering node previews with customizable views and controls.
253 lines
7.1 KiB
Python
253 lines
7.1 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 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()
|