Add data collection

app.py

@@ -1,5 +1,6 @@
 import json
 import os
+import shutil
 import subprocess
 import time
 import uuid
@@ -45,6 +46,10 @@ def main():
 
     # Define inputs and outputs
     in_prompt = gr.Textbox(label='Prompt', placeholder='Enter a prompt to generate a LEGO model.', max_length=500)
+    in_optout = gr.Checkbox(label='Do not save my data',
+                            info='We may collect model inputs and outputs to help us improve the model. '
+                                 'Your data will never be shared or used for any other purpose. '
+                                 'If you wish to opt out of data collection, please check the box below.')
     in_temperature = gr.Slider(0.01, 2.0, value=model_cfg.temperature, step=0.01,
                                label='Temperature', info=get_help_string('temperature'))
     in_seed = gr.Number(value=42, label='Seed', info='Random seed for generation.',
@@ -68,7 +73,7 @@ def main():
                     'The model is restricted to creating structures made of 1-unit-tall cuboid bricks on a 20x20x20 grid. It was trained on a dataset of 21 object categories: '
                     '*basket, bed, bench, birdhouse, bookshelf, bottle, bowl, bus, camera, car, chair, guitar, jar, mug, piano, pot, sofa, table, tower, train, vessel.* '
                     'Performance on prompts from outside these categories may be limited. This demo does not include texturing or coloring.',
-        inputs=[in_prompt],
+        inputs=[in_prompt, in_optout],
         additional_inputs=[in_temperature, in_seed, in_bricks, in_rejections, in_regenerations],
         outputs=[out_img, out_txt],
         flagging_mode='never',
@@ -100,12 +105,13 @@ class LegoGenerator:
     def generate_lego(
             self,
             prompt: str,
+            do_not_save_data: bool,
             temperature: float | None,
             seed: int | None,
             max_bricks: int | None,
             max_brick_rejections: int | None,
             max_regenerations: int | None,
-    ):
+    ) -> (str, str):
         # Set model parameters
         if temperature is not None: self.model.temperature = temperature
         if max_bricks is not None: self.model.max_bricks = max_bricks
@@ -125,19 +131,46 @@ class LegoGenerator:
         ldr_filename = os.path.join(output_dir, f'{output_uuid}.ldr')
         with open(ldr_filename, 'w') as f:
             f.write(output['lego'].to_ldr())
-        print(f'Finished generation in {time.time() - start_time:.1f}s!')
+        generation_time = time.time() - start_time
+        print(f'Finished generation in {generation_time:.1f}s!')
 
         # Render LEGO model to image
         print('Rendering image...')
-        start_time = time.time()
         img_filename = os.path.join(output_dir, f'{output_uuid}.png')
         subprocess.run(['python', 'render_lego.py', '--in_file', ldr_filename, '--out_file', img_filename],
                        check=True)  # Run render as a subprocess to prevent issues with Blender
-        print(f'Finished rendering in {time.time() - start_time:.1f}s!')
-
-        return img_filename, output['lego']
-
-    def generate_lego_subprocess(self, *args):
+        rendering_time = time.time() - start_time - generation_time
+        print(f'Finished rendering in {rendering_time:.1f}s!')
+
+        # Save data persistently
+        if not do_not_save_data:
+            data_dir = '/data/apun/legogpt_demo_out'
+            os.makedirs(data_dir, exist_ok=True)
+
+            # Copy output image to persistent storage
+            img_copy_filename = os.path.join(data_dir, f'{output_uuid}.png')
+            shutil.copy(img_filename, img_copy_filename)
+
+            # Save metadata
+            metadata_filename = os.path.join(data_dir, f'{output_uuid}.json')
+            with open(metadata_filename, 'w') as f:
+                json.dump({
+                    'prompt': prompt,
+                    'temperature': self.model.temperature,
+                    'seed': seed,
+                    'max_bricks': self.model.max_bricks,
+                    'max_brick_rejections': self.model.max_brick_rejections,
+                    'max_regenerations': self.model.max_regenerations,
+                    'start_timestamp': time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time)),
+                    'generation_time': generation_time,
+                    'rendering_time': rendering_time,
+                    'output_txt': output['lego'].to_txt(),
+                }, f)
+            print(f'Saved data to {metadata_filename}.')
+
+        return img_filename, output['lego'].to_txt()
+
+    def generate_lego_subprocess(self, *args) -> (str, str):
         """
         Run generation as a subprocess so that multiple requests can be handled concurrently.
         """