inflation and mask blur sliders fixed

app.py

@@ -1,11 +1,12 @@
+import cv2
 import random
-from typing import Tuple
+from typing import Tuple, Optional
 
 import gradio as gr
 import numpy as np
 import spaces
 import torch
-from PIL import Image, ImageFilter, ImageOps
+from PIL import Image, ImageFilter
 from diffusers import FluxInpaintPipeline
 from gradio_client import Client, handle_file
 
@@ -20,23 +21,24 @@ for taking it to the next level by enabling inpainting with the FLUX.
 MAX_SEED = np.iinfo(np.int32).max
 IMAGE_SIZE = 1024
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+PIPE = FluxInpaintPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(DEVICE)
+CLIENT = Client("SkalskiP/florence-sam-masking")
 
-client = Client("SkalskiP/florence-sam-masking")
-
-
-def remove_background(image: Image.Image, threshold: int = 50) -> Image.Image:
-    image = image.convert("RGBA")
-    data = image.getdata()
-    new_data = []
-    for item in data:
-        avg = sum(item[:3]) / 3
-        if avg < threshold:
-            new_data.append((0, 0, 0, 0))
-        else:
-            new_data.append(item)
 
-    image.putdata(new_data)
-    return image
+# def remove_background(image: Image.Image, threshold: int = 50) -> Image.Image:
+#     image = image.convert("RGBA")
+#     data = image.getdata()
+#     new_data = []
+#     for item in data:
+#         avg = sum(item[:3]) / 3
+#         if avg < threshold:
+#             new_data.append((0, 0, 0, 0))
+#         else:
+#             new_data.append(item)
+#
+#     image.putdata(new_data)
+#     return image
 
 
 # EXAMPLES = [
@@ -68,11 +70,8 @@ def remove_background(image: Image.Image, threshold: int = 50) -> Image.Image:
 #     ]
 # ]
 
-pipe = FluxInpaintPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(DEVICE)
-
 
-def resize_image_dimensions(
+def calculate_image_dimensions_for_flux(
     original_resolution_wh: Tuple[int, int],
     maximum_dimension: int = IMAGE_SIZE
 ) -> Tuple[int, int]:
@@ -92,12 +91,40 @@ def resize_image_dimensions(
     return new_width, new_height
 
 
-def is_image_empty(image: Image.Image) -> bool:
+def is_mask_empty(image: Image.Image) -> bool:
     gray_img = image.convert("L")
     pixels = list(gray_img.getdata())
     return all(pixel == 0 for pixel in pixels)
 
 
+def process_mask(
+    mask: Image.Image,
+    mask_inflation: Optional[int] = None,
+    mask_blur: Optional[int] = None
+) -> Image.Image:
+    """
+    Inflates and blurs the white regions of a mask.
+
+    Args:
+        mask (Image.Image): The input mask image.
+        mask_inflation (Optional[int]): The number of pixels to inflate the mask by.
+        mask_blur (Optional[int]): The radius of the Gaussian blur to apply.
+
+    Returns:
+        Image.Image: The processed mask with inflated and/or blurred regions.
+    """
+    if mask_inflation and mask_inflation > 0:
+        mask_array = np.array(mask)
+        kernel = np.ones((mask_inflation, mask_inflation), np.uint8)
+        mask_array = cv2.dilate(mask_array, kernel, iterations=1)
+        mask = Image.fromarray(mask_array)
+
+    if mask_blur and mask_blur > 0:
+        mask = mask.filter(ImageFilter.GaussianBlur(radius=mask_blur))
+
+    return mask
+
+
 @spaces.GPU(duration=100)
 def process(
     input_image_editor: dict,
@@ -125,26 +152,29 @@ def process(
         gr.Info("Please upload an image.")
         return None, None
 
-    if is_image_empty(mask) and not masking_prompt_text:
+    if is_mask_empty(mask) and not masking_prompt_text:
         gr.Info("Please draw a mask or enter a masking prompt.")
         return None, None
 
-    if not is_image_empty(mask) and masking_prompt_text:
+    if not is_mask_empty(mask) and masking_prompt_text:
         gr.Info("Both mask and masking prompt are provided. Please provide only one.")
         return None, None
 
-    if is_image_empty(mask):
-        mask = client.predict(
+    if is_mask_empty(mask):
+        mask = CLIENT.predict(
             image_input=handle_file(image_path),
             text_input=masking_prompt_text,
             api_name="/process_image")
         mask = Image.open(mask)
 
-    width, height = resize_image_dimensions(original_resolution_wh=image.size)
+    width, height = calculate_image_dimensions_for_flux(original_resolution_wh=image.size)
     image = image.resize((width, height), Image.LANCZOS)
     mask = mask.resize((width, height), Image.LANCZOS)
     if mask_inflation_slider:
-        mask = ImageOps.expand(mask, border=mask_inflation_slider, fill=255)
+        mask_array = np.array(mask)
+        kernel = np.ones((mask_inflation_slider, mask_inflation_slider), np.uint8)
+        mask_array = cv2.dilate(mask_array, kernel, iterations=1)
+        mask = Image.fromarray(mask_array)
     if mask_blur_slider:
         mask = mask.filter(ImageFilter.GaussianBlur(radius=mask_blur_slider))