Update app.py

app.py

@@ -5,14 +5,14 @@ import numpy as np
 import gradio as gr
 import spaces
 import torch
-from PIL import Image, UnidentifiedImageError
+from PIL import Image
 from diffusers import FluxInpaintPipeline
+from io import BytesIO
 
 MARKDOWN = """
 # FLUX.1 Inpainting 🔥
-Shoutout to [Black Forest Labs](https://huggingface.co/black-forest-labs) team for 
-creating this amazing model, and a big thanks to [Gothos](https://github.com/Gothos) 
-for taking it to the next level by enabling inpainting with the FLUX.
+
+Shoutout to [Black Forest Labs](https://huggingface.co/black-forest-labs) for creating this model and [Gothos](https://github.com/Gothos) for adding inpainting support to FLUX.
 """
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -20,8 +20,19 @@ IMAGE_SIZE = 1024
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 
+def load_image(url: str) -> Image.Image:
+    try:
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+        return Image.open(BytesIO(response.content)).convert("RGBA")
+    except requests.exceptions.RequestException as e:
+        print(f"Error loading image from {url}: {e}")
+        return None
+
+
 def remove_background(image: Image.Image, threshold: int = 50) -> Image.Image:
-    image = image.convert("RGBA")
+    if not image:
+        return None
     data = image.getdata()
     new_data = []
     for item in data:
@@ -30,28 +41,10 @@ def remove_background(image: Image.Image, threshold: int = 50) -> Image.Image:
             new_data.append((0, 0, 0, 0))
         else:
             new_data.append(item)
-
     image.putdata(new_data)
     return image
 
 
-def load_image(url: str) -> Image.Image:
-    try:
-        response = requests.get(url, stream=True)
-        response.raise_for_status()  # Raise an HTTPError for bad responses
-        image = Image.open(BytesIO(response.content))
-        return image
-    except requests.HTTPError as http_err:
-        print(f"HTTP error occurred: {http_err}")
-        return None
-    except UnidentifiedImageError:
-        print("Cannot identify image file")
-        return None
-    except Exception as err:
-        print(f"Other error occurred: {err}")
-        return None
-
-
 EXAMPLES = [
     [
         {
@@ -114,16 +107,19 @@ def process(
     progress=gr.Progress(track_tqdm=True)
 ):
     if not input_text:
-        return None, None, "Please enter a text prompt."
+        gr.Info("Please enter a text prompt.")
+        return None, None
 
     image = input_image_editor.get('background')
-    mask = input_image_editor.get('layers', [None])[0]
+    mask = input_image_editor.get('layers')[0] if input_image_editor.get('layers') else None
 
     if not image:
-        return None, None, "Please upload an image."
+        gr.Info("Please upload an image.")
+        return None, None
 
     if not mask:
-        return None, None, "Please draw a mask on the image."
+        gr.Info("Please draw a mask on the image.")
+        return None, None
 
     width, height = resize_image_dimensions(original_resolution_wh=image.size)
     resized_image = image.resize((width, height), Image.LANCZOS)
@@ -142,8 +138,8 @@ def process(
         generator=generator,
         num_inference_steps=num_inference_steps_slider
     ).images[0]
-    
-    return result, resized_mask, None
+    print('INFERENCE DONE')
+    return result, resized_mask
 
 
 with gr.Blocks() as demo: