app.py

import io
import os

import gradio as gr
from google import genai
from PIL import Image

# Initialize the Gemini client
client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))

GREEN_SCREEN_COLOR = (0, 255, 0)  # chroma key green for padding regions


def add_padding(image, side, padding_size):
    """Add green-screen padding to specified side of image"""
    width, height = image.size
    if side == "left":
        new_image = Image.new(
            "RGB", (width + padding_size, height), color=GREEN_SCREEN_COLOR
        )
        new_image.paste(image, (padding_size, 0))
    elif side == "right":
        new_image = Image.new(
            "RGB", (width + padding_size, height), color=GREEN_SCREEN_COLOR
        )
        new_image.paste(image, (0, 0))
    elif side == "top":
        new_image = Image.new(
            "RGB", (width, height + padding_size), color=GREEN_SCREEN_COLOR
        )
        new_image.paste(image, (0, padding_size))
    elif side == "bottom":
        new_image = Image.new(
            "RGB", (width, height + padding_size), color=GREEN_SCREEN_COLOR
        )
        new_image.paste(image, (0, 0))
    return new_image


def preview_padding(image, side, padding_size):
    """Add green-screen padding to image and show preview"""
    if image is None:
        return None, None, "No image to pad", ""
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)

    padded_image = add_padding(image, side, padding_size)
    width, height = padded_image.size
    info_text = (
        f"Padded size: {width} x {height} pixels ({padding_size}px added to {side})"
    )

    direction_text = {
        "left": "left side",
        "right": "right side",
        "top": "top",
        "bottom": "bottom",
    }
    default_prompt = (
        f"This image has a green-screen region on the {direction_text[side]}. "
        f"Please naturally extend and outpaint the image to fill in the green-screen area on the {direction_text[side]}. "
        f"Make it look seamless and consistent with the existing image content, matching the style, colors, and content naturally. "
        f"Do not leave any green regions - completely fill in the {direction_text[side]} area."
    )
    return padded_image, side, info_text, default_prompt


def outpaint_with_nano_banana(padded_image, side, custom_prompt):
    """Use Nano Banana to outpaint the padded image"""
    if padded_image is None:
        return None, "No padded image to outpaint", None
    if side is None:
        return None, "No side information available", None

    if custom_prompt and custom_prompt.strip():
        prompt = custom_prompt
    else:
        direction_text = {
            "left": "left side",
            "right": "right side",
            "top": "top",
            "bottom": "bottom",
        }
        prompt = (
            f"This image has a green-screen region on the {direction_text[side]}. "
            f"Please naturally extend and outpaint the image to fill in the green-screen area on the {direction_text[side]}."
        )

    try:
        response = client.models.generate_content(
            model="gemini-2.5-flash-image", contents=[padded_image, prompt]
        )
        for part in response.candidates[0].content.parts:
            if getattr(part, "inline_data", None):
                result_image = Image.open(io.BytesIO(part.inline_data.data))
                w, h = result_image.size
                return result_image, f"Output size: {w} x {h} pixels", result_image
        w, h = padded_image.size
        return padded_image, f"Output size: {w} x {h} pixels (no result from API)", padded_image
    except Exception as e:
        return padded_image, f"Error: {str(e)}", padded_image


def reuse_outpainted_image(outpainted_image):
    """Set the most recent outpainted image as the new input and reset padding state."""
    if outpainted_image is None:
        return (
            None,
            "No outpainted image available. Run an outpaint first.",
            None,
            "No padding applied",
            None,
            None,
        )
    if not isinstance(outpainted_image, Image.Image):
        outpainted_image = Image.fromarray(outpainted_image)
    w, h = outpainted_image.size
    return (
        outpainted_image,
        f"Image size: {w} x {h} pixels",
        None,
        "No padding applied",
        None,
        None,
    )


def get_image_info(image):
    """Get image size information"""
    if image is None:
        return "No image uploaded"
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    w, h = image.size
    return f"Image size: {w} x {h} pixels"


# Create Gradio interface
with gr.Blocks(title="Nano Banana Outpainting") as demo:
    gr.Markdown("# 🍌 Nano Banana Outpainting App")
    gr.Markdown(
        "Upload an image and click a direction to add green-screen padding. Review the padding, then click 'Outpaint' to fill it in with AI."
    )

    # State variables
    padded_state = gr.State(None)
    side_state = gr.State(None)
    outpainted_state = gr.State(None)

    # ─────────────────────────────
    # ROW 1: The three image panes
    # ─────────────────────────────
    with gr.Row():
        # LEFT: Upload + info
        with gr.Column():
            input_image = gr.Image(label="Upload Image", type="pil")
            image_info = gr.Textbox(
                label="Input Image Info", interactive=False, value="No image uploaded"
            )

        # MIDDLE: Padded preview + info
        with gr.Column():
            padded_preview = gr.Image(
                label="Padded Image (with green-screen pixels)",
                type="pil",
                interactive=False,
            )
            padded_info = gr.Textbox(
                label="Padded Image Info", interactive=False, value="No padding applied"
            )

        # RIGHT: Outpainted result + info
        with gr.Column():
            output_image = gr.Image(
                label="Outpainted Result", type="pil", interactive=False
            )
            output_info = gr.Textbox(
                label="Output Image Info", interactive=False, value="No output yet"
            )

    # ─────────────────────────────
    # ROW 2: Step 1 and Step 2 side-by-side
    # ─────────────────────────────
    with gr.Row():
        # STEP 1 PANEL (left)
        with gr.Column(scale=1, min_width=360):
            gr.Markdown("### Step 1: Choose direction to add green-screen padding")
            # Padding size belongs to Step 1
            gr.Markdown("**Padding Size:**")
            padding_slider = gr.Slider(
                minimum=50,
                maximum=500,
                value=100,
                step=10,
                label="Green-Screen Padding Size (pixels)",
                info="Adjust how many pixels to add",
            )
            with gr.Row():
                btn_top = gr.Button("⬆️ Top", size="lg")
            with gr.Row():
                btn_left = gr.Button("⬅️ Left", size="lg")
                btn_right = gr.Button("➡️ Right", size="lg")
            with gr.Row():
                btn_bottom = gr.Button("⬇️ Bottom", size="lg")

            gr.Markdown("### Step 3: Outpaint the green-screen region")
            btn_outpaint = gr.Button(
                "🎨 Outpaint with Nano Banana", size="lg", variant="primary"
            )
            gr.Markdown("### Optional: Reuse the outpainted result")
            btn_reuse = gr.Button("♻️ Reuse Outpainted Image")

        # STEP 2 PANEL (right)
        with gr.Column(scale=1, min_width=360):
            gr.Markdown("### Step 2: Customize Prompt (Optional)")
            prompt_textbox = gr.Textbox(
                label="Outpainting Prompt",
                placeholder="Click a direction button to see the default prompt, then edit if desired",
                value="",
                lines=6,
                info="Edit this prompt to control what Nano Banana generates.",
                show_label=True,
            )

    # Wire events
    input_image.change(fn=get_image_info, inputs=input_image, outputs=image_info)

    btn_top.click(
        fn=lambda img, pad_size: preview_padding(img, "top", pad_size),
        inputs=[input_image, padding_slider],
        outputs=[padded_preview, side_state, padded_info, prompt_textbox],
    ).then(fn=lambda img: img, inputs=padded_preview, outputs=padded_state)

    btn_bottom.click(
        fn=lambda img, pad_size: preview_padding(img, "bottom", pad_size),
        inputs=[input_image, padding_slider],
        outputs=[padded_preview, side_state, padded_info, prompt_textbox],
    ).then(fn=lambda img: img, inputs=padded_preview, outputs=padded_state)

    btn_left.click(
        fn=lambda img, pad_size: preview_padding(img, "left", pad_size),
        inputs=[input_image, padding_slider],
        outputs=[padded_preview, side_state, padded_info, prompt_textbox],
    ).then(fn=lambda img: img, inputs=padded_preview, outputs=padded_state)

    btn_right.click(
        fn=lambda img, pad_size: preview_padding(img, "right", pad_size),
        inputs=[input_image, padding_slider],
        outputs=[padded_preview, side_state, padded_info, prompt_textbox],
    ).then(fn=lambda img: img, inputs=padded_preview, outputs=padded_state)

    btn_outpaint.click(
        fn=outpaint_with_nano_banana,
        inputs=[padded_state, side_state, prompt_textbox],
        outputs=[output_image, output_info, outpainted_state],
    )

    btn_reuse.click(
        fn=reuse_outpainted_image,
        inputs=outpainted_state,
        outputs=[input_image, image_info, padded_preview, padded_info, padded_state, side_state],
    )

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