File size: 2,039 Bytes
8b4f646
 
ae727e3
8b4f646
 
 
 
 
 
 
 
 
 
4196028
 
 
8b4f646
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ae727e3
 
 
 
 
 
8b4f646
 
 
 
ae727e3
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
53
54
55
56
57
58
59
60
61
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)