Update app.py

app.py

@@ -24,6 +24,25 @@ repo_id = "black-forest-labs/FLUX.1-Fill-dev"
 if torch.cuda.is_available():
     pipe = FluxFillPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16).to("cuda")
 
+
+def create_mask_image(mask_array):
+    # Convert the mask to a numpy array if it's not already
+    if not isinstance(mask_array, np.ndarray):
+        mask_array = np.array(mask_array)
+    
+    # Create a new array with the same shape as the mask, but only for RGB channels
+    processed_mask = np.zeros((mask_array.shape[0], mask_array.shape[1], 3), dtype=np.uint8)
+    
+    # Set transparent parts (alpha=0) to black (0, 0, 0)
+    transparent_mask = mask_array[:, :, 3] == 0
+    processed_mask[transparent_mask] = [0, 0, 0]
+    
+    # Set black parts (RGB=0, 0, 0 and alpha=255) to white (255, 255, 255)
+    black_mask = (mask_array[:, :, :3] == [0, 0, 0]).all(axis=2) & (mask_array[:, :, 3] == 255)
+    processed_mask[black_mask] = [255, 255, 255]
+    
+    return Image.fromarray(processed_mask)
+
 @spaces.GPU()
 def inpaintGen(
         imgMask,
@@ -47,7 +66,9 @@ def inpaintGen(
         raise gr.Error("Please draw a mask on the image.")
 
     source_img = load_image(source_path)
-    mask_img = load_image(mask_path).split()[3]
+    mask_img = load_image(mask_path)
+
+    mask_img = create_mask_image(mask)
     
     width, height = source_img.size