Update app.py

app.py

@@ -1,6 +1,5 @@
-# app.py
 """
-Gradio Space: GPT‑Image‑1 – BYOT playground  
+Gradio Space: GPT‑Image‑1 – BYOT playground
 Generate · Edit (paint mask!) · Variations
 ==========================================
 Adds an **in‑browser paint tool** for the edit / inpaint workflow so users can
@@ -10,7 +9,7 @@ draw the mask directly instead of uploading one.
 * Upload an image.  
 * Use the *Mask* canvas to **paint the areas you’d like changed** (white =
   editable, black = keep).  
-  Gradio’s built‑in *sketch* tool captures your brush strokes.
+  The new `gr.ImageMask` component captures your brush strokes.
 * The painted mask is converted to a 1‑channel PNG and sent to the
   `images.edit()` endpoint.
 
@@ -22,7 +21,7 @@ from __future__ import annotations
 
 import io
 import os
-from typing import List, Optional
+from typing import List, Optional, Union, Dict, Any
 
 import gradio as gr
 import numpy as np
@@ -105,10 +104,33 @@ def _bytes_from_numpy(arr: np.ndarray) -> bytes:
     return out.getvalue()
 
 
+def _extract_mask_array(mask_value: Union[np.ndarray, Dict[str, Any], None]) -> Optional[np.ndarray]:
+    """Handle ImageMask / ImageEditor return formats and extract a numpy mask array."""
+    if mask_value is None:
+        return None
+
+    # If we already have a numpy array (ImageMask with type="numpy")
+    if isinstance(mask_value, np.ndarray):
+        return mask_value
+
+    # If it's an EditorValue dict coming from ImageEditor/ImageMask with type="numpy"
+    if isinstance(mask_value, dict):
+        # Prefer the composite (all layers merged) if present
+        comp = mask_value.get("composite")
+        if comp is not None:
+            return np.asarray(comp)
+        # Fallback to the topmost layer
+        layers = mask_value.get("layers")
+        if layers:
+            return np.asarray(layers[-1])
+    # Unknown format – ignore
+    return None
+
+
 def edit_image(
     api_key: str,
     image_numpy: np.ndarray,
-    mask_numpy: Optional[np.ndarray],
+    mask_value: Optional[Union[np.ndarray, Dict[str, Any]]],
     prompt: str,
     n: int,
     size: str,
@@ -122,11 +144,13 @@ def edit_image(
     img_bytes = _bytes_from_numpy(image_numpy)
 
     mask_bytes: Optional[bytes] = None
+    mask_numpy = _extract_mask_array(mask_value)
+
     if mask_numpy is not None:
-        # Convert painted area (alpha > 0) to white, else black; 1‑channel.
-        if mask_numpy.shape[-1] == 4:  # RGBA from gr.Image sketch
+        # Convert painted area (any non‑zero pixel) to white, else black; 1‑channel.
+        if mask_numpy.shape[-1] == 4:  # RGBA (has alpha channel)
             alpha = mask_numpy[:, :, 3]
-        else:  # RGB
+        else:  # RGB or grayscale
             alpha = np.any(mask_numpy != 0, axis=-1).astype(np.uint8) * 255
         bw = np.stack([alpha] * 3, axis=-1)  # 3‑channel white/black
         mask_bytes = _bytes_from_numpy(bw)
@@ -206,7 +230,7 @@ def build_ui():
             # ----- Edit Tab ----- #
             with gr.TabItem("Edit / Inpaint"):
                 img_edit = gr.Image(label="Image", type="numpy")
-                mask_canvas = gr.Image(label="Mask – paint white where the image should change", type="numpy", tool="sketch")
+                mask_canvas = gr.ImageMask(label="Mask – paint white where the image should change", type="numpy")
                 prompt_edit = gr.Textbox(label="Edit prompt", lines=2, placeholder="Replace the sky with a starry night")
                 btn_edit = gr.Button("Edit 🖌️")
                 gallery_edit = gr.Gallery(columns=2, height="auto")