Update app.py

app.py

@@ -158,38 +158,52 @@ with gr.Blocks(css=custom_css) as interface:
 if __name__ == "__main__":
     interface.launch(share=True)
 """
-
 import gradio as gr
-import onnxruntime
+import netron
 import os
-
-def visualize_onnx_model(onnx_model_path):
-  """
-  Visualizes the given ONNX model using Netron.
-
-  Args:
-    onnx_model_path: The path to the ONNX model file.
-  """
-  try:
-    # Save the ONNX model to a temporary file (optional, but can be helpful)
-    temp_model_path = "temp_model.onnx"
-    with open(temp_model_path, "wb") as f:
-      f.write(onnx_model_path.read())  # Assuming onnx_model_path is a file-like object
-
-    # Run Netron 
-    os.system(f"netron {temp_model_path}") 
-
-  except Exception as e:
-    print(f"Error visualizing model: {e}")
-
-# Create the Gradio interface
-iface = gr.Interface(
-    fn=visualize_onnx_model,
-    inputs="file",  # Accept the ONNX model as a file upload
-    outputs="text", 
-    title="Netron ONNX Model Visualization",
-    description="Upload an ONNX model file to visualize with Netron."
-)
+import threading
+import time
+
+# Function to start Netron server
+def start_netron_server(model_file, port=8080):
+    if not os.path.exists(model_file):
+        return "Error: Model file does not exist."
+    
+    # Start Netron in a thread
+    thread = threading.Thread(target=netron.start, args=(model_file,), kwargs={"port": port, "host": "0.0.0.0"}, daemon=True)
+    thread.start()
+    
+    # Wait for the server to initialize
+    time.sleep(2)
+
+# Function to return the Netron viewer iframe
+def visualize_model(file):
+    if file:
+        file_path = file.name
+        port = 8080  # Port for Netron server
+        start_netron_server(file_path, port)
+        
+        iframe_url = f"https://bhumikamak-neuralvista.hf.space:{port}/"
+        
+        # Embed the Netron viewer using iframe
+        iframe_html = f"""
+        <iframe 
+            src="{iframe_url}" 
+            width="100%" 
+            height="800" 
+            frameborder="0">
+        </iframe>
+        """
+        return iframe_html
+    
+    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")
+    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
-iface.launch(share=True) 
+demo.launch(server_name="0.0.0.0", server_port=7860)