add- inpainint

intruct.py

@@ -2,9 +2,9 @@ from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.responses import StreamingResponse
 import io
 import math
-from PIL import Image, ImageOps
+from PIL import Image, ImageOps, ImageDraw
 import torch
-from diffusers import StableDiffusionInstructPix2PixPipeline
+from diffusers import StableDiffusionInstructPix2PixPipeline, StableDiffusionInpaintPipeline
 from fastapi import FastAPI, Response
 from fastapi.responses import FileResponse
 import torch
@@ -13,27 +13,30 @@ from huggingface_hub import hf_hub_download, login
 from safetensors.torch import load_file
 from io import BytesIO
 import os
-import base64  # Added for encoding images as base64
-from typing import List  # Added for type hinting the list of prompts
-
-
+import base64
+from typing import List
 
 # Initialize FastAPI app
 app = FastAPI()
 
-# Load the pre-trained model once at startup
+# Load the pre-trained InstructPix2Pix model for editing
 model_id = "timbrooks/instruct-pix2pix"
-pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+pipe_edit = StableDiffusionInstructPix2PixPipeline.from_pretrained(
     model_id, torch_dtype=torch.float16, safety_checker=None
 ).to("cuda")
 
+# Load the pre-trained Inpainting model
+inpaint_model_id = "stabilityai/stable-diffusion-2-inpainting"
+pipe_inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+    inpaint_model_id, torch_dtype=torch.float16, safety_checker=None
+).to("cuda")
+
 # Default configuration values
 DEFAULT_STEPS = 50
 DEFAULT_TEXT_CFG = 7.5
 DEFAULT_IMAGE_CFG = 1.5
 DEFAULT_SEED = 1371
 
-
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 def load_model():
@@ -73,18 +76,16 @@ def load_model():
 
 # Load model at startup with error handling
 try:
-    pipe = load_model()
+    pipe_generate = load_model()
 except Exception as e:
     print(f"Model initialization failed: {str(e)}")
     raise
 
-
-
 @app.get("/generate")
 async def generate_image(prompt: str):
     try:
         # Generate image
-        image = pipe(
+        image = pipe_generate(
             prompt,
             num_inference_steps=4,
             guidance_scale=0
@@ -100,7 +101,6 @@ async def generate_image(prompt: str):
     except Exception as e:
         return {"error": str(e)}
 
-# New endpoint to handle a list of prompts
 @app.get("/generate_multiple")
 async def generate_multiple_images(prompts: List[str]):
     try:
@@ -109,7 +109,7 @@ async def generate_multiple_images(prompts: List[str]):
         
         # Generate an image for each prompt
         for prompt in prompts:
-            image = pipe(
+            image = pipe_generate(
                 prompt,
                 num_inference_steps=4,
                 guidance_scale=0
@@ -136,8 +136,6 @@ async def generate_multiple_images(prompts: List[str]):
 async def health_check():
     return {"status": "healthy"}
 
-
-
 def process_image(input_image: Image.Image, instruction: str, steps: int, text_cfg_scale: float, image_cfg_scale: float, seed: int):
     """
     Process the input image with the given instruction using InstructPix2Pix.
@@ -157,7 +155,7 @@ def process_image(input_image: Image.Image, instruction: str, steps: int, text_c
     generator = torch.manual_seed(seed)
 
     # Generate the edited image
-    edited_image = pipe(
+    edited_image = pipe_edit(
         instruction,
         image=input_image,
         guidance_scale=text_cfg_scale,
@@ -179,12 +177,6 @@ async def edit_image(
 ):
     """
     Endpoint to edit an image based on a text instruction.
-    - file: The input image to edit.
-    - instruction: The text instruction for editing the image.
-    - steps: Number of inference steps.
-    - text_cfg_scale: Text CFG weight.
-    - image_cfg_scale: Image CFG weight.
-    - seed: Random seed for reproducibility.
     """
     # Read and convert the uploaded image
     image_data = await file.read()
@@ -201,12 +193,86 @@ async def edit_image(
     # Return the image as a streaming response
     return StreamingResponse(img_byte_arr, media_type="image/png")
 
+# New endpoint for inpainting
+@app.post("/inpaint/")
+async def inpaint_image(
+    file: UploadFile = File(...),
+    prompt: str = Form(...),
+    mask_coordinates: str = Form(...),  # Format: "x1,y1,x2,y2" (top-left and bottom-right of the rectangle to inpaint)
+    steps: int = Form(default=DEFAULT_STEPS),
+    guidance_scale: float = Form(default=7.5),
+    seed: int = Form(default=DEFAULT_SEED)
+):
+    """
+    Endpoint to perform inpainting on an image.
+    - file: The input image to inpaint.
+    - prompt: The text prompt describing what to generate in the inpainted area.
+    - mask_coordinates: Coordinates of the rectangular area to inpaint (format: "x1,y1,x2,y2").
+    - steps: Number of inference steps.
+    - guidance_scale: Guidance scale for the inpainting process.
+    - seed: Random seed for reproducibility.
+    """
+    try:
+        # Read and convert the uploaded image
+        image_data = await file.read()
+        input_image = Image.open(io.BytesIO(image_data)).convert("RGB")
+
+        # Resize image to fit model requirements (must be divisible by 8 for inpainting)
+        width, height = input_image.size
+        factor = 512 / max(width, height)
+        factor = math.ceil(min(width, height) * factor / 8) * 8 / min(width, height)
+        width = int((width * factor) // 8) * 8
+        height = int((height * factor) // 8) * 8
+        input_image = ImageOps.fit(input_image, (width, height), method=Image.Resampling.LANCZOS)
+
+        # Create a mask for inpainting
+        mask = Image.new("L", (width, height), 0)  # Black image (0 = no inpainting)
+        draw = ImageDraw.Draw(mask)
+
+        # Parse the mask coordinates
+        try:
+            x1, y1, x2, y2 = map(int, mask_coordinates.split(","))
+            # Adjust coordinates based on resized image
+            x1 = int(x1 * factor)
+            y1 = int(y1 * factor)
+            x2 = int(x2 * factor)
+            y2 = int(y2 * factor)
+        except ValueError:
+            return {"error": "Invalid mask coordinates format. Use 'x1,y1,x2,y2'."}
+
+        # Draw a white rectangle on the mask (255 = area to inpaint)
+        draw.rectangle([x1, y1, x2, y2], fill=255)
+
+        # Set the random seed for reproducibility
+        generator = torch.manual_seed(seed)
+
+        # Perform inpainting
+        inpainted_image = pipe_inpaint(
+            prompt=prompt,
+            image=input_image,
+            mask_image=mask,
+            num_inference_steps=steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+        ).images[0]
+
+        # Convert the inpainted image to bytes
+        img_byte_arr = io.BytesIO()
+        inpainted_image.save(img_byte_arr, format="PNG")
+        img_byte_arr.seek(0)
+
+        # Return the image as a streaming response
+        return StreamingResponse(img_byte_arr, media_type="image/png")
+
+    except Exception as e:
+        return {"error": str(e)}
+
 @app.get("/")
 async def root():
     """
     Root endpoint for basic health check.
     """
-    return {"message": "InstructPix2Pix API is running. Use POST /edit-image/ to edit images."}
+    return {"message": "InstructPix2Pix API is running. Use POST /edit-image/ or /inpaint/ to edit images."}
 
 if __name__ == "__main__":
     import uvicorn