| |
| """MODNet portrait matting -> ExecuTorch WebGPU. |
| |
| The official code is on GitHub (ZHKKKe/MODNet) and the official checkpoints on Google |
| Drive; neither is reachable from the export host and the author has no Hugging Face |
| account. This pins a community Hugging Face mirror that carries a copy of the upstream |
| source tree and the photographic checkpoint (sha256 checked; four independent mirrors |
| publish the same checkpoint bytes). The webcam checkpoint comes from a second mirror. |
| |
| Input: [1, 3, 256, 256] float32 RGB normalized to [-1, 1] ((x / 255 - 0.5) / 0.5). |
| Output: [1, 1, 256, 256] alpha matte in [0, 1] (inference=True, pred_matte only). |
| """ |
|
|
| import argparse |
| import importlib |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| import _common |
|
|
| SOURCE = { |
| "repo": "boopathiraj/MODNet", |
| "revision": "a967175f8fd63de67a1c161d8daac966d2b57e3a", |
| "files": { |
| "modnet/__init__.py": None, |
| "modnet/modnet.py": None, |
| "modnet/src/__init__.py": None, |
| "modnet/src/models/__init__.py": None, |
| "modnet/src/models/backbones/__init__.py": None, |
| "modnet/src/models/backbones/wrapper.py": None, |
| "modnet/src/models/backbones/mobilenetv2.py": None, |
| "modnet/LICENSE": None, |
| }, |
| } |
| CHECKPOINTS = { |
| "photographic": { |
| "repo": "boopathiraj/MODNet", |
| "revision": "a967175f8fd63de67a1c161d8daac966d2b57e3a", |
| "file": "modnet/pretrained/modnet_photographic_portrait_matting.ckpt", |
| "sha256": "7c22235f0925deba15d4d63e53afcb654c47055bbcd98f56e393ab2584007ed8", |
| }, |
| "webcam": { |
| "repo": "DavG25/modnet-pretrained-models", |
| "revision": "903cc06311b3b12071edfa1b42b534e4a31c718a", |
| "file": "models/modnet_webcam_portrait_matting.ckpt", |
| "sha256": "913b82b66558db39b6286c150f809017d7528c872b156eb14333c9c6cb52108b", |
| }, |
| } |
|
|
|
|
| class InstanceNorm2dDecomposed(torch.nn.Module): |
| """InstanceNorm2d(affine=False) as mean/sub/mul/sqrt/div. |
| |
| The Vulkan/WebGPU partitioner does not take aten._native_batch_norm_legit.no_stats, |
| so the stock module would leave 16 CPU kernels splitting the graph into 13 delegates. |
| """ |
|
|
| def __init__(self, eps): |
| super().__init__() |
| self.eps = eps |
|
|
| def forward(self, x): |
| d = x - x.mean(dim=(2, 3), keepdim=True) |
| var = (d * d).mean(dim=(2, 3), keepdim=True) |
| return d / torch.sqrt(var + self.eps) |
|
|
|
|
| def build_model(variant): |
| root = None |
| files = {} |
| for f in SOURCE["files"]: |
| path, digest = _common.hf_file(SOURCE["repo"], f, SOURCE["revision"]) |
| files[f] = digest |
| root = Path(str(path)[: -len(f)]) |
| sys.path.insert(0, str(root)) |
| MODNet = importlib.import_module("modnet.modnet").MODNet |
| sys.path.remove(str(root)) |
|
|
| ck = CHECKPOINTS[variant] |
| path, _ = _common.hf_file(ck["repo"], ck["file"], ck["revision"], ck["sha256"]) |
| state = torch.load(path, map_location="cpu", weights_only=True) |
| |
| state = {k.removeprefix("module."): v for k, v in state.items()} |
| model = MODNet(backbone_pretrained=False) |
| model.load_state_dict(state, strict=True) |
| for parent in list(model.modules()): |
| for name, child in parent.named_children(): |
| if isinstance(child, torch.nn.InstanceNorm2d): |
| assert not child.affine and not child.track_running_stats |
| setattr(parent, name, InstanceNorm2dDecomposed(child.eps)) |
|
|
| class Matte(torch.nn.Module): |
| def __init__(self, m): |
| super().__init__() |
| self.m = m |
|
|
| def forward(self, img): |
| return self.m(img, True)[2] |
|
|
| return Matte(model).eval(), files |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) |
| ap.add_argument("--out", required=True, help="output directory") |
| ap.add_argument("--size", type=int, nargs=2, default=[256, 256], metavar=("H", "W")) |
| ap.add_argument("--variant", choices=sorted(CHECKPOINTS), default="photographic") |
| ap.add_argument( |
| "--golden", |
| help="optional goldens.manifest.json to compare against (key: bgremove)", |
| ) |
| args = ap.parse_args() |
|
|
| model, source_files = build_model(args.variant) |
| golden = _common.load_golden(args.golden, "bgremove") if args.golden else None |
| _common.export_vision( |
| name="modnet", |
| model=model, |
| example_input=torch.rand(1, 3, *args.size) * 2 - 1, |
| out_dir=args.out, |
| upstream={ |
| "model": f"ZHKKKe/MODNet ({args.variant} portrait matting checkpoint)", |
| "weights": { |
| **CHECKPOINTS[args.variant], |
| "official": False, |
| "note": "community mirror; official checkpoints are on Google Drive, linked from GitHub, neither reachable from the export host", |
| }, |
| "architecture_source": { |
| "repo": SOURCE["repo"], |
| "revision": SOURCE["revision"], |
| "files_sha256": source_files, |
| "official": False, |
| "note": "copy of the upstream src/models tree with modnet.py moved to the package root", |
| }, |
| "license": "Apache-2.0", |
| }, |
| extra_recipe={ |
| "variant": args.variant, |
| "preprocessing_outside_graph": "RGB / 255, then (x - 0.5) / 0.5", |
| "postprocessing_outside_graph": "alpha in [0, 1]", |
| }, |
| golden=golden, |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
