|
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: |
|
|
|
info = model_info(repo, token=oauth_token) |
|
|
|
|
|
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) |