chore: Update import statement for InpaintingRequest in painting.py and refactor code to use shared BaseModel for Painting and InpaintingRequest classes

Former-commit-id: a23f3f613d6636adcadd1b01d7f1d61993d91251 [formerly e5c83686f1238c8a9cb66b21590a5a1ee6597665]
Former-commit-id: 2252984afe4046eea2172d7ddd6d0096cf70b07b

api/endpoints.py

@@ -51,5 +51,5 @@ async def root():
 def check_health():
     return {"status": "ok"}
 
-uvicorn.run(app, host="0.0.0.0", port=8000)
+
 

api/models/painting.py

@@ -1,4 +1,5 @@
 from pydantic import BaseModel
+from fastapi import Form
 
 
 class InpaintingRequest(BaseModel):
@@ -7,3 +8,4 @@ class InpaintingRequest(BaseModel):
     num_inference_steps: int
     strength: float
     guidance_scale: float
+    

api/routers/painting.py

@@ -1,69 +1,34 @@
+from fastapi import  APIRouter, File, UploadFile, HTTPException, Form
+from PIL import Image
 import sys
 sys.path.append("../scripts")
-from fastapi import APIRouter, File, UploadFile, HTTPException
-from pydantic import BaseModel
-from PIL import Image
-from io import BytesIO
-from models.painting import InpaintingRequest
 import uuid
-from inpainting_pipeline import AutoPaintingPipeline
-from utils import pil_to_s3_json, ImageAugmentation
-from hydra import compose, initialize
 import lightning.pytorch as pl
-pl.seed_everything(42)
-
-router = APIRouter()
-
-
+from typing import List
+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 pydantic import BaseModel
+from async_batcher.batcher import AsyncBatcher
+from typing import Dict
 
-#class InpaintingRequest(BaseModel):
-   # prompt: str
-   # negative_prompt: str
-   # num_inference_steps: int
-   # strength: float
-   # guidance_scale: float
 
-def augment_image(image, target_width, target_height, roi_scale, segmentation_model_name, detection_model_name):
-    """
-    Augments an image with a given prompt, model, and other parameters.
+router = APIRouter()
+pl.seed_everything(42)
 
-    Parameters:
-    - image (str): The path to the image file.
-    - target_width (int): The desired width of the augmented image.
-    - target_height (int): The desired height of the augmented image.
-    - roi_scale (float): The scale factor for the region of interest.
+with initialize(version_base=None, config_path="../../configs"):
+    cfg = compose(config_name="inpainting")
+inpainting_pipeline = load_pipeline(cfg.model, accelerator(), enable_compile=True)
 
-    Returns:
-    - augmented_image (PIL.Image.Image): The augmented image.
-    - inverted_mask (PIL.Image.Image): The inverted mask generated from the augmented image.
-    """
-    image = Image.open(image)
+def augment_image(image_path, target_width, target_height, roi_scale, segmentation_model_name, detection_model_name):
+    image = Image.open(image_path)
     image_augmentation = ImageAugmentation(target_width, target_height, roi_scale)
     image = image_augmentation.extend_image(image)
     mask = image_augmentation.generate_mask_from_bbox(image, segmentation_model_name, detection_model_name)
     inverted_mask = image_augmentation.invert_mask(mask)
     return image, inverted_mask
 
-def run_inference(cfg: dict, image_path: str, prompt: str, negative_prompt: str, num_inference_steps: int, strength: float, guidance_scale: float):
-    """
-    Run inference using the provided configuration and input image.
-
-    Args:
-        cfg (dict): Configuration dictionary containing model parameters.
-        image_path (str): Path to the input image file.
-        prompt (str): Prompt for the inference process.
-        negative_prompt (str): Negative prompt for the inference process.
-        num_inference_steps (int): Number of inference steps to perform.
-        strength (float): Strength parameter for the inference.
-        guidance_scale (float): Guidance scale for the inference.
-
-    Returns:
-        dict: A JSON object containing the image ID and the signed URL.
-
-    Raises:
-        HTTPException: If an error occurs during the inference process.
-
-    """
+def run_inference(cfg: dict, image_path: str, prompt: str, negative_prompt: str, num_inference_steps: int, strength: float, guidance_scale: float, mode: str, num_images: int):
     image, mask_image = augment_image(image_path, 
                                       cfg['target_width'], 
                                       cfg['target_height'], 
@@ -71,25 +36,89 @@ def run_inference(cfg: dict, image_path: str, prompt: str, negative_prompt: str,
                                       cfg['segmentation_model'], 
                                       cfg['detection_model'])
     
-    pipeline = AutoPaintingPipeline(model_name=cfg['model'], 
-                                    image=image, 
-                                    mask_image=mask_image, 
-                                    target_height=cfg['target_height'], 
-                                    target_width=cfg['target_width'])
-    output = pipeline.run_inference(prompt=prompt, 
+    painting_pipeline = AutoPaintingPipeline(
+        pipeline=inpainting_pipeline,
+        image=image, 
+        mask_image=mask_image, 
+        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)
-    return pil_to_s3_json(output, file_name="output.png")
-
-@router.post("/kandinskyv2.2_inpainting")
-async def inpainting_inference(image: UploadFile = File(...), 
-                               prompt: str = "", 
-                               negative_prompt: str = "", 
-                               num_inference_steps: int = 50, 
-                               strength: float = 0.5, 
-                               guidance_scale: float = 7.5):
+    if mode == "s3_json":
+        return pil_to_s3_json(output, file_name="output.png")
+    elif mode == "b64_json":
+        return pil_to_b64_json(output)
+    else:
+        raise ValueError("Invalid mode. Supported modes are 'b64_json' and 's3_json'.")
+
+class InpaintingRequest(BaseModel):
+    prompt: str
+    negative_prompt: str
+    num_inference_steps: int
+    strength: float
+    guidance_scale: float
+    num_images: int = 1
+
+class InpaintingBatcher(AsyncBatcher[List[Dict], dict]):
+    def __init__(self, pipeline, cfg):
+        self.pipeline = pipeline
+        self.cfg = cfg
+
+    def process_batch(self, batch: List[Dict], image_paths: List[str]) -> List[dict]:
+        results = []
+        for data, image_path in zip(batch, image_paths):
+            try:
+                image, mask_image = augment_image(
+                    image_path,
+                    self.cfg['target_width'],
+                    self.cfg['target_height'],
+                    self.cfg['roi_scale'],
+                    self.cfg['segmentation_model'],
+                    self.cfg['detection_model']
+                )
+                
+                pipeline = AutoPaintingPipeline(
+                    image=image, 
+                    mask_image=mask_image, 
+                    target_height=self.cfg['target_height'], 
+                    target_width=self.cfg['target_width']
+                )
+                output = pipeline.run_inference(
+                    prompt=data['prompt'], 
+                    negative_prompt=data['negative_prompt'], 
+                    num_inference_steps=data['num_inference_steps'], 
+                    strength=data['strength'], 
+                    guidance_scale=data['guidance_scale']
+                )
+
+                if data['mode'] == "s3_json":
+                    result = pil_to_s3_json(output, 'inpainting_image')
+                elif data['mode'] == "b64_json":
+                    result = pil_to_b64_json(output)
+                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
+
+@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)
+):
     """
     Run the inpainting/outpainting inference pipeline.
 
@@ -100,6 +129,8 @@ async def inpainting_inference(image: UploadFile = File(...),
     - num_inference_steps: int - The number of inference steps to perform during the inpainting/outpainting process.
     - strength: float - The strength parameter for controlling the inpainting/outpainting process.
     - guidance_scale: float - The guidance scale parameter for controlling the inpainting/outpainting process.
+    - mode: str - The output mode, either "s3_json" or "b64_json".
+    - num_images: int - The number of images to generate.
 
     Returns:
     - result: The result of the inpainting/outpainting process.
@@ -113,14 +144,47 @@ async def inpainting_inference(image: UploadFile = File(...),
         with open(image_path, "wb") as f:
             f.write(image_bytes)
         
+        result = run_inference(
+            cfg, image_path, prompt, negative_prompt, num_inference_steps, strength, guidance_scale, mode, num_images
+        )
         
-        with initialize(version_base=None,config_path="../../configs"):
-            cfg = compose(config_name="inpainting")
+        return result
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@router.post("/inpainting_batch")
+async def inpainting_batch_inference(
+    batch: List[dict], 
+    images: List[UploadFile] = File(...)
+):
+    """
+    Run batch inpainting/outpainting inference pipeline.
 
-        result = run_inference(cfg, image_path, prompt, negative_prompt, num_inference_steps, strength, guidance_scale)
+    Parameters:
+    - batch: List[dict] - The batch of requests containing parameters for the inpainting/outpainting process.
+    - images: List[UploadFile] - The list of image files to be used for inpainting/outpainting.
+
+    Returns:
+    - results: The results of the inpainting/outpainting process for each request.
+
+    Raises:
+    - HTTPException: If an error occurs during the inpainting/outpainting process.
+    """
+    try:
+        image_paths = []
+        for image in images:
+            image_bytes = await image.read()
+            image_path = f"/tmp/{uuid.uuid4()}.png"
+            with open(image_path, "wb") as f:
+                f.write(image_bytes)
+            image_paths.append(image_path)
+
+        batcher = InpaintingBatcher(pipeline, cfg)
+        results = batcher.process_batch(batch, image_paths)
         
-        return result
+        return results
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
+    
 
 

configs/inpainting.yaml

@@ -1,5 +1,5 @@
-segmentation_model : 'facebook/sam-vit-huge'
-detection_model : 'yolov8l'
+segmentation_model : 'facebook/sam-vit-base'
+detection_model : 'yolov8s'
 model : 'kandinsky-community/kandinsky-2-2-decoder-inpaint'
 target_width : 2560
 target_height : 1472

scripts/config.py

@@ -9,6 +9,7 @@ CAPTIONING_MODEL_NAME = "Salesforce/blip-image-captioning-base"
 SEGMENTATION_MODEL_NAME = "facebook/sam-vit-large"
 DETECTION_MODEL_NAME = "yolov8l"
 ENABLE_COMPILE = False
+INPAINTING_MODEL_NAME = ''
 
 
 

scripts/inpainting_pipeline.py

@@ -1,81 +1,80 @@
 import torch
-from diffusers import AutoPipelineForInpainting,DiffusionPipeline
+from diffusers import AutoPipelineForInpainting
 from diffusers.utils import load_image
-from utils import (accelerator, ImageAugmentation, clear_memory)
+from utils import accelerator, ImageAugmentation
 import hydra
 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 = True):
+    pipeline = AutoPipelineForInpainting.from_pretrained(model_name, torch_dtype=torch.float16)
+    if enable_compile:
+        pipeline.unet.to(memory_format=torch.channels_last)
+        pipeline.unet = torch.compile(pipeline.unet, mode='reduce-overhead',fullgraph=True)
+    pipeline.to(device)
+    return pipeline
 
 class AutoPaintingPipeline:
-    """
-    AutoPaintingPipeline class represents a pipeline for auto painting using an inpainting model from diffusers.
-    
-    Args:
-        model_name (str): The name of the pretrained inpainting model.
-        image (Image): The input image to be processed.
-        mask_image (Image): The mask image indicating the areas to be inpainted.
-    """
-    
-    def __init__(self, model_name: str, image: Image, mask_image: Image,target_width: int, target_height: int):
-        self.model_name = model_name
-        self.device = accelerator()
-        self.pipeline = AutoPipelineForInpainting.from_pretrained(self.model_name, torch_dtype=torch.float16)
-        self.image = load_image(image)
-        self.mask_image = load_image(mask_image)
+    def __init__(self, pipeline, image: Image, mask_image: Image, target_width: int, target_height: int):
+        self.pipeline = pipeline
+        self.image = image
+        self.mask_image = mask_image
         self.target_width = target_width
         self.target_height = target_height
-        self.pipeline.to(self.device)
-        self.pipeline.unet = torch.compile(self.pipeline.unet,mode='max-autotune')
-        
-       
-        
-        
+
     def run_inference(self, prompt: str, negative_prompt: str, num_inference_steps: int, strength: float, guidance_scale: float):
-        """
-        Runs the inference on the input image using the inpainting pipeline.
-        
-        Returns:
-            Image: The output image after inpainting.
-        """
-       
-        image = load_image(self.image)
-        mask_image = load_image(self.mask_image)
-        output = self.pipeline(prompt=prompt,negative_prompt=negative_prompt,image=image,mask_image=mask_image,num_inference_steps=num_inference_steps,strength=strength,guidance_scale=guidance_scale, height = self.target_height ,width = self.target_width).images[0]
+        output = self.pipeline(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=self.image,
+            mask_image=self.mask_image,
+            num_inference_steps=num_inference_steps,
+            strength=strength,
+            guidance_scale=guidance_scale,
+            height=self.target_height,
+            width=self.target_width
+            
+        ).images[0]
         return output
-    
-    
-@hydra.main(version_base=None ,config_path="../configs", config_name="inpainting")
+
+@hydra.main(version_base=None, config_path="../configs", config_name="inpainting")
 def inference(cfg: DictConfig):
-    """
-    Load the configuration file for the inpainting pipeline.
-    
-    Args:
-        cfg (DictConfig): The configuration file for the inpainting pipeline.
-    """
+    # Load the pipeline once and cache it
+    pipeline = load_pipeline(cfg.model, accelerator(), True)
+
+    # Image augmentation and preparation
     augmenter = ImageAugmentation(target_width=cfg.target_width, target_height=cfg.target_height)
-    model_name = cfg.model
     image_path = "../sample_data/example3.jpg"
     image = Image.open(image_path)
     extended_image = augmenter.extend_image(image)
     mask_image = augmenter.generate_mask_from_bbox(extended_image, cfg.segmentation_model, cfg.detection_model)
     mask_image = augmenter.invert_mask(mask_image)
-    prompt = cfg.prompt
-    negative_prompt = cfg.negative_prompt
-    num_inference_steps = cfg.num_inference_steps
-    strength = cfg.strength
-    guidance_scale = cfg.guidance_scale
-    pipeline = AutoPaintingPipeline(model_name=model_name, image = extended_image, mask_image=mask_image, target_height=cfg.target_height, target_width=cfg.target_width)
-    output = pipeline.run_inference(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, strength=strength, guidance_scale=guidance_scale)
-    output.save(f'{cfg.output_path}/output.jpg')
-    mask_image.save(f'{cfg.output_path}/mask.jpg')
     
+    # Create AutoPaintingPipeline instance with cached pipeline
+    painting_pipeline = AutoPaintingPipeline(
+        pipeline=pipeline,
+        image=extended_image,
+        mask_image=mask_image,
+        target_height=cfg.target_height,
+        target_width=cfg.target_width
+    )
+    
+    # Run inference
+    output = painting_pipeline.run_inference(
+        prompt=cfg.prompt,
+        negative_prompt=cfg.negative_prompt,
+        num_inference_steps=cfg.num_inference_steps,
+        strength=cfg.strength,
+        guidance_scale=cfg.guidance_scale
+    )
     
+    # Save output and mask images
+    output.save(f'{cfg.output_path}/output.jpg')
+    mask_image.save(f'{cfg.output_path}/mask.jpg')
+
 if __name__ == "__main__":
     inference()
 
-        
-