Image_text_generation.py

import os
import io
import tempfile
from datetime import datetime
import PIL.Image
from google import genai
from google.genai import types
import httpx
from imgurpython import ImgurClient

genai_client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])

class Image_text_Generator:
    def __init__(self, imgur):
        # 初始化 Imgur 客戶端
        self.imgur_client = imgur

    def generate_image_with_gemini(self, prompt):
        """
        使用 Gemini 模型生成圖片。

        參數:
        prompt (str): 用於生成圖片的提示詞。

        返回:
        bytes: 生成的圖片的二進位資料，如果生成失敗則返回 None。
        """
        response = 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

    def generate_txt_with_gemini(self, prompt,doc_url):
        doc_data = httpx.get(doc_url).content

        response = client.models.generate_content(
            model="gemini-1.5-flash",
            contents=[
                types.Part.from_bytes(data=doc_data, mime_type='application/pdf',),prompt]
        )
        
        return response.text