+ gen mode
This commit is contained in:
@@ -1,15 +1,14 @@
|
||||
import os
|
||||
import io
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from typing import List, Union, Dict, Any
|
||||
from PIL import Image
|
||||
from typing import List, Union
|
||||
|
||||
# Импортируем из нового SDK
|
||||
from PIL import Image
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
|
||||
# Для настройки логгера
|
||||
from models.enums import AspectRatios, Quality
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -18,82 +17,100 @@ class GoogleAdapter:
|
||||
if not api_key:
|
||||
raise ValueError("API Key for Gemini is missing")
|
||||
self.client = genai.Client(api_key=api_key)
|
||||
# Укажите актуальную модель.
|
||||
# Если gemini-3-pro-image-preview недоступна, используйте gemini-2.0-flash-exp
|
||||
self.model_name = "gemini-3-pro-preview"
|
||||
|
||||
def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
image_bytes: bytes = None,
|
||||
generate_image: bool = False
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Универсальный метод:
|
||||
- Если generate_image=True: просим модель вернуть картинку (Image Generation).
|
||||
- Если image_bytes переданы + generate_image=False: это Vision (описание фото).
|
||||
- Если image_bytes + generate_image=True: это Image-to-Image (редактирование).
|
||||
"""
|
||||
if generate_image:
|
||||
self.model_name = "gemini-3-pro-image-preview"
|
||||
else :
|
||||
self.model_name = "gemini-3-pro-preview"
|
||||
# Константы моделей
|
||||
self.TEXT_MODEL = "gemini-3-pro-preview"
|
||||
self.IMAGE_MODEL = "gemini-3-pro-image-preview"
|
||||
|
||||
def _prepare_contents(self, prompt: str, images_list: List[bytes] = None) -> list:
|
||||
"""Вспомогательный метод для подготовки контента (текст + картинки)"""
|
||||
contents = [prompt]
|
||||
if images_list:
|
||||
for img_bytes in images_list:
|
||||
try:
|
||||
# Gemini API требует PIL Image на входе
|
||||
image = Image.open(io.BytesIO(img_bytes))
|
||||
contents.append(image)
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing input image: {e}")
|
||||
return contents
|
||||
|
||||
# Если есть входное изображение (для Vision или для редактирования)
|
||||
if image_bytes:
|
||||
try:
|
||||
image = Image.open(io.BytesIO(image_bytes))
|
||||
contents.append(image)
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing input image: {e}")
|
||||
return {"error": "Не удалось обработать входящее изображение."}
|
||||
|
||||
# Настраиваем конфигурацию
|
||||
# Для генерации картинок добавляем 'IMAGE' в response_modalities
|
||||
modalities = ['TEXT']
|
||||
if generate_image:
|
||||
modalities.append('IMAGE')
|
||||
def generate_text(self, prompt: str, images_list: List[bytes] = None) -> str:
|
||||
"""
|
||||
Генерация текста (Чат или Vision).
|
||||
Возвращает строку с ответом.
|
||||
"""
|
||||
contents = self._prepare_contents(prompt, images_list)
|
||||
|
||||
try:
|
||||
# Вызов API (синхронный метод в обертке, но aiogram вызывает его в треде,
|
||||
# либо используйте client.aio для асинхронности если поддерживается версией SDK)
|
||||
# В google-genai v0.3+ есть асинхронный клиент, но для простоты здесь стандартный вызов.
|
||||
# Чтобы не блокировать event loop, в main.py мы обернем это в to_thread при необходимости,
|
||||
# но пока используем стандартный вызов.
|
||||
|
||||
response = self.client.models.generate_content(
|
||||
model=self.model_name,
|
||||
model=self.TEXT_MODEL,
|
||||
contents=contents,
|
||||
config=types.GenerateContentConfig(
|
||||
response_modalities=modalities,
|
||||
temperature=0.7 if not generate_image else 1.0,
|
||||
response_modalities=['TEXT'],
|
||||
temperature=0.7,
|
||||
)
|
||||
)
|
||||
|
||||
result = {"text": "", "images": []}
|
||||
|
||||
# Парсим ответ (Text или Inline Data)
|
||||
# Собираем текст из всех частей ответа
|
||||
result_text = ""
|
||||
if response.parts:
|
||||
for part in response.parts:
|
||||
if part.text:
|
||||
result["text"] += part.text
|
||||
result_text += part.text
|
||||
|
||||
# Проверяем наличие сгенерированного изображения
|
||||
if part.inline_data:
|
||||
# ИСПРАВЛЕНИЕ: Берем "сырые" байты напрямую из ответа
|
||||
# Это работает быстрее и не вызывает ошибку с PIL
|
||||
|
||||
# part.inline_data.data — это уже bytes
|
||||
byte_arr = io.BytesIO(part.inline_data.data)
|
||||
now = datetime.now()
|
||||
# Имя файла для телеграма (формально)
|
||||
byte_arr.name = f'{now.timestamp()}.png'
|
||||
|
||||
result["images"].append(byte_arr)
|
||||
|
||||
return result
|
||||
return result_text
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Gemini API Error: {e}")
|
||||
return {"error": f"Ошибка API: {str(e)}"}
|
||||
logger.error(f"Gemini Text API Error: {e}")
|
||||
return f"Ошибка генерации текста: {e}"
|
||||
|
||||
def generate_image(self, prompt: str, aspect_ratio: AspectRatios, quality: Quality, images_list: List[bytes] = None, ) -> List[io.BytesIO]:
|
||||
"""
|
||||
Генерация изображений (Text-to-Image или Image-to-Image).
|
||||
Возвращает список байтовых потоков (готовых к отправке).
|
||||
"""
|
||||
contents = self._prepare_contents(prompt, images_list)
|
||||
|
||||
try:
|
||||
response = self.client.models.generate_content(
|
||||
model=self.IMAGE_MODEL,
|
||||
contents=contents,
|
||||
config=types.GenerateContentConfig(
|
||||
response_modalities=['IMAGE'],
|
||||
temperature=1.0,
|
||||
image_config=types.ImageConfig(
|
||||
aspect_ratio=aspect_ratio.value,
|
||||
image_size=quality.value
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
generated_images = []
|
||||
|
||||
if response.parts:
|
||||
for part in response.parts:
|
||||
# Ищем картинки (inline_data)
|
||||
if part.inline_data:
|
||||
try:
|
||||
# 1. Берем сырые байты
|
||||
raw_data = part.inline_data.data
|
||||
byte_arr = io.BytesIO(raw_data)
|
||||
|
||||
# 2. Нейминг (формально, для TG)
|
||||
timestamp = datetime.now().timestamp()
|
||||
byte_arr.name = f'{timestamp}.png'
|
||||
|
||||
# 3. Важно: сбросить курсор в начало
|
||||
byte_arr.seek(0)
|
||||
|
||||
generated_images.append(byte_arr)
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing output image: {e}")
|
||||
|
||||
return generated_images
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Gemini Image API Error: {e}")
|
||||
# В случае ошибки возвращаем пустой список (или можно рейзить исключение)
|
||||
return []
|
||||
Reference in New Issue
Block a user