commit

Former-commit-id: 522ceaeca59f5a0db3c0e1e300b6ac98b69656af

api/routers/painting.py

@@ -1,22 +1,20 @@
-from pathlib import Path
 import os
 import uuid
 from typing import List, Tuple, Any, Dict
-from fastapi import APIRouter, File, UploadFile, HTTPException, Form, Depends, Body
+from fastapi import APIRouter, File, UploadFile, HTTPException, Form
 from pydantic import BaseModel, Field
 from PIL import Image
 import lightning.pytorch as pl
 from scripts.api_utils import pil_to_s3_json, pil_to_b64_json, ImageAugmentation, accelerator
-from inpainting_pipeline import AutoPaintingPipeline, load_pipeline
+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
-import asyncio
 pl.seed_everything(42)
 router = APIRouter()
 
-# Initialize Hydra configuration
+
 with initialize(version_base=None, config_path="../../configs"):
     cfg = compose(config_name="inpainting")
 
@@ -31,7 +29,6 @@ def load_pipeline_wrapper():
     """
     pipeline = load_pipeline(cfg.model, accelerator(), enable_compile=True)
     return pipeline
-
 inpainting_pipeline = load_pipeline_wrapper()
 
 class InpaintingRequest(BaseModel):
@@ -45,7 +42,8 @@ class InpaintingRequest(BaseModel):
     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):
     """
     Model representing a batch request for inpainting inference.
@@ -68,14 +66,35 @@ async def save_image(image: UploadFile) -> str:
         f.write(await image.read())
     return file_path
 
-def run_inference(cfg, image_path: str, mask_image_path: str, request: InpaintingRequest):
+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.
-        mask_image_path (str): Path to the mask image file.
         request (InpaintingRequest): Pydantic model containing inference parameters.
 
     Returns:
@@ -84,8 +103,17 @@ def run_inference(cfg, image_path: str, mask_image_path: str, request: Inpaintin
     Raises:
         ValueError: If an invalid mode is provided.
     """
-    image = Image.open(image_path)
-    mask_image = Image.open(mask_image_path)
+    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, 
@@ -107,33 +135,26 @@ def run_inference(cfg, image_path: str, mask_image_path: str, request: Inpaintin
         raise ValueError("Invalid mode. Supported modes are 'b64_json' and 's3_json'.")
 
 class InpaintingBatcher(AsyncBatcher):
-    def __init__(self, max_batch_size: int):
-        super().__init__(max_batch_size)
-
-    async def process_batch(self, batch: Tuple[List[str], List[str], List[InpaintingRequest]]) -> List[Dict[str, Any]]:
+    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[str], List[InpaintingRequest]]): Tuple of image paths, mask image paths, and corresponding requests.
+            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, mask_image_paths, requests = batch
+        image_paths, requests = batch
         results = []
-        for image_path, mask_image_path, request in zip(image_paths, mask_image_paths, requests):
-            try:
-                result = run_inference(cfg, image_path, mask_image_path, request)
-                results.append(result)
-            except Exception as e:
-                results.append({"error": str(e)})
+        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(...),
-    mask_image: UploadFile = File(...),
     request_data: str = Form(...),
 ):
     """
@@ -141,7 +162,6 @@ async def inpainting_inference(
 
     Args:
         image (UploadFile): Uploaded image file.
-        mask_image (UploadFile): Uploaded mask image file.
         request_data (str): JSON string of the request parameters.
 
     Returns:
@@ -152,10 +172,9 @@ async def inpainting_inference(
     """
     try:
         image_path = await save_image(image)
-        mask_image_path = await save_image(mask_image)
         request_dict = json.loads(request_data)
         request = InpaintingRequest(**request_dict)
-        result = run_inference(cfg, image_path, mask_image_path, request)
+        result = run_inference(cfg, image_path, request)
         return result
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -163,7 +182,6 @@ async def inpainting_inference(
 @router.post("/inpainting/batch")
 async def inpainting_batch_inference(
     images: List[UploadFile] = File(...),
-    mask_images: List[UploadFile] = File(...),
     request_data: str = Form(...),
 ):
     """
@@ -171,7 +189,6 @@ async def inpainting_batch_inference(
 
     Args:
         images (List[UploadFile]): List of uploaded image files.
-        mask_images (List[UploadFile]): List of uploaded mask image files.
         request_data (str): JSON string of the request parameters.
 
     Returns:
@@ -185,13 +202,12 @@ async def inpainting_batch_inference(
         batch_request = InpaintingBatchRequestModel(**request_dict)
         requests = batch_request.requests
 
-        if len(images) != len(requests) or len(images) != len(mask_images):
-            raise HTTPException(status_code=400, detail="The number of images, mask images, and requests must match.")
+        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 = [save_image(image) for image in images]
-        mask_image_paths = [save_image(mask_image) for mask_image in mask_images]
-        results = await batcher.process_batch((image_paths, mask_image_paths, requests))
+        image_paths = [await save_image(image) for image in images]
+        results = batcher.process_batch((image_paths, requests))
 
         return results
     except Exception as e:

api/routers/sdxl_text_to_image.py

@@ -1,4 +1,4 @@
-import config
+import scripts.config as config
 from fastapi import APIRouter, HTTPException
 from typing import List
 from diffusers import DiffusionPipeline

api/yolov8l.pt.REMOVED.git-id

@@ -0,0 +1 @@
+85f4cb9e9454a5944aebcc4ac73fd5571f3ddcd3

picpilot.egg-info/SOURCES.txt

@@ -16,8 +16,9 @@ scripts/config.py
 scripts/inpainting_pipeline.py
 scripts/kandinsky3_inpainting.py
 scripts/load_pipeline.py
-scripts/logger.py
 scripts/products10k_captions.py
 scripts/s3_manager.py
 scripts/sdxl_lora_inference.py
-scripts/sdxl_lora_tuner.py
+scripts/sdxl_lora_tuner.py
+ui/__init__.py
+ui/ui.py

picpilot.egg-info/top_level.txt

@@ -1,2 +1,3 @@
 api
 scripts
+ui

scripts/api_utils.py

@@ -3,10 +3,10 @@ from ultralytics import YOLO
 from transformers import SamModel, SamProcessor
 import numpy as np
 from PIL import Image, ImageOps
-from config import SEGMENTATION_MODEL_NAME, DETECTION_MODEL_NAME
+from scripts.config import SEGMENTATION_MODEL_NAME, DETECTION_MODEL_NAME
 from diffusers.utils import load_image
 import gc
-from s3_manager import S3ManagerService
+from scripts.s3_manager import S3ManagerService
 import io
 from io import BytesIO
 import base64

scripts/inpainting_pipeline.py

@@ -1,7 +1,7 @@
 import torch
 from diffusers import AutoPipelineForInpainting,DiffusionPipeline
 from diffusers.utils import load_image
-from api_utils import accelerator, ImageAugmentation
+from scripts.api_utils import accelerator, ImageAugmentation
 import hydra
 from omegaconf import DictConfig
 from PIL import Image

scripts/kandinsky3_inpainting.py

@@ -1,12 +1,11 @@
-import sys
 import torch
-sys.path.append('../')
 from Kandinsky.kandinsky3 import get_inpainting_pipeline
-from api_utils import ImageAugmentation
+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('cuda:0')
+device_map = torch.device(accelerator())
 dtype_map = {
     'unet': torch.float16,
     'text_encoder': torch.float16,
@@ -18,12 +17,13 @@ pipe = get_inpainting_pipeline(
     device_map, dtype_map,
 )
 
-augmenter = ImageAugmentation(target_width=2560, target_height=1440)
-image = Image.open(image_path='/home/product_diffusion_api/sample_data/example1.jpg')
-extended_image = augmenter.extend_image(image)
-mask_image = augmenter.generate_mask_from_bbox(extended_image, segmentation_model='facebook/sam-vit-base', detection_model='yolov8s')
-mask_image = augmenter.invert_mask(mask_image)
+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", extended_image, mask_image)
+
+
+image = pipe( "Product on the Kitchen used for cooking", image, mask_image)
 image.save('output.jpg')

scripts/load_pipeline.py

@@ -2,7 +2,7 @@ from config import MODEL_NAME,ADAPTER_NAME
 import torch
 from diffusers import DiffusionPipeline
 from wandb.integration.diffusers import autolog
-from config import PROJECT_NAME
+from scripts.config import PROJECT_NAME
 autolog(init=dict(project=PROJECT_NAME))
 
 

ui/ui.py

@@ -1,92 +1,41 @@
 import gradio as gr
 import numpy as np
 import requests
-from io import BytesIO
 import json
 from PIL import Image
-from pydantic import BaseModel, Field
 from diffusers.utils import load_image
+from io import BytesIO
 
 # API endpoints
-sdxl_inference_endpoint = 'https://vikramsingh178-picpilot-server.hf.space/api/v1/product-diffusion/sdxl_v0_lora_inference'
-kandinsky_inpainting_inference = 'https://vikramsingh178-picpilot-server.hf.space/api/v1/product-diffusion/inpainting'
-
-class InputRequestSDXL(BaseModel):
-    prompt: str
-    num_inference_steps: int
-    guidance_scale: float
-    negative_prompt: str
-    num_images: int
-    mode: str
-
-class InpaintingRequestKandinsky(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")
+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'
 
 def generate_sdxl_lora_image(prompt, negative_prompt, num_inference_steps, guidance_scale, num_images, mode):
-    payload = InputRequestSDXL(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        num_inference_steps=num_inference_steps,
-        guidance_scale=guidance_scale,
-        num_images=num_images,
-        mode=mode
-    ).dict()
-    
-    try:
-        response = requests.post(sdxl_inference_endpoint, json=payload)
-        response.raise_for_status()
-        response_data = response.json()
-        url = response_data['url']
-        image = load_image(url)
-        return image
-    except requests.exceptions.RequestException as e:
-        print(f"Error in SDXL-Lora API request: {e}")
-        return None
-
-def process_masked_image(img):
-    if img is None or "composite" not in img:
-        return None, None
-    
-    base_image = Image.fromarray(img["composite"]).convert("RGB")
-    
-    if "layers" in img and len(img["layers"]) > 0:
-        alpha_channel = img["layers"][0][:, :, 3]
-        mask = np.where(alpha_channel == 0, 0, 255).astype(np.uint8)
-        mask = Image.fromarray(mask).convert("L")
-    else:
-        mask = Image.new("L", base_image.size, 0)
-    
-    return base_image, mask
-
-def generate_outpainting(prompt, negative_prompt, num_inference_steps, strength, guidance_scale, mode, num_images, masked_image, width, height):
-    base_image, mask = process_masked_image(masked_image)
-    
-    if base_image is None or mask is None:
-        return None, None
-    
-    # Resize base image and mask
-    base_image_resized = base_image.resize((width, height))
-    mask_resized = mask.resize((width, height))
+    payload = {
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "num_inference_steps": num_inference_steps,
+        "guidance_scale": guidance_scale,
+        "num_images": num_images,
+        "mode": mode
+    }
     
-    # Convert the resized images to bytes
+    response = requests.post(sdxl_inference_endpoint, json=payload)
+    response.raise_for_status()
+    response_data = response.json()
+    url = response_data['url']
+    image = load_image(url)
+    return image
+
+def generate_outpainting(prompt, negative_prompt, num_inference_steps, strength, guidance_scale, mode, num_images, image, width, height):
+    # Convert the image to bytes
     img_byte_arr = BytesIO()
-    base_image_resized.save(img_byte_arr, format='PNG')
+    image.save(img_byte_arr, format='PNG')
     img_byte_arr = img_byte_arr.getvalue()
 
-    mask_byte_arr = BytesIO()
-    mask_resized.save(mask_byte_arr, format='PNG')
-    mask_byte_arr = mask_byte_arr.getvalue()
-    
     # Prepare the files for multipart/form-data
     files = {
-        'image': ('image.png', img_byte_arr, 'image/png'),
-        'mask_image': ('mask.png', mask_byte_arr, 'image/png'),
+        'image': ('image.png', img_byte_arr, 'image/png')
     }
 
     # Prepare the request data
@@ -97,7 +46,9 @@ def generate_outpainting(prompt, negative_prompt, num_inference_steps, strength,
         "strength": strength,
         "guidance_scale": guidance_scale,
         "mode": mode,
-        "num_images": num_images
+        "num_images": num_images,
+        "width": width,
+        "height": height
     }
 
     # Convert request_data to a JSON string
@@ -108,30 +59,15 @@ def generate_outpainting(prompt, negative_prompt, num_inference_steps, strength,
         'request_data': request_data_json
     }
     
-    try:
-        response = requests.post(kandinsky_inpainting_inference, files=files, data=form_data)
-        response.raise_for_status()
-        response_data = response.json()
-        url = response_data['url']
-        outpainted_image = load_image(url)
-        return mask_resized, outpainted_image
-    except requests.exceptions.RequestException as e:
-        print(f"Error in Kandinsky Inpainting API request: {e}")
-        return None, None
-
-def generate_mask_preview(img):
-    base_image, mask = process_masked_image(img)
-    if mask is None:
-        return None
-    return mask
-
-def resize_image(img, width, height):
-    if img is None:
-        return None
-    resized_img = img.resize((width, height))
-    return resized_img
+    response = requests.post(kandinsky_inpainting_inference, files=files, data=form_data)
+    response.raise_for_status()
+    response_data = response.json()
+    image_url = response_data['image_url']
+    mask_url = response_data['mask_url']
+    outpainted_image = load_image(image_url)
+    mask_image = load_image(mask_url)
+    return outpainted_image, mask_image
 
-# Gradio interface setup
 with gr.Blocks(theme='VikramSingh178/Webui-Theme') as demo:
     with gr.Tab("SdxL-Lora"):
         with gr.Row():
@@ -149,32 +85,25 @@ with gr.Blocks(theme='VikramSingh178/Webui-Theme') as demo:
                 image_preview = gr.Image(label="Generated Image (SDXL-Lora)", show_download_button=True, show_share_button=True, container=True)
                 generate_button.click(generate_sdxl_lora_image, inputs=[prompt, negative_prompt, num_inference_steps, guidance_scale, num_images, mode], outputs=[image_preview])
                 
-
     with gr.Tab("Inpainting"):
         with gr.Row():
             with gr.Column():
                 with gr.Group():
-                    masked_image = gr.ImageMask(label="Upload Image and Draw Mask", format='png')
+                    input_image = gr.Image(label="Upload Image", type="pil")
                     prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here")
                     negative_prompt= gr.Textbox(label="Negative Prompt", placeholder="Enter negative prompt here")
                     num_inference_steps = gr.Slider(minimum=1, maximum=100, step=1, value=20, label="Inference Steps")
                     strength = gr.Slider(minimum=0.1, maximum=1, step=0.1, value=0.8, label="Strength")
                     guidance_scale = gr.Slider(minimum=1.0, maximum=10.0, step=0.1, value=7.5, label="Guidance Scale")
-                    num_images= gr.Slider(minimum=1, maximum=10, step=1, value=1, label="Number of Images")
+                    num_images = gr.Slider(minimum=1, maximum=10, step=1, value=1, label="Number of Images")
                     mode_kandinsky = gr.Dropdown(choices=["s3_json", "b64_json"], value="s3_json", label="Mode")
                     width_slider = gr.Slider(minimum=512, maximum=1024, step=1, value=800, label="Image Width")
                     height_slider = gr.Slider(minimum=512, maximum=1024, step=1, value=800, label="Image Height")
-                    resize_button = gr.Button("Resize Image", variant='secondary')
                     generate_button = gr.Button("Generate Inpainting", variant='primary')
-                    generate_mask_button_painting = gr.Button("Generate Mask", variant='primary')
 
             with gr.Column(scale=1):
-                mask_preview= gr.Image(label="Mask Preview", show_download_button=True, container=True)
                 outpainted_image_preview = gr.Image(label="Outpainted Image (Kandinsky)", show_download_button=True, show_share_button=True, container=True)
-                resize_button.click(resize_image, inputs=[masked_image, width_slider, height_slider], outputs=[masked_image])
-                generate_mask_button_painting.click(generate_mask_preview, inputs=masked_image, outputs=[mask_preview])
-                generate_button.click(generate_outpainting, 
-                                      inputs=[prompt, negative_prompt, num_inference_steps, strength, guidance_scale, mode_kandinsky, num_images, masked_image, width_slider, height_slider], 
-                                      outputs=[mask_preview, outpainted_image_preview])
+                mask_image_preview = gr.Image(label="Generated Mask", show_download_button=True, show_share_button=True, container=True)
+                generate_button.click(generate_outpainting, inputs=[prompt, negative_prompt, num_inference_steps, strength, guidance_scale, mode_kandinsky, num_images, input_image, width_slider, height_slider], outputs=[outpainted_image_preview, mask_image_preview])
 
-demo.launch()
+demo.launch()