netron-ui / app.py
brian-xetdata's picture
Skipping api calls if the repo string is invalid
4196028
import gradio as gr
from huggingface_hub import model_info, hf_hub_url
from gradio_huggingfacehub_search import HuggingfaceHubSearch
from netron import build_netron_iframe, is_netron_file
current_repo = None
def repo_change(repo: str, oauth_token: str) -> dict[str, any] | None:
global current_repo
current_repo = repo
choices = []
if len(repo) == 0 or repo.find("/") == -1:
return gr.update(choices=choices, value=None, interactive=False)
try:
# get the files from the model repo
info = model_info(repo, token=oauth_token)
# filter only netron capable files
for sibling in info.siblings:
if is_netron_file(sibling.rfilename):
choices.append(sibling.rfilename)
return gr.update(choices=choices, value=None, interactive=True)
except Exception as e:
print("Exception", e)
return None
def viewer_change(file):
if current_repo is None:
return gr.update(value="<p>Repo not selected</p>")
file_url = hf_hub_url(current_repo, file)
if file_url is None:
return gr.update(value="<p>File not found</p>")
return gr.update(value=build_netron_iframe(file_url))
with gr.Blocks() as viewer_app:
with gr.Row():
gr.Markdown("# Netron Model Viewer")
gr.LoginButton(min_width=250)
with gr.Column():
with gr.Row():
model_in = HuggingfaceHubSearch(
label="Model ID",
placeholder="Search for models on Huggingface",
search_type="model",
)
files = gr.Dropdown(choices=[], interactive=False, label="File")
viewer = gr.HTML("""<p>Enter the model ID and select the compatible model file to view the model's graph view.</p>""")
gr.on([model_in.submit, model_in.change], fn=repo_change, inputs=[model_in], outputs=[files], trigger_mode="always_last")
files.select(fn=viewer_change, inputs=[files], outputs=[viewer])
if __name__ == "__main__":
viewer_app.launch(share=True)