Update app.py

app.py

@@ -4,7 +4,11 @@ import os
 import matplotlib.pyplot as plt # Kept for potential future gr.Plot use or if figures are handled directly
 from simulation import RFTSimulation
 from utils import export_json, seed_everything
-from IPython.display import Image, display # For Colab dev mode inline image display
+
+# Conditionally import IPython.display only if in Colab
+IN_COLAB = 'COLAB_GPU' in os.environ or 'Google' in os.environ.get('COLAB_CLOUD_PLATFORM', '')
+if IN_COLAB:
+    from IPython.display import Image, display # For Colab dev mode inline image display
 
 # Helper function to run simulation and generate plots/export data
 def run_simulation_and_plot(num_agents, steps, mutation_rate, drift_rate, seed=None):
@@ -67,8 +71,7 @@ iface = gr.Interface(
 )
 
 # Colab-friendly "dev mode" or Gradio launch logic
-# Check if running in Colab environment (simple heuristic)
-IN_COLAB = 'COLAB_GPU' in os.environ or 'Google' in os.environ.get('COLAB_CLOUD_PLATFORM', '')
+# The IN_COLAB check is now done at the top for conditional import
 
 if IN_COLAB:
     print("Running in Colab development mode (skipping Gradio launch).")