Update app.py

app.py

@@ -3,17 +3,19 @@ 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():
+def generate_image(prompt, output_filename="generated_image"):
+    # Initialize client with the API key
     client = genai.Client(
-        api_key=os.environ.get("GEMINI_API_KEY"),
+        api_key="AIzaSyAQcy3LfrkMy6DqS_8MqftAXu1Bx_ov_E8",
     )
 
     model = "gemini-2.0-flash-exp-image-generation"
@@ -21,7 +23,7 @@ def generate():
         types.Content(
             role="user",
             parts=[
-                types.Part.from_text(text="""INSERT_INPUT_HERE"""),
+                types.Part.from_text(text=prompt),
             ],
         ),
     ]
@@ -37,33 +39,66 @@ def generate():
         safety_settings=[
             types.SafetySetting(
                 category="HARM_CATEGORY_CIVIC_INTEGRITY",
-                threshold="OFF",  # Off
+                threshold="OFF",
             ),
         ],
         response_mime_type="text/plain",
     )
 
-    for chunk in client.models.generate_content_stream(
+    # 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:
-            file_name = "ENTER_FILE_NAME"
             inline_data = chunk.candidates[0].content.parts[0].inline_data
             file_extension = mimetypes.guess_extension(inline_data.mime_type)
-            save_binary_file(
-                f"{file_name}{file_extension}", inline_data.data
+            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"
+
+# Create Gradio interface
+with gr.Blocks(title="Image Generation App") as demo:
+    gr.Markdown("# Image Generation with Gemini")
+    gr.Markdown("Enter a prompt to generate an image using Google's Gemini model")
+    
+    with gr.Row():
+        with gr.Column():
+            prompt_input = gr.Textbox(
+                label="Prompt",
+                placeholder="Enter your image description here...",
+                lines=3
             )
-            print(
-                "File of mime type"
-                f" {inline_data.mime_type} saved"
-                f"to: {file_name}"
+            filename_input = gr.Textbox(
+                label="Output Filename",
+                value="generated_image",
+                placeholder="Enter desired filename (without extension)"
             )
-        else:
-            print(chunk.text)
+            generate_btn = gr.Button("Generate Image")
+        
+        with gr.Column():
+            image_output = gr.Image(label="Generated Image")
+            text_output = gr.Textbox(label="Status")
+
+    # Connect the button to the generation function
+    generate_btn.click(
+        fn=generate_image,
+        inputs=[prompt_input, filename_input],
+        outputs=[image_output, text_output]
+    )
 
 if __name__ == "__main__":
-    generate()
+    demo.launch()