app.py

import base64
import os
import mimetypes
from google import genai
from google.genai import types
import gradio as gr
import io
from PIL import Image

def save_binary_file(file_name, data):
    f = open(file_name, "wb")
    f.write(data)
    f.close()

def generate_image(prompt, output_filename="generated_image"):
    # Initialize client with the API key
    client = genai.Client(
        api_key="AIzaSyAQcy3LfrkMy6DqS_8MqftAXu1Bx_ov_E8",
    )

    model = "gemini-2.0-flash-exp-image-generation"
    contents = [
        types.Content(
            role="user",
            parts=[
                types.Part.from_text(text=prompt),
            ],
        ),
    ]
    generate_content_config = types.GenerateContentConfig(
        temperature=1,
        top_p=0.95,
        top_k=40,
        max_output_tokens=8192,
        response_modalities=[
            "image",
            "text",
        ],
        safety_settings=[
            types.SafetySetting(
                category="HARM_CATEGORY_CIVIC_INTEGRITY",
                threshold="OFF",
            ),
        ],
        response_mime_type="text/plain",
    )

    # Generate the content
    response = client.models.generate_content_stream(
        model=model,
        contents=contents,
        config=generate_content_config,
    )

    # Process the response
    for chunk in response:
        if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
            continue
        if chunk.candidates[0].content.parts[0].inline_data:
            inline_data = chunk.candidates[0].content.parts[0].inline_data
            file_extension = mimetypes.guess_extension(inline_data.mime_type)
            filename = f"{output_filename}{file_extension}"
            save_binary_file(filename, inline_data.data)
            
            # Convert binary data to PIL Image for Gradio display
            img = Image.open(io.BytesIO(inline_data.data))
            return img, f"Image saved as {filename}"
        else:
            return None, chunk.text

    return None, "No image generated"

# Function to handle chat interaction
def chat_handler(user_input, chat_history):
    # Add user message to chat history
    chat_history.append({"role": "user", "content": user_input})
    
    # Generate image based on user input
    img, status = generate_image(user_input)
    
    # Add AI response to chat history
    if img:
        chat_history.append({"role": "assistant", "content": img})
    chat_history.append({"role": "assistant", "content": status})
    
    return chat_history, ""

# Create Gradio interface with chatbot layout
with gr.Blocks(title="Image Editing Chatbot") as demo:
    gr.Markdown("# Image Editing Chatbot")
    gr.Markdown("Type a prompt to generate or edit an image using Google's Gemini model")
    
    # Chatbot display area
    chatbot = gr.Chatbot(
        label="Chat",
        height=400,
        type="messages",  # Explicitly set to 'messages' format
        avatar_images=(None, None)  # No avatars for simplicity
    )
    
    # Input area
    with gr.Row():
        prompt_input = gr.Textbox(
            label="",
            placeholder="Type something",
            show_label=False,
            container=False,
            scale=4
        )
        run_btn = gr.Button("Run", scale=1)
    
    # State to maintain chat history
    chat_state = gr.State([])

    # Connect the button to the chat handler
    run_btn.click(
        fn=chat_handler,
        inputs=[prompt_input, chat_state],
        outputs=[chatbot, prompt_input]
    )

    # Also allow Enter key to submit
    prompt_input.submit(
        fn=chat_handler,
        inputs=[prompt_input, chat_state],
        outputs=[chatbot, prompt_input]
    )

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