Spaces:
Running
Running
| import base64 | |
| import os | |
| import mimetypes | |
| import tempfile | |
| import time | |
| from google import genai | |
| from google.genai import types | |
| import gradio as gr | |
| def generate_image(api_key, prompt, file_name): | |
| # Validate inputs | |
| if not api_key: | |
| return None, "Please enter your Gemini API key" | |
| if not prompt: | |
| return None, "Please enter a prompt for image generation" | |
| if not file_name: | |
| file_name = f"gemini_image_{int(time.time())}" | |
| try: | |
| # Set up the client | |
| client = genai.Client(api_key=api_key) | |
| model = "gemini-2.0-flash-exp-image-generation" | |
| # Create content | |
| contents = [ | |
| types.Content( | |
| role="user", | |
| parts=[types.Part.from_text(text=prompt)], | |
| ), | |
| ] | |
| # Configure generation settings | |
| generate_content_config = types.GenerateContentConfig( | |
| response_modalities=["image", "text"], | |
| safety_settings=[ | |
| types.SafetySetting(category="HARM_CATEGORY_HARASSMENT", threshold="BLOCK_NONE"), | |
| types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="BLOCK_NONE"), | |
| types.SafetySetting(category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="BLOCK_NONE"), | |
| types.SafetySetting(category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="BLOCK_NONE"), | |
| types.SafetySetting(category="HARM_CATEGORY_CIVIC_INTEGRITY", threshold="BLOCK_NONE"), | |
| ], | |
| response_mime_type="text/plain", | |
| ) | |
| # Create a temporary directory to store the image | |
| temp_dir = tempfile.mkdtemp() | |
| generated_image_path = None | |
| generation_text = "" | |
| # Stream the content | |
| for chunk in client.models.generate_content_stream( | |
| model=model, | |
| contents=contents, | |
| config=generate_content_config, | |
| ): | |
| 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) or '.jpg' | |
| generated_image_path = os.path.join(temp_dir, f"{file_name}{file_extension}") | |
| with open(generated_image_path, "wb") as f: | |
| f.write(inline_data.data) | |
| generation_text += f"Image of type {inline_data.mime_type} generated successfully." | |
| else: | |
| generation_text += chunk.text | |
| if generated_image_path: | |
| return generated_image_path, generation_text | |
| else: | |
| return None, "No image was generated. Try a different prompt." | |
| except Exception as e: | |
| return None, f"Error: {str(e)}" | |
| def create_ui(): | |
| with gr.Blocks(title="Gemini Image Generator") as demo: | |
| gr.Markdown("# Gemini Image Generator") | |
| gr.Markdown("Generate images using Google's Gemini 2.0 Flash Image Generation model") | |
| with gr.Row(): | |
| with gr.Column(): | |
| api_key = gr.Textbox( | |
| label="Gemini API Key", | |
| placeholder="Enter your Gemini API key here", | |
| type="password" | |
| ) | |
| prompt = gr.Textbox( | |
| label="Prompt", | |
| placeholder="Describe the image you want to generate", | |
| lines=3 | |
| ) | |
| file_name = gr.Textbox( | |
| label="Output File Name (optional)", | |
| placeholder="Enter a file name (without extension)" | |
| ) | |
| generate_btn = gr.Button("Generate Image", variant="primary") | |
| with gr.Column(): | |
| output_image = gr.Image(label="Generated Image", type="filepath") | |
| output_text = gr.Textbox(label="Generation Info", lines=2) | |
| generate_btn.click( | |
| fn=generate_image, | |
| inputs=[api_key, prompt, file_name], | |
| outputs=[output_image, output_text] | |
| ) | |
| gr.Markdown(""" | |
| ## How to use | |
| 1. Enter your Gemini API key (get one from https://ai.google.dev/) | |
| 2. Write a detailed prompt describing the image you want to generate | |
| 3. (Optional) Provide a file name for your generated image | |
| 4. Click "Generate Image" and wait for the result | |
| ## Notes | |
| - The model used is `gemini-2.0-flash-exp-image-generation` | |
| - All safety filters are set to "BLOCK_NONE" - use responsibly | |
| - Image generation may take a few seconds to complete | |
| """) | |
| return demo | |
| if __name__ == "__main__": | |
| demo = create_ui() | |
| demo.launch() |