Create app.py

app.py

@@ -0,0 +1,273 @@
+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")
+