| """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: |
| |
| |
| |
| |
| 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) |
|
|
