Update app.py

app.py

@@ -9,6 +9,7 @@ import numpy as np
 from PIL import Image
 import openai
 
+# --- Constants and Helper Functions (Keep as before) ---
 MODEL = "gpt-image-1"
 SIZE_CHOICES = ["auto", "1024x1024", "1536x1024", "1024x1536"]
 QUALITY_CHOICES = ["auto", "low", "medium", "high"]
@@ -46,37 +47,24 @@ def _common_kwargs(
     kwargs: Dict[str, Any] = dict(
         model=MODEL,
         n=n,
-        # REMOVED: response_format="b64_json", # This parameter caused the BadRequestError
     )
-
-    # Use API defaults if 'auto' is selected
     if size != "auto":
         kwargs["size"] = size
     if quality != "auto":
         kwargs["quality"] = quality
-
-    # Prompt is optional for variations
     if prompt is not None:
         kwargs["prompt"] = prompt
-
-    # Output format specific settings (API default is png)
     if out_fmt != "png":
         kwargs["output_format"] = out_fmt
-
-    # Transparency via background parameter (png & webp only)
     if transparent_bg and out_fmt in {"png", "webp"}:
         kwargs["background"] = "transparent"
-
-    # Compression for lossy formats (API expects integer 0-100)
     if out_fmt in {"jpeg", "webp"}:
-        # Ensure compression is an integer as expected by the API
         kwargs["output_compression"] = int(compression)
-
     return kwargs
 
+# --- API Call Functions (Keep as corrected before) ---
 
 # ---------- Generate ---------- #
-
 def generate(
     api_key: str,
     prompt: str,
@@ -90,7 +78,7 @@ def generate(
     """Calls the OpenAI image generation endpoint."""
     if not prompt:
         raise gr.Error("Please enter a prompt.")
-    client = _client(api_key)
+    client = _client(api_key) # API key used here
     try:
         common_args = _common_kwargs(prompt, n, size, quality, out_fmt, compression, transparent_bg)
         resp = client.images.generate(**common_args)
@@ -101,18 +89,16 @@ def generate(
     except openai.RateLimitError:
         raise gr.Error("Rate limit exceeded. Please try again later.")
     except openai.BadRequestError as e:
-        # Extract the specific error message if possible
         error_message = str(e)
         try:
-            # Attempt to parse the error body if it's JSON-like
             import json
-            body = json.loads(str(e.body)) # e.body might be bytes
+            body = json.loads(str(e.body))
             if isinstance(body, dict) and 'error' in body and 'message' in body['error']:
                 error_message = f"OpenAI Bad Request: {body['error']['message']}"
             else:
                  error_message = f"OpenAI Bad Request: {e}"
         except:
-             error_message = f"OpenAI Bad Request: {e}" # Fallback
+             error_message = f"OpenAI Bad Request: {e}"
         raise gr.Error(error_message)
     except Exception as e:
         raise gr.Error(f"An unexpected error occurred: {e}")
@@ -120,7 +106,6 @@ def generate(
 
 
 # ---------- Edit / Inpaint ---------- #
-
 def _bytes_from_numpy(arr: np.ndarray) -> bytes:
     """Convert RGBA/RGB uint8 numpy array to PNG bytes."""
     img = Image.fromarray(arr.astype(np.uint8))
@@ -128,43 +113,21 @@ def _bytes_from_numpy(arr: np.ndarray) -> bytes:
     img.save(out, format="PNG")
     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):
-        mask_arr = mask_value
-    # If it's an EditorValue dict coming from ImageEditor/ImageMask with type="file" or "pil"
-    elif isinstance(mask_value, dict):
-        # Prefer the composite (all layers merged) if present
+    if mask_value is None: return None
+    if isinstance(mask_value, np.ndarray): return mask_value
+    if isinstance(mask_value, dict):
         comp = mask_value.get("composite")
         if comp is not None and isinstance(comp, (Image.Image, np.ndarray)):
-             mask_arr = np.array(comp) if isinstance(comp, Image.Image) else comp
-        # Fallback to the mask if present (often from ImageMask)
+             return np.array(comp) if isinstance(comp, Image.Image) else comp
         elif mask_value.get("mask") is not None and isinstance(mask_value["mask"], (Image.Image, np.ndarray)):
-             mask_arr = np.array(mask_value["mask"]) if isinstance(mask_value["mask"], Image.Image) else mask_value["mask"]
-        # Fallback to the topmost layer
+             return np.array(mask_value["mask"]) if isinstance(mask_value["mask"], Image.Image) else mask_value["mask"]
         elif mask_value.get("layers"):
             top_layer = mask_value["layers"][-1]
             if isinstance(top_layer, (Image.Image, np.ndarray)):
-                 mask_arr = np.array(top_layer) if isinstance(top_layer, Image.Image) else top_layer
-            else:
-                 return None # Cannot process layer format
-        else:
-            return None # No usable image data found in dict
-    else:
-        # Unknown format – ignore
-        return None
-
-    # Ensure mask_arr is a numpy array now
-    if not isinstance(mask_arr, np.ndarray):
-        return None # Should not happen after above checks, but safeguard
-
-    return mask_arr
-
+                 return np.array(top_layer) if isinstance(top_layer, Image.Image) else top_layer
+    return None
 
 def edit_image(
     api_key: str,
@@ -179,99 +142,52 @@ def edit_image(
     transparent_bg: bool,
 ):
     """Calls the OpenAI image edit endpoint."""
-    if image_numpy is None:
-        raise gr.Error("Please upload an image.")
-    if not prompt:
-        raise gr.Error("Please enter an edit prompt.")
+    if image_numpy is None: raise gr.Error("Please upload an image.")
+    if not prompt: raise gr.Error("Please enter an edit prompt.")
 
     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:
-        # Check if the mask seems empty (all black or fully transparent)
         is_empty = False
-        if mask_numpy.ndim == 2: # Grayscale
-            is_empty = np.all(mask_numpy == 0)
-        elif mask_numpy.shape[-1] == 4: # RGBA
-            is_empty = np.all(mask_numpy[:, :, 3] == 0)
-        elif mask_numpy.shape[-1] == 3: # RGB
-            is_empty = np.all(mask_numpy == 0)
+        if mask_numpy.ndim == 2: is_empty = np.all(mask_numpy == 0)
+        elif mask_numpy.shape[-1] == 4: is_empty = np.all(mask_numpy[:, :, 3] == 0)
+        elif mask_numpy.shape[-1] == 3: is_empty = np.all(mask_numpy == 0)
 
         if is_empty:
-             gr.Warning("The provided mask appears empty (all black/transparent). The API might edit the entire image or ignore the mask.")
-             # Pass None if the mask is effectively empty, as per API docs (transparent areas are edited)
+             gr.Warning("Mask appears empty. API might edit entire image or ignore mask.")
              mask_bytes = None
         else:
-            # Convert the mask to the format required by the API:
-            # A PNG image where TRANSPARENT areas indicate where the image should be edited.
-            # Our Gradio mask uses WHITE to indicate the edit area.
-            # So, we need to create an alpha channel where white pixels in the input mask become transparent (0),
-            # and black/other pixels become opaque (255).
-
-            if mask_numpy.ndim == 2: # Grayscale input mask
-                # Assume white (255) means edit -> make transparent (0 alpha)
-                # Assume black (0) means keep -> make opaque (255 alpha)
-                alpha = (mask_numpy == 0).astype(np.uint8) * 255
-            elif mask_numpy.shape[-1] == 4:  # RGBA input mask (from gr.ImageMask)
-                # Use the alpha channel directly if it exists and seems meaningful,
-                # otherwise, treat non-black RGB as edit area.
-                # gr.ImageMask often returns RGBA where painted area is white [255,255,255,255] and background is [0,0,0,0]
-                # We want the painted (white) area to be transparent in the final mask.
-                # We want the unpainted (transparent black) area to be opaque in the final mask.
-                alpha = (mask_numpy[:, :, 3] == 0).astype(np.uint8) * 255
-            elif mask_numpy.shape[-1] == 3: # RGB input mask
-                # Assume white [255, 255, 255] means edit -> make transparent (0 alpha)
-                # Assume black [0, 0, 0] or other colors mean keep -> make opaque (255 alpha)
+            if mask_numpy.ndim == 2: alpha = (mask_numpy == 0).astype(np.uint8) * 255
+            elif mask_numpy.shape[-1] == 4: alpha = (mask_numpy[:, :, 3] == 0).astype(np.uint8) * 255
+            elif mask_numpy.shape[-1] == 3:
                 is_white = np.all(mask_numpy == [255, 255, 255], axis=-1)
                 alpha = (~is_white).astype(np.uint8) * 255
-            else:
-                 raise gr.Error("Unsupported mask format.")
+            else: raise gr.Error("Unsupported mask format.")
 
-            # Create a single-channel L mode image (grayscale/alpha) for the mask
             mask_img = Image.fromarray(alpha, mode='L')
-
-            # The API expects an RGBA PNG where the alpha channel defines the mask.
-            # Create a black image with the calculated alpha channel.
             rgba_mask = Image.new("RGBA", mask_img.size, (0, 0, 0, 0))
-            black_opaque = Image.new("L", mask_img.size, 0) # Black base
-            rgba_mask.putalpha(mask_img) # Use the calculated alpha
-
+            rgba_mask.putalpha(mask_img)
             out = io.BytesIO()
             rgba_mask.save(out, format="PNG")
             mask_bytes = out.getvalue()
-
-            # Debug: Save mask locally to check
-            # rgba_mask.save("debug_mask_sent_to_api.png")
-
     else:
-        gr.Info("No mask provided. The API will attempt to edit the image based on the prompt without a specific mask.")
-        mask_bytes = None # Explicitly pass None if no mask is usable
+        gr.Info("No mask provided. Editing without specific mask.")
+        mask_bytes = None
 
-    client = _client(api_key)
+    client = _client(api_key) # API key used here
     try:
         common_args = _common_kwargs(prompt, n, size, quality, out_fmt, compression, transparent_bg)
-        # The edit endpoint requires the prompt
-        if "prompt" not in common_args:
-             common_args["prompt"] = prompt # Should always be there via _common_kwargs, but safeguard
-
-        # Ensure image and mask are passed correctly
-        api_kwargs = {
-             "image": img_bytes,
-             **common_args
-        }
-        if mask_bytes is not None:
-             api_kwargs["mask"] = mask_bytes
-
+        api_kwargs = {"image": img_bytes, **common_args}
+        if mask_bytes is not None: api_kwargs["mask"] = mask_bytes
         resp = client.images.edit(**api_kwargs)
-
     except openai.AuthenticationError:
          raise gr.Error("Invalid OpenAI API key.")
     except openai.PermissionDeniedError:
-        raise gr.Error("Permission denied. Check your API key permissions or complete required verification for gpt-image-1.")
+        raise gr.Error("Permission denied. Check API key permissions/verification.")
     except openai.RateLimitError:
-        raise gr.Error("Rate limit exceeded. Please try again later.")
+        raise gr.Error("Rate limit exceeded.")
     except openai.BadRequestError as e:
         error_message = str(e)
         try:
@@ -279,15 +195,10 @@ def edit_image(
             body = json.loads(str(e.body))
             if isinstance(body, dict) and 'error' in body and 'message' in body['error']:
                 error_message = f"OpenAI Bad Request: {body['error']['message']}"
-                # Add specific advice based on common mask errors
-                if "mask" in error_message.lower():
-                     error_message += " (Ensure mask is a valid PNG with an alpha channel and matches the image dimensions.)"
-                elif "size" in error_message.lower():
-                     error_message += " (Ensure image and mask dimensions match and are supported.)"
-            else:
-                 error_message = f"OpenAI Bad Request: {e}"
-        except:
-             error_message = f"OpenAI Bad Request: {e}" # Fallback
+                if "mask" in error_message.lower(): error_message += " (Check mask format/dimensions)"
+                elif "size" in error_message.lower(): error_message += " (Check image/mask dimensions)"
+            else: error_message = f"OpenAI Bad Request: {e}"
+        except: error_message = f"OpenAI Bad Request: {e}"
         raise gr.Error(error_message)
     except Exception as e:
         raise gr.Error(f"An unexpected error occurred: {e}")
@@ -295,7 +206,6 @@ def edit_image(
 
 
 # ---------- Variations ---------- #
-
 def variation_image(
     api_key: str,
     image_numpy: np.ndarray,
@@ -307,27 +217,19 @@ def variation_image(
     transparent_bg: bool,
 ):
     """Calls the OpenAI image variations endpoint."""
-    # NOTE: Variations are only supported for DALL-E 2 according to docs.
-    # This might fail with gpt-image-1. Consider adding a check or using DALL-E 2.
-    gr.Warning("Note: Image variations are officially supported for DALL·E 2, not gpt-image-1. This may not work as expected.")
-
-    if image_numpy is None:
-        raise gr.Error("Please upload an image.")
+    gr.Warning("Note: Variations may not work with gpt-image-1 (use DALL·E 2).")
+    if image_numpy is None: raise gr.Error("Please upload an image.")
     img_bytes = _bytes_from_numpy(image_numpy)
-    client = _client(api_key)
+    client = _client(api_key) # API key used here
     try:
-        # Prompt is None for variations
         common_args = _common_kwargs(None, n, size, quality, out_fmt, compression, transparent_bg)
-        resp = client.images.variations(
-            image=img_bytes,
-            **common_args,
-        )
+        resp = client.images.variations(image=img_bytes, **common_args)
     except openai.AuthenticationError:
          raise gr.Error("Invalid OpenAI API key.")
     except openai.PermissionDeniedError:
-        raise gr.Error("Permission denied. Check your API key permissions.")
+        raise gr.Error("Permission denied.")
     except openai.RateLimitError:
-        raise gr.Error("Rate limit exceeded. Please try again later.")
+        raise gr.Error("Rate limit exceeded.")
     except openai.BadRequestError as e:
         error_message = str(e)
         try:
@@ -336,11 +238,9 @@ def variation_image(
             if isinstance(body, dict) and 'error' in body and 'message' in body['error']:
                  error_message = f"OpenAI Bad Request: {body['error']['message']}"
                  if "model does not support variations" in error_message.lower():
-                      error_message += " (gpt-image-1 does not support variations, use DALL·E 2 instead)."
-            else:
-                 error_message = f"OpenAI Bad Request: {e}"
-        except:
-             error_message = f"OpenAI Bad Request: {e}" # Fallback
+                      error_message += " (gpt-image-1 does not support variations)."
+            else: error_message = f"OpenAI Bad Request: {e}"
+        except: error_message = f"OpenAI Bad Request: {e}"
         raise gr.Error(error_message)
     except Exception as e:
         raise gr.Error(f"An unexpected error occurred: {e}")
@@ -359,11 +259,12 @@ def build_ui():
         )
 
         with gr.Accordion("🔐 API key", open=False):
+            # API key input component
             api = gr.Textbox(label="OpenAI API key", type="password", placeholder="sk-…")
 
         # Common controls
         with gr.Row():
-             n_slider = gr.Slider(1, 4, value=1, step=1, label="Number of images (n)", info="Max 4 for this demo.") # Limit n for stability/cost
+             n_slider = gr.Slider(1, 4, value=1, step=1, label="Number of images (n)", info="Max 4 for this demo.")
              size = gr.Dropdown(SIZE_CHOICES, value="auto", label="Size", info="API default if 'auto'.")
              quality = gr.Dropdown(QUALITY_CHOICES, value="auto", label="Quality", info="API default if 'auto'.")
         with gr.Row():
@@ -376,7 +277,8 @@ def build_ui():
 
         out_fmt.change(_toggle_compression, inputs=out_fmt, outputs=compression)
 
-        common_inputs = [api, n_slider, size, quality, out_fmt, compression, transparent]
+        # Define the list of common controls *excluding* the API key
+        common_controls = [n_slider, size, quality, out_fmt, compression, transparent]
 
         with gr.Tabs():
             # ----- Generate Tab ----- #
@@ -385,9 +287,11 @@ def build_ui():
                     prompt_gen = gr.Textbox(label="Prompt", lines=3, placeholder="A photorealistic ginger cat astronaut on Mars", scale=4)
                     btn_gen = gr.Button("Generate 🚀", variant="primary", scale=1)
                 gallery_gen = gr.Gallery(label="Generated Images", columns=2, height="auto", preview=True)
+
+                # CORRECTED inputs list for generate
                 btn_gen.click(
                     generate,
-                    inputs=[prompt_gen] + common_inputs, # Prepend specific inputs
+                    inputs=[api, prompt_gen] + common_controls, # API key first
                     outputs=gallery_gen,
                     api_name="generate"
                 )
@@ -397,19 +301,20 @@ def build_ui():
                 gr.Markdown("Upload an image, then **paint the area to change** in the mask canvas below (white = edit area). The API requires the mask and image to have the same dimensions.")
                 with gr.Row():
                     img_edit = gr.Image(label="Source Image", type="numpy", height=400)
-                    # Use ImageMask component for interactive painting
                     mask_canvas = gr.ImageMask(
                          label="Mask – Paint White Where Image Should Change",
-                         type="numpy", # Get mask as numpy array
+                         type="numpy",
                          height=400
                     )
                 with gr.Row():
                     prompt_edit = gr.Textbox(label="Edit prompt", lines=2, placeholder="Replace the sky with a starry night", scale=4)
                     btn_edit = gr.Button("Edit 🖌️", variant="primary", scale=1)
                 gallery_edit = gr.Gallery(label="Edited Images", columns=2, height="auto", preview=True)
+
+                # CORRECTED inputs list for edit_image
                 btn_edit.click(
                     edit_image,
-                    inputs=[img_edit, mask_canvas, prompt_edit] + common_inputs, # Prepend specific inputs
+                    inputs=[api, img_edit, mask_canvas, prompt_edit] + common_controls, # API key first
                     outputs=gallery_edit,
                     api_name="edit"
                 )
@@ -421,9 +326,11 @@ def build_ui():
                     img_var = gr.Image(label="Source Image", type="numpy", height=400, scale=4)
                     btn_var = gr.Button("Create Variations ✨", variant="primary", scale=1)
                 gallery_var = gr.Gallery(label="Variations", columns=2, height="auto", preview=True)
+
+                # CORRECTED inputs list for variation_image
                 btn_var.click(
                     variation_image,
-                    inputs=[img_var] + common_inputs, # Prepend specific inputs
+                    inputs=[api, img_var] + common_controls, # API key first
                     outputs=gallery_var,
                     api_name="variations"
                 )
@@ -432,6 +339,4 @@ def build_ui():
 
 if __name__ == "__main__":
     app = build_ui()
-    # Set share=True to create a public link (useful for Spaces)
-    # Set debug=True for more detailed logs in the console
     app.launch(share=os.getenv("GRADIO_SHARE") == "true", debug=True)