Update gemini.py

gemini.py

@@ -1,13 +1,12 @@
 import os
-import uuid
 from io import BytesIO
 from PIL import Image
 import gradio as gr
-from google import genai
-from google.genai import types
+import google.generativeai as genai
+from google.generativeai import types
 import logging
 
-# 設定logging
+# 設定 logging
 logging.basicConfig(
     filename='app.log',
     level=logging.INFO,
@@ -16,18 +15,12 @@ logging.basicConfig(
 
 # 設定 Gemini API 金鑰
 GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
-client = genai.Client(api_key=GEMINI_API_KEY)
-
-# 設定圖片儲存目錄
-STATIC_IMAGE_PATH = "static/images"
-os.makedirs(STATIC_IMAGE_PATH, exist_ok=True)
-
-# 取得 Hugging Face Space 的主機名稱
-SPACE_HOST = os.environ.get("SPACE_HOST", "your-space-name.hf.space")
+genai.configure(api_key=GEMINI_API_KEY)
+client = genai.Client()
 
 def generate_image(prompt):
     """
-    使用 Gemini API 根據提示詞生成圖片，並返回圖片的公開 URL。
+    使用 Gemini API 根據提示詞生成圖片，並返回 PIL 圖像物件。
     """
     response = client.models.generate_content(
         model="gemini-2.0-flash-exp-image-generation",
@@ -41,16 +34,11 @@ def generate_image(prompt):
     for part in response.candidates[0].content.parts:
         if part.inline_data is not None:
             image = Image.open(BytesIO(part.inline_data.data))
-            filename = f"{uuid.uuid4().hex}.png"
-            image_path = os.path.join(STATIC_IMAGE_PATH, filename)
-            image.save(image_path)
-
-            # 建立圖片的公開 URL
-            image_url = f"https://{SPACE_HOST}/static/images/{filename}"
-            logging.info(f"生成的圖片 URL: {image_url}")
-            return image_url
+            logging.info("成功生成圖片。")
+            return image
 
-    return "未能生成圖片，請嘗試其他提示詞。"
+    logging.warning("未能生成圖片，請嘗試其他提示詞。")
+    return None
 
 # 建立 Gradio 介面
 with gr.Blocks() as demo:
@@ -60,8 +48,8 @@ with gr.Blocks() as demo:
     image_output = gr.Image(label="生成的圖片")
 
     def on_generate(prompt):
-        image_url = generate_image(prompt)
-        return image_url
+        image = generate_image(prompt)
+        return image
 
     generate_button.click(fn=on_generate, inputs=prompt_input, outputs=image_output)