Initial commit of CoreMLaMa: LaMa for Core ML

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+.DS_Store

CoreMLaMa.py

@@ -0,0 +1,14 @@
+import torch
+
+# A "wrapper" module around LaMa which handles some input/output
+# pre- and post-processing beyond CoreML's built-in capabilities
+class CoreMLaMa(torch.nn.Module):
+    def __init__(self, lama):
+        super(CoreMLaMa, self).__init__()
+        self.lama = lama
+
+    def forward(self, image, mask):
+        normalized_mask = ((mask > 0) * 1).byte()
+        lama_out = self.lama(image, normalized_mask)
+        output = torch.clamp(lama_out * 255, min=0, max=255)
+        return output

README.md

@@ -0,0 +1,30 @@
+## CoreMLaMa: LaMa for Core ML
+
+This repo contains a script for converting a [LaMa](https://advimman.github.io/lama-project/) (aka cute, fuzzy 🦙) model to Apple's Core ML model format.
+
+### Conversion Instructions
+
+1. Create a Conda environment for CoreMLaMa:
+    ```sh
+    conda create -n coremllama python=3.10 # works with mamba, too
+    conda activate coremllama
+    pip install -r requirements.txt
+    ```
+
+2. Run the conversion script:
+    ```sh
+    python convert_lama.py
+    ```
+
+This script will download and convert [Big LaMa](https://github.com/advimman/lama#models-options) to a Core ML package named `LaMa.mlpackage`.
+
+### Acknowledgements and Thanks
+
+Thanks to the authors of LaMa:
+
+[[Project page](https://advimman.github.io/lama-project/)] [[arXiv](https://arxiv.org/abs/2109.07161)] [[Supplementary](https://ashukha.com/projects/lama_21/lama_supmat_2021.pdf)] [[BibTeX](https://senya-ashukha.github.io/projects/lama_21/paper.txt)] [[Casual GAN Papers Summary](https://www.casualganpapers.com/large-masks-fourier-convolutions-inpainting/LaMa-explained.html)]
+
+
+CoreMLaMa uses the LaMa model and supporting code from [LaMa Cleaner](https://lama-cleaner-docs.vercel.app). LaMa Cleaner makes this project much simpler.
+
+

convert_lama.py

@@ -0,0 +1,49 @@
+import coremltools as ct
+import torch
+
+from lama_cleaner.model.lama import LaMa
+from CoreMLaMa import CoreMLaMa
+
+model_manager = LaMa("cpu")
+
+# Fixed image/mask size
+# Flexible input shapes are not (currently) supported, for various reasons
+size = (800, 800) # pixel width x height
+
+# Image/mask shapes in PyTorch format
+image_shape=(1, 3, size[1], size[0])
+mask_shape=(1, 1, size[1], size[0])
+
+lama_inpaint_model = model_manager.model
+model = CoreMLaMa(lama_inpaint_model).eval()
+
+print("Scripting CoreMLaMa")
+jit_model = torch.jit.script(model)
+
+print("Converting model")
+# Note that ct.ImageType assumes an 8 bpp image, while LaMa
+# uses 32-bit FP math internally. Creating a CoreML model
+# that can work with 32-bit FP image inputs is on the "To Do"
+# list
+coreml_model = ct.convert(
+    jit_model,
+    convert_to="mlprogram",
+    compute_precision=ct.precision.FLOAT32,
+    compute_units=ct.ComputeUnit.CPU_AND_GPU,
+    inputs=[
+        ct.ImageType(name="image",
+                     shape=image_shape,
+                     scale=1/255.0),
+        ct.ImageType(
+            name="mask",
+            shape=mask_shape,
+            color_layout=ct.colorlayout.GRAYSCALE)
+    ],
+    outputs=[ct.ImageType(name="output")],
+    skip_model_load=True
+)
+
+coreml_model_file_name = "LaMa.mlpackage"
+print(f"Saving model to {coreml_model_file_name}")
+coreml_model.save(coreml_model_file_name)
+print("Done!")

requirements.txt

@@ -0,0 +1,3 @@
+torch>=2
+lama-cleaner>=1.1
+coremltools>=6.3