Update app.py

app.py

@@ -8,6 +8,7 @@ from analysis.spectre_utils import CrossDomainSamplingMetrics
 import networkx as nx
 import numpy as np
 import matplotlib.pyplot as plt
+import torch
 
 
 cfg = OmegaConf.load('./config.yaml')
@@ -22,9 +23,9 @@ input_dims, output_dims = compute_input_output_dims(data_loaders['train'], extra
 
 sampling_metrics = CrossDomainSamplingMetrics(data_loaders)
 
-model = LGGMText2Graph_Demo.load_from_checkpoint('last-v1.ckpt')
+model = LGGMText2Graph_Demo.load_from_checkpoint('last-v1.ckpt', map_location=torch.device("cpu"))
 
-model.init_prompt_encoder()
+model.init_prompt_encoder_pretrained()
 
 def calculate_average_degree(graph):
     num_nodes = graph.number_of_nodes()
@@ -34,7 +35,7 @@ def calculate_average_degree(graph):
 
 def predict(text, num_nodes = None):
     # Assuming model.generate and other processes are defined as before
-    graphs = model.generate(text, int(num_nodes))
+    graphs = model.generate_pretrained(text, int(num_nodes))
     ccs = []
     degs = []
     images = []
@@ -51,8 +52,11 @@ def predict(text, num_nodes = None):
         plt.close(fig)
         
         images.append(image)
+    
+    avg_deg = np.mean(degs)
+    avg_cc = np.mean(ccs)
 
-    return images[0], images[1], images[2], images[3], images[4], ccs[0], ccs[1], ccs[2], ccs[3], ccs[4], degs[0], degs[1], degs[2], degs[3], degs[4]
+    return images[0], images[1], images[2], images[3], images[4], ccs[0], ccs[1], ccs[2], ccs[3], ccs[4], degs[0], degs[1], degs[2], degs[3], degs[4], avg_cc, avg_deg
 
 def clear(input_text):
     return None, None
@@ -67,7 +71,7 @@ with gr.Blocks() as demo:
             input_num = gr.Slider(5, 200, value=10, label="Count", info="Number of nodes in the graph to be generated")
         with gr.Column():
             gr.Markdown("### Suggested Prompts")
-            gr.Markdown("1. Create a complex network with high clustering coefficient, exhibiting a very dense connection.\n2. Create a graph with extremely low number of triangles, which means it has very low clustering coefficient.")
+            gr.Markdown("1. Create a complex network with high clustering coefficient.\n2. Create a graph with extremely low number of triangles.\n 3. Please give me a Power Network with extremely low number of triangles but with medium level of average degree.")
 
     with gr.Row() as output_row:
         output_images = [gr.Image(label = f"Generated Network #{_}") for _ in range(5)]
@@ -75,15 +79,19 @@ with gr.Blocks() as demo:
         output_texts_cc = [gr.Textbox(label=f"CC #{_}") for _ in range(5)]
     with gr.Row():
         output_texts_deg = [gr.Textbox(label=f"DEG #{_}") for _ in range(5)]
+    
+    with gr.Row():
+        avg_cc_text = gr.Textbox(label="Average Clustering Coefficient")
+        avg_deg_text = gr.Textbox(label="Average Degree")
 
     with gr.Row():
         submit_button = gr.Button("Submit")
         clear_button = gr.Button("Clear")
 
     # Change function is linked to the submit button
-    submit_button.click(fn=predict, inputs=[input_text, input_num], outputs=output_images + output_texts_cc + output_texts_deg)
+    submit_button.click(fn=predict, inputs=[input_text, input_num], outputs=output_images + output_texts_cc + output_texts_deg + [avg_cc_text, avg_deg_text])
 
     # Clear function resets the text input and clears the outputs
-    clear_button.click(fn=clear, inputs=input_text, outputs=output_images + output_texts_cc + output_texts_deg)
+    clear_button.click(fn=clear, inputs=input_text, outputs=output_images + output_texts_cc + output_texts_deg + [avg_cc_text, avg_deg_text])
 
 demo.launch()