Skip to content

Commit 4d6209b

Browse files
authored
Expose Model Explorer's visualize_from_config (#14005)
This PR exposes the Model Explorer visualisation API, which provides the foundation to enable us to load custom file types and plugins.
1 parent d4b58c9 commit 4d6209b

File tree

2 files changed

+13
-3
lines changed

2 files changed

+13
-3
lines changed

devtools/visualization/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,4 +9,5 @@
99
SingletonModelExplorerServer,
1010
visualize,
1111
visualize_graph,
12+
visualize_model_explorer,
1213
)

devtools/visualization/visualization_utils.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ def visualize(
108108
**kwargs,
109109
):
110110
"""Wraps the visualize_from_config call from model_explorer.
111-
For convenicence, figures out how to find the exported_program
111+
For convenience, figures out how to find the exported_program
112112
from EdgeProgramManager and ExecutorchProgramManager for you.
113113
114114
See https://github.com/google-ai-edge/model-explorer/wiki/4.-API-Guide#visualize-pytorch-models
@@ -123,13 +123,22 @@ def visualize(
123123
)
124124
if reuse_server:
125125
cur_config.set_reuse_server()
126-
visualize_from_config(
127-
cur_config,
126+
visualize_model_explorer(
127+
config=kwargs.pop("config", cur_config),
128128
no_open_in_browser=no_open_in_browser,
129129
**kwargs,
130130
)
131131

132132

133+
def visualize_model_explorer(
134+
**kwargs,
135+
):
136+
"""Wraps the visualize_from_config call from model_explorer."""
137+
visualize_from_config(
138+
**kwargs,
139+
)
140+
141+
133142
def visualize_graph(
134143
graph_module: GraphModule,
135144
exported_program: ExportedProgram | EdgeProgramManager | ExecutorchProgramManager,

0 commit comments

Comments
 (0)