Spaces:
Running
Running
File size: 3,812 Bytes
ac3fd77 19ffd31 ac3fd77 19ffd31 ac3fd77 19ffd31 ac3fd77 19ffd31 ac3fd77 19ffd31 ac3fd77 19ffd31 ac3fd77 19ffd31 ac3fd77 19ffd31 ed4625b 31bf755 19ffd31 ed4625b 31bf755 ed4625b 31bf755 ed4625b 19ffd31 ed4625b 31bf755 ed4625b 19ffd31 ed4625b 19ffd31 ac3fd77 19ffd31 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
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() |