Update app.py

app.py

@@ -1,3 +1,4 @@
+"""
 import gradio as gr
 import os
 from PIL import Image
@@ -44,13 +45,13 @@ def process_image(sample_choice, uploaded_image, yolo_versions):
 
 # Custom CSS for styling (optional)
 custom_css = """
- run_button {
-    background-color: purple;
-    color: white;
-    width: 120px;
-    border-radius: 5px;
-    font-size: 14px;
-}
+ #run_button {
+  #  background-color: purple;
+   # color: white;
+   # width: 120px;
+    #border-radius: 5px;
+    #font-size: 14px;
+#}
 """
 
 # Gradio UI
@@ -113,3 +114,67 @@ with gr.Blocks(css=custom_css) as interface:
 # Launch Gradio app
 if __name__ == "__main__":
     interface.launch(share=True)
+"""
+
+import gradio as gr
+import netron
+import threading
+import time
+import tempfile
+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:
+        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="https://netron.app/?url=https://huggingface.co/FFusion/FFusionXL-BASE/blob/main/vae_encoder/model.onnx" 
+                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 app
+with gr.Blocks() as demo:
+    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
+demo.launch()