File size: 2,457 Bytes
6013524 | 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 | # Copyright 2026 Spawn Lau
# SPDX-License-Identifier: Apache-2.0
"""Refine the proven 630-entry candidate to the native 70-token grid bound."""
import argparse
import hashlib
import json
from pathlib import Path
from ai_edge_litert import schema_py_generated as schema
from litert_lm_builder import LitertLmFileBuilder
from build_size_candidates import crop_positions, pack_model
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--candidate-dir", type=Path, required=True)
args = parser.parse_args()
unpacked = args.candidate_dir / "unpacked"
report_path = args.candidate_dir / "build-report.json"
report = json.loads(report_path.read_text(encoding="utf-8"))
assert not any(x["name"] == "position210" for x in report["candidates"])
model_path, = unpacked.glob("*TFLiteModel*tf_lite_vision_encoder*.tflite")
weight_path, = unpacked.glob("*TFLiteWeights*tf_lite_vision_encoder*.weight")
builder = LitertLmFileBuilder.from_toml_file(str(unpacked / "model.toml"))
model = schema.ModelT.InitFromObj(schema.Model.GetRootAsModel(model_path.read_bytes(), 0))
refinement = crop_positions(model, weight_path, 210)
refinement["native_bound_argument"] = {
"patch_size": 16, "pooling_kernel_size": 3, "max_subpatches": 630,
"preprocessed_dimensions_are_multiples_of_pixels": 48,
"minimum_patch_grid_side": 3,
"maximum_patch_grid_side": 210,
"valid_position_coordinate_range": [0, 209],
"source": "https://github.com/google-ai-edge/LiteRT-LM/blob/v0.18.0/support/preprocessor/image_preprocessor_utils.cc",
"scope": "Current app's native image preprocessing with vision_tokens_per_image=70 only.",
}
model_path.write_bytes(pack_model(model))
output_path = args.candidate_dir / "eg2-gallery-position210.litertlm"
with output_path.open("wb") as destination:
builder.build(destination)
report["position210_refinement"] = refinement
candidate = {"name": "position210", "path": str(output_path), "bytes": output_path.stat().st_size,
"saved_bytes": report["source_bytes"] - output_path.stat().st_size,
"sha256": hashlib.file_digest(output_path.open("rb"), "sha256").hexdigest()}
report["candidates"].append(candidate)
report_path.write_text(json.dumps(report, indent=2), encoding="utf-8")
print(json.dumps(candidate, indent=2))
if __name__ == "__main__":
main()
|