test new `masking_prompt_text` mode

app.py

@@ -1,13 +1,15 @@
 from typing import Tuple
 
+import os
 import requests
 import random
 import numpy as np
 import gradio as gr
 import spaces
 import torch
-from PIL import Image
+from PIL import Image, ImageFilter
 from diffusers import FluxInpaintPipeline
+from gradio_client import Client, handle_file
 
 MARKDOWN = """
 # FLUX.1 Inpainting 🔥
@@ -21,6 +23,9 @@ MAX_SEED = np.iinfo(np.int32).max
 IMAGE_SIZE = 1024
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+client = Client("SkalskiP/florence-sam-masking", hf_token=HF_TOKEN)
+
 
 def remove_background(image: Image.Image, threshold: int = 50) -> Image.Image:
     image = image.convert("RGBA")
@@ -45,6 +50,7 @@ EXAMPLES = [
             "composite": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-composite-2.png", stream=True).raw),
         },
         "little lion",
+        None,
         42,
         False,
         0.85,
@@ -57,6 +63,7 @@ EXAMPLES = [
             "composite": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-composite-3.png", stream=True).raw),
         },
         "tattoos",
+        None,
         42,
         False,
         0.85,
@@ -74,11 +81,6 @@ def resize_image_dimensions(
 ) -> Tuple[int, int]:
     width, height = original_resolution_wh
 
-    # if width <= maximum_dimension and height <= maximum_dimension:
-    #     width = width - (width % 32)
-    #     height = height - (height % 32)
-    #     return width, height
-
     if width > height:
         scaling_factor = maximum_dimension / width
     else:
@@ -93,31 +95,53 @@ def resize_image_dimensions(
     return new_width, new_height
 
 
+def is_image_empty(image: Image.Image) -> bool:
+    gray_img = image.convert("L")
+    pixels = list(gray_img.getdata())
+    return all(pixel == 0 for pixel in pixels)
+
+
 @spaces.GPU(duration=100)
 def process(
     input_image_editor: dict,
-    input_text: str,
+    inpainting_prompt_text: str,
+    masking_prompt_text: str,
     seed_slicer: int,
     randomize_seed_checkbox: bool,
     strength_slider: float,
     num_inference_steps_slider: int,
     progress=gr.Progress(track_tqdm=True)
 ):
-    if not input_text:
+    if not inpainting_prompt_text:
         gr.Info("Please enter a text prompt.")
         return None, None
 
-    image = input_image_editor['background']
-    mask = input_image_editor['layers'][0]
+    image_path = input_image_editor['background']
+    mask_path = input_image_editor['layers'][0]
+
+    image = Image.open(image_path)
+    mask = Image.open(mask_path)
 
     if not image:
         gr.Info("Please upload an image.")
         return None, None
 
-    if not mask:
-        gr.Info("Please draw a mask on the image.")
+    if is_image_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:
+        gr.Info("Both mask and masking prompt are provided. Please provide only one.")
         return None, None
 
+    if is_image_empty(mask):
+        mask = client.predict(
+            image_input=handle_file(image_path),
+            text_input=masking_prompt_text,
+            api_name="/process_image")
+        mask = Image.open(mask)
+
+    mask = mask.filter(ImageFilter.GaussianBlur(radius=5))
     width, height = resize_image_dimensions(original_resolution_wh=image.size)
     resized_image = image.resize((width, height), Image.LANCZOS)
     resized_mask = mask.resize((width, height), Image.LANCZOS)
@@ -126,7 +150,7 @@ def process(
         seed_slicer = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed_slicer)
     result = pipe(
-        prompt=input_text,
+        prompt=inpainting_prompt_text,
         image=resized_image,
         mask_image=resized_mask,
         width=width,
@@ -145,24 +169,31 @@ with gr.Blocks() as demo:
         with gr.Column():
             input_image_editor_component = gr.ImageEditor(
                 label='Image',
-                type='pil',
+                type='filepath',
                 sources=["upload", "webcam"],
                 image_mode='RGB',
                 layers=False,
                 brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed"))
 
             with gr.Row():
-                input_text_component = gr.Text(
+                inpainting_prompt_text_component = gr.Text(
                     label="Prompt",
                     show_label=False,
                     max_lines=1,
-                    placeholder="Enter your prompt",
+                    placeholder="Enter text to generate inpainting",
                     container=False,
                 )
                 submit_button_component = gr.Button(
                     value='Submit', variant='primary', scale=0)
 
             with gr.Accordion("Advanced Settings", open=False):
+                masking_prompt_text_component = gr.Text(
+                    label="Prompt",
+                    show_label=False,
+                    max_lines=1,
+                    placeholder="Enter text to generate masking",
+                    container=False,
+                )
                 seed_slicer_component = gr.Slider(
                     label="Seed",
                     minimum=0,
@@ -207,7 +238,8 @@ with gr.Blocks() as demo:
             examples=EXAMPLES,
             inputs=[
                 input_image_editor_component,
-                input_text_component,
+                inpainting_prompt_text_component,
+                masking_prompt_text_component,
                 seed_slicer_component,
                 randomize_seed_checkbox_component,
                 strength_slider_component,
@@ -225,7 +257,8 @@ with gr.Blocks() as demo:
         fn=process,
         inputs=[
             input_image_editor_component,
-            input_text_component,
+            inpainting_prompt_text_component,
+            masking_prompt_text_component,
             seed_slicer_component,
             randomize_seed_checkbox_component,
             strength_slider_component,