ImageGen / gemini.py
chen0857's picture
Update gemini.py
e5a77e5 verified
raw
history blame
1.73 kB
import os
from io import BytesIO
from PIL import Image
import gradio as gr
from google import genai
from google.genai import types
import logging
# 設定 logging
logging.basicConfig(
filename='app.log',
level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s'
)
# 設定 Gemini API 金鑰
# === 初始化 Google Gemini ===
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
client = genai.Client(api_key=GEMINI_API_KEY)
def generate_image(prompt):
"""
使用 Gemini API 根據提示詞生成圖片,並返回 PIL 圖像物件。
"""
response = client.models.generate_content(
model="gemini-2.0-flash-exp-image-generation",
contents=prompt,
config=types.GenerateContentConfig(
response_modalities=["TEXT", "IMAGE"]
),
)
# 處理回應中的圖片
for part in response.candidates[0].content.parts:
if part.inline_data is not None:
image = Image.open(BytesIO(part.inline_data.data))
logging.info("成功生成圖片。")
return image
logging.warning("未能生成圖片,請嘗試其他提示詞。")
return None
# 建立 Gradio 介面
with gr.Blocks() as demo:
gr.Markdown("## 🖼️ B11090054 陳緯綸 Gemini 圖片生成器")
prompt_input = gr.Textbox(label="輸入提示詞", placeholder="例如:一隻戴著墨鏡的貓在沙灘上")
generate_button = gr.Button("生成圖片")
image_output = gr.Image(label="生成的圖片")
def on_generate(prompt):
image = generate_image(prompt)
return image
generate_button.click(fn=on_generate, inputs=prompt_input, outputs=image_output)
if __name__ == "__main__":
demo.launch()