commit

Former-commit-id: 6e67c162ab28997e4b3581e2482f192da255a4be [formerly 522ceaeca59f5a0db3c0e1e300b6ac98b69656af]

Former-commit-id: 62721f17c8bc87e55be13bfc8363e09454779126

api/routers/painting.py

@@ -1,49 +1,64 @@
 import os
 import uuid
-import json
-from typing import List
+from typing import List, Tuple, Any, Dict
 from fastapi import APIRouter, File, UploadFile, HTTPException, Form
-from pydantic import BaseModel, Field, ValidationError
+from pydantic import BaseModel, Field
+from PIL import Image
 import lightning.pytorch as pl
-from scripts.api_utils import pil_to_b64_json
-from scripts.outpainting import ControlNetZoeDepthOutpainting
+from scripts.api_utils import pil_to_s3_json, pil_to_b64_json, ImageAugmentation, accelerator
+from scripts.inpainting_pipeline import AutoPaintingPipeline, load_pipeline
+from hydra import compose, initialize
 from async_batcher.batcher import AsyncBatcher
+import json
 from functools import lru_cache
-
 pl.seed_everything(42)
-
 router = APIRouter()
 
-@lru_cache(maxsize=1)
-def load_pipeline():
-    outpainting_pipeline = ControlNetZoeDepthOutpainting(target_size=(1024, 1024))
-    return outpainting_pipeline
 
-class OutpaintingRequest(BaseModel):
+with initialize(version_base=None, config_path="../../configs"):
+    cfg = compose(config_name="inpainting")
+
+# Load the inpainting pipeline
+@lru_cache(maxsize=1)
+def load_pipeline_wrapper():
     """
-    Model representing a request for outpainting inference.
+    Load the inpainting pipeline with the specified configuration.
+
+    Returns:
+        pipeline: The loaded inpainting pipeline.
     """
-    controlnet_prompt: str = Field(...)
-    controlnet_negative_prompt: str = Field(...)
-    controlnet_conditioning_scale: float = Field(...)
-    controlnet_guidance_scale: float = Field(...)
-    controlnet_num_inference_steps: int = Field(...)
-    controlnet_guidance_end: float = Field(...)
-    inpainting_prompt: str = Field(...)
-    inpainting_negative_prompt: str = Field(...)
-    inpainting_guidance_scale: float = Field(...)
-    inpainting_strength: float = Field(...)
-    inpainting_num_inference_steps: int = Field(...)
+    pipeline = load_pipeline(cfg.model, accelerator(), enable_compile=True)
+    return pipeline
+inpainting_pipeline = load_pipeline_wrapper()
 
-class OutpaintingBatchRequestModel(BaseModel):
+class InpaintingRequest(BaseModel):
+    """
+    Model representing a request for inpainting inference.
     """
-    Model representing a batch request for outpainting 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")
+    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[OutpaintingRequest]
+    Model representing a batch request for inpainting inference.
+    """
+    requests: List[InpaintingRequest]
 
 async def save_image(image: UploadFile) -> str:
     """
     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)
@@ -51,75 +66,149 @@ async def save_image(image: UploadFile) -> str:
         f.write(await image.read())
     return file_path
 
-def run_inference(image_path: str, request: OutpaintingRequest):
-    pipeline = load_pipeline()
-    result = pipeline.run_pipeline(
-        image_path,
-        controlnet_prompt=request.controlnet_prompt,
-        controlnet_negative_prompt=request.controlnet_negative_prompt,
-        controlnet_conditioning_scale=request.controlnet_conditioning_scale,
-        controlnet_guidance_scale=request.controlnet_guidance_scale,
-        controlnet_num_inference_steps=request.controlnet_num_inference_steps,
-        controlnet_guidance_end=request.controlnet_guidance_end,
-        inpainting_prompt=request.inpainting_prompt,
-        inpainting_negative_prompt=request.inpainting_negative_prompt,
-        inpainting_guidance_scale=request.inpainting_guidance_scale,
-        inpainting_strength=request.inpainting_strength,
-        inpainting_num_inference_steps=request.inpainting_num_inference_steps
+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)
+    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, 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'], 
+                                          cfg['roi_scale'], 
+                                          cfg['segmentation_model'], 
+                                          cfg['detection_model'])
+    else:
+        image = Image.open(image_path)
+        mask_image = None  
+    
+    painting_pipeline = AutoPaintingPipeline(
+        pipeline=inpainting_pipeline,
+        image=image, 
+        mask_image=mask_image, 
+        target_height=cfg['target_height'], 
+        target_width=cfg['target_width']
     )
-    return result
+    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,
+                                    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'.")
 
-@router.post("/outpaint")
-async def outpaint(
+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):
+            result = run_inference(cfg, image_path, request)
+            results.append(result)
+        return results
+
+@router.post("/inpainting")
+async def inpainting_inference(
     image: UploadFile = File(...),
-    request: str = Form(...)
+    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:
-        request_dict = json.loads(request)
-        outpainting_request = OutpaintingRequest(**request_dict)
-        
         image_path = await save_image(image)
-        result = run_inference(image_path, outpainting_request)
-        
-        result_json = pil_to_b64_json(result)
-        
-        os.remove(image_path)
-        
-        return {"result": result_json}
-    except json.JSONDecodeError:
-        raise HTTPException(status_code=400, detail="Invalid JSON in request data")
-    except ValidationError as e:
-        raise HTTPException(status_code=422, detail=str(e))
+        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))
 
-class OutpaintingBatcher(AsyncBatcher):
-    async def process_batch(self, batch):
-        results = []
-        for image, request in batch:
-            image_path = await save_image(image)
-            try:
-                result = run_inference(image_path, request)
-                results.append(result)
-            finally:
-                os.remove(image_path)
-        return results
+@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').
 
-@router.post("/batch_outpaint")
-async def batch_outpaint(images: List[UploadFile] = File(...), batch_request: str = Form(...)):
+    Raises:
+        HTTPException: If there is an error during image processing.
+    """
     try:
-        batch_request_dict = json.loads(batch_request)
-        batch_outpainting_request = OutpaintingBatchRequestModel(**batch_request_dict)
-        
-        batcher = OutpaintingBatcher(max_queue_size=64)
-        results = await batcher.process_batch(list(zip(images, batch_outpainting_request.requests)))
-        
-        result_jsons = [pil_to_b64_json(result) for result in results]
-
-        return {"results": result_jsons}
-    except json.JSONDecodeError:
-        raise HTTPException(status_code=400, detail="Invalid JSON in batch request data")
-    except ValidationError as e:
-        raise HTTPException(status_code=422, detail=str(e))
+        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 = batcher.process_batch((image_paths, requests))
+
+        return results
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))

scripts/kandinsky3_inpainting.py

@@ -0,0 +1,29 @@
+import torch
+from Kandinsky.kandinsky3 import get_inpainting_pipeline
+from scripts.api_utils import ImageAugmentation,accelerator
+from diffusers.utils import load_image
+import numpy as np
+from PIL import Image
+
+device_map = torch.device(accelerator())
+dtype_map = {
+    'unet': torch.float16,
+    'text_encoder': torch.float16,
+    'movq': torch.float32,
+}
+
+
+pipe = get_inpainting_pipeline(
+    device_map, dtype_map,
+)
+
+image = Image.open('/home/PicPilot/sample_data/image.png')
+mask_image = Image.open('/home/PicPilot/sample_data/mask_image.png')
+image = load_image(image=image)
+mask_image = np.array(mask_image)
+
+
+
+
+image = pipe( "Product on the Kitchen used for cooking", image, mask_image)
+image.save('output.jpg')

ui/ui.py

@@ -5,11 +5,18 @@ import json
 from PIL import Image
 from diffusers.utils import load_image
 from io import BytesIO
+<<<<<<< HEAD
 from vars import base_url
 
 # API endpoints
 sdxl_inference_endpoint = f'{base_url}/api/v1/product-diffusion/sdxl_v0_lora_inference'
 kandinsky_inpainting_inference = f'{base_url}/api/v1/product-diffusion/inpainting'
+=======
+
+# API endpoints
+sdxl_inference_endpoint = 'http://localhost:7860/api/v1/product-diffusion/sdxl_v0_lora_inference'
+kandinsky_inpainting_inference = 'http://localhost:7860/api/v1/product-diffusion/inpainting'
+>>>>>>> 41b06fc (commit)
 
 def generate_sdxl_lora_image(prompt, negative_prompt, num_inference_steps, guidance_scale, num_images, mode):
     payload = {