immich/machine-learning/export/ann/run.py

import os
import platform
import subprocess
from tempfile import TemporaryDirectory
from typing import Callable, ClassVar

import onnx
import torch
from onnx2torch import convert
from onnx2torch.node_converters.registry import add_converter
from onnxruntime.tools.onnx_model_utils import fix_output_shapes, make_input_shape_fixed
from tinynn.converter import TFLiteConverter
from huggingface_hub import snapshot_download
from onnx2torch.onnx_graph import OnnxGraph
from onnx2torch.onnx_node import OnnxNode
from onnx2torch.utils.common import OperationConverterResult, onnx_mapping_from_node
from onnx.shape_inference import infer_shapes_path
from huggingface_hub import login, upload_file

# egregious hacks:
# changed `Clip`'s min/max logic to skip empty strings
# changed OnnxSqueezeDynamicAxes to use `sorted` instead of `torch.sort``
# commented out shape inference in `fix_output_shapes``


class ArgMax(torch.nn.Module):
    def __init__(self, dim: int = -1, keepdim: bool = False):
        super().__init__()
        self.dim = dim
        self.keepdim = keepdim

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return torch.argmax(input, dim=self.dim, keepdim=self.keepdim)


class Erf(torch.nn.Module):
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return torch.erf(input)


@add_converter(operation_type="ArgMax", version=13)
def _(node: OnnxNode, graph: OnnxGraph) -> OperationConverterResult:
    return OperationConverterResult(
        torch_module=ArgMax(),
        onnx_mapping=onnx_mapping_from_node(node=node),
    )


class ExportBase(torch.nn.Module):
    task: ClassVar[str]

    def __init__(
        self,
        name: str,
        input_shape: tuple[int, ...],
        pretrained: str | None = None,
        optimization_level: int = 5,
    ):
        super().__init__()
        self.name = name
        self.optimize = optimization_level
        self.nchw_transpose = False
        self.input_shape = input_shape
        self.pretrained = pretrained
        self.dummy_param = torch.nn.Parameter(torch.empty(0))
        self.model = self.load().eval()
        for param in self.parameters():
            param.requires_grad_(False)
        self.eval()

    def load(self) -> torch.nn.Module:
        cache_dir = os.path.join(os.environ["CACHE_DIR"], self.model_name)
        task_path = os.path.join(cache_dir, self.task)
        model_path = os.path.join(task_path, "model.onnx")
        if not os.path.isfile(model_path):
            snapshot_download(self.repo_name, cache_dir=cache_dir, local_dir=cache_dir)
        infer_shapes_path(model_path, check_type=True, strict_mode=True, data_prop=True)
        onnx_model = onnx.load_model(model_path)
        make_input_shape_fixed(onnx_model.graph, onnx_model.graph.input[0].name, self.input_shape)
        fix_output_shapes(onnx_model)
        # try:
        # onnx.save(onnx_model, model_path)
        # except:
        # onnx.save(onnx_model, model_path, save_as_external_data=True, all_tensors_to_one_file=False)
        # infer_shapes_path(model_path, check_type=True, strict_mode=True, data_prop=True)
        # onnx_model = onnx.load_model(model_path)
        # onnx_model = infer_shapes(onnx_model, check_type=True, strict_mode=True, data_prop=True)
        return convert(onnx_model)

    def forward(self, *inputs: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor]:
        if self.precision == "fp16":
            inputs = tuple(i.half() for i in inputs)

        out = self._forward(*inputs)
        if self.precision == "fp16":
            if isinstance(out, tuple):
                return tuple(o.float() for o in out)
            return out.float()
        return out

    def _forward(self, *inputs: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor]:
        return self.model(*inputs)

    def to_armnn(self, output_path: str) -> None:
        output_dir = os.path.dirname(output_path)
        os.makedirs(output_dir, exist_ok=True)
        self(*self.dummy_inputs)
        print(f"Exporting {self.model_name} ({self.task}) with {self.precision} precision")
        jit = torch.jit.trace(self, self.dummy_inputs).eval()
        with TemporaryDirectory() as tmpdir:
            tflite_model_path = os.path.join(tmpdir, "model.tflite")
            converter = TFLiteConverter(
                jit,
                self.dummy_inputs,
                tflite_model_path,
                optimize=self.optimize,
                nchw_transpose=self.nchw_transpose,
            )
            # segfaults on ARM, must run on x86_64 / AMD64
            converter.convert()

            subprocess.run(
                [
                    "./armnnconverter",
                    "-f",
                    "tflite-binary",
                    "-m",
                    tflite_model_path,
                    "-i",
                    "input_tensor",
                    "-o",
                    "output_tensor",
                    "-p",
                    output_path,
                ],
                capture_output=True,
            )
        print(f"Finished exporting {self.name} ({self.task}) with {self.precision} precision")

    @property
    def dummy_inputs(self) -> tuple[torch.FloatTensor]:
        return (torch.rand(self.input_shape, device=self.device, dtype=self.dtype),)

    @property
    def model_name(self) -> str:
        return f"{self.name}__{self.pretrained}" if self.pretrained else self.name

    @property
    def repo_name(self) -> str:
        return f"immich-app/{self.model_name}"

    @property
    def device(self) -> torch.device:
        return self.dummy_param.device

    @property
    def dtype(self) -> torch.dtype:
        return self.dummy_param.dtype

    @property
    def precision(self) -> str:
        match self.dtype:
            case torch.float32:
                return "fp32"
            case torch.float16:
                return "fp16"
            case _:
                raise ValueError(f"Unsupported dtype {self.dtype}")


class ArcFace(ExportBase):
    task = "recognition"


class RetinaFace(ExportBase):
    task = "detection"


class OpenClipVisual(ExportBase):
    task = "visual"


class OpenClipTextual(ExportBase):
    task = "textual"

    @property
    def dummy_inputs(self) -> tuple[torch.LongTensor]:
        return (torch.randint(0, 5000, self.input_shape, device=self.device, dtype=torch.int32),)


class MClipTextual(ExportBase):
    task = "textual"

    @property
    def dummy_inputs(self) -> tuple[torch.LongTensor]:
        return (
            torch.randint(0, 5000, self.input_shape, device=self.device, dtype=torch.int32),
            torch.randint(0, 1, self.input_shape, device=self.device, dtype=torch.int32),
        )


def main() -> None:
    if platform.machine() not in ("x86_64", "AMD64"):
        raise RuntimeError(f"Can only run on x86_64 / AMD64, not {platform.machine()}")
    login(token=os.environ["HF_AUTH_TOKEN"])
    os.environ["LD_LIBRARY_PATH"] = "armnn"
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    failed: list[Callable[[], ExportBase]] = [
        lambda: OpenClipVisual("ViT-H-14-378-quickgelu", (1, 3, 378, 378), pretrained="dfn5b"), # flatbuffers: cannot grow buffer beyond 2 gigabytes (will probably work with fp16)
        lambda: OpenClipVisual("ViT-H-14-quickgelu", (1, 3, 224, 224), pretrained="dfn5b"), # flatbuffers: cannot grow buffer beyond 2 gigabytes (will probably work with fp16)
        lambda: OpenClipTextual("nllb-clip-base-siglip", (1, 77), pretrained="v1"), # ERROR (tinynn.converter.base) Unsupported ops: aten::logical_not
        lambda: OpenClipTextual("nllb-clip-large-siglip", (1, 77), pretrained="v1"), # ERROR (tinynn.converter.base) Unsupported ops: aten::logical_not
        lambda: OpenClipVisual("ViT-B-32", (1, 3, 224, 224), pretrained="laion2b_e16"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipTextual("ViT-B-32", (1, 77), pretrained="laion2b_e16"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipVisual("ViT-B-32", (1, 3, 224, 224), pretrained="laion400m_e31"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipTextual("ViT-B-32", (1, 77), pretrained="laion400m_e31"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipVisual("ViT-B-32", (1, 3, 224, 224), pretrained="laion400m_e32"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipTextual("ViT-B-32", (1, 77), pretrained="laion400m_e32"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipVisual("ViT-B-32", (1, 3, 224, 224), pretrained="laion2b-s34b-b79k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipTextual("ViT-B-32", (1, 77), pretrained="laion2b-s34b-b79k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipVisual("ViT-B-16", (1, 3, 224, 224), pretrained="laion400m_e31"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipTextual("ViT-B-16", (1, 77), pretrained="laion400m_e31"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipVisual("ViT-B-16", (1, 3, 224, 224), pretrained="laion400m_e32"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipTextual("ViT-B-16", (1, 77), pretrained="laion400m_e32"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipVisual("ViT-B-16-plus-240", (1, 3, 224, 224), pretrained="laion400m_e31"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipTextual("ViT-B-16-plus-240", (1, 77), pretrained="laion400m_e31"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipVisual("ViT-L-14", (1, 3, 224, 224), pretrained="laion400m_e31"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipTextual("ViT-L-14", (1, 77), pretrained="laion400m_e31"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipVisual("ViT-L-14", (1, 3, 224, 224), pretrained="laion400m_e32"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipTextual("ViT-L-14", (1, 77), pretrained="laion400m_e32"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipVisual("ViT-L-14", (1, 3, 224, 224), pretrained="laion2b-s32b-b82k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipTextual("ViT-L-14", (1, 77), pretrained="laion2b-s32b-b82k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipVisual("ViT-H-14", (1, 3, 224, 224), pretrained="laion2b-s32b-b79k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipTextual("ViT-H-14", (1, 77), pretrained="laion2b-s32b-b79k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipVisual("ViT-g-14", (1, 3, 224, 224), pretrained="laion2b-s12b-b42k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipTextual("ViT-g-14", (1, 77), pretrained="laion2b-s12b-b42k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipVisual("XLM-Roberta-Large-Vit-B-16Plus", (1, 3, 240, 240)), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipVisual("XLM-Roberta-Large-ViT-H-14", (1, 3, 224, 224), pretrained="frozen_laion5b_s13b_b90k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipVisual("nllb-clip-base-siglip", (1, 3, 384, 384), pretrained="v1"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipVisual("nllb-clip-large-siglip", (1, 3, 384, 384), pretrained="v1"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
        lambda: OpenClipVisual("RN50", (1, 3, 224, 224), pretrained="yfcc15m"), # BatchNorm operation with mean/var output is not implemented
        lambda: OpenClipTextual("RN50", (1, 77), pretrained="yfcc15m"), # BatchNorm operation with mean/var output is not implemented
        lambda: OpenClipVisual("RN50", (1, 3, 224, 224), pretrained="cc12m"), # BatchNorm operation with mean/var output is not implemented
        lambda: OpenClipTextual("RN50", (1, 77), pretrained="cc12m"), # BatchNorm operation with mean/var output is not implemented
        lambda: MClipTextual("XLM-Roberta-Large-Vit-L-14", (1, 77)), # Expected normalized_shape to be at least 1-dimensional, i.e., containing at least one element, but got normalized_shape = []
        lambda: MClipTextual("XLM-Roberta-Large-Vit-B-16Plus", (1, 77)), # Expected normalized_shape to be at least 1-dimensional, i.e., containing at least one element, but got normalized_shape = []
        lambda: MClipTextual("LABSE-Vit-L-14", (1, 77)), # Expected normalized_shape to be at least 1-dimensional, i.e., containing at least one element, but got normalized_shape = []
        lambda: OpenClipTextual("XLM-Roberta-Large-ViT-H-14", (1, 77), pretrained="frozen_laion5b_s13b_b90k"), # Expected normalized_shape to be at least 1-dimensional, i.e., containing at least one element, but got normalized_shape = []
    ]
    
    succeeded: list[Callable[[], ExportBase]] = [
        lambda: OpenClipVisual("ViT-B-32", (1, 3, 224, 224), pretrained="openai"),
        lambda: OpenClipTextual("ViT-B-32", (1, 77), pretrained="openai"),
        lambda: OpenClipVisual("ViT-B-16", (1, 3, 224, 224), pretrained="openai"),
        lambda: OpenClipTextual("ViT-B-16", (1, 77), pretrained="openai"),
        lambda: OpenClipVisual("ViT-L-14", (1, 3, 224, 224), pretrained="openai"),
        lambda: OpenClipTextual("ViT-L-14", (1, 77), pretrained="openai"),
        lambda: OpenClipVisual("ViT-L-14-336", (1, 3, 336, 336), pretrained="openai"),
        lambda: OpenClipTextual("ViT-L-14-336", (1, 77), pretrained="openai"),
        lambda: OpenClipVisual("RN50", (1, 3, 224, 224), pretrained="openai"),
        lambda: OpenClipTextual("RN50", (1, 77), pretrained="openai"),
        lambda: OpenClipTextual("ViT-H-14-quickgelu", (1, 77), pretrained="dfn5b"),
        lambda: OpenClipTextual("ViT-H-14-378-quickgelu", (1, 77), pretrained="dfn5b"),
        lambda: OpenClipVisual("XLM-Roberta-Large-Vit-L-14", (1, 3, 224, 224)),
        lambda: OpenClipVisual("XLM-Roberta-Large-Vit-B-32", (1, 3, 224, 224)),
        lambda: ArcFace("buffalo_s", (1, 3, 112, 112), optimization_level=3),
        lambda: RetinaFace("buffalo_s", (1, 3, 640, 640), optimization_level=3),
        lambda: ArcFace("buffalo_m", (1, 3, 112, 112), optimization_level=3),
        lambda: RetinaFace("buffalo_m", (1, 3, 640, 640), optimization_level=3),
        lambda: ArcFace("buffalo_l", (1, 3, 112, 112), optimization_level=3),
        lambda: RetinaFace("buffalo_l", (1, 3, 640, 640), optimization_level=3),
        lambda: ArcFace("antelopev2", (1, 3, 112, 112), optimization_level=3),
        lambda: RetinaFace("antelopev2", (1, 3, 640, 640), optimization_level=3),
    ]

    models: list[Callable[[], ExportBase]] = [*failed, *succeeded]
    for _model in succeeded:
        model = _model().to(device)
        try:
            relative_path = os.path.join(model.task, "model.armnn")
            output_path = os.path.join("output", model.model_name, relative_path)
            model.to_armnn(output_path)
            upload_file(path_or_fileobj=output_path, path_in_repo=relative_path, repo_id=model.repo_name)
            if device == torch.device("cuda"):
                model.half()
                relative_path = os.path.join(model.task, "fp16", "model.armnn")
                output_path = os.path.join("output", model.model_name, relative_path)
                model.to_armnn(output_path)
                upload_file(path_or_fileobj=output_path, path_in_repo=relative_path, repo_id=model.repo_name)
            
        except Exception as exc:
            print(f"Failed to export {model.model_name} ({model.task}): {exc}")


if __name__ == "__main__":
    with torch.no_grad():
        main()