Update app.py

app.py

@@ -159,49 +159,63 @@ if __name__ == "__main__":
     interface.launch(share=True)
 """
 import gradio as gr
-import os
+import netron
+import threading
+import time
 import tempfile
-import subprocess
-
-
-# Function to generate a Netron static visualization
-def generate_netron_static(model_file):
-    if not os.path.exists(model_file):
-        return "<p>Error: Model file does not exist.</p>"
-
-    # Create a temporary directory to save the static Netron files
-    output_dir = tempfile.mkdtemp()
-    
-    # Run the Netron CLI to export the visualization to static files
-    subprocess.run(["netron", "--export", output_dir, model_file], check=True)
-    
-    # Create a static HTML viewer from the output directory
-    static_html_path = os.path.join(output_dir, "index.html")
-    if os.path.exists(static_html_path):
-        # Embed the static HTML in an iframe
-        return f"""
-        <iframe 
-            srcdoc="{open(static_html_path).read()}" 
-            width="100%" 
-            height="800" 
-            frameborder="0">
-        </iframe>
-        """
-    return "<p>Error: Unable to generate static visualization.</p>"
-
-
-# Gradio function to handle file upload
+import os
+
+# Function to start a Netron server and get the file-based URL
+def start_netron_server(file_path):
+    if not os.path.exists(file_path):
+        return None
+
+    # Temporary directory to store Netron visualizations
+    temp_dir = tempfile.mkdtemp()
+    output_path = os.path.join(temp_dir, "model_view.html")
+
+    # Run Netron in a separate thread with export option
+    def run_netron():
+        netron.start(file_path, browse=False, output=output_path)
+
+    thread = threading.Thread(target=run_netron, daemon=True)
+    thread.start()
+
+    # Wait for Netron to finish initializing
+    time.sleep(2)
+
+    # Ensure the exported file exists
+    if os.path.exists(output_path):
+        return f"file://{output_path}"
+    else:
+        return None
+
+
+# Gradio function to display the model viewer
 def visualize_model(file):
     if file:
-        return generate_netron_static(file.name)
+        file_path = file.name
+        netron_url = start_netron_server(file_path)
+        if netron_url:
+            # Embed the Netron viewer using an iframe with the generated URL
+            iframe_html = f"""
+            <iframe 
+                src="{netron_url}" 
+                width="100%" 
+                height="800" 
+                frameborder="0">
+            </iframe>
+            """
+            return iframe_html
+        return "<p>Error: Unable to generate Netron visualization.</p>"
     return "<p>Please upload a valid model file.</p>"
 
 
-# Gradio Interface
+# Gradio app
 with gr.Blocks() as demo:
-    gr.Markdown("## Netron Model Visualization with Gradio")
-    file_input = gr.File(label="Upload Model File (ONNX, TFLite, etc.)")
-    output_html = gr.HTML(label="Netron Viewer")
+    gr.Markdown("## Netron Model Viewer with Gradio")
+    file_input = gr.File(label="Upload a Model File (e.g., ONNX, TFLite, etc.)")
+    output_html = gr.HTML(label="Model Visualization")
     file_input.upload(visualize_model, file_input, output_html)
 
 # Launch the Gradio app