chore: Update dependencies in requirements.txt and fix import statement in painting.py

Former-commit-id: 4655fd747d81318b35f11a66885e8c4e5ce78e8b [formerly 076e97d0b56f534dc5908d0998b7e195c94f8f8a]
Former-commit-id: d6f4bb04acf087d0c30c84af9c3e0b8f9903becd

api/endpoints.py

@@ -52,4 +52,4 @@ def check_health():
 
 
 
-#uvicorn.run(app, host='127.0.0.1', port=8000)
+uvicorn.run(app, host='127.0.0.1', port=8000)

api/requirements.txt

@@ -1,12 +1,12 @@
+async_batcher==0.2.0
 config==0.5.1
 diffusers==0.27.2
 fastapi==0.111.0
 hydra-core==1.3.2
 lightning==2.2.3
-logfire==0.41.0
+logfire==0.42.0
 Pillow==10.3.0
-pydantic==2.7.3
+pydantic==2.7.4
 torch==2.2.0
 utils==1.0.2
 uvicorn==0.30.1
-

api/routers/painting.py

@@ -2,7 +2,7 @@ import sys
 sys.path.append('../scripts')
 import os
 import uuid
-from typing import List, Tuple, Any,Dict
+from typing import List, Tuple, Any, Dict
 from fastapi import APIRouter, File, UploadFile, HTTPException, Form, Depends, Body
 from pydantic import BaseModel, Field
 from PIL import Image
@@ -11,9 +11,7 @@ from utils import pil_to_s3_json, pil_to_b64_json, ImageAugmentation, accelerato
 from inpainting_pipeline import AutoPaintingPipeline, load_pipeline
 from hydra import compose, initialize
 from async_batcher.batcher import AsyncBatcher
-from concurrent.futures import Executor
 import json
-import asyncio
 from functools import lru_cache
 
 pl.seed_everything(42)
@@ -26,11 +24,20 @@ with initialize(version_base=None, config_path="../../configs"):
 # Load the inpainting pipeline
 @lru_cache(maxsize=1)
 def load_pipeline_wrapper():
+    """
+    Load the inpainting pipeline with the specified configuration.
+
+    Returns:
+        pipeline: The loaded inpainting pipeline.
+    """
     pipeline = load_pipeline(cfg.model, accelerator(), enable_compile=True)
     return pipeline
 inpainting_pipeline = load_pipeline_wrapper()
 
 class InpaintingRequest(BaseModel):
+    """
+    Model representing a request for inpainting inference.
+    """
     prompt: str = Field(..., description="Prompt text for inference")
     negative_prompt: str = Field(..., description="Negative prompt text for inference")
     num_inference_steps: int = Field(..., description="Number of inference steps")
@@ -41,6 +48,9 @@ class InpaintingRequest(BaseModel):
     use_augmentation: bool = Field(True, description="Whether to use image augmentation")
     
 class InpaintingBatchRequestModel(BaseModel):
+    """
+    Model representing a batch request for inpainting inference.
+    """
     requests: List[InpaintingRequest]
 
 async def save_image(image: UploadFile) -> str:
@@ -105,7 +115,7 @@ def run_inference(cfg, image_path: str, request: InpaintingRequest):
                                           cfg['detection_model'])
     else:
         image = Image.open(image_path)
-        mask_image = None  # Assume mask_image is provided or generated separately
+        mask_image = None  
     
     painting_pipeline = AutoPaintingPipeline(
         pipeline=inpainting_pipeline,
@@ -118,16 +128,26 @@ def run_inference(cfg, image_path: str, request: InpaintingRequest):
                                     negative_prompt=request.negative_prompt, 
                                     num_inference_steps=request.num_inference_steps, 
                                     strength=request.strength, 
-                                    guidance_scale=request.guidance_scale)
+                                    guidance_scale=request.guidance_scale,
+                                    num_images=request.num_images)
     if request.mode == "s3_json":
         return pil_to_s3_json(output, file_name="output.png")
     elif request.mode == "b64_json":
         return pil_to_b64_json(output)
     else:
         raise ValueError("Invalid mode. Supported modes are 'b64_json' and 's3_json'.")
-    
+
 class InpaintingBatcher(AsyncBatcher):
     async def process_batch(self, batch: Tuple[List[str], List[InpaintingRequest]]) -> List[Dict[str, Any]]:
+        """
+        Process a batch of images and requests for inpainting inference.
+
+        Args:
+            batch (Tuple[List[str], List[InpaintingRequest]]): Tuple of image paths and corresponding requests.
+
+        Returns:
+            List[Dict[str, Any]]: List of resulting images in the specified mode ('b64_json' or 's3_json').
+        """
         image_paths, requests = batch
         results = []
         for image_path, request in zip(image_paths, requests):
@@ -190,7 +210,7 @@ async def inpainting_batch_inference(
 
         batcher = InpaintingBatcher(max_batch_size=64)
         image_paths = [await save_image(image) for image in images]
-        results = await batcher.process_batch((image_paths, requests))
+        results = batcher.process_batch((image_paths, requests))
 
         return results
     except Exception as e:

api/routers/sdxl_text_to_image.py

@@ -184,7 +184,7 @@ async def sdxl_v0_lora_inference(data: InputFormat):
 async def sdxl_v0_lora_inference_batch(data: List[InputFormat]):
     batcher = SDXLLoraBatcher(max_batch_size=-1)
     try:
-        predictions = await batcher.process_batch(data)
+        predictions =  batcher.process_batch(data)
         return predictions
     except Exception as e:
         print(f"Error in /sdxl_v0_lora_inference/batch: {e}")

scripts/inpainting_pipeline.py

@@ -24,7 +24,7 @@ class AutoPaintingPipeline:
         self.target_width = target_width
         self.target_height = target_height
 
-    def run_inference(self, prompt: str, negative_prompt: str, num_inference_steps: int, strength: float, guidance_scale: float):
+    def run_inference(self, prompt: str, negative_prompt: str, num_inference_steps: int, strength: float, guidance_scale: float,num_images):
         output = self.pipeline(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -33,6 +33,7 @@ class AutoPaintingPipeline:
             num_inference_steps=num_inference_steps,
             strength=strength,
             guidance_scale=guidance_scale,
+            num_images_per_prompt = num_images,
             height=self.target_height,
             width=self.target_width