Add application file5

app.py

@@ -31,6 +31,8 @@ ROOM_CATEGORIES = {
     'Storage': 10, 'Front Door': 11, 'Other': 13
 }
 
+rooms_data = []
+edges_data = []
 
 def save_samples(
         sample, ext, model_kwargs,
@@ -346,8 +348,7 @@ def create_layout(graphs, corners, room_type, metrics=False, use_ddim=True, ddim
     print(graphs, corners, room_type)
     model_kwargs = function_test(graphs, corners, room_type)
     for key in model_kwargs:
-        # model_kwargs[key] = th.from_numpy(np.array([model_kwargs[key]])).cuda()
-        model_kwargs[key] = th.from_numpy(np.array([model_kwargs[key]])).cpu()
+        model_kwargs[key] = th.from_numpy(np.array([model_kwargs[key]])).cuda()
 
     png_paths = []
     svg_paths = []
@@ -388,11 +389,8 @@ def create_layout(graphs, corners, room_type, metrics=False, use_ddim=True, ddim
     return png_paths, svg_paths
 
 
-rooms_data = []
-edges_data = []
-
-
 def generate_layout(metrics: bool, ddim_steps: int, num_samples: int):
+
     room_list = []
     room_corners = []
     living_room = 0
@@ -447,7 +445,13 @@ def generate_layout(metrics: bool, ddim_steps: int, num_samples: int):
 
     png_color_guide = './color_guide.png'
 
-    return png_paths, svg_paths, png_color_guide
+    rooms_data.clear()
+    edges_data.clear()
+
+    rooms_df = pd.DataFrame(columns=["room_id", "room_type", "num_corners"])
+    edges_df = pd.DataFrame(columns=["edge_id", "room1_id", "room2_id"])
+
+    return png_paths, svg_paths, png_color_guide, rooms_df, edges_df
 
 
 with gr.Blocks() as demo:
@@ -503,8 +507,18 @@ with gr.Blocks() as demo:
 
 
     generate_button = gr.Button("Generate Layout")
-    generate_button.click(generate_layout, inputs=[metrics_toggle, ddim_input, num_sample], outputs=[png_gallery, svg_files, png_color_guide])
-
+    # generate_button.click(generate_layout, inputs=[metrics_toggle, ddim_input, num_sample], outputs=[png_gallery, svg_files, png_color_guide])
+    generate_button.click(
+        generate_layout,
+        inputs=[metrics_toggle, ddim_input, num_sample],
+        outputs=[
+            png_gallery,
+            svg_files,
+            png_color_guide,
+            rooms_table,
+            edges_table
+        ]
+    )
     add_room_button.click(add_room, inputs=[room_type, num_corners],
                           outputs=[rooms_table, edges_table, room1_id, room2_id])
     add_edge_button.click(add_edge, inputs=[room1_id, room2_id], outputs=[rooms_table, edges_table, room1_id, room2_id])