From cca4325afe5045e6cf748daf09ea09c9d512bd8e Mon Sep 17 00:00:00 2001
From: Mergen Nachin <mnachin@meta.com>
Date: Wed, 15 Jan 2025 13:45:39 -0500
Subject: [PATCH] Move ai-edge-model-explorer into
 devtools/install_requirements.sh

Summary:

From top-level requirements, move to devtools/install_requirement.sh instead.

ai-edge-model-explorer is too new and currently requires numpy<2 which conflicts with our recent upgrade.

Let's not take core dependency on this tool yet.
---
 .lintrunner.toml                                   |  1 +
 devtools/install_requirements.sh                   | 11 +++++++++++
 devtools/visualization/visualization_utils.py      |  9 ++++++++-
 devtools/visualization/visualization_utils_test.py |  9 ++++++++-
 install_requirements.py                            |  1 -
 pytest.ini                                         |  1 +
 6 files changed, 29 insertions(+), 3 deletions(-)
 create mode 100755 devtools/install_requirements.sh

diff --git a/.lintrunner.toml b/.lintrunner.toml
index fe8ecad1fc..dd75ea8f32 100644
--- a/.lintrunner.toml
+++ b/.lintrunner.toml
@@ -294,6 +294,7 @@ include_patterns = [
     'build/**/*.py',
     'codegen/**/*.py',
     # 'devtools/**/*.py',
+    'devtools/visualization/**/*.py',
     'docs/**/*.py',
     # 'examples/**/*.py',
     # 'exir/**/*.py',
diff --git a/devtools/install_requirements.sh b/devtools/install_requirements.sh
new file mode 100755
index 0000000000..242bc70257
--- /dev/null
+++ b/devtools/install_requirements.sh
@@ -0,0 +1,11 @@
+#!/bin/bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Conflict: this requires numpy<2 whereas ExecuTorch core requires numpy>=2
+# Follow https://github.com/google-ai-edge/model-explorer/issues/277 for potential
+# resolution
+pip install ai-edge-model-explorer>=0.1.16
diff --git a/devtools/visualization/visualization_utils.py b/devtools/visualization/visualization_utils.py
index a2ee4c6050..4d520a6636 100644
--- a/devtools/visualization/visualization_utils.py
+++ b/devtools/visualization/visualization_utils.py
@@ -8,9 +8,16 @@
 import time
 
 from executorch.exir import EdgeProgramManager, ExecutorchProgramManager
-from model_explorer import config, consts, visualize_from_config  # type: ignore
 from torch.export.exported_program import ExportedProgram
 
+try:
+    from model_explorer import config, consts, visualize_from_config  # type: ignore
+except ImportError:
+    print(
+        "Error: 'model_explorer' is not installed. Install using devtools/install_requirement.sh"
+    )
+    raise
+
 
 class SingletonModelExplorerServer:
     """Singleton context manager for starting a model-explorer server.
diff --git a/devtools/visualization/visualization_utils_test.py b/devtools/visualization/visualization_utils_test.py
index 89781ab4f4..dafefa7dfd 100644
--- a/devtools/visualization/visualization_utils_test.py
+++ b/devtools/visualization/visualization_utils_test.py
@@ -17,7 +17,14 @@
     visualize,
 )
 from executorch.exir import ExportedProgram
-from model_explorer.config import ModelExplorerConfig  # type: ignore
+
+try:
+    from model_explorer.config import ModelExplorerConfig  # type: ignore
+except ImportError:
+    print(
+        "Error: 'model_explorer' is not installed. Install using devtools/install_requirement.sh"
+    )
+    raise
 
 
 @pytest.fixture
diff --git a/install_requirements.py b/install_requirements.py
index 26093cab84..adb26170cd 100644
--- a/install_requirements.py
+++ b/install_requirements.py
@@ -170,7 +170,6 @@ def python_is_compatible():
     "tomli",  # Imported by extract_sources.py when using python < 3.11.
     "wheel",  # For building the pip package archive.
     "zstd",  # Imported by resolve_buck.py.
-    "ai-edge-model-explorer>=0.1.16",  # For visualizing ExportedPrograms
 ]
 
 # Assemble the list of requirements to actually install.
diff --git a/pytest.ini b/pytest.ini
index b7e9afb9b9..d0c27fdfab 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -14,6 +14,7 @@ addopts =
     # explicitly list out tests that are running successfully in oss
     examples/models/test
     devtools/
+    --ignore=devtools/visualization/visualization_utils_test.py
     # examples
     examples/models/llama/tests
     examples/models/llama3_2_vision/preprocess