Update Image_generation.py

Image_generation.py

@@ -1,47 +1,76 @@
 import os
 import io
-import PIL.Image
 import tempfile
-from google import genai
-from google.genai import types
 from datetime import datetime
+import PIL.Image
+import google.generativeai as genai
+from google.generativeai import types
+from imgurpython import ImgurClient
+
+class ImageGenerator:
+    def __init__(self, google_api_key, imgur_client_id, imgur_client_secret, imgur_access_token, imgur_refresh_token):
+        # 初始化 Gemini 客戶端
+        self.genai_client = genai.Client(api_key=google_api_key)
+        # 初始化 Imgur 客戶端
+        self.imgur_client = ImgurClient(imgur_client_id, imgur_client_secret, imgur_access_token, imgur_refresh_token)
+
+    def generate_image_with_gemini(self, prompt):
+        """
+        使用 Gemini 模型生成圖片。
+
+        參數:
+        prompt (str): 用於生成圖片的提示詞。
+
+        返回:
+        bytes: 生成的圖片的二進位資料，如果生成失敗則返回 None。
+        """
+        response = self.genai_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.text is not None:
+                print(part.text)
+            elif part.inline_data is not None:
+                return part.inline_data.data
+        return None
+
+    def upload_image_to_imgur(self, image_binary, album=None, name="gemini-image", title="Gemini Generated Image"):
+        """
+        將圖片上傳到 Imgur。
+
+        參數:
+        image_binary (bytes): 要上傳的圖片的二進位資料。
+        album (str): 要上傳到的相簿 ID，預設為 None。
+        name (str): 圖片的名稱，預設為 "gemini-image"。
+        title (str): 圖片的標題，預設為 "Gemini Generated Image"。
+
+        返回:
+        str: 上傳後的圖片 URL，如果上傳失敗則返回 None。
+        """
+        try:
+            # 將二進位資料轉換為 PIL Image
+            image = PIL.Image.open(io.BytesIO(image_binary))
+
+            # 建立暫存檔案以便上傳
+            with tempfile.NamedTemporaryFile(suffix=".png", delete=True) as tmp:
+                image.save(tmp.name, format='PNG')
+
+                # 準備上傳資訊
+                config = {
+                    'album': album,
+                    'name': name,
+                    'title': title,
+                    'description': f'Generated by Gemini - {datetime.now()}'
+                }
+
+                # 使用 Imgur 客戶端上傳圖片
+                uploaded_image = self.imgur_client.upload_from_path(tmp.name, config=config, anon=False)
 
-# 使用 Gemini 生成圖片
-client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])
-
-def generate_image_with_gemini(prompt):
-    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.text is not None:
-            print(part.text)
-        elif part.inline_data is not None:
-            return part.inline_data.data
-
-    return
-
-def upload_image_to_imgur(client, image_binary, album=None, name="gemini-image", title="gemini Generated Image"):
-    # 將 binary 資料轉為 PIL Image
-    image = PIL.Image.open(io.BytesIO(image_binary))
-
-    # 建立暫存檔案來上傳 (因為 ImgurClient 需要檔案路徑)
-    with tempfile.NamedTemporaryFile(suffix=".png", delete=True) as tmp:
-        image.save(tmp.name, format='PNG')
-
-        # 準備上傳資訊
-        config = {
-            'album': album,
-            'name': name,
-            'title': title,
-            'description': f'Generated by gemini - {datetime.now()}'
-        }
-
-        # 使用 client 進行圖片上傳
-        uploaded_image = client.upload_from_path(tmp.name, config=config, anon=False)
-
-    # 回傳圖片網址
-    return uploaded_image['link']
+            # 返回圖片的連結
+            return uploaded_image['link']
+        except Exception as e:
+            print(f"圖片上傳失敗: {e}")
+            return None