File size: 1,728 Bytes
8b4f646 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import gradio as gr
from huggingface_hub import model_info, hf_hub_url
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 = []
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():
repo = gr.Textbox(placeholder="owner/model name (Sign in above to view private models)", label="Model Repository")
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>""")
repo.change(fn=repo_change, inputs=[repo], outputs=[files], trigger_mode='always_last')
files.select(fn=viewer_change, inputs=[files], outputs=[viewer])
if __name__ == "__main__":
viewer_app.launch(share=True) |