Initial public release

tools/models/__init__.py

@@ -0,0 +1 @@
+"""Model preparation, export, and upload tools."""

tools/models/birefnet_export.py

@@ -0,0 +1,108 @@
+# /// script
+# requires-python = ">=3.11,<3.12"
+# dependencies = [
+#   "torch==2.13.0", "torchvision==0.28.0", "onnx==1.22.0",
+#   "numpy==2.4.6", "transformers==4.57.6", "timm==1.0.25",
+#   "kornia==0.8.2", "einops==0.8.2", "safetensors==0.7.0",
+# ]
+# [tool.uv.sources]
+# torch = { index = "pytorch" }
+# torchvision = { index = "pytorch" }
+# [[tool.uv.index]]
+# name = "pytorch"
+# url = "https://download.pytorch.org/whl/cu130"
+# explicit = true
+# ///
+"""Export pinned official BiRefNet weights to fixed-resolution ONNX alpha."""
+
+import argparse
+import hashlib
+import json
+import sys
+from pathlib import Path
+from urllib.request import urlopen
+
+import onnx
+import torch
+from huggingface_hub import hf_hub_download
+from transformers import AutoModelForImageSegmentation
+
+from onnx_weights import compress_weights, compress_exact_constants
+from deform_sampling import deform_conv, validate_deform
+
+
+SOURCES = {
+    "lite": ("ZhengPeng7/BiRefNet_lite", "aa62cd87eafb9cc43056d08ef3615a14628b831d", 1024),
+    "hr-matting": ("ZhengPeng7/BiRefNet_HR-matting", "5d6b6f8adcb5b417c871b1d84ceaae9871355b7f", 2048),
+}
+WEIGHT_HASHES = {
+    "lite": "4417d89795250e698c3cb0ae8df15743810065f646f48a694fdfa7ca052d0815",
+    "hr-matting": "a5a4de698739ea5e0e8bbab28e1b293dde95092b87a442d566cbc585c53cef55",
+}
+LICENSE_REVISION = "ebcc0bc8ec7fe919cec829f2dea656b3078acddc"
+
+
+class AlphaModel(torch.nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, image):
+        return self.model(image)[-1].sigmoid()
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--variant", choices=SOURCES, required=True)
+    parser.add_argument("--destination", type=Path, required=True)
+    parser.add_argument("--checkpoint", type=Path)
+    parser.add_argument("--fp32-output", type=Path)
+    args = parser.parse_args()
+    repository, revision, size = SOURCES[args.variant]
+    checkpoint = args.checkpoint / "model.safetensors" if args.checkpoint else Path(
+        hf_hub_download(repository, "model.safetensors", revision=revision)
+    )
+    with checkpoint.open("rb") as stream:
+        if hashlib.file_digest(stream, "sha256").hexdigest() != WEIGHT_HASHES[args.variant]:
+            raise RuntimeError("BiRefNet checkpoint SHA-256 mismatch")
+    torch.set_num_threads(8)
+    model = AutoModelForImageSegmentation.from_pretrained(
+        str(args.checkpoint) if args.checkpoint else repository,
+        revision=revision,
+        trust_remote_code=True,
+    ).eval()
+    validate_deform()
+    sys.modules[type(model).__module__].deform_conv2d = deform_conv
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = AlphaModel(model).to(device)
+    args.destination.mkdir(parents=True, exist_ok=True)
+    output = args.destination / "model.onnx"
+    with torch.inference_mode():
+        torch.onnx.export(
+            model, torch.zeros(1, 3, size, size, device=device),
+            str(output),
+            input_names=["image"], output_names=["alpha"],
+            opset_version=17, dynamo=False, do_constant_folding=False,
+        )
+    graph = onnx.load(output)
+    if args.fp32_output:
+        onnx.save(graph, args.fp32_output)
+    compress_weights(graph)
+    compress_exact_constants(graph)
+    onnx.checker.check_model(graph)
+    onnx.save(graph, output)
+    license_url = f"https://raw.githubusercontent.com/ZhengPeng7/BiRefNet/{LICENSE_REVISION}/LICENSE"
+    license_text = urlopen(license_url, timeout=60).read()
+    if hashlib.sha256(license_text).hexdigest() != "92a7089e0915fc32bc40067560b398f1e6a7a5958abd7d04eda393629a5acefb":
+        raise RuntimeError("BiRefNet license SHA-256 mismatch")
+    (args.destination / "LICENSE.txt").write_bytes(license_text)
+    with output.open("rb") as stream:
+        checksum = hashlib.file_digest(stream, "sha256").hexdigest()
+    record = {"repository": repository, "revision": revision, "input_size": size,
+              "bytes": output.stat().st_size, "sha256": checksum,
+              "checkpoint_sha256": WEIGHT_HASHES[args.variant]}
+    print(json.dumps(record), flush=True)
+
+
+if __name__ == "__main__":
+    main()

tools/models/cli.py

@@ -0,0 +1,262 @@
+"""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()

tools/models/deform_sampling.py

@@ -0,0 +1,47 @@
+"""Express BiRefNet's single-group deformable convolution with GridSample."""
+
+import torch
+from torch.nn import functional as F
+
+
+def deform_conv(input, offset, weight, bias=None, stride=(1, 1), padding=0, mask=None):
+    batch, channels, height, width = map(int, input.shape)
+    output_height, output_width = map(int, offset.shape[-2:])
+    kernel_height, kernel_width = map(int, weight.shape[-2:])
+    count = kernel_height * kernel_width
+    y, x = torch.meshgrid(
+        torch.arange(output_height, dtype=input.dtype, device=input.device),
+        torch.arange(output_width, dtype=input.dtype, device=input.device), indexing="ij",
+    )
+    ky, kx = torch.meshgrid(
+        torch.arange(kernel_height, dtype=input.dtype, device=input.device),
+        torch.arange(kernel_width, dtype=input.dtype, device=input.device), indexing="ij",
+    )
+    offsets = offset.reshape(batch, count, 2, output_height, output_width)
+    sy = y[None, None] * stride[0] - padding + ky.reshape(1, count, 1, 1) + offsets[:, :, 0]
+    sx = x[None, None] * stride[1] - padding + kx.reshape(1, count, 1, 1) + offsets[:, :, 1]
+    coordinates = torch.stack((sx * (2.0 / (width - 1)) - 1, sy * (2.0 / (height - 1)) - 1), -1)
+    grid = coordinates.permute(0, 2, 3, 1, 4).reshape(batch, output_height, output_width * count, 2)
+    sampled = F.grid_sample(input, grid, mode="bilinear", padding_mode="zeros", align_corners=True)
+    sampled = sampled.reshape(batch, channels, output_height, output_width, count).permute(0, 1, 4, 2, 3)
+    sampled = (sampled * mask[:, None]).reshape(batch, channels * count, output_height, output_width)
+    return F.conv2d(sampled, weight.reshape(int(weight.shape[0]), channels * count, 1, 1), bias)
+
+
+def validate_deform():
+    from torchvision.ops import deform_conv2d
+
+    torch.manual_seed(12)
+    for kernel in (1, 3, 7):
+        image = torch.randn(1, 4, 17, 19)
+        offset = torch.randn(1, kernel * kernel * 2, 17, 19) * 3
+        mask = torch.rand(1, kernel * kernel, 17, 19) * 2
+        weights = torch.randn(5, 4, kernel, kernel)
+        bias = torch.randn(5)
+        reference = deform_conv2d(image, offset, weights, bias, padding=kernel // 2, mask=mask)
+        actual = deform_conv(image, offset, weights, bias, padding=kernel // 2, mask=mask)
+        torch.testing.assert_close(actual, reference, atol=1e-4, rtol=1e-4)
+
+
+if __name__ == "__main__":
+    validate_deform()

tools/models/exports.py

@@ -0,0 +1,109 @@
+"""Export official PyTorch super-resolution models to ONNX."""
+
+import numpy as np
+
+
+OPSET_VERSION = 17
+
+
+def _require_build_dependencies():
+    try:
+        import onnx
+        import onnxruntime
+        import torch
+    except ImportError as error:
+        raise RuntimeError(
+            "Model build dependencies are missing. Run with: "
+            "uv sync --group models"
+        ) from error
+    return onnx, onnxruntime, torch
+
+
+def _load_weights(torch, checkpoint_path, key):
+    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+    return checkpoint[key]
+
+
+def _validate_export(onnx, onnxruntime, torch, model, sample, destination):
+    onnx.checker.check_model(onnx.load(destination))
+    expected = model(sample).detach().numpy()
+    session = onnxruntime.InferenceSession(
+        str(destination),
+        providers=["CPUExecutionProvider"],
+    )
+    actual = session.run(None, {"input": sample.numpy()})[0]
+    np.testing.assert_allclose(actual, expected, rtol=1e-4, atol=1e-5)
+
+
+def _export(torch, model, sample, destination, dynamic_axes=None):
+    model.eval()
+    torch.onnx.export(
+        model,
+        sample,
+        destination,
+        export_params=True,
+        opset_version=OPSET_VERSION,
+        do_constant_folding=True,
+        input_names=["input"],
+        output_names=["output"],
+        dynamic_axes=dynamic_axes,
+        dynamo=False,
+    )
+
+
+def _restore_realesrgan_scope_names(onnx, destination):
+    model = onnx.load(destination)
+    for node in model.graph.node:
+        node.name = node.name.replace("/body.", "/body/body.")
+        for values in (node.input, node.output):
+            for index, value in enumerate(values):
+                values[index] = value.replace("/body.", "/body/body.")
+    onnx.save(model, destination)
+
+
+def export_realesrgan_wdn_x4(checkpoint_path, destination):
+    onnx, onnxruntime, torch = _require_build_dependencies()
+    nn = torch.nn
+
+    class SRVGGNetCompact(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.upscale = 4
+            self.body = nn.ModuleList([nn.Conv2d(3, 64, 3, 1, 1), nn.PReLU(64)])
+            for _ in range(32):
+                self.body.extend([nn.Conv2d(64, 64, 3, 1, 1), nn.PReLU(64)])
+            self.body.append(nn.Conv2d(64, 3 * self.upscale**2, 3, 1, 1))
+            self.upsampler = nn.PixelShuffle(self.upscale)
+
+        def forward(self, value):
+            result = value
+            for layer in self.body:
+                result = layer(result)
+            result = self.upsampler(result)
+            return result + torch.nn.functional.interpolate(
+                value,
+                scale_factor=self.upscale,
+                mode="nearest",
+            )
+
+    model = SRVGGNetCompact()
+    model.load_state_dict(_load_weights(torch, checkpoint_path, "params"), strict=True)
+    sample = torch.rand(1, 3, 16, 16)
+    dynamic_axes = {
+        "input": {2: "height", 3: "width"},
+        "output": {2: "height4", 3: "width4"},
+    }
+    _export(torch, model, sample, destination, dynamic_axes)
+    _restore_realesrgan_scope_names(onnx, destination)
+    _validate_export(onnx, onnxruntime, torch, model, sample, destination)
+
+
+def export_model(exporter, checkpoint_path, destination):
+    exporters = {
+        "realesrgan_wdn_x4": export_realesrgan_wdn_x4,
+    }
+    try:
+        export = exporters[exporter]
+    except KeyError as error:
+        raise ValueError(f"Unknown model exporter: {exporter}") from error
+    export(checkpoint_path, destination)

tools/models/hat_export.py

@@ -0,0 +1,95 @@
+# /// script
+# requires-python = ">=3.11,<3.12"
+# dependencies = ["torch==2.13.0", "onnx==1.19.0", "numpy==1.26.4", "einops==0.8.2", "gdown==6.2.0"]
+# [tool.uv.sources]
+# torch = { index = "pytorch" }
+# [[tool.uv.index]]
+# name = "pytorch"
+# url = "https://download.pytorch.org/whl/cpu"
+# explicit = true
+# ///
+"""Export the official HAT Sharper checkpoint with FP16 weight storage."""
+
+import argparse
+import hashlib
+import json
+from pathlib import Path
+from urllib.request import urlopen
+
+import gdown
+import onnx
+import torch
+
+from onnx_weights import compress_weights
+
+
+REVISION = "1638a9a822581657811867bf670717f8371fc3e5"
+CHECKPOINT_SHA256 = "5800b67136006eb8cab3b4ed7c8d73b6a195bb18e6cc709b674f9aa069c00271"
+SOURCE_URL = f"https://raw.githubusercontent.com/XPixelGroup/HAT/{REVISION}"
+
+
+def load_model(checkpoint, cache):
+    architecture = cache / "hat_arch.py"
+    if not architecture.exists():
+        architecture.write_bytes(urlopen(SOURCE_URL + "/hat/archs/hat_arch.py", timeout=60).read())
+    if hashlib.sha256(architecture.read_bytes()).hexdigest() != "81d8cecf491975246c9ebb20480898c656f59de020f38b05daa68741e426117f":
+        raise RuntimeError("HAT architecture SHA-256 mismatch")
+    # Export only: remove BasicSR's training registry and use the equivalent
+    # PyTorch initialization utilities. The official forward remains unchanged.
+    source = architecture.read_text(encoding="utf-8")
+    source = source.replace("from basicsr.utils.registry import ARCH_REGISTRY", "")
+    source = source.replace("@ARCH_REGISTRY.register()", "")
+    source = source.replace(
+        "from basicsr.archs.arch_util import to_2tuple, trunc_normal_",
+        "from torch.nn.modules.utils import _pair as to_2tuple\nfrom torch.nn.init import trunc_normal_",
+    )
+    namespace = {"__name__": "hat_arch"}
+    exec(compile(source, str(architecture), "exec"), namespace)
+    model = namespace["HAT"](
+        upscale=4, in_chans=3, img_size=64, window_size=16,
+        compress_ratio=3, squeeze_factor=30, conv_scale=0.01,
+        overlap_ratio=0.5, img_range=1., depths=[6] * 6,
+        embed_dim=180, num_heads=[6] * 6, mlp_ratio=2,
+        upsampler="pixelshuffle", resi_connection="1conv",
+    ).eval()
+    with checkpoint.open("rb") as stream:
+        if hashlib.file_digest(stream, "sha256").hexdigest() != CHECKPOINT_SHA256:
+            raise RuntimeError("HAT Sharper checkpoint SHA-256 mismatch")
+    model.load_state_dict(torch.load(checkpoint, map_location="cpu", weights_only=True)["params_ema"], strict=True)
+    return model
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--destination", type=Path, required=True)
+    parser.add_argument("--checkpoint", type=Path)
+    args = parser.parse_args()
+    cache = Path(__file__).resolve().parents[2] / ".model-cache" / "hat-sharper"
+    cache.mkdir(parents=True, exist_ok=True)
+    checkpoint = args.checkpoint or cache / "Real_HAT_GAN_sharper.pth"
+    if not checkpoint.exists():
+        gdown.download(id="1EioFq5-mKmv1uqta_Byd9cgXp9SU3zjj", output=str(checkpoint))
+    torch.set_num_threads(8)
+    torch.manual_seed(0)
+    model = load_model(checkpoint, cache)
+    args.destination.mkdir(parents=True, exist_ok=True)
+    output = args.destination / "model.onnx"
+    with torch.inference_mode():
+        torch.onnx.export(model, torch.zeros(1, 3, 256, 256), str(output),
+                          opset_version=17, input_names=["input"], output_names=["output"],
+                          do_constant_folding=False, dynamo=False)
+    graph = onnx.load(output)
+    compress_weights(graph)
+    onnx.checker.check_model(graph)
+    onnx.save(graph, output)
+    (args.destination / "LICENSE.txt").write_bytes(urlopen(SOURCE_URL + "/LICENSE", timeout=60).read())
+    if sum(path.stat().st_size for path in args.destination.iterdir()) >= 60_000_000:
+        raise RuntimeError("HAT asset exceeds 60 MB")
+    with output.open("rb") as stream:
+        digest = hashlib.file_digest(stream, "sha256").hexdigest()
+    print(json.dumps({"bytes": output.stat().st_size, "sha256": digest,
+                      "revision": REVISION, "checkpoint_sha256": CHECKPOINT_SHA256}), flush=True)
+
+
+if __name__ == "__main__":
+    main()

tools/models/moge3_export.py

@@ -0,0 +1,167 @@
+# /// script
+# requires-python = ">=3.11,<3.12"
+# dependencies = [
+#   "torch==2.13.0", "torchvision==0.28.0", "onnx==1.22.0", "numpy==2.4.6",
+#   "triton-windows==3.7.1.post27; sys_platform == 'win32'",
+#   "moge @ git+https://github.com/microsoft/MoGe.git@74fbce054ebed49800de42d0ad0e83495065719a",
+# ]
+# [tool.uv.sources]
+# torch = { index = "pytorch" }
+# torchvision = { index = "pytorch" }
+# [[tool.uv.index]]
+# name = "pytorch"
+# url = "https://download.pytorch.org/whl/cu130"
+# explicit = true
+# ///
+"""Export the dynamic MoGe-3 backbone and sparse refinement ONNX graphs."""
+
+import argparse
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from moge.model.v3 import MoGeModel
+from moge.utils.geometry_torch import normalized_view_plane_uv
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src" / "anyimage"))
+from server.models.onnx_moge3 import build_sparse_inputs
+
+
+class Moge3Backbone(torch.nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, image, num_tokens):
+        m = self.model
+        aspect = image.shape[-1] / image.shape[-2]
+        bh = (num_tokens / aspect).sqrt().round().long()
+        bw = (num_tokens * aspect).sqrt().round().long()
+        feat, cls = m.encoder(image, bh, bw, return_class_token=True)
+        features = []
+        for level in range(5):
+            uv = normalized_view_plane_uv(
+                width=bw * 2**level,
+                height=bh * 2**level,
+                aspect_ratio=aspect,
+                dtype=image.dtype,
+                device=image.device,
+            ).permute(2, 0, 1)[None]
+            features.append(torch.cat([feat, uv], 1) if level == 0 else uv)
+        neck = m.neck(features)
+        return (
+            m.points_head(neck)[-1].permute(0, 2, 3, 1).float(),
+            features[0],
+            m.normal_head(neck)[-1],
+            m.mask_head(neck)[-1],
+            m.scale_head(cls),
+        )
+
+
+class Moge3Refiner(torch.nn.Module):
+    def __init__(self, r):
+        super().__init__()
+        self.r = r
+        self.n = len(r.down_stages)
+
+    def conv(self, m, x, idx):
+        padded = torch.cat([x, torch.zeros_like(x[:1])], 0)
+        weights = m.weight.reshape(m.out_channels, 27, m.in_channels).permute(1, 2, 0)
+        result = padded[idx[:, 0]] @ weights[0]
+        for k in range(1, 27):
+            result = result + padded[idx[:, k]] @ weights[k]
+        return result + m.bias
+
+    def stage(self, stage, x, idx):
+        for b in stage:
+            y = F.silu(b.norm1(x))
+            y = F.silu(self.conv(b.conv1, y, idx))
+            y = self.conv(b.conv2, y, idx)
+            x = y + b.skip_connection(x)
+        return x
+
+    def forward(self, coord, encoder, *args):
+        n = self.n
+        ns = args[:n]
+        ps = args[n : 2 * n - 1]
+        us = args[2 * n - 1 : 3 * n - 2]
+        ei = args[-1]
+        r = self.r
+        x = r.input_proj(coord.reshape(-1, 3))
+        skips = []
+        for i, stage in enumerate(r.down_stages):
+            x = self.stage(stage, x, ns[i])
+            if i < n - 1:
+                skips.append(x)
+                padded = torch.cat([x, torch.zeros_like(x[:1])], 0)
+                count = (ps[i] < x.shape[0]).sum(1, keepdim=True).to(x.dtype)
+                x = r.downsample_blocks[i].linear(padded[ps[i]].sum(1) / count)
+        ef = encoder.permute(0, 2, 3, 1).reshape(-1, encoder.shape[1])[ei]
+        x = r.fuse_proj(torch.cat([x, r.encoder_fuse(ef)], -1))
+        x = self.stage(r.bottleneck_stage, x, ns[-1])
+        for i, (up, stage) in enumerate(zip(r.upsample_blocks, r.up_stages)):
+            level = n - 2 - i
+            x = up.linear(x)[us[level]] + skips[level]
+            x = self.stage(stage, x, ns[level])
+        return r.out_proj(x).reshape(coord.shape[0], coord.shape[1], coord.shape[2])
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--destination", type=Path, required=True)
+    parser.add_argument("--checkpoint", default="Ruicheng/moge-3-vitl")
+    args = parser.parse_args()
+    args.destination.mkdir(parents=True, exist_ok=True)
+    checkpoint_options = (
+        {"revision": "184008f877d7ad1ad4c2cd2182a9bd1f63d0e5be"}
+        if args.checkpoint == "Ruicheng/moge-3-vitl"
+        else {}
+    )
+    model = (
+        MoGeModel.from_pretrained(args.checkpoint, **checkpoint_options).eval().cuda()
+    )
+    model.onnx_compatible_mode = True
+    backbone = Moge3Backbone(model).eval()
+    image = torch.zeros(1, 3, 192, 256, device="cuda")
+    with torch.no_grad():
+        coord, encoder, *_ = backbone(image, torch.tensor(64, device="cuda"))
+    sparse = build_sparse_inputs(coord.cpu().numpy())
+    indices = [torch.from_numpy(value).cuda() for value in sparse.values()]
+    names = ["coord", "encoder", *sparse]
+    axes = {name: {0: "count_" + name} for name in sparse}
+    axes.update(
+        coord={1: "height", 2: "width"},
+        encoder={2: "token_height", 3: "token_width"},
+        delta={1: "height", 2: "width"},
+    )
+    torch.onnx.export(
+        Moge3Refiner(model.refiner).eval(),
+        (coord, encoder, *indices),
+        str(args.destination / "refiner.onnx"),
+        input_names=names,
+        output_names=["delta"],
+        dynamic_axes=axes,
+        opset_version=18,
+        dynamo=False,
+    )
+    torch.onnx.export(
+        backbone,
+        (image, torch.tensor(1200, device="cuda")),
+        str(args.destination / "backbone.onnx"),
+        input_names=["image", "num_tokens"],
+        output_names=["coord", "encoder", "normal", "mask", "scale"],
+        dynamic_axes={
+            "image": {2: "height", 3: "width"},
+            "coord": {1: "grid_height", 2: "grid_width"},
+            "encoder": {2: "token_height", 3: "token_width"},
+            "normal": {2: "grid_height", 3: "grid_width"},
+            "mask": {2: "grid_height", 3: "grid_width"},
+        },
+        opset_version=18,
+        dynamo=False,
+    )
+
+
+if __name__ == "__main__":
+    main()

tools/models/onnx_weights.py

@@ -0,0 +1,49 @@
+"""Compress stored ONNX weights without changing graph arithmetic."""
+
+import numpy as np
+from onnx import TensorProto, helper, numpy_helper
+
+
+def compress_weights(model):
+    casts = []
+    for initializer in model.graph.initializer:
+        if initializer.data_type != TensorProto.FLOAT:
+            continue
+        name = initializer.name
+        values = numpy_helper.to_array(initializer)
+        if not np.isfinite(values).all() or np.max(np.abs(values), initial=0) > np.finfo(np.float16).max:
+            continue
+        array = values.astype(np.float16)
+        stored_name = name + "__fp16_storage"
+        initializer.CopyFrom(numpy_helper.from_array(array, stored_name))
+        casts.append(helper.make_node(
+            "Cast", [stored_name], [name], to=TensorProto.FLOAT,
+            name=stored_name + "_restore",
+        ))
+    nodes = list(model.graph.node)
+    del model.graph.node[:]
+    model.graph.node.extend(casts + nodes)
+
+
+def compress_exact_constants(model):
+    """Compress only constant tensors that round-trip without any value change."""
+    nodes = []
+    for node in model.graph.node:
+        nodes.append(node)
+        if node.op_type != "Constant":
+            continue
+        attribute = next((item for item in node.attribute if item.name == "value"), None)
+        if attribute is None or attribute.t.data_type != TensorProto.FLOAT:
+            continue
+        values = numpy_helper.to_array(attribute.t)
+        if not np.isfinite(values).all() or np.max(np.abs(values), initial=0) > np.finfo(np.float16).max:
+            continue
+        stored = values.astype(np.float16)
+        if values.nbytes < 64 or not np.array_equal(values, stored.astype(np.float32)):
+            continue
+        name = node.output[0]
+        node.output[0] = name + "__fp16_constant"
+        attribute.t.CopyFrom(numpy_helper.from_array(stored))
+        nodes.append(helper.make_node("Cast", [node.output[0]], [name], to=TensorProto.FLOAT))
+    del model.graph.node[:]
+    model.graph.node.extend(nodes)