bug fixes

configs/tti_settings.py

@@ -7,7 +7,7 @@ class TTI_SETTINGS(BaseSettings):
     ENABLE_COMPILE: bool = False
     DEVICE: str = "cuda"
     TRITON_MODEL_NAME: str = "PICPILOT_PRODUCTION_SERVER"
-    MAX_BATCH_SIZE: int = 8
+    MAX_BATCH_SIZE: int = 32
     MAX_QUEUE_DELAY_MICROSECONDS: int = 100
     TORCH_INDUCTOR_CONFIG: dict = {
         "conv_1x1_as_mm": True,

scripts/s3_manager.py

@@ -20,12 +20,21 @@ class S3ManagerService:
             region_name=settings.AWS_REGION,
         )
     
-    def generate_signed_url(self, file_name: str, exp: int = 43200) -> str:  # 43200 seconds = 12 hours
-        return self.s3.generate_presigned_url(
-            "get_object",
-            Params={"Bucket": settings.AWS_BUCKET_NAME, "Key": file_name},
-            ExpiresIn=exp,
-        )
+    def generate_signed_url(self, file_name: str, exp: int = 43200) -> str:
+        try:
+            url = self.s3.generate_presigned_url(
+                ClientMethod='get_object',
+                Params={
+                    'Bucket': settings.AWS_BUCKET_NAME,
+                    'Key': file_name
+                },
+                ExpiresIn=exp,
+                HttpMethod='GET'
+            )
+            return url
+        except Exception as e:
+            print(f"Error generating presigned URL: {e}")
+            return None
 
     def generate_unique_file_name(self, file_name: str) -> str:
         random_string = "".join(

triton-api/client.py

@@ -128,7 +128,7 @@ async def generate_image(request: ImageGenerationRequest) -> Dict[str, Any]:
         inputs = _prepare_inference_inputs(request)
         result_dict = await triton_client.infer_sample(**inputs)
 
-        output = result_dict["output"][0]
+        output = result_dict["output"]
         return json.loads(output.decode("utf-8"))
     except Exception as e:
         logger.error(f"Error generating image: {e}")

triton-api/text-to-image.py

@@ -130,18 +130,20 @@ class SDXLLoraInference:
         Raises:
             ValueError: If an invalid output mode is specified.
         """
-        image = self.pipe(
+        images = self.pipe(
             prompt=self.prompt,
             num_inference_steps=self.num_inference_steps,
             guidance_scale=self.guidance_scale,
             negative_prompt=self.negative_prompt,
             num_images_per_prompt=self.num_images,
-        ).images[0]
+        ).images
 
         if self.mode == "s3_json":
-            return pil_to_s3_json(image, "sdxl_image")
+            for image in images:
+                return pil_to_s3_json(image, "sdxl_image")
         elif self.mode == "b64_json":
-            return pil_to_b64_json(image)
+            for image in images:
+                return pil_to_b64_json(image)
         else:
             raise ValueError(
                 "Invalid mode. Supported modes are 'b64_json' and 's3_json'."