#!/usr/bin/env python3 """Export Apple SHARP to a browser-friendly ONNX predictor graph. This exporter emits the raw SHARP predictor outputs (Gaussians in SHARP's pre-unprojection space / NDC-aligned coordinates). The web worker in this repository performs the final NDC->metric conversion and covariance decomposition in JavaScript, which avoids relying on ONNX `SVD` support in browser runtimes. Expected exported inputs: - image: float32 [1, 3, 1536, 1536] in [0, 1] - disparity_factor: float32 [1] (f_px / image_width) Expected exported outputs: - mean_vectors_ndc: float32 [1, N, 3] - singular_values_ndc: float32 [1, N, 3] - quaternions_ndc: float32 [1, N, 4] - colors: float32 [1, N, 3] - opacities: float32 [1, N] Notes: - The released Apple SHARP model weights are licensed separately for research-only use. - Export success depends on local Python / PyTorch / ONNX package versions. """ from __future__ import annotations import argparse import sys from pathlib import Path def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--sharp-repo", type=Path, required=True, help="Path to a local clone of https://github.com/apple/ml-sharp", ) parser.add_argument( "--checkpoint", type=Path, default=None, help="Path to SHARP .pt checkpoint (optional; downloads default if omitted).", ) parser.add_argument( "--output", type=Path, default=Path("public/models/sharp_web_predictor.onnx"), help="Output ONNX path.", ) parser.add_argument( "--device", type=str, default="cpu", choices=["cpu", "cuda", "mps"], help="Device for export (cpu recommended).", ) parser.add_argument( "--opset", type=int, default=20, help="ONNX opset version.", ) parser.add_argument( "--verbose", action="store_true", help="Print extra diagnostics.", ) return parser.parse_args() def import_sharp_modules(sharp_repo: Path): src_path = sharp_repo / "src" if not src_path.exists(): raise FileNotFoundError( "Could not find %s. Expected a local clone of apple/ml-sharp with a src/ directory." % src_path ) sys.path.insert(0, str(src_path)) import torch # noqa: WPS433 from sharp.cli.predict import DEFAULT_MODEL_URL # noqa: WPS433 from sharp.models import PredictorParams, create_predictor # noqa: WPS433 return torch, DEFAULT_MODEL_URL, PredictorParams, create_predictor def load_predictor(torch, create_predictor, predictor_params_cls, checkpoint_path, default_url, device, verbose): predictor = create_predictor(predictor_params_cls()) if checkpoint_path is None: if verbose: print("Downloading checkpoint from %s" % default_url) state_dict = torch.hub.load_state_dict_from_url(default_url, progress=True) else: if verbose: print("Loading checkpoint from %s" % checkpoint_path) try: state_dict = torch.load(checkpoint_path, weights_only=True, map_location=device) except TypeError: state_dict = torch.load(checkpoint_path, map_location=device) predictor.load_state_dict(state_dict) predictor.eval() predictor.to(device) return predictor def make_export_module(torch, predictor): class SharpPredictorExport(torch.nn.Module): def __init__(self): super().__init__() self.predictor = predictor def forward(self, image, disparity_factor): gaussians = self.predictor(image, disparity_factor) return ( gaussians.mean_vectors, gaussians.singular_values, gaussians.quaternions, gaussians.colors, gaussians.opacities, ) return SharpPredictorExport() def export_onnx() -> None: args = parse_args() sharp_repo = args.sharp_repo.expanduser().resolve() output_path = args.output.expanduser().resolve() checkpoint_path = args.checkpoint.expanduser().resolve() if args.checkpoint else None torch, default_url, predictor_params_cls, create_predictor = import_sharp_modules(sharp_repo) device = torch.device(args.device) predictor = load_predictor( torch=torch, create_predictor=create_predictor, predictor_params_cls=predictor_params_cls, checkpoint_path=checkpoint_path, default_url=default_url, device=device, verbose=args.verbose, ) export_module = make_export_module(torch, predictor).eval().to(device) dummy_image = torch.rand((1, 3, 1536, 1536), dtype=torch.float32, device=device) dummy_disparity_factor = torch.tensor([1.0], dtype=torch.float32, device=device) output_path.parent.mkdir(parents=True, exist_ok=True) if args.verbose: print("Exporting ONNX to %s" % output_path) preexisting_files = {p.name for p in output_path.parent.iterdir()} if output_path.parent.exists() else set() with torch.no_grad(): torch.onnx.export( export_module, (dummy_image, dummy_disparity_factor), str(output_path), export_params=True, do_constant_folding=True, opset_version=args.opset, input_names=["image", "disparity_factor"], output_names=[ "mean_vectors_ndc", "singular_values_ndc", "quaternions_ndc", "colors", "opacities", ], ) # Consolidate PyTorch-exported external tensor shards into a single .onnx.data file when # the model exceeds the 2GB protobuf limit. This is much easier to statically host than # hundreds of shard files. import onnx # noqa: WPS433 model_proto = onnx.load(str(output_path), load_external_data=True) has_external_data = any( initializer.data_location == onnx.TensorProto.EXTERNAL for initializer in model_proto.graph.initializer ) if has_external_data: data_filename = output_path.name + ".data" onnx.save_model( model_proto, str(output_path), save_as_external_data=True, all_tensors_to_one_file=True, location=data_filename, size_threshold=1024, ) keep = {output_path.name, data_filename} for path in output_path.parent.iterdir(): if path.name in keep or path.name in preexisting_files: continue if path.is_file(): path.unlink() print("Wrote %s" % output_path) if has_external_data: print("Wrote external tensor data %s" % (output_path.parent / (output_path.name + ".data"))) print( "Next: start the web app with `bun dev` (or `bun run build` + static serve) and point it to this ONNX file." ) if __name__ == "__main__": export_onnx()