| import os, sys, resource |
| os.environ['XFORMERS_DISABLED'] = '1' |
| sys.path.insert(0, os.path.abspath('mogesrc/MoGe-main')); import glob; sys.path.insert(0, os.path.abspath(glob.glob('u3d/*')[0])) |
| import numpy as np, torch |
| from moge.model.v2 import MoGeModel |
| name, H, W = sys.argv[1], int(sys.argv[2]), int(sys.argv[3]) |
| tokens = (H // 14) * (W // 14) |
| model = MoGeModel.from_pretrained(f'Ruicheng/moge-2-{name}-normal').eval() |
| model.onnx_compatible_mode = True |
| import torch.nn.functional as F |
| from moge.utils.geometry_torch import normalized_view_plane_uv |
| rows, cols = H // 14, W // 14 |
| vit = model.encoder.backbone |
| with torch.no_grad(): |
| pos = vit.interpolate_pos_encoding(torch.zeros(1, rows * cols + 1, vit.pos_embed.shape[-1]), H, W).clone() |
| vit.interpolate_pos_encoding = lambda x, h, w: pos |
| uvs = [normalized_view_plane_uv(width=cols * 2 ** l, height=rows * 2 ** l, aspect_ratio=W / H, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0) for l in range(5)] |
| class Wrap(torch.nn.Module): |
| def __init__(s, m): super().__init__(); s.m = m |
| def forward(s, image): |
| m = s.m |
| features = m.encoder(image, rows, cols, return_class_token=True)[0] |
| feats = [torch.concat([features, uvs[0]], dim=1)] + uvs[1:] |
| feats = m.neck(feats) |
| points = F.interpolate(m.points_head(feats)[-1], (H, W), mode='bilinear', align_corners=False) |
| mask = F.interpolate(m.mask_head(feats)[-1], (H, W), mode='bilinear', align_corners=False) |
| return m._remap_points(points.permute(0, 2, 3, 1)), mask.squeeze(1).sigmoid() |
| x = torch.rand(1, 3, H, W) |
| with torch.no_grad(): |
| w = Wrap(model); ref = w(x) |
| print('forward ok', [tuple(r.shape) for r in ref], 'rss', resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e9, flush=True) |
| traced = torch.jit.trace(w, x, strict=False, check_trace=False) |
| print('traced', resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e9, flush=True) |
| import coremltools as ct |
| from coremltools.converters.mil.frontend.torch import ops as _ops |
| from coremltools.converters.mil import Builder as _mb |
| def _cast(context, node, dtype, dtype_name): |
| x = _ops._get_inputs(context, node, expected=1)[0] |
| if x.can_be_folded_to_const(): res = _mb.const(val=dtype(np.asarray(x.val).reshape(-1)[0]), name=node.name) |
| elif len(x.shape) > 0: res = _mb.cast(x=_mb.squeeze(x=x, name=node.name + '_item'), dtype=dtype_name, name=node.name) |
| else: res = _mb.cast(x=x, dtype=dtype_name, name=node.name) |
| context.add(res, node.name) |
| _ops._cast = _cast |
| ml = ct.convert(traced, inputs=[ct.ImageType(name='image', shape=(1, 3, H, W), scale=1 / 255.0, color_layout=ct.colorlayout.RGB)], |
| outputs=[ct.TensorType(name='points'), ct.TensorType(name='mask')], |
| convert_to='mlprogram', compute_precision=ct.precision.FLOAT16, minimum_deployment_target=ct.target.macOS15) |
| out = f'moge2_{name}_{W}x{H}.mlpackage'; ml.save(out) |
| print('saved', out, 'rss', resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e9, flush=True) |
|
|
