Add multiprocessing

app.py

@@ -8,6 +8,7 @@ from dataclasses import fields
 from urllib.request import urlretrieve
 
 import gradio as gr
+import torch.multiprocessing as mp
 import transformers
 from legogpt.models import LegoGPT, LegoGPTConfig
 
@@ -39,47 +40,8 @@ def main():
         print('Running in Hugging Face Space, setting up environment...')
         setup()
 
-    model_cfg = LegoGPTConfig(max_regenerations=10)
-    model = LegoGPT(model_cfg)
-
-    def generate_lego(
-            prompt: str,
-            temperature: float | None,
-            seed: int | None,
-            max_bricks: int | None,
-            max_brick_rejections: int | None,
-            max_regenerations: int | None,
-    ):
-        # Set model parameters
-        if temperature is not None: model.temperature = temperature
-        if max_bricks is not None: model.max_bricks = max_bricks
-        if max_brick_rejections is not None: model.max_brick_rejections = max_brick_rejections
-        if max_regenerations is not None: model.max_regenerations = max_regenerations
-        if seed is not None: transformers.set_seed(seed)
-
-        # Generate LEGO
-        print(f'Generating LEGO for prompt: "{prompt}"')
-        start_time = time.time()
-        output = model(prompt)
-
-        # Write output LDR to file
-        output_dir = os.path.abspath('out')
-        output_uuid = str(uuid.uuid4())
-        os.makedirs(output_dir, exist_ok=True)
-        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!')
-
-        # 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']
+    model_cfg = LegoGPTConfig(max_regenerations=5)
+    generator = LegoGenerator(LegoGPT(model_cfg))
 
     # Define inputs and outputs
     in_prompt = gr.Textbox(label='Prompt', placeholder='Enter a prompt to generate a LEGO model.')
@@ -99,7 +61,7 @@ def main():
 
     # Define Gradio interface
     demo = gr.Interface(
-        fn=generate_lego,
+        fn=generator.generate_lego_subprocess,
         title='LegoGPT Demo',
         description='Official demo for [LegoGPT](https://avalovelace1.github.io/LegoGPT/), the first approach for generating physically stable LEGO brick models from text prompts.\n\n'
                     '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: '
@@ -123,9 +85,65 @@ def main():
                 fn=lambda *args: (args[-1], examples[args[0]]['output_txt']),
                 run_on_click=True,
             )
+
+    concurrency_limit = 2 if os.environ.get('CONCURRENCY_LIMIT') is None else int(os.environ.get('CONCURRENCY_LIMIT'))
+    demo.queue(default_concurrency_limit=concurrency_limit)
     demo.launch(share=True)
 
 
+class LegoGenerator:
+    def __init__(self, model: LegoGPT):
+        self.model = model
+        self.ctx = mp.get_context('spawn')
+
+    def generate_lego(
+            self,
+            prompt: str,
+            temperature: float | None,
+            seed: int | None,
+            max_bricks: int | None,
+            max_brick_rejections: int | None,
+            max_regenerations: int | None,
+    ):
+        # Set model parameters
+        if temperature is not None: self.model.temperature = temperature
+        if max_bricks is not None: self.model.max_bricks = max_bricks
+        if max_brick_rejections is not None: self.model.max_brick_rejections = max_brick_rejections
+        if max_regenerations is not None: self.model.max_regenerations = max_regenerations
+        if seed is not None: transformers.set_seed(seed)
+
+        # Generate LEGO
+        print(f'Generating LEGO for prompt: "{prompt}"')
+        start_time = time.time()
+        output = self.model(prompt)
+
+        # Write output LDR to file
+        output_dir = os.path.abspath('out')
+        output_uuid = str(uuid.uuid4())
+        os.makedirs(output_dir, exist_ok=True)
+        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!')
+
+        # 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):
+        """
+        Run generation as a subprocess so that multiple requests can be handled concurrently.
+        """
+        with self.ctx.Pool(1) as pool:
+            return pool.starmap(self.generate_lego, [args])[0]
+
+
 def get_help_string(field_name: str) -> str:
     """
     :param field_name: Name of a field in LegoGPTConfig.