use CUDA

app.py

@@ -43,7 +43,7 @@ def main():
         setup()
 
     model_cfg = BrickGPTConfig(max_regenerations=5, device='cuda')
-    generator = BrickGenerator(BrickGPT(model_cfg))
+    generator = BrickGenerator(model_cfg)
 
     # Define inputs and outputs
     in_prompt = gr.Textbox(label='Prompt', placeholder='Enter a prompt to generate a brick model.', max_length=500)
@@ -100,9 +100,9 @@ def main():
 
 
 class BrickGenerator:
-    def __init__(self, model: BrickGPT):
-        self.model = model
-        self.ctx = mp.get_context('spawn')
+    def __init__(self, model_cfg: BrickGPTConfig):
+        self.model_cfg = model_cfg
+        self.model = None
 
     @spaces.GPU
     def generate_bricks(
@@ -115,6 +115,8 @@ class BrickGenerator:
             max_brick_rejections: int | None,
             max_regenerations: int | None,
     ) -> (str, str):
+        self.model = BrickGPT(self.model_cfg)
+
         # Set model parameters
         if temperature is not None: self.model.temperature = temperature
         if max_bricks is not None: self.model.max_bricks = max_bricks
@@ -173,13 +175,6 @@ class BrickGenerator:
 
         return img_filename, output['bricks'].to_txt()
 
-    def generate_bricks_subprocess(self, *args) -> (str, str):
-        """
-        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_bricks, [args])[0]
-
 
 def get_help_string(field_name: str) -> str:
     """