Deadmon commited on
Commit
25b4ca5
·
verified ·
1 Parent(s): 2150980

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -129
app.py CHANGED
@@ -1,117 +1,4 @@
1
- import base64
2
- import os
3
- import mimetypes
4
- from google import genai
5
- from google.genai import types
6
- import gradio as gr
7
- import io
8
- from PIL import Image
9
-
10
- def save_binary_file(file_name, data):
11
- f = open(file_name, "wb")
12
- f.write(data)
13
- f.close()
14
-
15
- def generate_image(prompt, image=None, output_filename="generated_image"):
16
- # Initialize client with the API key
17
- client = genai.Client(
18
- api_key="AIzaSyAQcy3LfrkMy6DqS_8MqftAXu1Bx_ov_E8",
19
- )
20
-
21
- model = "gemini-2.0-flash-exp-image-generation"
22
- parts = [types.Part.from_text(text=prompt)]
23
-
24
- # If an image is provided, add it to the content
25
- if image:
26
- # Convert PIL Image to bytes
27
- img_byte_arr = io.BytesIO()
28
- image.save(img_byte_arr, format="PNG")
29
- img_bytes = img_byte_arr.getvalue()
30
- # Add the image as a Part with inline_data
31
- parts.append({
32
- "inline_data": {
33
- "mime_type": "image/png",
34
- "data": img_bytes
35
- }
36
- })
37
-
38
- contents = [
39
- types.Content(
40
- role="user",
41
- parts=parts,
42
- ),
43
- ]
44
- generate_content_config = types.GenerateContentConfig(
45
- temperature=1,
46
- top_p=0.95,
47
- top_k=40,
48
- max_output_tokens=8192,
49
- response_modalities=[
50
- "image",
51
- "text",
52
- ],
53
- safety_settings=[
54
- types.SafetySetting(
55
- category="HARM_CATEGORY_CIVIC_INTEGRITY",
56
- threshold="OFF",
57
- ),
58
- ],
59
- response_mime_type="text/plain",
60
- )
61
-
62
- # Generate the content
63
- response = client.models.generate_content_stream(
64
- model=model,
65
- contents=contents,
66
- config=generate_content_config,
67
- )
68
-
69
- # Process the response
70
- for chunk in response:
71
- if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
72
- continue
73
- if chunk.candidates[0].content.parts[0].inline_data:
74
- inline_data = chunk.candidates[0].content.parts[0].inline_data
75
- file_extension = mimetypes.guess_extension(inline_data.mime_type)
76
- filename = f"{output_filename}{file_extension}"
77
- save_binary_file(filename, inline_data.data)
78
-
79
- # Convert binary data to PIL Image for Gradio display
80
- img = Image.open(io.BytesIO(inline_data.data))
81
- return img, f"Image saved as {filename}"
82
- else:
83
- return None, chunk.text
84
-
85
- return None, "No image generated"
86
-
87
- # Function to handle chat interaction
88
- def chat_handler(prompt, user_image, chat_history, image_history, output_filename="generated_image"):
89
- # Add the prompt to the chat history
90
- if prompt:
91
- chat_history.append({"role": "user", "content": prompt})
92
-
93
- # Add the uploaded image to the image history
94
- if user_image:
95
- image_history.append(user_image)
96
-
97
- # If no input, return early
98
- if not prompt and not user_image:
99
- chat_history.append({"role": "assistant", "content": "Please provide a prompt or an image."})
100
- return chat_history, user_image, None, image_history, ""
101
-
102
- # Generate image based on user input
103
- img, status = generate_image(prompt or "Generate an image", user_image, output_filename)
104
-
105
- # Add the status message to the chat history
106
- chat_history.append({"role": "assistant", "content": status})
107
-
108
- # Add the generated image to the image history
109
- if img:
110
- image_history.append(img)
111
-
112
- return chat_history, user_image, img, image_history, ""
113
-
114
- # Create Gradio interface
115
  with gr.Blocks(title="Image Editing Chatbot") as demo:
116
  gr.Markdown("# Image Editing Chatbot")
117
  gr.Markdown("Upload an image and/or type a prompt to generate or edit an image using Google's Gemini model")
@@ -124,16 +11,10 @@ with gr.Blocks(title="Image Editing Chatbot") as demo:
124
  avatar_images=(None, None)
125
  )
126
 
127
- # Gallery to display image history as thumbnails
128
- image_history_display = gr.Gallery(
129
- label="Image History",
130
- height=200,
131
- columns=4,
132
- object_fit="contain",
133
- preview=True
134
- )
135
 
136
- # Separate image outputs
137
  with gr.Row():
138
  uploaded_image_output = gr.Image(label="Uploaded Image")
139
  generated_image_output = gr.Image(label="Generated Image")
@@ -159,22 +40,48 @@ with gr.Blocks(title="Image Editing Chatbot") as demo:
159
  )
160
  generate_btn = gr.Button("Generate Image")
161
 
162
- # State to maintain chat history and image history
163
  chat_state = gr.State([])
164
- image_history_state = gr.State([])
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
165
 
166
  # Connect the button to the chat handler
167
  generate_btn.click(
168
  fn=chat_handler,
169
- inputs=[prompt_input, image_input, chat_state, image_history_state, filename_input],
170
- outputs=[chatbot, uploaded_image_output, generated_image_output, image_history_display, prompt_input]
171
  )
172
 
173
  # Also allow Enter key to submit
174
  prompt_input.submit(
175
  fn=chat_handler,
176
- inputs=[prompt_input, image_input, chat_state, image_history_state, filename_input],
177
- outputs=[chatbot, uploaded_image_output, generated_image_output, image_history_display, prompt_input]
178
  )
179
 
180
  if __name__ == "__main__":
 
1
+ # Fallback: Use a gr.Gallery component for thumbnails
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
  with gr.Blocks(title="Image Editing Chatbot") as demo:
3
  gr.Markdown("# Image Editing Chatbot")
4
  gr.Markdown("Upload an image and/or type a prompt to generate or edit an image using Google's Gemini model")
 
11
  avatar_images=(None, None)
12
  )
13
 
14
+ # Gallery for thumbnails
15
+ thumbnail_gallery = gr.Gallery(label="Conversation Thumbnails", columns=4, height=150)
 
 
 
 
 
 
16
 
17
+ # Separate image outputs (gallery)
18
  with gr.Row():
19
  uploaded_image_output = gr.Image(label="Uploaded Image")
20
  generated_image_output = gr.Image(label="Generated Image")
 
40
  )
41
  generate_btn = gr.Button("Generate Image")
42
 
43
+ # State to maintain chat history and thumbnail list
44
  chat_state = gr.State([])
45
+ thumbnail_list = gr.State([])
46
+
47
+ def chat_handler(prompt, user_image, chat_history, thumbnails, output_filename="generated_image"):
48
+ # Add the prompt to the chat history
49
+ if prompt:
50
+ chat_history.append({"role": "user", "content": prompt})
51
+
52
+ # Add the uploaded image to the thumbnail list
53
+ if user_image:
54
+ thumbnails.append(user_image)
55
+
56
+ # If no input, return early
57
+ if not prompt and not user_image:
58
+ chat_history.append({"role": "assistant", "content": "Please provide a prompt or an image."})
59
+ return chat_history, user_image, None, thumbnails, ""
60
+
61
+ # Generate image based on user input
62
+ img, status = generate_image(prompt or "Generate an image", user_image, output_filename)
63
+
64
+ # Add the status message to the chat history
65
+ chat_history.append({"role": "assistant", "content": status})
66
+
67
+ # Add the generated image to the thumbnail list
68
+ if img:
69
+ thumbnails.append(img)
70
+
71
+ return chat_history, user_image, img, thumbnails, ""
72
 
73
  # Connect the button to the chat handler
74
  generate_btn.click(
75
  fn=chat_handler,
76
+ inputs=[prompt_input, image_input, chat_state, thumbnail_list, filename_input],
77
+ outputs=[chatbot, uploaded_image_output, generated_image_output, thumbnail_gallery, prompt_input]
78
  )
79
 
80
  # Also allow Enter key to submit
81
  prompt_input.submit(
82
  fn=chat_handler,
83
+ inputs=[prompt_input, image_input, chat_state, thumbnail_list, filename_input],
84
+ outputs=[chatbot, uploaded_image_output, generated_image_output, thumbnail_gallery, prompt_input]
85
  )
86
 
87
  if __name__ == "__main__":