import sys, torch, torch.nn as nn
import coremltools as ct
sys.path.insert(0, "/tmp/clothseg/repo")
from network import U2NET

CKPT = "/tmp/clothseg/cloth_segm_u2net_latest.pth"
OUT  = "/tmp/clothseg/ClothSegmentation.mlpackage"

# Build + load (the checkpoint was saved as a plain state_dict, possibly with a module. prefix).
net = U2NET(in_ch=3, out_ch=4)
sd = torch.load(CKPT, map_location="cpu")
if isinstance(sd, dict) and "model_state_dict" in sd:
    sd = sd["model_state_dict"]
sd = { (k[7:] if k.startswith("module.") else k): v for k, v in sd.items() }
net.load_state_dict(sd)
net.eval()

# Wrapper: return only the fused output d0, as per-class softmax probabilities (1,4,768,768).
class Wrap(nn.Module):
    def __init__(self, m): super().__init__(); self.m = m
    def forward(self, x):
        d0 = self.m(x)[0]                 # (1,4,768,768) logits
        return torch.softmax(d0, dim=1)   # per-class probabilities

wrap = Wrap(net).eval()
dummy = torch.randn(1, 3, 768, 768)
with torch.no_grad():
    traced = torch.jit.trace(wrap, dummy)

# CoreML image input: pixels [0,255] -> model wants (p/255-0.5)/0.5 = p/127.5 - 1.
mlmodel = ct.convert(
    traced,
    inputs=[ct.ImageType(name="image", shape=(1, 3, 768, 768),
                         scale=1.0/127.5, bias=[-1.0, -1.0, -1.0],
                         color_layout=ct.colorlayout.RGB)],
    outputs=[ct.TensorType(name="probs")],
    minimum_deployment_target=ct.target.iOS16,
    compute_precision=ct.precision.FLOAT16,
)
mlmodel.short_description = "U2NET cloth segmentation (bg/upper/lower/full) — 768x768, softmax probs"
mlmodel.save(OUT)
print("SAVED", OUT)