app.py

import os
import mimetypes
import tempfile
import urllib.parse
import urllib.request

import gradio as gr
from PIL import Image  # kept for compatibility / future UI use

from google import genai
from google.genai import types

# Config
MODEL_ID = os.environ.get("GEMINI_MODEL", "gemini-2.5-flash-image")


def _load_bytes(path_or_url: str) -> tuple[bytes, str]:
    """Return (data, mime_type) from either a local path or a http(s) URL."""
    parsed = urllib.parse.urlparse(path_or_url)
    is_url = parsed.scheme in ("http", "https")

    if is_url:
        req = urllib.request.Request(path_or_url, headers={"User-Agent": "Mozilla/5.0"})
        with urllib.request.urlopen(req, timeout=30) as resp:
            data = resp.read()
            mime_type = resp.headers.get_content_type() or "application/octet-stream"
        if mime_type in ("application/octet-stream", "binary/octet-stream"):
            guess, _ = mimetypes.guess_type(path_or_url)
            if guess:
                mime_type = guess
        return data, mime_type

    with open(path_or_url, "rb") as f:
        data = f.read()
    mime_type, _ = mimetypes.guess_type(path_or_url)
    return data, (mime_type or "application/octet-stream")


def generate_content(prompt: str, image1: str | None, image2: str | None, image3: str | None, api_key: str) -> tuple[str | None, str | None]:
    """Generate content using the current Google GenAI SDK (google-genai)."""
    api_key = (api_key or "").strip() or os.environ.get("GEMINI_API_KEY")
    if not api_key:
        raise gr.Error("Please set Space secret GEMINI_API_KEY (recommended) or paste your key into the UI field.")

    try:
        client = genai.Client(api_key=api_key)

        # Prepare request contents: images first, then instruction text
        contents: list[types.Part | str] = []
        for img_path in [image1, image2, image3]:
            if img_path:
                data, mime_type = _load_bytes(img_path)
                contents.append(types.Part.from_bytes(data=data, mime_type=mime_type))

        contents.append((prompt or "").strip() or "Describe the provided images.")

        # Ask for TEXT+IMAGE when supported; retry as TEXT-only if needed.
        try:
            response = client.models.generate_content(
                model=MODEL_ID,
                contents=contents,
                config=types.GenerateContentConfig(response_modalities=["TEXT", "IMAGE"]),
            )
        except Exception:
            response = client.models.generate_content(model=MODEL_ID, contents=contents)

        image_output: str | None = None
        text_chunks: list[str] = []

        if getattr(response, "candidates", None):
            for part in response.candidates[0].content.parts:
                if getattr(part, "text", None):
                    text_chunks.append(part.text)

                inline = getattr(part, "inline_data", None)
                if inline and getattr(inline, "data", None) and image_output is None:
                    filename = os.path.join(tempfile.gettempdir(), f"gemini_output_{os.urandom(4).hex()}.png")
                    with open(filename, "wb") as f:
                        f.write(inline.data)
                    image_output = filename

        if not text_chunks and getattr(response, "text", None):
            text_chunks.append(response.text)

        text_output = "\n".join([t for t in text_chunks if t]).strip() or None

        if image_output:
            return image_output, text_output or "Image generated successfully!"
        if text_output:
            return None, text_output
        return None, "No content generated. Try a different prompt or model."

    except Exception as e:
        raise gr.Error(f"Error generating content: {type(e).__name__}: {str(e)}")


# Gradio app
with gr.Blocks() as demo:
    gr.HTML('''        <div style="text-align: center; margin-bottom: 20px;">
      <h1>✨ Gemini Flash Studio</h1>
      <p>Image Generation & Editing powered by the Gemini API</p>
      <p style="opacity:0.8;">Model: <code>{model}</code> (override with <code>GEMINI_MODEL</code>)</p>
      <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank"
         style="color: #007bff; text-decoration: none; font-weight: bold;">Built with anycoder</a>
    </div>
    '''.format(model=MODEL_ID))

    gr.Markdown("### Configuration")
    with gr.Row():
        api_key_input = gr.Textbox(
            label="Gemini API Key (optional if GEMINI_API_KEY secret is set)",
            type="password",
            placeholder="If you set GEMINI_API_KEY in Space Secrets, you can leave this empty.",
            scale=4,
        )

    gr.Markdown("### Inputs")
    with gr.Row():
        img_input_1 = gr.Image(label="Image 1 (Drag & Drop)", type="filepath", sources=["upload", "clipboard"], height=300)
        img_input_2 = gr.Image(label="Image 2 (Drag & Drop)", type="filepath", sources=["upload", "clipboard"], height=300)
        img_input_3 = gr.Image(label="Image 3 (Drag & Drop)", type="filepath", sources=["upload", "clipboard"], height=300)

    prompt_input = gr.Textbox(
        label="✍️ Prompt",
        placeholder="Describe the image you want to generate or how you want to edit the uploaded images...",
        lines=3,
        show_copy_button=True,
    )

    generate_btn = gr.Button("✨ Generate", variant="primary", size="lg")

    gr.Markdown("### Output")
    with gr.Row():
        with gr.Column():
            output_image = gr.Image(label="Generated Image", type="filepath", height=400)
        with gr.Column():
            output_text = gr.Textbox(label="Response / Description", lines=10)

    gr.Examples(
        examples=[
            ["Generate a futuristic cyberpunk city with neon lights", None, None, None],
            ["Describe what is in this image", "https://gradio-builds.s3.amazonaws.com/assets/cheetah-003.jpg", None, None],
            ["Edit this image: make it look like a painting", None, None, None],
        ],
        inputs=[prompt_input, img_input_1, img_input_2, img_input_3],
    )

    generate_btn.click(
        fn=generate_content,
        inputs=[prompt_input, img_input_1, img_input_2, img_input_3, api_key_input],
        outputs=[output_image, output_text],
        api_visibility="public",
    )

    demo.launch(
        theme=gr.themes.Soft(
            primary_hue="indigo",
            secondary_hue="cyan",
            neutral_hue="slate",
            font=gr.themes.GoogleFont("Inter"),
            text_size="lg",
            spacing_size="lg",
            radius_size="md",
        ).set(
            button_primary_background_fill="*primary_500",
            button_primary_background_fill_hover="*primary_600",
            button_primary_text_color="white",
            block_border_width="0px",
            block_shadow="*shadow_drop_lg",
        ),
        footer_links=[{"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}],
    )