Update app.py

app.py

@@ -62,8 +62,12 @@ def parse_input(num_nodes, edges_str, node_features_str):
 # ---------- VISUALIZACIÓN ----------
 
 def draw_graph(G, pred_label):
+    import matplotlib.pyplot as plt
+    import io
+
     pos = nx.spring_layout(G)
-    node_colors = 'lightgreen' if pred_label == 1 else 'lightcoral'
+    node_colors = ['lightgreen' if pred_label == 1 else 'lightcoral'] * G.number_of_nodes()
+
     plt.figure(figsize=(4, 4))
     nx.draw(G, pos, with_labels=True, node_color=node_colors, edge_color='gray', node_size=800)
     plt.title("Grafo de entrada")
@@ -71,10 +75,9 @@ def draw_graph(G, pred_label):
     buf = io.BytesIO()
     plt.savefig(buf, format='png')
     buf.seek(0)
-    img_base64 = base64.b64encode(buf.read()).decode('utf-8')
     plt.close()
 
-    return f"data:image/png;base64,{img_base64}"
+    return buf  # Devuelve un objeto tipo archivo
 
 # ---------- FUNCIÓN DE PREDICCIÓN ----------
 
@@ -92,8 +95,7 @@ def predict_graph(num_nodes, edges_str, node_features_str):
         pred = out.argmax(dim=1).item()
 
     label_text = "Mutagénico ✅" if pred == 1 else "No mutagénico ❌"
-    img = draw_graph(G, pred)
-    return label_text, img
+    return label_text, draw_graph(G, pred)
 
 # ---------- INTERFAZ GRADIO ----------