Spaces:
Sleeping
Sleeping
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) |