gimlet-api 0.0.7__py3-none-any.whl → 0.0.9__py3-none-any.whl

{gimlet_api-0.0.7.dist-info → gimlet_api-0.0.9.dist-info}/METADATA

@@ -14,6 +14,6 @@ Requires-Dist: torch-mlir-gml
 Requires-Dist: numpy<2.0.0
 Requires-Dist: transformers>=4.43.3
 Requires-Dist: safetensors-mlir
-Version: 0.0.7
+Version: 0.0.9
 
 UNKNOWN

{gimlet_api-0.0.7.dist-info → gimlet_api-0.0.9.dist-info}/RECORD

@@ -1,14 +1,14 @@
 gml/__init__.py,sha256=H3WQZ_RaN7VNeb__qeHEbKLEwkaG7gpL5FQ8s1IotUA,773
 gml/_utils.py,sha256=mSCWHhCdzcUvHqmJIB2FS215K1LMgJCWcZ6e6FWK3hQ,1184
 gml/asset_manager.py,sha256=VnbqUZHPOgPrAh6ri9C0EuNhS8tAHIrbUyJPAJuD9po,2053
-gml/client.py,sha256=5QDKljltBeBTCd2hH38--fTSP0bVVcAvSnWsA9YEFQc,13819
-gml/compile.py,sha256=K4WdC01WkyLlbcoSJzdF2LaVmOgxIkIdc3YNjRFRw9s,10849
-gml/device.py,sha256=sMILurG02aDjL8wrdBW3ftC44WoAPUeZ4Y0yQ0DtaBk,2665
-gml/hf.py,sha256=hi0Af0Q3FM7VvfLB1PkrNai1j7siH6Ouwc_sXf8QE8c,17900
-gml/model.py,sha256=cHFjIEplWCDeSCSl_IPHYzNyv-KOenx-OAsEd-5TpTs,7260
+gml/client.py,sha256=jPA71PTbv-4DX2FsfFcj1VPN-eMGdF4tKCj2NnOW7ZE,13862
+gml/compile.py,sha256=Ih43r_zU07p91w9aiA0lrPJfmACpAWg0x_HFddMSy7Q,8346
+gml/device.py,sha256=Iw71NnuLcgjY32ZMXHlnlPkosTuHEmL9E98utmNChlM,2650
+gml/hf.py,sha256=pp215wNmaPyCVy4DqFJbe_vEe1BRJ1GAJEURZnLuU0g,28220
+gml/model.py,sha256=xESdD7tlqn93ym67Lyyk7TZdM3wUqyn7qWdP2AbgdkI,7261
 gml/model_utils.py,sha256=vZvE5cHZIDkUkeZ4Pk4hhV-zOYMiREluv4b8kdqQ3Ig,1375
-gml/pipelines.py,sha256=6tujvMpAACwmEmUfZGFuRn0N8zKvjSVIsIX7FKcPvEU,4301
-gml/preprocessing.py,sha256=STQDSA1_jXPTenJotNtsNMXOc9h1x_wJyQ100LXS6-g,3209
+gml/pipelines.py,sha256=LKj_lh5I5HzyUUIPG4CImiqBnQPrJsj0CHPKhLiOOGo,8374
+gml/preprocessing.py,sha256=MaKkEW4ZP9fjpkJQfpc0X3rCUuSuSmJnGMClHamKmZU,3210
 gml/proto/gogoproto/gogo_pb2.py,sha256=WVMIAR8K--mCUkTPM7mEeeXGpQlRRtt_kco10iP3CZs,15728
 gml/proto/mediapipe/framework/calculator_contract_test_pb2.py,sha256=hNjyZCBz3RYa6rN4xR3FOCZKA24gq_LsJ3EMegl5wK4,2031
 gml/proto/mediapipe/framework/calculator_options_pb2.py,sha256=Nq1BQRtLdsIgfkw7ymD3eg2p2_RSlZhiHS7YbDhNHR0,1563
@@ -25,16 +25,16 @@ gml/proto/mediapipe/framework/thread_pool_executor_pb2.py,sha256=9TJ66fqSo1BiJmE
 gml/proto/opentelemetry/proto/common/v1/common_pb2.py,sha256=wQjeDti-C8JiNwRn-z5M5p-Fqxm-SmnbPaoitJcSK-4,2860
 gml/proto/opentelemetry/proto/metrics/v1/metrics_pb2.py,sha256=t2Far6oVcUFQIimzgAkZ8vQd0asMIlvECp4osC0ujgg,9735
 gml/proto/opentelemetry/proto/resource/v1/resource_pb2.py,sha256=cbNmE12Nm3PjW4NXU7-Z-9m_0Zs3Ab8R1xLkDnvclCg,1730
-gml/proto/src/api/corepb/v1/controlplane_pb2.py,sha256=ChS3EbDhgtywVrI1fU-CqqynhRpRU-YrdK7gFCaN46w,6008
-gml/proto/src/api/corepb/v1/cp_edge_pb2.py,sha256=O0YhmiQrwy-bFuqYt5Vno2M8m99-qlwyhkfeLBB14WQ,14903
-gml/proto/src/api/corepb/v1/device_info_pb2.py,sha256=j4gvzhM1MsDLMHCAwah1X5xxpfG_5nODs7K83mTY0zI,3425
-gml/proto/src/api/corepb/v1/gem_config_pb2.py,sha256=yyEqUqq3-YiX-ByAhbTbZfdh09KuNzEtIYhgk_noJVM,3367
-gml/proto/src/api/corepb/v1/mediastream_pb2.py,sha256=_WV7zav0uaPHzP-yvjRtUtrwexWiz4eqVclIZmhqdcY,7193
-gml/proto/src/api/corepb/v1/model_exec_pb2.py,sha256=1DM58lSFgfoHk0ui3ZTjDfifgp4dhE7nHvhMwmInpsA,27103
+gml/proto/src/api/corepb/v1/controlplane_pb2.py,sha256=BosvQ6GYaUGsNTkRZH7osP2dZGWP6U9WyxItIQ_QS-8,9769
+gml/proto/src/api/corepb/v1/cp_edge_pb2.py,sha256=oIpxq13C1ynK3alzDNZTOL5URxz5qzbDLD9NOM5xxjE,14511
+gml/proto/src/api/corepb/v1/device_info_pb2.py,sha256=5orIOJAkvtH9pWBSXveDASFi4Rn59YWdOSnVLdj891A,5356
+gml/proto/src/api/corepb/v1/gem_config_pb2.py,sha256=2ljfF16Xeqgj9TM3gHN54BqRHqS3SQNhOCenEY9K9qU,4718
+gml/proto/src/api/corepb/v1/mediastream_pb2.py,sha256=Un9OwDUmWdqv92QP66K-WVOAzxP_4hMoz33JI4W1G5Y,7868
+gml/proto/src/api/corepb/v1/model_exec_pb2.py,sha256=o0drstrDssejqCFo8Cmm9F0zDw_bmzeOUHiYFrruOqE,29877
 gml/proto/src/common/typespb/jwt_pb2.py,sha256=lxy-bqbyg96i9n_xr2JbkuWX-ldnoJavXPMnApzVSio,5580
 gml/proto/src/common/typespb/status_pb2.py,sha256=IbBJnbsAlvsuTtyT285ZuW6k5VaPfl5kRSOnBxD_H8M,2109
 gml/proto/src/common/typespb/uuid_pb2.py,sha256=5Fm3jYpCPX7sMrP6RhRYsF0SnuZNIBEQJk9f0jwZ2Rw,1188
-gml/proto/src/controlplane/compiler/cpb/v1/cpb_pb2.py,sha256=tkJFPWpndKZy19TFuLKlBfWW1fUQPj0lJLiQ9HfugZU,3213
+gml/proto/src/controlplane/compiler/cpb/v1/cpb_pb2.py,sha256=R8jcxOlR1iz4Y7MnxIKoJ2RaNayqWPiBSt0W496QT-c,3262
 gml/proto/src/controlplane/compiler/cpb/v1/cpb_pb2_grpc.py,sha256=l-gTK9nYpTlVb7QGAckSQXlHhkRdKe2-nrxXc8NQavY,2912
 gml/proto/src/controlplane/directory/directorypb/v1/directory_pb2.py,sha256=KgoUT8ccF-yJPe1r4otQjAPQoKBaQzdBlHoIUSkk0yE,11445
 gml/proto/src/controlplane/directory/directorypb/v1/directory_pb2_grpc.py,sha256=p3OpT8-hfNHu4-29qr-ZahRwO-LoCYM9Q4jomAHTXGA,24572
@@ -44,7 +44,8 @@ gml/proto/src/controlplane/logicalpipeline/lppb/v1/lppb_pb2.py,sha256=wvLQvoh2UA
 gml/proto/src/controlplane/logicalpipeline/lppb/v1/lppb_pb2_grpc.py,sha256=-snjW7n6JveUzJVPFcm25XlL19kowPSKgd61l_jPnHA,9541
 gml/proto/src/controlplane/model/mpb/v1/mpb_pb2.py,sha256=RVedXkNYu2iF5OHiXoYyRw9AGRCUWG7qNyY-5QY71Go,3762
 gml/proto/src/controlplane/model/mpb/v1/mpb_pb2_grpc.py,sha256=KSdb6V04qUHDsb1R2o3wixwTyZgrhwnPYobjnRgWX4I,4735
-gml/tensor.py,sha256=753IsMFYZD7p_f0cQPt4nTIBo5p5S5ELqwCuoHORdMk,14823
-gimlet_api-0.0.7.dist-info/WHEEL,sha256=sobxWSyDDkdg_rinUth-jxhXHqoNqlmNMJY3aTZn2Us,91
-gimlet_api-0.0.7.dist-info/METADATA,sha256=IhN5ODabyw5hfMJDD-LCeGoiZSmqn1kWCKDNay8or5s,531
-gimlet_api-0.0.7.dist-info/RECORD,,
+gml/register_submodules.py,sha256=U8IwjVygX2vxNi_aK6ljHOD4mmrOhbyVczvy4wwulqU,5027
+gml/tensor.py,sha256=aPLm3I3qkYNDcJmntaUycqqN5rsZmcj8ql0EkupJudY,14977
+gimlet_api-0.0.9.dist-info/WHEEL,sha256=sobxWSyDDkdg_rinUth-jxhXHqoNqlmNMJY3aTZn2Us,91
+gimlet_api-0.0.9.dist-info/METADATA,sha256=P5wOKzPZyJroiZTPRpWsDdjiS5XQL21GK-heo5Set_E,531
+gimlet_api-0.0.9.dist-info/RECORD,,

gml/client.py

@@ -19,6 +19,8 @@ import uuid
 from pathlib import Path
 from typing import BinaryIO, List, Optional, TextIO, Union
 
+import grpc
+
 import gml.proto.src.api.corepb.v1.model_exec_pb2 as modelexecpb
 import gml.proto.src.common.typespb.uuid_pb2 as uuidpb
 import gml.proto.src.controlplane.compiler.cpb.v1.cpb_pb2 as cpb
@@ -31,7 +33,6 @@ import gml.proto.src.controlplane.logicalpipeline.lppb.v1.lppb_pb2 as lppb
 import gml.proto.src.controlplane.logicalpipeline.lppb.v1.lppb_pb2_grpc as lppb_grpc
 import gml.proto.src.controlplane.model.mpb.v1.mpb_pb2 as mpb
 import gml.proto.src.controlplane.model.mpb.v1.mpb_pb2_grpc as mpb_grpc
-import grpc
 from gml._utils import chunk_file, sha256sum
 from gml.device import DeviceCapabilities
 from gml.model import Model
@@ -252,7 +253,7 @@ class Client:
                 raise Exception("file status is deleted or unknown, cannot re-upload")
         return file_info
 
-    def _get_model_if_exists(self, name: str) -> Optional[cpb.Model]:
+    def _get_model_if_exists(self, name: str) -> Optional[modelexecpb.Model]:
         req = mpb.GetModelRequest(
             name=name,
             org_id=self._get_org_id(),
@@ -260,13 +261,13 @@ class Client:
         stub = self._ms_stub()
         try:
             resp = stub.GetModel(req, metadata=self._get_request_metadata())
-            return cpb.Model(id=resp.id, info=resp.model_info)
+            return modelexecpb.Model(id=resp.id, info=resp.model_info)
         except grpc.RpcError as e:
             if e.code() != grpc.StatusCode.NOT_FOUND:
                 raise e
             return None
 
-    def _create_model(self, model_info: modelexecpb.ModelInfo) -> cpb.Model:
+    def _create_model(self, model_info: modelexecpb.ModelInfo) -> modelexecpb.Model:
         req = mpb.CreateModelRequest(
             org_id=self._get_org_id(),
             name=model_info.name,
@@ -276,9 +277,9 @@ class Client:
         resp = stub.CreateModel(
             req, metadata=self._get_request_metadata(idempotent=True)
         )
-        return cpb.Model(id=resp.id, info=model_info)
+        return modelexecpb.Model(id=resp.id, info=model_info)
 
-    def create_model(self, model: Model) -> cpb.Model:
+    def create_model(self, model: Model) -> modelexecpb.Model:
         existing_model = self._get_model_if_exists(model.name)
         if existing_model is not None:
             print(
@@ -350,8 +351,8 @@ class Client:
         models: List[Model],
         pipeline_file: Optional[Path] = None,
         pipeline: Optional[Union[str, Pipeline]] = None,
-        runtimes=List[str],
-        cameras=List[str],
+        runtimes: List[str],
+        cameras: List[str],
     ):
         model_with_assets = []
         for model in models:

gml/compile.py

@@ -21,7 +21,6 @@ from typing import Any, Dict, List, Optional, Sequence, Union
 import safetensors_mlir
 import torch
 import torch_mlir
-from gml.asset_manager import AssetManager
 from mlir.ir import (
     BF16Type,
     ComplexType,
@@ -39,9 +38,12 @@ from torch._decomp import remove_decompositions
 from torch.export._trace import _export
 from torch_mlir.dialects import torch as torch_d
 from torch_mlir.extras.fx_decomp_util import get_decomposition_table
-from torch_mlir.extras.fx_importer import FxImporter, FxImporterHooks
+from torch_mlir.extras.fx_importer import FxImporter, FxImporterHooks, InputInfo
 from torch_mlir.fx import export_and_import
 
+from gml.asset_manager import AssetManager
+from gml.register_submodules import submodule_registration_workarounds
+
 
 def _default_decomposition_denylist():
     """These ops will not be decomposed by default."""
@@ -130,7 +132,10 @@ class SafetensorImporterHooks(FxImporterHooks):
         self._tensors = TensorSet()
 
     def resolve_literal(
-        self, gni: "torch_mlir.extras.fx_importer.GraphNodeImporter", literal: Any
+        self,
+        gni: "torch_mlir.extras.fx_importer.GraphNodeImporter",
+        literal: Any,
+        info: Optional[InputInfo],
     ) -> Optional[Value]:
         if not isinstance(literal, torch.Tensor):
             return None
@@ -166,92 +171,6 @@ class SafetensorImporterHooks(FxImporterHooks):
         save_file(tensors, file_path)
 
 
-def _to_module_list(val):
-    if isinstance(val, torch.nn.Module):
-        return val
-
-    converted = []
-    for item in val:
-        c = _to_module_container(item)
-        if c is None:
-            return None
-        converted.append(c)
-    if not converted:
-        return None
-    return torch.nn.ModuleList(converted)
-
-
-def _to_module_dict(val):
-    if isinstance(val, torch.nn.Module):
-        return val
-
-    converted = dict()
-    for k, v in val.items():
-        c = _to_module_container(v)
-        if c is None:
-            return None
-        converted[k] = v
-    if not converted:
-        return None
-    return torch.nn.ModuleDict(converted)
-
-
-def _to_module_container(val, root=False):
-    if isinstance(val, torch.nn.Module) and not root:
-        return val
-    if isinstance(val, dict):
-        return _to_module_dict(val)
-    if isinstance(val, list) or isinstance(val, tuple):
-        return _to_module_list(val)
-
-    return None
-
-
-def _replace_containers_with_torch_containers(mod: torch.nn.Module):
-    """Replaces any lists, dict, or nested combinations of lists/dicts that are attributes of `mod` with torch.nn.ModuleList/torch.nn.ModuleDict
-
-    This fixes some `module is not installed as a submodule` errors.
-    ."""
-    _excludes = set(["_modules"])
-    replacements = dict()
-    for name, val in mod.__dict__.items():
-        if name in _excludes:
-            continue
-        c = _to_module_container(val, root=True)
-        if c is None:
-            continue
-        replacements[name] = c
-
-    for name, repl in replacements.items():
-        setattr(mod, name, repl)
-
-
-def _ensure_submodules_accessed_through_getattr(mod: torch.nn.Module):
-    """This removes any registered modules from `mod.__dict__`.
-
-    This ensures that all accesses of submodules go through torch's __getattr__ infra,
-    preventing certain cases of `module is not installed as a submodule` errors.
-    """
-    if not hasattr(mod, "_modules"):
-        return
-    for name in mod._modules:
-        if name in mod.__dict__:
-            del mod.__dict__[name]
-
-
-def _submodule_registration_workarounds(mod: torch.nn.Module):
-    """Apply submodule registration workarounds recursively to all submodules of `mod`."""
-    _ensure_submodules_accessed_through_getattr(mod)
-    _replace_containers_with_torch_containers(mod)
-    # We intentionally don't use `mod.modules()` (which returns all recursive submodules) here because we want only
-    # the direct dependencies of `mod`. So that we get a pre-order traversal, ensuring the workarounds are applied
-    # before we check for submodules.
-    for submod in mod._modules.values():
-        if submod is mod:
-            continue
-        _submodule_registration_workarounds(submod)
-
-
 def to_torch_mlir(
     model: torch.nn.Module,
     example_inputs: Sequence[torch.Tensor],
@@ -266,16 +185,17 @@ def to_torch_mlir(
             if not isinstance(shape, dict):
                 continue
             for idx in shape:
-                if isinstance(shape[idx], torch.export.dynamic_shapes._Dim):
+                # Assign the value so that pyright understands the type.
+                value = shape[idx]
+                if isinstance(value, torch.export.dynamic_shapes._Dim):
                     continue
-                shape[idx] = torch.export.Dim(shape[idx])
-
+                shape[idx] = torch.export.Dim(value)
     if decomposition_denylist is None:
         decomposition_denylist = _default_decomposition_denylist()
 
     model = model.eval().to("cpu")
 
-    _submodule_registration_workarounds(model)
+    submodule_registration_workarounds(model)
 
     try:
         # Running the model a few times on the inputs, leads to more consistent compiled results.

gml/device.py

@@ -51,16 +51,14 @@ class DeviceCapabilities:
 
 def _runtime_str_to_runtime_protos(
     runtime: str,
-) -> deviceinfopb.ModelRuntimeInfo.ModelRuntimeType:
+) -> deviceinfopb.ModelRuntimeType:
     match runtime.lower():
         case "tensorrt":
-            return (
-                deviceinfopb.ModelRuntimeInfo.ModelRuntimeType.MODEL_RUNTIME_TYPE_TENSORRT
-            )
+            return deviceinfopb.ModelRuntimeType.MODEL_RUNTIME_TYPE_TENSORRT
         case "openvino":
-            return (
-                deviceinfopb.ModelRuntimeInfo.ModelRuntimeType.MODEL_RUNTIME_TYPE_OPENVINO
-            )
+            return deviceinfopb.ModelRuntimeType.MODEL_RUNTIME_TYPE_OPENVINO
+        case "hailort":
+            return deviceinfopb.ModelRuntimeType.MODEL_RUNTIME_TYPE_HAILORT
         case _:
             raise ValueError("invalid runtime: {}".format(runtime))