chore: Update Inpainting model name in config.py

Former-commit-id: 22fa0b9f756523fb282ef6ae1b1d691a810b21ca [formerly 98fe0a0084679efccb8eb55bc9b061f439126cce]
Former-commit-id: e181cfba0759e658521596a5031cce2e1a5b1f5b

.vscode/settings.json

@@ -2,5 +2,8 @@
     "python.REPL.enableREPLSmartSend": false,
     "circleci.persistedProjectSelection": [
         "circleci/BQyaZDyai3ejufsymJgXtx/G7GFeA574Ga7r2VtdHXBbu"
+    ],
+    "python.analysis.extraPaths": [
+        "./scripts"
     ]
 }

api/endpoints.py

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

api/routers/batch_painting.py

@@ -0,0 +1,16 @@
+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 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 load_pipeline
+from async_batcher.batcher import AsyncBatcher
+from config import INPAINTING_MODEL_NAME
+router = APIRouter()
+
+
+inpainting_pipeline = load_pipeline(model_name = INPAINTING_MODEL_NAME,enable_)

api/routers/painting.py

@@ -1,25 +1,32 @@
-from fastapi import  APIRouter, File, UploadFile, HTTPException, Form
-from PIL import Image
 import sys
-sys.path.append("../scripts")
+sys.path.append('../scripts')
+import os
 import uuid
+from typing import List, Tuple, Any
+from fastapi import APIRouter, File, UploadFile, HTTPException, Form
+from PIL import Image
 import lightning.pytorch as pl
-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
-
 
-router = APIRouter()
 pl.seed_everything(42)
+router = APIRouter()
 
+# Initialize the configuration and pipeline
 with initialize(version_base=None, config_path="../../configs"):
     cfg = compose(config_name="inpainting")
 inpainting_pipeline = load_pipeline(cfg.model, accelerator(), enable_compile=True)
 
+async def save_image(image: UploadFile) -> str:
+    """Save an uploaded image to a temporary file and return the file path."""
+    file_name = f"{uuid.uuid4()}.png"
+    file_path = os.path.join("/tmp", file_name)
+    with open(file_path, "wb") as f:
+        f.write(await image.read())
+    return file_path
+
 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)
@@ -28,13 +35,26 @@ 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: 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'], 
-                                      cfg['roi_scale'], 
-                                      cfg['segmentation_model'], 
-                                      cfg['detection_model'])
+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:
+        image, mask_image = augment_image(image_path, 
+                                          cfg['target_width'], 
+                                          cfg['target_height'], 
+                                          cfg['roi_scale'], 
+                                          cfg['segmentation_model'], 
+                                          cfg['detection_model'])
+    else:
+        image = Image.open(image_path)
+        mask_image = None  # Assume mask_image is provided or generated separately
     
     painting_pipeline = AutoPaintingPipeline(
         pipeline=inpainting_pipeline,
@@ -55,58 +75,7 @@ def run_inference(cfg: dict, image_path: str, prompt: str, negative_prompt: str,
     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(
@@ -117,74 +86,14 @@ async def inpainting_inference(
     strength: float = Form(...),
     guidance_scale: float = Form(...),
     mode: str = Form(...),
-    num_images: int = Form(1)
+    num_images: int = Form(1),
+    
 ):
-    """
-    Run the inpainting/outpainting inference pipeline.
-
-    Parameters:
-    - image: UploadFile - The image file to be used for inpainting/outpainting.
-    - prompt: str - The prompt text for guiding the inpainting/outpainting process.
-    - negative_prompt: str - The negative prompt text for guiding the inpainting/outpainting process.
-    - 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.
-
-    Raises:
-    - HTTPException: If an error occurs during the inpainting/outpainting process.
-    """
     try:
-        image_bytes = await image.read()
-        image_path = f"/tmp/{uuid.uuid4()}.png"
-        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
-        )
-        
+        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)
         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.
-
-    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 results
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(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=64)
     try:
-        predictions = batcher.process_batch(data)
+        predictions = await batcher.process_batch(data)
         return predictions
     except Exception as e:
         print(f"Error in /sdxl_v0_lora_inference/batch: {e}")

scripts/config.py

@@ -9,7 +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 = ''
+INPAINTING_MODEL_NAME = 'kandinskyv2.2'
 
 
 

scripts/inpainting_pipeline.py

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