Converting models such as super-resolution and brightness correction from spandrel to CoreML

You can convert various models in the same way.

To convert a machine learning model to CoreML for use on iOS, you previously had to initialize the model architecture from the model repository and load the weights.

Spandrel is a library that can restore the structure of a model from a pytorch checkpoint file (pth), and supports a variety of models including super-resolution and inpainting. All you
need is the weights (checkpoint file).

It would be convenient to convert these to CoreML and use them on iOS and macOS devices, so we will convert them.

This time, let’s try using RetinexFormer to enhance low-light images.

Example: RetinexFormer conversion procedure

install

Install with pip.

pip install spandrel spandrel_extra_arches
# git clone https://github.com/chaiNNer-org/spandrel.git
# import shutil
# shutil.copytree("spandrel/libs/spandrel/spandrel/","/usr/local/lib/python3.10/dist-packages/spandrel")

Basically, you can use just pip, but for some reason when you install with pip, some packages don’t include the model architecture, so if you have trouble initializing the model later, clone the repository and replace it with the package.

Downloading checkpoint files

The checkpoint file can be downloaded from the link in the spandrel readme.

https://github.com/chaiNNer-org/spandrel

Initializing the model

Load the pth file into ModelLoader and initialize the model.
Normally, that’s all you need to do, but if you don’t have a commercial license, you need to register the model from spandrel_extra_arches. The same goes for codeformer.

from spandrel import MAIN_REGISTRY, ModelLoader, ImageModelDescriptor
from spandrel_extra_arches import EXTRA_REGISTRY

# add extra architectures before `ModelLoader` is used
MAIN_REGISTRY.add(*EXTRA_REGISTRY)
# load a model from disk
model = ModelLoader().load_from_file(r"LOL_v2_real.pth")
# make sure it's an image to image model
assert isinstance(model, ImageModelDescriptor)
# send it to the GPU and put it in inference mode
model.cuda().eval()

Let’s check the operation of the model with torch.
See below for the inference method.

Make a wrapped model

Create an nn.module that adds post-processing to the model.

import torch

torch_model = model.model.eval().cpu()

class CoreMLModel(torch.nn.Module):
    def __init__(self, m):
        super(CoreMLModel, self).__init__()
        self.m = m

    def forward(self, image):
        pred = self.m(image)
        output = torch.clamp(pred * 255, min=0, max=255)
        return output

coremlmodel = CoreMLModel(torch_model)

Convert your wrapped model.

import coremltools as ct
import torch

ex = torch.randn(1,3,512,512).cpu()
ts = torch.jit.trace(coremlmodel, ex)
mlmodel = ct.convert(ts, inputs=[ct.ImageType(shape=ex.shape,scale=1/255)],outputs=[ct.ImageType(name="output")])
mlmodel.save("retiinexformerNTIRE.mlpackage")

🐣

🐣

I’m a freelance engineer.
Work consultation
Please feel free to contact us with a brief development description.
rockyshikoku@gmail.com

I am creating applications using machine learning and AR technology.

I send machine learning / AR related information.

GitHub

Twitter
Medium