Image_generation.py

import os
import io
import tempfile
from datetime import datetime
import PIL.Image
from google import genai
from google.genai 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)

            # 返回圖片的連結
            return uploaded_image['link']
        except Exception as e:
            print(f"圖片上傳失敗: {e}")
            return None