Spaces:
Running
Running
File size: 4,999 Bytes
e2ff448 |
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 131 |
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() |