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):   # coremltools 9.0 calls int() on a one-element array and numpy 2 refuses
    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)