diff --git a/adapters/google_adapter.py b/adapters/google_adapter.py index ddd6434..52dcfb9 100644 --- a/adapters/google_adapter.py +++ b/adapters/google_adapter.py @@ -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)}"} \ No newline at end of file + 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 [] \ No newline at end of file diff --git a/keyboards.py b/keyboards.py index 1179cf7..ac92f4d 100644 --- a/keyboards.py +++ b/keyboards.py @@ -1,5 +1,11 @@ +from aiogram.fsm.context import FSMContext from aiogram.types import InlineKeyboardMarkup, InlineKeyboardButton +from models.enums import AspectRatios, Quality, GenType +from repos.dao import DAO + + + def get_request_kb(): return InlineKeyboardMarkup(inline_keyboard=[ [InlineKeyboardButton(text="🔐 Запросить доступ", callback_data="req_access")] @@ -12,4 +18,21 @@ def get_admin_decision_kb(user_id: int): InlineKeyboardButton(text="✅ Разрешить", callback_data=f"access_allow_{user_id}"), InlineKeyboardButton(text="🚫 Запретить", callback_data=f"access_deny_{user_id}") ] + ]) + +async def get_gen_mode_kb(state: FSMContext, dao: DAO): + data = await state.get_data() + char = await dao.chars.get_character(character_id=data['char_id']) + return InlineKeyboardMarkup(inline_keyboard=[ + [ + InlineKeyboardButton(text=f'Перс: {char.name}', callback_data=f'gen_mode_change_char'), + ], + [ + InlineKeyboardButton(text=f"🔁{AspectRatios[data['aspect_ratio']].value}", callback_data=f'gen_mode_change_aspect_ratio'), + InlineKeyboardButton(text=f"🔁{Quality[data['quality']].value}", callback_data=f'gen_mode_change_quality'), + InlineKeyboardButton(text=f"🔁{GenType[data['type']].value}",callback_data=f'gen_mode_change_type') + ], + [ + InlineKeyboardButton(text="❌ Выйти из режима генерации", callback_data=f'gen_mode_off'), + ] ]) \ No newline at end of file diff --git a/main.py b/main.py index 4a976a9..a117e91 100644 --- a/main.py +++ b/main.py @@ -13,6 +13,7 @@ from motor.motor_asyncio import AsyncIOMotorClient # Импорты from adapters.google_adapter import GoogleAdapter +from middlewares.album import AlbumMiddleware from middlewares.auth import AuthMiddleware from middlewares.dao import DaoMiddleware from repos.char_repo import CharacterRepo @@ -64,7 +65,8 @@ dp.include_router(gen_router) # Вешаем защиту ТОЛЬКО на основной роутер main_router.message.middleware(AuthMiddleware(repo=users_repo, admin_id=ADMIN_ID)) gen_router.message.middleware(AuthMiddleware(repo=users_repo, admin_id=ADMIN_ID)) -char_router.message.middleware(DaoMiddleware(dao=DAO(client=mongo_client))) +gen_router.message.middleware(AlbumMiddleware(latency=0.8)) +dp.update.middleware(DaoMiddleware(dao=DAO(client=mongo_client))) diff --git a/middlewares/album.py b/middlewares/album.py new file mode 100644 index 0000000..b471316 --- /dev/null +++ b/middlewares/album.py @@ -0,0 +1,51 @@ +import asyncio +from typing import Any, Dict, List, Union, Callable, Awaitable +from aiogram import BaseMiddleware +from aiogram.types import Message + + +class AlbumMiddleware(BaseMiddleware): + def __init__(self, latency: float = 0.5): + # latency - задержка в секундах для сбора частей альбома + self.latency = latency + self.album_data: Dict[str, List[Message]] = {} + + async def __call__( + self, + handler: Callable[[Message, Dict[str, Any]], Awaitable[Any]], + event: Message, + data: Dict[str, Any] + ) -> Any: + # Если у сообщения нет media_group_id, это не альбом -> пропускаем дальше как обычно + if not event.media_group_id: + return await handler(event, data) + + group_id = event.media_group_id + + try: + # Если этот альбом мы еще не видели (первое сообщение из пачки) + if group_id not in self.album_data: + self.album_data[group_id] = [event] # Создаем список + await asyncio.sleep(self.latency) # Ждем остальные части + + # После ожидания кладем собранный список в data + # Теперь в хендлере будет доступен аргумент 'album' + data["album"] = self.album_data[group_id] + + # Вызываем хендлер ОДИН раз + return await handler(event, data) + + else: + # Если альбом уже собирается, просто добавляем сообщение в список + # и НЕ вызываем хендлер (прерываем цепочку для этого сообщения) + self.album_data[group_id].append(event) + return + + finally: + # Чистим память после обработки, если это был "главный" поток обработки + if group_id in self.album_data and len(self.album_data[group_id]) > 1: + # Маленький хак: удаляем только если обработчик завершился + # Проверка len нужна, чтобы не удалить раньше времени в параллельных тасках, + # но корректнее просто удалять в блоке первого сообщения. + if event == self.album_data[group_id][0]: + del self.album_data[group_id] \ No newline at end of file diff --git a/models/Character.py b/models/Character.py index 2b5a055..f73cd9c 100644 --- a/models/Character.py +++ b/models/Character.py @@ -2,7 +2,8 @@ from pydantic import BaseModel class Character(BaseModel): - id: int | None + id: str name: str character_image: bytes - character_bio: str \ No newline at end of file + character_bio: str + diff --git a/models/enums.py b/models/enums.py new file mode 100644 index 0000000..10d5c49 --- /dev/null +++ b/models/enums.py @@ -0,0 +1,19 @@ +from enum import Enum + + +class AspectRatios(Enum): + NINESIXTEEN = '9:16' + SIXTEENNINE = '16:9' + THREEFOUR = '3:4' + FOURTHREE = '4:3' + + +class Quality(Enum): + ONEK = '1K' + TWOK = '2K' + FOURK = '4K' + + +class GenType(Enum): + TEXT = 'Text' + IMAGE = 'Image' \ No newline at end of file diff --git a/repos/char_repo.py b/repos/char_repo.py index 5eef2ea..89598aa 100644 --- a/repos/char_repo.py +++ b/repos/char_repo.py @@ -1,3 +1,6 @@ +from typing import List + +from bson import ObjectId from motor.motor_asyncio import AsyncIOMotorClient from models.Character import Character @@ -12,5 +15,23 @@ class CharacterRepo: character.id = op.inserted_id return character - async def get_character(self, character_id: int) -> Character: - return await self.collection.find_one({"id": character_id}) \ No newline at end of file + async def get_character(self, character_id: str) -> Character | None: + res = await self.collection.find_one({"_id": ObjectId(character_id)}) + if res is None: + return None + else: + res["id"] = str(res.pop("_id")) + return Character(**res) + + async def get_all_characters(self) -> List[Character]: + docs = await self.collection.find().to_list(None) + + characters = [] + for doc in docs: + # Конвертируем ObjectId в строку и кладем в поле id + doc["id"] = str(doc.pop("_id")) + + # Создаем объект + characters.append(Character(**doc)) + + return characters \ No newline at end of file diff --git a/routers/char_router.py b/routers/char_router.py index 6110a72..30c043e 100644 --- a/routers/char_router.py +++ b/routers/char_router.py @@ -25,9 +25,11 @@ async def add_char(message: Message, state: FSMContext, dao: DAO): await message.answer("Кайф, теперь напиши ее имя") -@router.callback_query(States.char_wait_name) +@router.message(States.char_wait_name) async def new_char_name(message: Message, state: FSMContext, dao: DAO): - await state.set_data({"name": message.text}) + data = await state.get_data() + data["name"] = message.text + await state.set_data(data) await state.set_state(States.char_wait_bio) await message.answer("А теперь напиши био. Хоть чуть чуть.") @@ -37,7 +39,7 @@ async def new_char_name(message: Message, state: FSMContext, dao: DAO): async def new_char_bio(message: Message, state: FSMContext, dao: DAO, bot: Bot): # Получаем все накопленные данные data = await state.get_data() - file_id = data["file_id"] + file_id = data["photo"] name = data["name"] bio = message.text @@ -64,9 +66,9 @@ async def new_char_bio(message: Message, state: FSMContext, dao: DAO, bot: Bot): await message.answer_photo( photo=BufferedInputFile(photo_bytes, filename="char.png"), caption=( - "🎉 **Персонаж создан!**\n\n" - f"👤 **Имя:** {char.name}\n" - f"📝 **Био:** {char.character_bio}" + "🎉 Персонаж создан!\n\n" + f"👤 Имя: {char.name}\n" + f"📝 Био: {char.character_bio}" ) ) await wait_msg.delete() @@ -79,10 +81,51 @@ async def new_char_bio(message: Message, state: FSMContext, dao: DAO, bot: Bot): # Не сбрасываем стейт, даем возможность попробовать ввести био снова или начать заново +@router.message(Command("chars")) +async def get_chars(message: Message, state: FSMContext, dao: DAO): + wait_msg = await message.answer("Ищем персонажей") + chars = await dao.chars.get_all_characters() + keyboards = [] + if len(chars) > 0: + for char in chars: + keyboards.append(InlineKeyboardButton(text=char.name, callback_data=f'char_info_{char.id}')) + keyboard = InlineKeyboardMarkup(inline_keyboard=[keyboards]) + else: + keyboard = InlineKeyboardMarkup( + inline_keyboard=[[InlineKeyboardButton(text="Персонажей нет", callback_data=f'no_chars')]]) + await message.answer("Сейчас есть такие персонажи:", reply_markup=keyboard) + await wait_msg.delete() + + +@router.callback_query(F.data.startswith("char_info_")) +async def get_char_info(callback_query: CallbackQuery, state: FSMContext, dao: DAO): + await callback_query.message.delete() + wait_msg = await callback_query.message.answer("Ищем инфу о персонаже") + char = await dao.chars.get_character(callback_query.data.split("_")[-1]) + if char is None: + await callback_query.message.answer("Информация о персонаже не найдена") + await get_chars(callback_query.message, state, dao) + await wait_msg.delete() + return + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [InlineKeyboardButton(text="Запросить фото в документе", callback_data=f'char_photo_file_{char.id}')]]) + await callback_query.message.answer_photo(photo=BufferedInputFile(char.character_image, f"photo_{char.id}.png"), caption=f"👤 Имя: {char.name}\n" + f"📝 Био: {char.character_bio}", + reply_markup=keyboard) + + await wait_msg.delete() + + +@router.callback_query(F.data.startswith("char_photo_file")) +async def get_char_info_photo_file(callback_query: CallbackQuery, state: FSMContext, dao: DAO): + char = await dao.chars.get_character(callback_query.data.split("_")[-1]) + await callback_query.message.answer_document(BufferedInputFile(char.character_image, f"photo_{char.id}.png")) + + # 4. Хендлер-помощник (если отправили команду без файла) @router.message(Command("add_char")) async def add_char_help(message: Message): await message.answer( "ℹ️ **Как добавить персонажа:**\n" "Прикрепите фото (файлом/документом) и добавьте в подпись команду `/add_char`." - ) \ No newline at end of file + ) diff --git a/routers/gen_router.py b/routers/gen_router.py index 27a1ffd..8c68f5b 100644 --- a/routers/gen_router.py +++ b/routers/gen_router.py @@ -1,48 +1,342 @@ import asyncio +import random +from enum import Enum +from typing import List from aiogram import Router, Bot, F from aiogram.enums import ParseMode +from aiogram.exceptions import TelegramBadRequest from aiogram.filters import * +from aiogram.fsm.context import FSMContext +from aiogram.fsm.state import StatesGroup, State from aiogram.types import * +from aiogram.types import message +import keyboards from adapters.google_adapter import GoogleAdapter +from models.Character import Character +from models.enums import AspectRatios, Quality, GenType +from repos.dao import DAO router = Router() -@router.message(Command("image")) -async def cmd_image_gen(message: Message, command: CommandObject, gemini: GoogleAdapter, bot: Bot): - # ... ваш код ... - # Обратите внимание: gemini теперь прилетает аргументом, так как мы сделали dp["gemini"] - prompt = command.args - if not prompt: - await message.answer("⚠️ Напиши промпт.") + +class States(StatesGroup): + gen_mode_wait_char = State() + gen_mode = State() + + + + +async def init_gen_mode(state: FSMContext, dao: DAO): + data = await state.get_data() + data['aspect_ratio'] = AspectRatios.NINESIXTEEN.name + data['quality'] = Quality.ONEK.name + data['type'] = GenType.TEXT.name + await state.update_data(data) + + +@router.message(Command("gen_mode")) +async def gen_mode(message: Message, state: FSMContext, dao: DAO): + state_on = await state.get_state() + if state_on is None and state_on is not States.gen_mode: + await message.answer("Включить режим генерации?", reply_markup=InlineKeyboardMarkup( + inline_keyboard=[[InlineKeyboardButton(text="✅ Включить!", callback_data="gen_mode_on")]])) + else: + await gen_mode_base_msg(message, state, dao, call_type="start") + + +@router.callback_query(F.data == "gen_mode_on") +async def gen_mode_on(callback_query: CallbackQuery, state: FSMContext, dao: DAO): + await callback_query.message.delete() + chars = await dao.chars.get_all_characters() + if len(chars) == 0: + await callback_query.message.answer( + "Персонажи не найдены! Сперва создайте персонажа отправив его фото ФАЙЛОМ с командой /add_char") + keyboards = [] + for char in chars: + keyboards.append(InlineKeyboardButton(text=char.name, callback_data=f'select_char_{char.id}')) + await state.set_state(States.gen_mode_wait_char) + await callback_query.message.answer("Выбери персонажа", + reply_markup=InlineKeyboardMarkup(inline_keyboard=[keyboards])) + + +@router.callback_query(States.gen_mode_wait_char, F.data.startswith("select_char_")) +async def select_char(call: CallbackQuery, state: FSMContext, dao: DAO): + await state.update_data({"char_id": call.data.split("_")[-1]}) + await init_gen_mode(state=state, dao=dao) + await state.set_state(States.gen_mode) + await gen_mode_base_msg(call.message, state=state, dao=dao, call_type="start") + + +@router.callback_query(States.gen_mode, F.data == 'gen_mode_off') +async def gen_mode_off(call: CallbackQuery, state: FSMContext, dao: DAO): + await state.clear() + await state.set_data({}) + await call.message.delete() + await call.message.answer("Режим генерации выключен. Нажмите /start для продолжения!") + + +@router.callback_query(States.gen_mode, F.data == 'gen_mode_change_char') +async def gen_mode_change_char(call: CallbackQuery, state: FSMContext, dao: DAO): + chars = await dao.chars.get_all_characters() + if len(chars) == 0: + await call.message.edit_caption( + "Персонажи не найдены! Сперва создайте персонажа отправив его фото ФАЙЛОМ с командой /add_char", + reply_markup=None) return + else: + keyboards = [] + for char in chars: + keyboards.append(InlineKeyboardButton(text=char.name, callback_data=f'select_new_char_{char.id}')) + keyboards.append(InlineKeyboardButton(text="⬅️ Назад", callback_data="gen_mode_cancel_char_change")) + await call.message.edit_caption("Выбери персонажа", reply_markup=InlineKeyboardMarkup(inline_keyboard=[keyboards])) + + +@router.callback_query(States.gen_mode, F.data.startswith('select_new_char_')) +async def change_char(call: CallbackQuery, state: FSMContext, dao: DAO): + await state.update_data({"char_id": call.data.split("_")[-1]}) + await gen_mode_base_msg(call.message, state=state, dao=dao) + + +@router.callback_query(States.gen_mode, F.data == 'gen_mode_change_aspect_ratio') +async def gen_mode_change_aspect_ratio(call: CallbackQuery, state: FSMContext, dao: DAO): + keyboards = [] + for ratio in AspectRatios: + keyboards.append(InlineKeyboardButton(text=ratio.value, callback_data=f'select_ratio_{ratio.name}')) + await call.message.edit_caption(caption="Выбери соотношение сторон", + reply_markup=InlineKeyboardMarkup(inline_keyboard=[keyboards, [InlineKeyboardButton( + text="⬅️ Назад", callback_data="gen_mode_cancel_ratio_change")]])) + + +@router.callback_query(States.gen_mode, F.data.startswith('select_ratio_')) +async def change_aspect_ratio(call: CallbackQuery, state: FSMContext, dao: DAO): + await state.update_data({"aspect_ratio": call.data.split("_")[-1]}) + await gen_mode_base_msg(call.message, state=state, dao=dao) + + +@router.callback_query(States.gen_mode, F.data == 'gen_mode_change_quality') +async def gen_mode_change_quality(call: CallbackQuery, state: FSMContext, dao: DAO): + keyboards = [] + for quality in Quality: + keyboards.append(InlineKeyboardButton(text=quality.value, callback_data=f'select_quality_{quality.name}')) + await call.message.edit_caption(caption="Выбери качество", + reply_markup=InlineKeyboardMarkup(inline_keyboard=[keyboards, [InlineKeyboardButton( + text="⬅️ Назад", callback_data="gen_mode_cancel_quality_change")]])) + + +@router.callback_query(States.gen_mode, F.data.startswith('select_quality_')) +async def change_quality(call: CallbackQuery, state: FSMContext, dao: DAO): + await state.update_data({"quality": call.data.split("_")[-1]}) + await gen_mode_base_msg(call.message, state=state, dao=dao) + + +@router.callback_query(States.gen_mode, F.data == 'gen_mode_change_type') +async def gen_mode_change_type(call: CallbackQuery, state: FSMContext, dao: DAO): + keyboards = [] + for gen_type in GenType: + keyboards.append(InlineKeyboardButton(text=gen_type.value, callback_data=f'select_type_{gen_type.name}')) + await call.message.edit_caption(caption="Выбери тип генерации", reply_markup=InlineKeyboardMarkup( + inline_keyboard=[keyboards, + [InlineKeyboardButton(text="⬅️ Назад", callback_data="gen_mode_cancel_type_change")]])) + + +@router.callback_query(States.gen_mode, F.data.startswith('select_type_')) +async def change_quality(call: CallbackQuery, state: FSMContext, dao: DAO): + await state.update_data({"type": call.data.split("_")[-1]}) + await gen_mode_base_msg(call.message, state=state, dao=dao) + + +@router.callback_query(States.gen_mode, F.data == 'gen_mode_cancel_char_change') +@router.callback_query(States.gen_mode, F.data == 'gen_mode_cancel_ratio_change') +@router.callback_query(States.gen_mode, F.data == 'gen_mode_cancel_type_change') +@router.callback_query(States.gen_mode, F.data == 'gen_mode_cancel_quality_change') +async def cancel_gen_mode_change(call: CallbackQuery, state: FSMContext, dao: DAO): + await gen_mode_base_msg(call.message, state=state, dao=dao) + + +async def gen_mode_base_msg(message: Message, state: FSMContext, dao: DAO, call_type='continue'): + data = await state.get_data() + char: Character = await dao.chars.get_character(data["char_id"]) + if call_type == "start": + await message.answer_photo(BufferedInputFile(char.character_image, f'{char.id}.png'), + caption="🎉 Режим генерации включен! Просто пиши мне промпт и я отправлю в генерацию по указанным настройкам.\n\n" + "Фото девушки грузить не надо, оно загрузится по дефолту\n\n" + "Но дополнительные фото можно загрузить.", + reply_markup=await keyboards.get_gen_mode_kb(state=state, dao=dao)) + else: + try: + await message.edit_caption( + caption="🎉 Режим генерации включен!", + reply_markup=await keyboards.get_gen_mode_kb(state=state, dao=dao)) + except TelegramBadRequest as tbr: + await message.edit_text( + text="🎉 Режим генерации включен!", + reply_markup=await keyboards.get_gen_mode_kb(state=state, dao=dao)) + + +@router.message(States.gen_mode, F.media_group_id) +async def handle_album( + message: Message, + album: List[Message], + state: FSMContext, + dao: DAO, + gemini: GoogleAdapter, + bot: Bot +): + """ + Обработка альбома (группы фото). + """ + # 1. Ищем промпт (подпись) в любом из сообщений альбома + # message.text в альбомах обычно None + prompt = next((msg.caption for msg in album if msg.caption), None) + + if not prompt: + await message.answer("⚠️ Вы отправили альбом, но не добавили описание (промпт).") + return + + # 2. Собираем file_id всех фото + file_ids = [] + for msg in album: + if msg.photo: + file_ids.append(msg.photo[-1].file_id) + elif msg.video: + # Если нужно, можно добавить обработку видео (пока пропускаем) + pass + + await message.answer(f"📥 Принято {len(album)} файлов. Начинаю генерацию...") + wait_msg = await message.answer("🎨 Генерирую...") + + # 3. Вызываем генерацию + try: + generated_files = await generate_image( + prompt=prompt, + media=file_ids, + state=state, + dao=dao, + bot=bot, + gemini=gemini + ) + + await wait_msg.delete() + + # 4. Отправляем результат + if generated_files: + for file in generated_files: + await message.answer_document(file, caption="✨ Generated result") + else: + await message.answer("❌ Генерация не вернула изображений.") + await gen_mode_base_msg(message=message, state=state, dao=dao,call_type="start" ) + + except Exception as e: + await wait_msg.edit_text(f"❌ Ошибка: {e}") + + +@router.message(States.gen_mode) +async def gen_mode_start( + message: Message, + state: FSMContext, + dao: DAO, + gemini: GoogleAdapter, + bot: Bot +): + """ + Обработка одиночного сообщения (Текст или Фото+Подпись) + """ + # 1. Получаем промпт (Текст или Подпись) + prompt = message.text or message.caption + + if not prompt: + await message.answer("⚠️ Напиши промпт (или отправь фото с подписью).") + return + + # 2. Проверяем, есть ли прикрепленное фото + media_ids = [] + if message.photo: + media_ids.append(message.photo[-1].file_id) + elif message.reply_to_message and message.reply_to_message.photo: + # Поддержка реплая на фото + media_ids.append(message.reply_to_message.photo[-1].file_id) wait_msg = await message.answer("🎨 Генерирую...") - # Получение байтов фото (логика та же) - image_bytes = None - if message.photo: - file_io = await bot.download(message.photo[-1].file_id) - image_bytes = file_io.getvalue() - elif message.reply_to_message and message.reply_to_message.photo: - file_io = await bot.download(message.reply_to_message.photo[-1].file_id) - image_bytes = file_io.getvalue() + try: + generated_files = await generate_image( + prompt=prompt, + media=media_ids, # Передаем список (пустой или с 1 фото) + state=state, + dao=dao, + bot=bot, + gemini=gemini + ) - result = await asyncio.to_thread( - gemini.generate, prompt=prompt, image_bytes=image_bytes, generate_image=True + await wait_msg.delete() + + if generated_files: + for file in generated_files: + await message.answer_document(file, caption="✨ Generated result") + else: + await message.answer("❌ Ничего не сгенерировалось.") + await gen_mode_base_msg(message=message, state=state, dao=dao,call_type="start" ) + + except Exception as e: + await wait_msg.edit_text(f"❌ Ошибка: {e}") + + +async def generate_image( + prompt: str, + media: List[str] | None, + state: FSMContext, + dao: DAO, + bot: Bot, + gemini: GoogleAdapter +) -> List[BufferedInputFile]: + # 1. Получаем данные персонажа + data = await state.get_data() + char_id = data.get("char_id") + + if not char_id: + raise ValueError("Character ID not found in state") + + char: Character = await dao.chars.get_character(char_id) + + # Начинаем список с фото персонажа + media_group_bytes = [char.character_image] + + # 2. Скачиваем дополнительные файлы (если переданы) + if media: + # Создаем задачи для скачивания + tasks = [bot.download(file_id) for file_id in media] + + # Запускаем все скачивания параллельно + downloaded_files = await asyncio.gather(*tasks) + + # Добавляем байты скачанных файлов в общий список + for f in downloaded_files: + media_group_bytes.append(f.getvalue()) + + # 3. Генерация в Gemini + generated_images_io = await asyncio.to_thread( + gemini.generate_image, + prompt=prompt, + images_list=media_group_bytes, + aspect_ratio=AspectRatios[data['aspect_ratio']], + quality=Quality[data['quality']], ) - await wait_msg.delete() - - if result.get("images"): - for img in result["images"]: - await message.answer_document(BufferedInputFile(img.read(), "img.png")) - elif result.get("text"): - await message.answer(result["text"]) - else: - await message.answer(f"Ошибка: {result.get('error', 'Unknown')}") + # 4. Упаковка результата + images = [] + if generated_images_io: + for i, img_io in enumerate(generated_images_io): + # Важно: img_io.read() работает корректно, если курсор в начале (adapter это делает) + images.append( + BufferedInputFile( + img_io.read(), + filename=f"img_{random.randint(1000, 9999)}.png" + ) + ) + return images @router.message(F.text) async def handle_text(message: Message, gemini: GoogleAdapter, bot: Bot): @@ -52,4 +346,3 @@ async def handle_text(message: Message, gemini: GoogleAdapter, bot: Bot): await message.answer(result["text"], parse_mode=ParseMode.MARKDOWN) else: await message.answer("Ошибка генерации") -