fix mask generation

app.py

@@ -3,7 +3,12 @@ import numpy as np
 
 def create_mask(img):
     alpha_channel = img["layers"][0][:, :, 3]
-    mask = np.where(alpha_channel == 0, 0, 255)
+    # Create an output image with 4 channels (RGBA)
+    mask = np.zeros_like(img["layers"][0])
+    # Set RGB channels to white where alpha is not zero
+    mask[:, :, :3] = np.where(alpha_channel[:, :, None] == 0, 0, 255)
+    # Set alpha channel to 0 where original alpha is zero, else 255
+    mask[:, :, 3] = np.where(alpha_channel == 0, 255, 0)
     return mask
 
 with gr.Blocks(title="In-Paint Mask Creator") as demo: