MertandGitHub
23f1f87dd0
* add immich_model * graph per shape * decouple ocr from ort * prepared graphs * use f16 for accelerators * preload all shapes * warm-up pass after preload * keep precision for fp32 graphs * ensure padding for text recognition * optimize postprocessing for fp16 outputs * handle raw ctc logits * use f16 output for rknn * mmap token embedding * read device metadata
108 lines
4.4 KiB
Python
108 lines
4.4 KiB
Python
import fcntl
|
|
import pickle
|
|
import sys
|
|
from pathlib import Path
|
|
from shutil import rmtree
|
|
from typing import Any
|
|
|
|
import onnx
|
|
from google.protobuf.message import DecodeError
|
|
from immich_model.onnx.f16 import derive
|
|
from immich_model.runtime import apply_rewrites
|
|
from onnxruntime.capi.onnxruntime_pybind11_state import InvalidProtobuf, NoSuchFile
|
|
|
|
from ..config import settings
|
|
from .ort import UNREADABLE, GraphSpec, Manifest, fresh
|
|
|
|
|
|
def main() -> None:
|
|
spec: GraphSpec = pickle.loads(sys.stdin.buffer.read())
|
|
with spec.model_path.with_suffix(".lock").open("w") as lock:
|
|
fcntl.flock(lock, fcntl.LOCK_EX) # several workers may want the same graph
|
|
if fresh(spec) is None:
|
|
try:
|
|
prepare(spec)
|
|
except (InvalidProtobuf, NoSuchFile, DecodeError):
|
|
sys.exit(UNREADABLE)
|
|
|
|
|
|
def prepare(spec: GraphSpec) -> None:
|
|
rmtree(spec.directory, ignore_errors=True) # whatever is here was made under other facts, or not finished
|
|
spec.directory.mkdir(parents=True)
|
|
graph = rewritten(spec, narrowed(spec))
|
|
if spec.provider == "CPUExecutionProvider":
|
|
graph = prepacked(spec, graph) # allows mmap'ing the file instead of prepacking to dirty heap memory
|
|
elif spec.directory.as_posix() in spec.provider_options[0].values():
|
|
spec.session(graph) # the provider keeps what it compiles, in the directory it was given
|
|
manifest = Manifest(facts=spec.facts, graph=graph.relative_to(spec.model_path.parent).as_posix())
|
|
spec.manifest.write_text(manifest.model_dump_json()) # last, so that it means complete
|
|
|
|
|
|
def narrowed(spec: GraphSpec) -> Path:
|
|
if not spec.half:
|
|
return spec.model_path
|
|
destination = spec.directory / "model_fp16.onnx"
|
|
derive(spec.model_path, destination)
|
|
return destination
|
|
|
|
|
|
def rewritten(spec: GraphSpec, graph: Path) -> Path:
|
|
written: Path = graph if settings.legacy_models else apply_rewrites(graph, spec.plan)
|
|
return written
|
|
|
|
|
|
def prepacked(spec: GraphSpec, graph: Path) -> Path:
|
|
destination = spec.directory / "model.onnx"
|
|
options = spec.sess_options()
|
|
options.optimized_model_filepath = destination.as_posix()
|
|
options.add_session_config_entry("session.optimized_model_external_initializers_file_name", "model.data")
|
|
options.add_session_config_entry("session.save_external_prepacked_constant_initializers", "1")
|
|
options.add_session_config_entry("session.optimized_model_external_initializers_min_size_in_bytes", "1024")
|
|
spec.session(graph, options)
|
|
align(destination)
|
|
return destination
|
|
|
|
|
|
def align(destination: Path) -> None:
|
|
"""Re-lays the data file so every region starts aligned. ORT pads only large regions, so one odd length skews
|
|
all that follow, and MLAS reads a packed weight with aligned loads: it faults, and reads the rest slower."""
|
|
model = onnx.load(destination.as_posix(), load_external_data=False)
|
|
length_at: dict[int, int] = {}
|
|
offsets: list[Any] = []
|
|
packs: list[tuple[Any, str, list[list[str]]]] = []
|
|
for tensor in model.graph.initializer:
|
|
fields = {entry.key: entry for entry in tensor.external_data}
|
|
for key, entry in fields.items():
|
|
if key == "offset":
|
|
length_at[int(entry.value)] = int(fields["length"].value)
|
|
offsets.append(entry)
|
|
elif key.startswith("prepacked_"):
|
|
blob_key, *rest = entry.value.split("|")
|
|
triples = [item.split(";") for item in rest]
|
|
length_at.update((int(offset), int(length)) for offset, length, _ in triples)
|
|
packs.append((entry, blob_key, triples))
|
|
|
|
moved, cursor = {}, 0
|
|
for offset in sorted(length_at):
|
|
cursor += -cursor % (4096 if length_at[offset] > 1 << 20 else 64)
|
|
moved[offset] = cursor
|
|
cursor += length_at[offset]
|
|
|
|
data = destination.with_suffix(".data")
|
|
with data.open("rb") as source, data.with_suffix(".aligned").open("wb") as target:
|
|
for offset in sorted(length_at):
|
|
source.seek(offset)
|
|
target.seek(moved[offset]) # what it skips reads as zeros
|
|
target.write(source.read(length_at[offset]))
|
|
data.with_suffix(".aligned").replace(data)
|
|
|
|
for entry in offsets:
|
|
entry.value = str(moved[int(entry.value)])
|
|
for entry, blob_key, triples in packs:
|
|
entry.value = "|".join([blob_key, *(f"{moved[int(o)]};{n};{c}" for o, n, c in triples)])
|
|
onnx.save(model, destination.as_posix())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|