Refactor code to use shared BaseModel for Painting and InpaintingRequest classes

Former-commit-id: 4484440577c279f069f9e6c8c11d9ab9685f26e8 [formerly b52b26e6289548705fa59307a387beb19edec4f6]
Former-commit-id: 60b102b660610226910f419923c877a820891aab

api/endpoints.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
-from routers import (sdxl_text_to_image,painting,batch_painting)
+from routers import (sdxl_text_to_image,painting)
 import logfire
 import uvicorn 
 
@@ -26,7 +26,6 @@ app.add_middleware(
 
 app.include_router(sdxl_text_to_image.router, prefix='/api/v1/product-diffusion')
 app.include_router(painting.router,prefix='/api/v1/product-diffusion')
-app.include_router(batch_painting.router,prefix='/api/v1/product-diffusion')
 logfire.instrument_fastapi(app)
 
 

api/routers/painting.py

@@ -2,25 +2,57 @@ import sys
 sys.path.append('../scripts')
 import os
 import uuid
-from typing import List, Tuple, Any
-from fastapi import APIRouter, File, UploadFile, HTTPException, Form
+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
 import lightning.pytorch as pl
 from utils import pil_to_s3_json, pil_to_b64_json, ImageAugmentation, accelerator
 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)
 router = APIRouter()
 
-# Initialize the configuration and pipeline
+# Initialize Hydra configuration
 with initialize(version_base=None, config_path="../../configs"):
     cfg = compose(config_name="inpainting")
-inpainting_pipeline = load_pipeline(cfg.model, accelerator(), enable_compile=True)
+
+# Load the inpainting pipeline
+@lru_cache(maxsize=1)
+def load_pipeline_wrapper():
+    pipeline = load_pipeline(cfg.model, accelerator(), enable_compile=True)
+    return pipeline
+inpainting_pipeline = load_pipeline_wrapper()
+
+class InpaintingRequest(BaseModel):
+    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")
+    strength: float = Field(..., description="Strength of the inference")
+    guidance_scale: float = Field(..., description="Guidance scale for inference")
+    mode: str = Field(..., description="Mode for output ('b64_json' or 's3_json')")
+    num_images: int = Field(..., description="Number of images to generate")
+    use_augmentation: bool = Field(True, description="Whether to use image augmentation")
+    
+class InpaintingBatchRequestModel(BaseModel):
+    requests: List[InpaintingRequest]
 
 async def save_image(image: UploadFile) -> str:
-    """Save an uploaded image to a temporary file and return the file path."""
+    """
+    Save an uploaded image to a temporary file and return the file path.
+
+    Args:
+        image (UploadFile): The uploaded image file.
+
+    Returns:
+        str: File path where the image is saved.
+    """
     file_name = f"{uuid.uuid4()}.png"
     file_path = os.path.join("/tmp", file_name)
     with open(file_path, "wb") as f:
@@ -28,6 +60,20 @@ async def save_image(image: UploadFile) -> str:
     return file_path
 
 def augment_image(image_path, target_width, target_height, roi_scale, segmentation_model_name, detection_model_name):
+    """
+    Augment an image by extending its dimensions and generating masks.
+
+    Args:
+        image_path (str): Path to the image file.
+        target_width (int): Target width for augmentation.
+        target_height (int): Target height for augmentation.
+        roi_scale (float): Scale factor for region of interest.
+        segmentation_model_name (str): Name of the segmentation model.
+        detection_model_name (str): Name of the detection model.
+
+    Returns:
+        Tuple[Image.Image, Image.Image]: Augmented image and inverted mask.
+    """
     image = Image.open(image_path)
     image_augmentation = ImageAugmentation(target_width, target_height, roi_scale)
     image = image_augmentation.extend_image(image)
@@ -35,17 +81,22 @@ def augment_image(image_path, target_width, target_height, roi_scale, segmentati
     inverted_mask = image_augmentation.invert_mask(mask)
     return image, inverted_mask
 
-def run_inference(cfg,
-                  image_path: str, 
-                  prompt: str, 
-                  negative_prompt: str, 
-                  num_inference_steps: int, 
-                  strength: float, 
-                  guidance_scale: float, 
-                  mode: str, 
-                  num_images: int,
-                  use_augmentation: bool):
-    if use_augmentation:
+def run_inference(cfg, image_path: str, request: InpaintingRequest):
+    """
+    Run inference using an inpainting pipeline on an image.
+
+    Args:
+        cfg (dict): Configuration dictionary.
+        image_path (str): Path to the image file.
+        request (InpaintingRequest): Pydantic model containing inference parameters.
+
+    Returns:
+        dict: Resulting image in the specified mode ('b64_json' or 's3_json').
+    
+    Raises:
+        ValueError: If an invalid mode is provided.
+    """
+    if request.use_augmentation:
         image, mask_image = augment_image(image_path, 
                                           cfg['target_width'], 
                                           cfg['target_height'], 
@@ -63,37 +114,84 @@ def run_inference(cfg,
         target_height=cfg['target_height'], 
         target_width=cfg['target_width']
     )
-    output = painting_pipeline.run_inference(prompt=prompt, 
-                                    negative_prompt=negative_prompt, 
-                                    num_inference_steps=num_inference_steps, 
-                                    strength=strength, 
-                                    guidance_scale=guidance_scale)
-    if mode == "s3_json":
+    output = painting_pipeline.run_inference(prompt=request.prompt, 
+                                    negative_prompt=request.negative_prompt, 
+                                    num_inference_steps=request.num_inference_steps, 
+                                    strength=request.strength, 
+                                    guidance_scale=request.guidance_scale)
+    if request.mode == "s3_json":
         return pil_to_s3_json(output, file_name="output.png")
-    elif mode == "b64_json":
+    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]]:
+        image_paths, requests = batch
+        results = []
+        for image_path, request in zip(image_paths, requests):
+            result = run_inference(cfg, image_path, request)
+            results.append(result)
+        return results
 
 @router.post("/inpainting")
 async def inpainting_inference(
     image: UploadFile = File(...),
-    prompt: str = Form(...),
-    negative_prompt: str = Form(...),
-    num_inference_steps: int = Form(...),
-    strength: float = Form(...),
-    guidance_scale: float = Form(...),
-    mode: str = Form(...),
-    num_images: int = Form(1),
-    
+    request_data: str = Form(...),
 ):
+    """
+    Handle POST request for inpainting inference.
+
+    Args:
+        image (UploadFile): Uploaded image file.
+        request_data (str): JSON string of the request parameters.
+
+    Returns:
+        dict: Resulting image in the specified mode ('b64_json' or 's3_json').
+
+    Raises:
+        HTTPException: If there is an error during image processing.
+    """
     try:
         image_path = await save_image(image)
-        result = run_inference(cfg, image_path, prompt, negative_prompt, num_inference_steps, strength, guidance_scale, mode, num_images, use_augmentation=True)
+        request_dict = json.loads(request_data)
+        request = InpaintingRequest(**request_dict)
+        result = run_inference(cfg, image_path, request)
         return result
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+@router.post("/inpainting/batch")
+async def inpainting_batch_inference(
+    images: List[UploadFile] = File(...),
+    request_data: str = Form(...),
+):
+    """
+    Handle POST request for batch inpainting inference.
+
+    Args:
+        images (List[UploadFile]): List of uploaded image files.
+        request_data (str): JSON string of the request parameters.
+
+    Returns:
+        List[dict]: List of resulting images in the specified mode ('b64_json' or 's3_json').
 
+    Raises:
+        HTTPException: If there is an error during image processing.
+    """
+    try:
+        request_dict = json.loads(request_data)
+        batch_request = InpaintingBatchRequestModel(**request_dict)
+        requests = batch_request.requests
+
+        if len(images) != len(requests):
+            raise HTTPException(status_code=400, detail="The number of images and requests must match.")
+
+        batcher = InpaintingBatcher(max_batch_size=64)
+        image_paths = [await save_image(image) for image in images]
+        results = await batcher.process_batch((image_paths, requests))
+
+        return results
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

api/routers/sdxl_text_to_image.py

@@ -32,7 +32,7 @@ router = APIRouter()
 
 
 # Load the diffusion pipeline
-@lru_cache(maxsize=1)
+
 def load_pipeline(model_name, adapter_name,enable_compile:bool):
     """
     Load the diffusion pipeline with the specified model and adapter names.
@@ -182,7 +182,7 @@ async def sdxl_v0_lora_inference(data: InputFormat):
 
 @router.post("/sdxl_v0_lora_inference/batch")
 async def sdxl_v0_lora_inference_batch(data: List[InputFormat]):
-    batcher = SDXLLoraBatcher(max_batch_size=64)
+    batcher = SDXLLoraBatcher(max_batch_size=-1)
     try:
         predictions = await batcher.process_batch(data)
         return predictions

scripts/inpainting_pipeline.py

@@ -7,8 +7,8 @@ from omegaconf import DictConfig
 from PIL import Image
 from functools import lru_cache
 
-@lru_cache(maxsize=1)
-def load_pipeline(model_name: str, device, enable_compile: bool = False):
+
+def load_pipeline(model_name: str, device, enable_compile: bool = True):
     pipeline = AutoPipelineForInpainting.from_pretrained(model_name, torch_dtype=torch.float16)
     if enable_compile:
         pipeline.unet.to(memory_format=torch.channels_last)