chore: Refactor code to use shared BaseModel for Painting and InputFormat classes

Former-commit-id: 407eeff325cb35937072ae4d08ae116b0c938b52 [formerly aca93c4c8e33ae484649dafdb75471c4cd35c4b7]
Former-commit-id: 62ead1acffc613e477d909cb2c553fd5da68ba9f

api/models/painting.py

@@ -0,0 +1,11 @@
+from pydantic import BaseModel
+
+
+class Painting(BaseModel):
+    prompt: str
+    num_inference_steps: int
+    guidance_scale: float
+    negative_prompt: str
+    num_images: int
+    mode: str
+    

api/models/sdxl_input.py

@@ -0,0 +1,12 @@
+
+from pydantic import BaseModel
+
+
+
+class InputFormat(BaseModel):
+    prompt: str
+    num_inference_steps: int
+    guidance_scale: float
+    negative_prompt: str
+    num_images: int
+    mode: str

api/routers/sdxl_text_to_image.py

@@ -14,6 +14,8 @@ from s3_manager import S3ManagerService
 from PIL import Image
 import io
 from utils import accelerator
+from models.sdxl_input import InputFormat
+from async_batcher.batcher import AsyncBatcher
 
 device = accelerator()
 torch._inductor.config.conv_1x1_as_mm = True
@@ -101,7 +103,7 @@ loaded_pipeline = load_pipeline(config.MODEL_NAME, config.ADAPTER_NAME, config.E
 
 
 # SDXLLoraInference class for running inference
-class SDXLLoraInference:
+class SDXLLoraInference(AsyncBatcher):
     """
     Class for performing SDXL Lora inference.
 
@@ -169,21 +171,37 @@ class SDXLLoraInference:
             return pil_to_b64_json(image)
         else:
             raise ValueError("Invalid mode. Supported modes are 'b64_json' and 's3_json'.")
+        
+    
+class SDXLLoraBatcher(AsyncBatcher[InputFormat, dict]):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.pipe = loaded_pipeline
 
-
-# Input format for single request
-class InputFormat(BaseModel):
-    prompt: str
-    num_inference_steps: int
-    guidance_scale: float
-    negative_prompt: str
-    num_images: int
-    mode: str
-
-
-# Input format for batch requests
-class BatchInputFormat(BaseModel):
-    batch_input: List[InputFormat]
+    def process_batch(self, batch: List[InputFormat]) -> List[dict]:
+        results = []
+        for data in batch:
+            try:
+                images = self.pipe(
+                    prompt=data.prompt,
+                    num_inference_steps=data.num_inference_steps,
+                    guidance_scale=data.guidance_scale,
+                    negative_prompt=data.negative_prompt,
+                    num_images_per_prompt=data.num_images,
+                ).images
+                
+                for image in images:
+                    if data.mode == "s3_json":
+                        result = pil_to_s3_json(image, 'sdxl_image')
+                    elif data.mode == "b64_json":
+                        result = pil_to_b64_json(image)
+                    else:
+                        raise ValueError("Invalid mode. Supported modes are 'b64_json' and 's3_json'.")
+                    results.append(result)
+            except Exception as e:
+                print(f"Error in process_batch: {e}")
+                raise HTTPException(status_code=500, detail="Batch inference failed")
+        return results
 
 
 # Endpoint for single request
@@ -195,47 +213,24 @@ async def sdxl_v0_lora_inference(data: InputFormat):
         data.num_images,
         data.num_inference_steps,
         data.guidance_scale,
-        data.mode
+        data.mode,
+        
+        
     )
     output_json = inference.run_inference()
     return output_json
 
 
+# Endpoint for batch requests
 
 @router.post("/sdxl_v0_lora_inference/batch")
-async def sdxl_v0_lora_inference_batch(data: BatchInputFormat):
-    """
-    Perform batch inference for SDXL V0 LoRa model.
-
-    Args:
-        data (BatchInputFormat): The input data containing a batch of requests.
-
-    Returns:
-        dict: A dictionary containing the message and processed requests data.
-
-    Raises:
-        HTTPException: If the number of requests exceeds the maximum queue size.
-    """
-    MAX_QUEUE_SIZE = 64
-
-    if len(data.batch_input) > MAX_QUEUE_SIZE:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Number of requests exceeds maximum queue size ({MAX_QUEUE_SIZE})",
-        )
-
-    processed_requests = []
-    for item in data.batch_input:
-        inference = SDXLLoraInference(
-            item.prompt,
-            item.negative_prompt,
-            item.num_images,
-            item.num_inference_steps,
-            item.guidance_scale,
-            item.mode,
-        )
-        output_json = inference.run_inference()
-        processed_requests.append(output_json)
-
-    return {"message": "Requests processed successfully", "data": processed_requests}
+async def sdxl_v0_lora_inference_batch(data: List[InputFormat]):
+    batcher = SDXLLoraBatcher(max_batch_size=64, max_queue_time=0.001)
+    try:
+        predictions = await batcher.process(batch=data)
+        return predictions
+    except Exception as e:
+        print(f"Error in /sdxl_v0_lora_inference/batch: {e}")
+        raise HTTPException(status_code=500, detail="Batch inference endpoint failed")
+