app.py

import base64
import io
import os
import random

import gradio as gr
import openai
import requests
from PIL import Image


def generate_image(
    prompt: str,
    api_key: str,
    base_url: str,
    model: str,
    size: str,
    quality: str,
    style: str,
    progress=gr.Progress(track_tqdm=True),
):
    # Validate required parameters
    if not prompt:
        raise gr.Error("Please provide a prompt for the image generation")
    if not api_key:
        raise gr.Error("API key is required")
    if not base_url:
        raise gr.Error("Base URL is required")
    if not model:
        raise gr.Error("Model name is required")

    try:
        # Initialize client (fast operation)
        progress(0.1, desc="Initializing client...")
        client = openai.OpenAI(api_key=api_key, base_url=base_url)

        # Generate image (slowest operation)
        progress(0.2, desc="Sending request to API...")
        response = client.images.generate(model=model, prompt=prompt, size=size, quality=quality, style=style, n=1)
        progress(0.6, desc="Generating image... This may take 10-30 seconds")

        if hasattr(response, "data") and hasattr(response.data[0], "url") and response.data[0].url:
            image_url = response.data[0].url
            progress(0.8, desc="Downloading generated image...")
            image_response = requests.get(image_url)
            if image_response.status_code != 200:
                raise gr.Error("Failed to download the generated image")

            progress(0.9, desc="Processing image...")
            img = Image.open(io.BytesIO(image_response.content))
            progress(1.0, desc="Complete!")
            return img

        elif hasattr(response, "data") and hasattr(response.data[0], "b64_json") and response.data[0].b64_json:
            progress(0.8, desc="Decoding base64 image...")
            b64_data = response.data[0].b64_json
            img_data = base64.b64decode(b64_data)

            progress(0.9, desc="Processing image...")
            img = Image.open(io.BytesIO(img_data))
            progress(1.0, desc="Complete!")
            return img
        else:
            raise gr.Error("No image data received from the API")

    except openai.AuthenticationError:
        raise gr.Error("Invalid API key. Please check your OpenAI API key and try again.")
    except openai.APIError as e:
        error_msg = str(e)
        if "api_key" in error_msg.lower():
            raise gr.Error("API authentication error. Please check your credentials.")
        else:
            raise gr.Error("OpenAI API error occurred. Please try again.")
    except requests.RequestException:
        raise gr.Error("Network error occurred while downloading the image. Please try again.")
    except Exception as e:
        # Log the full error internally but show a generic message to the user
        print(f"Internal error: {type(e).__name__}")  # Log error type without sensitive data
        raise gr.Error("An unexpected error occurred. Please try again.")


css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}

/* Control image size */
.generated-image {
    max-height: 512px;
    width: auto;
    object-fit: contain;
}
"""

examples = [
    "A serene lake surrounded by mountains at sunset",
    "A futuristic cityscape at night",
    "A watercolor painting of a blooming cherry tree",
    "A majestic lion resting on a rocky outcrop in the African savanna",
    "A cozy cottage nestled in a snowy forest during winter",
    "An astronaut floating in space with Earth in the background",
    "A bustling marketplace in a Moroccan city",
    "A vibrant coral reef teeming with marine life",
    "A steampunk-inspired robot tending a garden",
    "A minimalist abstract painting with bold colors",
    "A hyperrealistic close-up of a dew-covered spiderweb",
    "A fantasy landscape with floating islands and waterfalls",
    "A vintage photograph of a jazz band in a smoky club",
    "A serene beach with crystal-clear water and palm trees",
    "A vibrant street market in a Southeast Asian city",
    "A futuristic laboratory with advanced technology",
    "A dense jungle with exotic plants and animals",
    "A medieval castle on a hilltop overlooking a village",
    "A bustling coffee shop in a rainy city",
    "A peaceful Zen garden with carefully raked gravel",
    "A majestic dragon soaring through a stormy sky",
    "A bioluminescent forest at twilight",
    "A portrait of a wise old wizard with a long beard",
    "A group of penguins waddling across the Antarctic ice",
    "A stack of pancakes with syrup and berries",
    "A close-up of a blooming sunflower in a field",
    "A cityscape reflected in a rain puddle",
    "A cup of coffee with latte art",
    "A snowy mountain range under a starry sky",
    "A field of lavender in Provence, France",
    "A plate of sushi with various types of fish",
    "A hot air balloon floating over a valley",
    "A lighthouse on a rocky cliff overlooking the ocean",
    "A bowl of ramen with chopsticks",
    "A tropical beach with a hammock and palm trees",
    "A grand library with towering bookshelves",
    "A cobblestone street in a European village",
    "A waterfall cascading into a clear pool",
    "A field of tulips in the Netherlands",
    "A campfire under a starry night sky",
    "A slice of pizza with pepperoni and cheese",
    "A dense bamboo forest in Japan",
    "A plate of pasta with tomato sauce and basil",
    "A serene Japanese garden with a koi pond",
    "A vibrant carnival with colorful lights and rides",
    "A cozy fireplace in a log cabin",
    "A field of wildflowers in the spring",
    "A bustling train station in a major city",
    "A quiet countryside road with rolling hills",
    "A modern art museum with abstract sculptures",
]

with gr.Blocks(css=css, title="OpenAI Image Generator") as demo:
    gr.Markdown("# OpenAI Compatible Image Generator")
    gr.Markdown("Generate images using OpenAI's DALL-E or compatible APIs")

    # Initialize browser state with default values in a list
    settings_state = gr.BrowserState(
        [
            "",  # api_key
            "https://api.openai.com/v1",  # base_url
            "dall-e-3",  # model
            "1024x1024",  # size
            "standard",  # quality
            "vivid",  # style
        ],
        storage_key="openai_image_generator_settings",
        secret=os.getenv("STORAGE_SECRET", "secret"),
    )

    with gr.Row():
        # Left column for settings
        with gr.Column(scale=1):
            gr.Markdown("### Settings")
            api_key = gr.Textbox(label="API Key", placeholder="Your OpenAI API key", type="password", value="")
            base_url = gr.Textbox(label="Base URL", placeholder="API base URL", value="https://api.openai.com/v1")
            model = gr.Textbox(label="Model", placeholder="dall-e-3", value="dall-e-3")
            size = gr.Dropdown(
                label="Size", choices=["256x256", "512x512", "1024x1024", "1792x1024", "1024x1792"], value="1024x1024"
            )
            quality = gr.Dropdown(label="Quality", choices=["standard", "hd"], value="standard")
            style = gr.Dropdown(label="Style", choices=["vivid", "natural"], value="vivid")

        # Right column for prompt and image
        with gr.Column(scale=2):
            with gr.Row():
                prompt = gr.Text(
                    label="Prompt",
                    show_label=False,
                    max_lines=1,
                    placeholder=random.choice(examples),
                    value=random.choice(examples),
                    container=False,
                )
                run_button = gr.Button("Run", scale=0, variant="primary")

            result = gr.Image(
                label="Result",
                show_label=False,
                elem_classes=["generated-image"],  # Add the CSS class
            )
            gr.Examples(examples=examples, inputs=[prompt])

    # Load settings from browser storage
    @demo.load(inputs=[settings_state], outputs=[api_key, base_url, model, size, quality, style])
    def load_from_local_storage(saved_values):
        return (
            saved_values[0],  # api_key
            saved_values[1],  # base_url
            saved_values[2],  # model
            saved_values[3],  # size
            saved_values[4],  # quality
            saved_values[5],  # style
        )

    # Save settings to browser storage
    @gr.on(
        inputs=[api_key, base_url, model, size, quality, style],
        outputs=[settings_state],
        triggers=[api_key.change, base_url.change, model.change, size.change, quality.change, style.change],
    )
    def save_to_local_storage(api_key, base_url, model, size, quality, style):
        return [api_key, base_url, model, size, quality, style]

    # Main generation event
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=generate_image,
        inputs=[
            prompt,
            api_key,
            base_url,
            model,
            size,
            quality,
            style,
        ],
        outputs=result,
    )

if __name__ == "__main__":
    demo.launch()