Add MI-GAN 512 Places2 for Core ML

migan/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Picsart AI Research (PAIR)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

migan/MIGAN.mlpackage.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c279edffd937627477b716376cea8e51aa8ac42c1691c8b47365020b915903c8
+size 24529826

migan/convert.py

@@ -0,0 +1,34 @@
+import sys, torch, torch.nn as nn, coremltools as ct
+src, ckpt, out = sys.argv[1], sys.argv[2], sys.argv[3]
+sys.path.insert(0, src)
+from lib.model_zoo.migan_inference import Generator
+
+class Wrapped(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.g = Generator(resolution=512)
+        self.g.load_state_dict(torch.load(ckpt, map_location="cpu"))
+        self.g.eval()
+    def forward(self, image, mask):
+        # image: -1..1 (from the ImageType scale); mask: 0..1, 1 = the hole.
+        known = (mask < 0.5).float()
+        x = torch.cat([known - 0.5, image * known], dim=1)
+        y = self.g(x)
+        return ((y * 0.5 + 0.5) * 255).clamp(0, 255)
+
+m = Wrapped().eval()
+ex = (torch.zeros(1, 3, 512, 512), torch.zeros(1, 1, 512, 512))
+traced = torch.jit.trace(m, ex)
+mlm = ct.convert(
+    traced,
+    inputs=[ct.ImageType(name="image", shape=(1, 3, 512, 512), scale=2 / 255, bias=[-1, -1, -1], color_layout=ct.colorlayout.RGB),
+            ct.ImageType(name="mask", shape=(1, 1, 512, 512), scale=1 / 255, color_layout=ct.colorlayout.GRAYSCALE)],
+    outputs=[ct.ImageType(name="output", color_layout=ct.colorlayout.RGB)],
+    minimum_deployment_target=ct.target.iOS17,
+    compute_precision=ct.precision.FLOAT32,
+)
+mlm.author = "Picsart AI Research (MI-GAN, MIT); Core ML conversion for PhotoWagon"
+mlm.license = "MIT"
+mlm.short_description = "MI-GAN 512 Places2: image and mask (white = hole) in, filled image out."
+mlm.save(out)
+print("saved", out)