app.py

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 requests
import json


class ImageAgent:
    def __init__(self, model: str='gpt-image-1', api_key: str=None):
        self.model = model
        self.origin = 'bj'
        self.api_key = api_key or os.getenv('API_KEY')
        self.gen_url = f'https://gpt-{self.origin}.singularity-ai.com/gpt-proxy/azure/imagen'
        self.edit_url = f'https://gpt-{self.origin}.singularity-ai.com/gpt-proxy/azure/imagen/edit'
        
    def image_generate(self, data: dict = None, prompt: str = None, retries: int = 3, **kwargs):
        headers = {
            "app_key": self.api_key,
            "Content-Type": "application/json"
        }

        request_data = data or {
            "model": self.model or "gpt-image-1",
            "prompt": prompt or "a red fox in a snowy forest",
            "size": kwargs.get("size", "auto"),
            "quality": kwargs.get("quality", "high"),
            "n": kwargs.get("n", 1)
        }

        print("***** Request Data - Image Generate *****")
        print(json.dumps(request_data, indent=2))

        for i in range(retries):
            try:
                print(f"第 {i+1} 次发送图像生成请求...")
                response = requests.post(self.gen_url, json=request_data, headers=headers, stream=True)
                print(f"响应状态码: {response.status_code}")

                if response.status_code != 200:
                    raise Exception(f"请求失败：{response.text}")

                try:
                    response_json = json.loads(response.text)
                except json.JSONDecodeError:
                    raise Exception(f"响应内容无法解析为 JSON：{response.text}")

                if not response_json.get("data"):
                    raise Exception(f"响应内容 data 字段为空，准备重试...")
                
                return response_json
            except Exception as e:
                raise Exception(f"发生错误：{e}")
        
        raise Exception(f"重试超过{retries}次，图像生成失败!")

    def image_edit(self, image: tuple, data: dict = None, prompt: str = None, retries: int = 3, **kwargs):
        
        request_data = data or {
            "model": self.model or "gpt-image-1",
            "prompt": prompt,
            "size": kwargs.get("size", "auto"),
            "quality": kwargs.get("quality", "high"),
            "n": kwargs.get("n", 1)
        }

        print("***** Request Data - Image Edit *****")
        print(json.dumps(request_data, indent=2))

        for i in range(retries):
            try:
                print(f"第 {i+1} 次发送图像编辑请求...")
                headers = {
                    "app_key": self.api_key
                }
                response = requests.post(self.edit_url, headers=headers, files={"image": image}, data=request_data, timeout=180)
                print(f"响应状态码: {response.status_code}")

                if response.status_code != 200:
                    raise Exception(f"请求失败：{response.text}")

                try:
                    response_json = json.loads(response.text)
                except json.JSONDecodeError:
                    raise Exception(f"响应内容无法解析为JSON: {response.text}")

                if not response_json.get("data"):
                    raise Exception(f"响应内容 data 字段为空: {response.text}, 准备重试...")

                return response_json
            except Exception as e:
                raise Exception(f"发生错误：{e}")
        
        raise Exception(f"重试超过{retries}次，图像编辑失败!")


# --- Constants ---
MODEL = "gpt-image-1"
SIZE_CHOICES = ["auto", "1024x1024", "1536x1024", "1024x1536"]
QUALITY_CHOICES = ["auto", "low", "medium", "high"]
FORMAT_CHOICES = ["png"]


def _client(key: str) -> ImageAgent:
    """Initializes the Image Agent with the provided API key."""
    api_key = key.strip() or os.getenv("API_KEY", "")
    if not api_key:
        raise gr.Error("Please enter your API key")
    return ImageAgent(api_key=api_key)


def _img_list(resp: Dict[str, Any]) -> 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.get("data", []):
        if d.get("b64_json", None):
            data = base64.b64decode(d.get("b64_json"))
            img = Image.open(io.BytesIO(data))
            imgs.append(np.array(img))
        elif d.get("url", None):
            imgs.append(d.get("url"))
    return imgs


def _common_kwargs(
    prompt: Optional[str],
    n: int,
    size: str,
    quality: str,
) -> Dict[str, Any]:
    """Prepare keyword args for Images API."""
    kwargs: Dict[str, Any] = {
        "model": MODEL,
        "n": n,
    }
    if size != "auto":
        kwargs["size"] = size
    if quality != "auto":
        kwargs["quality"] = quality
    if prompt is not None:
        kwargs["prompt"] = prompt
    return kwargs


def convert_to_format(
    img_array: np.ndarray,
    target_fmt: str,
    quality: int = 75,
) -> np.ndarray:
    """
    Convert a PIL numpy array to target_fmt (JPEG/WebP) and return as numpy array.
    """
    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)

# ---------- Generate ---------- #
def generate(
    api_key: str,
    prompt: str,
    n: int,
    size: str,
    quality: str,
):
    if not prompt:
        raise gr.Error("Please enter a prompt.")
    try:
        agent = _client(api_key)
        common_args = _common_kwargs(prompt, n, size, quality)
        api_kwargs = {"retries": 3, **common_args}
        resp = agent.image_generate(**api_kwargs)
        imgs = _img_list(resp)
        # if out_fmt in {"jpeg", "webp"}:
        #     imgs = [convert_to_format(img, out_fmt) for img in imgs]
        return imgs
    except Exception as e:
        raise gr.Error(str(e))


# ---------- Edit ---------- #
def _bytes_from_numpy(arr: np.ndarray) -> bytes:
    img = Image.fromarray(arr.astype(np.uint8))
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()

def edit_image(
        api_key: str,
        image_numpy: Optional[np.ndarray],
        prompt: str,
        n: int,
        size: str,
        quality: str
    ):
        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)

        try:
            agent = _client(api_key)
            common_args = _common_kwargs(prompt, n, size, quality)

            image_tuple = ("image.png", img_bytes, "image/png")
            api_kwargs = {"image": image_tuple, "retries": 3, **common_args}

            resp = agent.image_edit(**api_kwargs)
            imgs = _img_list(resp)
            # if out_fmt in {"jpeg", "webp"}:
            #     imgs = [convert_to_format(img, out_fmt) for img in imgs]
            return imgs
        except Exception as e:
            raise gr.Error(str(e))


# ---------- UI ---------- #
def build_ui():
    with gr.Blocks(title="GPT-Image-1 (BYOT)") as demo:
        gr.Markdown("""# 🐍 GPT-Image-1 Playground""")
        with gr.Accordion("🔐 API key", open=False):
            api = gr.Textbox(label="OpenAI API key", type="password", placeholder="gpt-...")

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

        common_controls = [n_slider, size, quality]

        with gr.Tabs():
            with gr.TabItem("Generate"):
                prompt_gen = gr.Textbox(
                    label="Prompt",
                    lines=3,
                    placeholder="Write down your prompt here",
                    autofocus=True,
                    container=False
                )
                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
                )

            with gr.TabItem("Edit"):
                img_edit = gr.Image(type="numpy", label="Image to edit", height=400)
                prompt_edit = gr.Textbox(label="Edit prompt", lines=2, placeholder="Write down your prompt here")
                btn_edit = gr.Button("Edit 🖌️")
                gallery_edit = gr.Gallery(columns=2, height="auto")
                btn_edit.click(edit_image, inputs=[api, img_edit, prompt_edit] + common_controls, outputs=gallery_edit)

    return demo


if __name__ == "__main__":
    app = build_ui()
    app.launch(share=os.getenv("GRADIO_SHARE") == "true", debug=os.getenv("GRADIO_DEBUG") == "true")