Merge pull request #232 from kartynnik/main

Make the dependency on PyTorch optional

src/server/package/pyproject.toml

@@ -23,7 +23,6 @@ dependencies = [
   "requests",
   "termcolor",
   "typing-extensions",
-  "torch >= 2.2",
   "numpy < 2",
 ]
 

src/server/package/src/model_explorer/apis.py

@@ -15,7 +15,6 @@
 
 from typing import TypedDict, Union
 
-import torch
 from typing_extensions import NotRequired
 
 from . import server
@@ -27,6 +26,11 @@
     DEFAULT_SETTINGS,
 )
 
+try:
+  import torch
+except ImportError:
+  torch = None
+
 NodeDataInfo = TypedDict(
     'NodeDataInfo',
     {
@@ -107,7 +111,7 @@ def visualize(
 
 def visualize_pytorch(
     name: str,
-    exported_program: torch.export.ExportedProgram,
+    exported_program: 'torch.export.ExportedProgram',
     host=DEFAULT_HOST,
     port=DEFAULT_PORT,
     extensions: list[str] = [],

src/server/package/src/model_explorer/builtin_pytorch_exportedprogram_adapter.py

@@ -15,12 +15,16 @@
 
 from typing import Dict
 
-import torch
-import torch.fx
+try:
+  import torch
+except ImportError:
+  torch = None
 
 from .adapter import Adapter, AdapterMetadata
 from .types import ModelExplorerGraphs
-from .pytorch_exported_program_adater_impl import PytorchExportedProgramAdapterImpl
+
+if torch is not None:
+  from .pytorch_exported_program_adater_impl import PytorchExportedProgramAdapterImpl
 
 
 class BuiltinPytorchExportedProgramAdapter(Adapter):
@@ -40,5 +44,10 @@ def __init__(self):
     super().__init__()
 
   def convert(self, model_path: str, settings: Dict) -> ModelExplorerGraphs:
+    if torch is None:
+      raise ImportError(
+          'Please install the `torch` package, e.g. via `pip install torch`, '
+          'and restart the Model Explorer server.'
+      )
     ep = torch.export.load(model_path)
     return PytorchExportedProgramAdapterImpl(ep, settings).convert()

src/server/package/src/model_explorer/config.py

@@ -19,14 +19,20 @@
 from urllib.parse import quote
 
 import requests
-import torch
 from typing_extensions import NotRequired
 
+try:
+  import torch
+except ImportError:
+  torch = None
+
 from .consts import DEFAULT_HOST, DEFAULT_SETTINGS
 from .node_data_builder import NodeData
-from .pytorch_exported_program_adater_impl import PytorchExportedProgramAdapterImpl
 from .types import ModelExplorerGraphs
 
+if torch is not None:
+  from .pytorch_exported_program_adater_impl import PytorchExportedProgramAdapterImpl
+
 ModelSource = TypedDict(
     'ModelSource', {'url': str, 'adapterId': NotRequired[str]}
 )
@@ -79,7 +85,7 @@ def add_model_from_path(
   def add_model_from_pytorch(
       self,
       name: str,
-      exported_program: torch.export.ExportedProgram,
+      exported_program: 'torch.export.ExportedProgram',
       settings=DEFAULT_SETTINGS,
   ) -> 'ModelExplorerConfig':
     """Adds the given pytorch model.