wip

supervisely-ecosystem · Sep 24, 2024 · 7aa5f45 · 7aa5f45
1 parent 7488161
commit 7aa5f45
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 1 deletion.
diff --git a/src/ui/dtl/actions/neural_networks/deploy/deploy.py b/src/ui/dtl/actions/neural_networks/deploy/deploy.py
@@ -126,6 +126,8 @@ def agent_selector_sidebar_save_btn_cb():
             model_selector_sidebar_custom_model_table,
             # public options
             model_selector_sidebar_public_model_table,
+            model_selector_runtime_selector_sidebar,
+            model_selector_sidebar_public_container,
             # sidebar
             model_selector_sidebar_model_source_tabs,
             model_selector_sidebar_save_btn,
@@ -150,6 +152,7 @@ def model_selector_sidebar_save_btn_cb():
                 saved_settings,
                 model_selector_sidebar_model_source_tabs,
                 model_selector_sidebar_public_model_table,
+                model_selector_runtime_selector_sidebar,
                 model_selector_sidebar_custom_model_table,
                 model_selector_stop_model_after_pipeline_checkbox,
             )
@@ -260,6 +263,7 @@ def _save_settings() -> dict:
                 agent_selector_sidebar_device_selector,
                 model_selector_sidebar_model_source_tabs,
                 model_selector_sidebar_public_model_table,
+                model_selector_runtime_selector_sidebar,
                 model_selector_sidebar_custom_model_table,
                 model_selector_stop_model_after_pipeline_checkbox,
             )

diff --git a/src/ui/dtl/actions/neural_networks/deploy/layout/model_selector.py b/src/ui/dtl/actions/neural_networks/deploy/layout/model_selector.py
@@ -6,6 +6,7 @@
     CustomModelsSelector,
     PretrainedModelsSelector,
     Checkbox,
+    Select,
 )
 
 
@@ -46,6 +47,13 @@ def create_model_selector_widgets(
     )
     if "object detection" in pretrained_model_selector_task_types:
         model_selector_sidebar_public_model_table.set_active_task_type("object detection")
+    available_runtimes = ["PyTorch", "ONNXRuntime", "TensorRT"]
+    model_selector_runtime_selector_sidebar = Select(
+        [Select.Item(runtime, runtime) for runtime in available_runtimes]
+    )
+    model_selector_sidebar_public_container = Container(
+        model_selector_sidebar_public_model_table, model_selector_runtime_selector_sidebar
+    )
     # ------------------------------
 
     # CUSTOM /PUBLIC TABS
@@ -54,7 +62,7 @@ def create_model_selector_widgets(
         descriptions=["Models trained by you", f"Models trained by {framework_name} team"],
         contents=[
             model_selector_sidebar_custom_model_table,
-            model_selector_sidebar_public_model_table,
+            model_selector_sidebar_public_container,
         ],
     )
     if len(custom_models) == 0:
@@ -107,6 +115,8 @@ def create_model_selector_widgets(
         model_selector_sidebar_custom_model_table,
         # public options
         model_selector_sidebar_public_model_table,
+        model_selector_runtime_selector_sidebar,
+        model_selector_sidebar_public_container,
         # sidebar
         model_selector_sidebar_model_source_tabs,
         model_selector_sidebar_save_btn,

diff --git a/src/ui/dtl/actions/neural_networks/deploy/layout/utils.py b/src/ui/dtl/actions/neural_networks/deploy/layout/utils.py
@@ -64,6 +64,7 @@ def save_model_settings(
     settings: dict,
     model_selector_sidebar_model_source_tabs: RadioGroup,
     model_selector_sidebar_public_model_table: PretrainedModelsSelector,
+    model_selector_runtime_selector_sidebar: Select,
     model_selector_sidebar_custom_model_table: CustomModelsSelector,
     model_selector_stop_model_after_pipeline_checkbox: Checkbox,
 ):
@@ -72,6 +73,7 @@ def save_model_settings(
     if model_source == "Pretrained public models":
         model_source = "Pretrained models"
         model_params = model_selector_sidebar_public_model_table.get_selected_model_params()
+        model_params["runtime"] = model_selector_runtime_selector_sidebar.get_value()
 
     elif model_source == "Custom models":
         model_source = "Custom models"
@@ -93,6 +95,9 @@ def save_model_settings(
     arch_type = model_params.get("arch_type", None)
     if arch_type is not None:
         settings["arch_type"] = arch_type
+    runtime = model_params.get("runtime", None)
+    if runtime is not None:
+        settings["runtime"] = runtime
 
     return settings
 
@@ -108,6 +113,7 @@ def save_settings(
     agent_selector_sidebar_device_selector: Select,
     model_selector_sidebar_model_source_tabs: RadioTabs,
     model_selector_sidebar_public_model_table: PretrainedModelsSelector,
+    model_selector_runtime_selector_sidebar: Select,
     model_selector_sidebar_custom_model_table: CustomModelsSelector,
     model_selector_stop_model_after_pipeline_checkbox: Checkbox,
 ):
@@ -118,6 +124,7 @@ def save_settings(
         settings,
         model_selector_sidebar_model_source_tabs,
         model_selector_sidebar_public_model_table,
+        model_selector_runtime_selector_sidebar,
         model_selector_sidebar_custom_model_table,
         model_selector_stop_model_after_pipeline_checkbox,
     )
@@ -195,6 +202,9 @@ def deploy_model(api: Api, session_id: int, saved_settings: dict):
     arch_type = saved_settings.get("arch_type", None)
     if arch_type is not None:
         deploy_params["arch_type"] = arch_type
+    runtime = saved_settings.get("runtime", None)
+    if runtime is not None:
+        deploy_params["runtime"] = runtime
 
     api.task.send_request(session_id, "deploy_from_api", data={"deploy_params": {**deploy_params}})