Update app.py

app.py

@@ -159,59 +159,50 @@ if __name__ == "__main__":
     interface.launch(share=True)
 """
 import gradio as gr
-import netron
 import os
-import threading
-import time
+import tempfile
+import subprocess
 
-# Function to start Netron server and get its port
-def start_netron_server(model_file):
+
+# Function to generate a Netron static visualization
+def generate_netron_static(model_file):
     if not os.path.exists(model_file):
-        return "Error: Model file does not exist."
-    
-    # Start Netron in a thread
-    def run_netron():
-        netron.start(model_file, browse=False)  # Prevent opening in a local browser
-    
-    thread = threading.Thread(target=run_netron, daemon=True)
-    thread.start()
+        return "<p>Error: Model file does not exist.</p>"
+
+    # Create a temporary directory to save the static Netron files
+    output_dir = tempfile.mkdtemp()
     
-    # Wait for the server to start
-    time.sleep(2)
+    # Run the Netron CLI to export the visualization to static files
+    subprocess.run(["netron", "--export", output_dir, model_file], check=True)
     
-    # Netron typically uses port 8080 by default, but you may need to confirm this dynamically
-    port = 8080
-    return port
-
-# Gradio function to visualize the model
-def visualize_model(file):
-    if file:
-        file_path = file.name
-        port = start_netron_server(file_path)
-        
-        # Construct the iframe URL using Hugging Face's public domain
-        space_url = os.getenv("SPACE_ID", "your-space-id")  # Replace 'your-space-id' with actual ID if testing locally
-        iframe_url = f"https://bhumikamak-neuralvista.hf.space:{port}/"
-        
-        # Embed the iframe into the Gradio app
-        iframe_html = f"""
+    # 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 
-            src="{iframe_url}" 
+            srcdoc="{open(static_html_path).read()}" 
             width="100%" 
             height="800" 
             frameborder="0">
         </iframe>
         """
-        return iframe_html
-    
+    return "<p>Error: Unable to generate static visualization.</p>"
+
+
+# Gradio function to handle file upload
+def visualize_model(file):
+    if file:
+        return generate_netron_static(file.name)
     return "<p>Please upload a valid model file.</p>"
 
+
 # Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("## Netron Model Visualization on Hugging Face Spaces")
+    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")
     file_input.upload(visualize_model, file_input, output_html)
 
 # Launch the Gradio app
-demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.launch()