(core) redistribute optional imports around modelExporters

PhasesResearchLab · Mar 8, 2024 · 8c863bf · 8c863bf
1 parent fe8c990
commit 8c863bf
Showing 1 changed file with 17 additions and 14 deletions.
diff --git a/pysipfenn/core/modelExporters.py b/pysipfenn/core/modelExporters.py
@@ -24,14 +24,7 @@ class ONNXExporter:
     """
 
     def __init__(self, calculator: Calculator):
-        """Initialize the ``ONNXExporter`` using a calculator object."""
-        try:
-            from onnxconverter_common import float16
-            from onnxsim import simplify
-        except ModuleNotFoundError as e:
-            print('\n\nNote: Export Dependencies are not installed by default. If you need them, you have to install pySIPFENN in '
-                  '`dev` mode like: `pip install "pysipfenn[dev]"`, or like `pip install -e ".[dev]"` (see pysipfenn.org)')
-
+        """Initialize the ``ONNXExporter`` using a calculator object."""    
         self.simplifiedDict = {model: False for model in calculator.loadedModels.keys()}
         self.fp16Dict = {model: False for model in calculator.loadedModels.keys()}
         self.calculator = calculator
@@ -86,6 +79,11 @@ def simplify(self, model: str) -> None:
         Returns:
             None
         """
+        try:
+            from onnxsim import simplify
+        except ModuleNotFoundError as e:
+            print('\n\nNote: Export Dependencies are not installed by default. If you need them, you have to install pySIPFENN in '
+                  '`dev` mode like: `pip install "pysipfenn[dev]"`, or like `pip install -e ".[dev]"` (see pysipfenn.org)')
         print(f'Simplifying {model}')
         assert model in self.calculator.loadedModels, f'{model} not loaded in calculator. Nothing to simplify.'
         loadedModel = self.calculator.loadedModels[model]
@@ -110,6 +108,11 @@ def toFP16(self, model: str) -> None:
         Returns:
             None
         """
+        try:
+            from onnxconverter_common import float16
+        except ModuleNotFoundError as e:
+            print('\n\nNote: Export Dependencies are not installed by default. If you need them, you have to install pySIPFENN in '
+                  '`dev` mode like: `pip install "pysipfenn[dev]"`, or like `pip install -e ".[dev]"` (see pysipfenn.org)')
         print(f'Converting {model} to FP16')
         assert model in self.calculator.loadedModels, f'{model} not loaded in calculator. Nothing to convert to FP16.'
         loadedModel = self.calculator.loadedModels[model]
@@ -239,12 +242,6 @@ class CoreMLExporter:
     """
 
     def __init__(self, calculator: Calculator):
-        try:
-            import coremltools as ct
-        except ModuleNotFoundError as e:
-            print('\n\nNote: Export Dependencies are not installed by default. If you need them, you have to install pySIPFENN in '
-                  '`dev` mode like: `pip install "pysipfenn[dev]"`, or like `pip install -e ".[dev]"` (see pysipfenn.org)')
-
         self.calculator = calculator
         assert len(self.calculator.loadedModels)>0, 'No models loaded in calculator. Nothing to export.'
         print(f'Initialized CoreMLExporter with models: {list(self.calculator.loadedModels.keys())}')
@@ -265,6 +262,12 @@ def export(self, model: str, append: str = '') -> None:
         Returns:
             None
         """
+        try:
+            import coremltools as ct
+        except ModuleNotFoundError as e:
+            print('\n\nNote: Export Dependencies are not installed by default. If you need them, you have to install pySIPFENN in '
+                  '`dev` mode like: `pip install "pysipfenn[dev]"`, or like `pip install -e ".[dev]"` (see pysipfenn.org)')
+
         print(f'Exporting {model} to CoreML')
         assert model in self.calculator.loadedModels, f'{model} not loaded in calculator. Nothing to export.'
         loadedModel = self.calculator.loadedModels[model]