add-blend

runway.py

@@ -26,34 +26,48 @@ async def root():
     """
     return {"message": "InstructPix2Pix API is running. Use POST /inpaint/ or /inpaint-with-reference/ to edit images."}
 
-def blend_images(original_image: Image, reference_image: Image, mask_image: Image) -> Image:
+def prepare_guided_image(original_image: Image, reference_image: Image, mask_image: Image) -> Image:
     """
-    Blend the original image with a reference image using the mask.
-    - Unmasked areas (white in mask) take pixels from the original image.
-    - Masked areas (black in mask) take pixels from the reference image as a starting point.
+    Prepare an initial image by softly blending the reference image into the masked area.
+    - Unmasked areas (white in mask, 255) remain fully from the original image.
+    - Masked areas (black in mask, 0) take content from the reference image with soft blending.
     
     Args:
         original_image (Image): The original image (RGB).
-        reference_image (Image): The reference image to blend with (RGB).
+        reference_image (Image): The reference image to copy from (RGB).
         mask_image (Image): The mask image (grayscale, L mode).
     
     Returns:
-        Image: The blended image.
+        Image: The blended image to guide inpainting.
     """
     # Convert images to numpy arrays
     original_array = np.array(original_image)
     reference_array = np.array(reference_image)
-    mask_array = np.array(mask_image) / 255.0  # Normalize mask to [0, 1]
+    mask_array = np.array(mask_image) / 255.0  # Normalize to [0, 1] for soft blending
 
-    # Ensure mask is broadcastable to RGB channels
+    # Expand mask to RGB channels
     mask_array = mask_array[:, :, np.newaxis]
 
-    # Blend: unmasked areas (mask=1) keep original, masked areas (mask=0) use reference
+    # Softly blend: unmasked areas (1) keep original, masked areas (0) use reference
     blended_array = original_array * mask_array + reference_array * (1 - mask_array)
     blended_array = blended_array.astype(np.uint8)
 
     return Image.fromarray(blended_array)
 
+def soften_mask(mask_image: Image, softness: int = 5) -> Image:
+    """
+    Soften the edges of the mask for smoother transitions.
+    
+    Args:
+        mask_image (Image): The original mask (grayscale, L mode).
+        softness (int): Size of the Gaussian blur kernel for softening edges.
+    
+    Returns:
+        Image: The softened mask.
+    """
+    from PIL import ImageFilter
+    return mask_image.filter(ImageFilter.GaussianBlur(radius=softness))
+
 @app.post("/inpaint/")
 async def inpaint_image(
     image: UploadFile = File(...),
@@ -104,17 +118,17 @@ async def inpaint_with_reference(
     image: UploadFile = File(...),
     mask: UploadFile = File(...),
     reference_image: UploadFile = File(...),
-    prompt: str = "Fill the masked area with appropriate content."
+    prompt: str = "Integrate the reference content naturally into the masked area, matching style and lighting."
 ):
     """
-    Endpoint for image inpainting using both a text prompt and a reference image.
+    Endpoint for replacing masked areas with reference image content, refined to look natural.
     - `image`: Original image file (PNG/JPG).
-    - `mask`: Mask file indicating areas to inpaint (black for masked areas, white for unmasked).
-    - `reference_image`: Reference image to guide the inpainting (PNG/JPG).
-    - `prompt`: Text prompt describing the desired output.
+    - `mask`: Mask file (black for areas to replace, white for areas to keep).
+    - `reference_image`: Reference image to guide the replacement (PNG/JPG).
+    - `prompt`: Text prompt for inpainting refinement.
     
     Returns:
-    - The inpainted image as a PNG file.
+    - The resulting image as a PNG file.
     """
     try:
         # Load the uploaded image, mask, and reference image
@@ -132,11 +146,20 @@ async def inpaint_with_reference(
         if original_image.size != reference_image.size:
             reference_image = reference_image.resize(original_image.size, Image.Resampling.LANCZOS)
 
-        # Blend the original and reference images using the mask
-        blended_image = blend_images(original_image, reference_image, mask_image)
+        # Soften the mask for smoother transitions
+        softened_mask = soften_mask(mask_image, softness=5)
 
-        # Perform inpainting using the pipeline
-        result = pipe(prompt=prompt, image=blended_image, mask_image=mask_image).images[0]
+        # Prepare the initial guided image by blending reference content into the masked area
+        guided_image = prepare_guided_image(original_image, reference_image, softened_mask)
+
+        # Perform inpainting to refine the result and make it look natural
+        result = pipe(
+            prompt=prompt,
+            image=guided_image,
+            mask_image=softened_mask,  # Use softened mask for inpainting
+            strength=0.75,  # Control how much inpainting modifies the image (0.0 to 1.0)
+            guidance_scale=7.5  # Control how closely the prompt is followed
+        ).images[0]
 
         # Convert result to bytes for response
         result_bytes = io.BytesIO()
@@ -147,11 +170,11 @@ async def inpaint_with_reference(
         return StreamingResponse(
             result_bytes,
             media_type="image/png",
-            headers={"Content-Disposition": "attachment; filename=inpainted_with_reference_image.png"}
+            headers={"Content-Disposition": "attachment; filename=natural_inpaint_image.png"}
         )
     
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error during inpainting with reference: {e}")
+        raise HTTPException(status_code=500, detail=f"Error during natural inpainting: {e}")
 
 if __name__ == "__main__":
     import uvicorn