File size: 2,989 Bytes
ebec118 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | """Export LiquidInbetween to Core ML.
Images in and out (RGB, 512 × 512): pass two pixel buffers and a moment
t ∈ (0, 1), get the frame at that moment. Weights are FP16 (about 8.6 MB).
python export_coreml.py <weights.pt | teacher.safetensors> <out.mlpackage> [--int8]
"""
import sys
from pathlib import Path
import coremltools as ct
import coremltools.optimize.coreml as cto
import numpy as np
import torch
sys.path.insert(0, str(Path(__file__).parent))
from liquid_model import IFNet, load_teacher, student_from_teacher
SIZE = 512
class Wrapper(torch.nn.Module):
"""Pixel values 0...255 in and out, as Core ML image types expect."""
def __init__(self, net):
super().__init__()
self.net = net
def forward(self, first, second, t):
out = self.net(first / 255.0, second / 255.0, t)
return (out.clamp(0, 1) * 255.0)
weights, out = Path(sys.argv[1]), Path(sys.argv[2])
if weights.suffix == ".safetensors":
net = student_from_teacher(load_teacher(weights))
else:
net = IFNet(timestep=True)
net.load_state_dict(torch.load(weights, map_location="cpu"))
model = Wrapper(net.eval()).eval()
example = (torch.rand(1, 3, SIZE, SIZE) * 255, torch.rand(1, 3, SIZE, SIZE) * 255, torch.full((1, 1, 1, 1), 0.5))
with torch.no_grad():
traced = torch.jit.trace(model, example)
mlmodel = ct.convert(
traced,
inputs=[
ct.ImageType(name="first", shape=(1, 3, SIZE, SIZE), color_layout=ct.colorlayout.RGB),
ct.ImageType(name="second", shape=(1, 3, SIZE, SIZE), color_layout=ct.colorlayout.RGB),
ct.TensorType(name="t", shape=(1, 1, 1, 1), dtype=np.float32),
],
outputs=[ct.ImageType(name="frame", color_layout=ct.colorlayout.RGB)],
convert_to="mlprogram",
compute_precision=ct.precision.FLOAT16,
minimum_deployment_target=ct.target.iOS17,
)
if "--int8" in sys.argv:
# Experimental: 8-bit weights halve the size, but measured 24.7 dB against
# PyTorch (FP16 alone: 42.9 dB), visibly worse, so FP16 is the default.
# The transposed convolutions are the flow and mask heads: a tiny weight
# error there becomes pixels of motion error, so they stay FP16.
config = cto.OptimizationConfig(
global_config=cto.OpLinearQuantizerConfig(mode="linear_symmetric", dtype="int8", granularity="per_channel"),
op_type_configs={"conv_transpose": None},
)
mlmodel = cto.linear_quantize_weights(mlmodel, config)
mlmodel.short_description = "LiquidInbetween: the frame at moment t between two photos."
mlmodel.author = "Haplo LLC (jc_builds). Based on RIFE (Huang et al., ECCV 2022, MIT)."
mlmodel.license = "MIT"
mlmodel.version = "1.0"
mlmodel.input_description["first"] = "The first photo, 512 × 512 RGB."
mlmodel.input_description["second"] = "The second photo, 512 × 512 RGB."
mlmodel.input_description["t"] = "The moment between them, 0...1."
mlmodel.output_description["frame"] = "The frame at moment t."
mlmodel.save(str(out))
print("saved", out)
|