BotVkAI / main.py
HollyCorn's picture
Upload main.py
c5c50ce verified
Raw
History Blame Contribute Delete
31.1 kB
import os
import threading
import hashlib
import asyncio
import io
import aiohttp
import random
import uuid
from flask import Flask, request
from vkbottle import Keyboard, KeyboardButtonColor, Text, OpenLink, PhotoMessageUploader
from vkbottle.bot import Bot, Message
from supabase import create_client, Client
from loguru import logger
from dotenv import load_dotenv
load_dotenv()
# Импортируем вспомогательные файлы
import prompts
import generator
# --- КОНФИГУРАЦИЯ ---
TOKEN = os.environ.get("VK_TOKEN")
YOOMONEY_SECRET = os.environ.get("YOOMONEY_SECRET")
SUPABASE_URL = os.environ.get("SUPABASE_URL")
SUPABASE_KEY = os.environ.get("SUPABASE_KEY")
YOOMONEY_WALLET = "4100118637309832"
GROUP_ID = 171587553
ADMIN_ID = 139711548
BUCKET_NAME = "photos"
REFERRAL_BONUS = 2 # Бонус за первую покупку друга
# Инициализация бота
bot = Bot(token=TOKEN)
uploader = PhotoMessageUploader(bot.api)
supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)
app = Flask(__name__)
loop = None
logger.add("bot_log.log", rotation="10 MB", level="INFO")
# --- ВСПОМОГАТЕЛЬНЫЕ ФУНКЦИИ ---
async def upload_to_bucket(url, user_id, prefix="temp"):
try:
async with aiohttp.ClientSession() as session:
async with session.get(url, timeout=30) as resp:
if resp.status == 200:
img_data = await resp.read()
file_name = f"{user_id}/{prefix}_{uuid.uuid4()}.png"
supabase.storage.from_(BUCKET_NAME).upload(
path=file_name,
file=img_data,
file_options={"content-type": "image/png"}
)
res = supabase.storage.from_(BUCKET_NAME).get_public_url(file_name)
return file_name, res
except Exception as e:
logger.error(f"Ошибка загрузки в бакет ({prefix}): {e}")
return None, None
async def download_file(url, retries=5, delay=3):
"""Пытается скачать файл из бакета несколько раз с задержкой и проверкой размера"""
for i in range(retries):
try:
async with aiohttp.ClientSession() as session:
async with session.get(url, timeout=15) as resp:
if resp.status == 200:
data = await resp.read()
# Проверка: файл должен существовать и быть больше 500 байт (минимум для картинки)
if data and len(data) > 500:
return io.BytesIO(data)
else:
logger.warning(
f"Попытка {i + 1}: Файл пустой или неполный ({len(data) if data else 0} байт)")
logger.warning(f"Попытка {i + 1} скачивания не удалась, ждем {delay} сек...")
await asyncio.sleep(delay)
except Exception as e:
logger.error(f"Ошибка на попытке {i + 1} скачивания файла для ВК: {e}")
await asyncio.sleep(delay)
return None
def get_user(user_id, ref_id=None):
res = supabase.table("users").select("*").eq("id", user_id).execute()
if not res.data:
# Логика определения пригласителя
inviter = None
if ref_id and str(ref_id).isdigit() and int(ref_id) != user_id:
inviter = int(ref_id)
new_user = {
"id": user_id, "balance": 0, "free_used": False, "status": "main",
"gender": None, "primary_photo": None, "last_style": None,
"invited_by": inviter, "ref_bonus_done": False
}
supabase.table("users").insert(new_user).execute()
return new_user
return res.data[0]
async def is_subscribed(user_id):
try:
res = await bot.api.groups.is_member(group_id=GROUP_ID, user_id=user_id)
return bool(res)
except Exception as e:
logger.error(f"Ошибка проверки подписки: {e}")
return True
# --- КЛАВИАТУРЫ ---
def get_main_menu(free_used=False):
kb = Keyboard(one_time=False)
if not free_used:
kb.add(Text("🎁 Одна фотография в подарок"), color=KeyboardButtonColor.POSITIVE).row()
kb.add(Text("🖤 Стили фото"), color=KeyboardButtonColor.PRIMARY).row()
kb.add(Text("✨ Творческий режим"), color=KeyboardButtonColor.PRIMARY).row()
kb.add(Text("🔄 Замена фото"), color=KeyboardButtonColor.PRIMARY).row()
kb.add(Text("💰 Кошелек"), color=KeyboardButtonColor.PRIMARY).row()
kb.add(Text("🆘 Поддержка"), color=KeyboardButtonColor.PRIMARY)
return kb.get_json()
def get_photo_change_menu():
kb = Keyboard(inline=True)
kb.add(Text("🔄 Сменить пол"), color=KeyboardButtonColor.PRIMARY).row()
kb.add(Text("🔄 Сменить фото"), color=KeyboardButtonColor.PRIMARY).row()
kb.add(Text("⬅️ В главное меню"), color=KeyboardButtonColor.PRIMARY)
return kb.get_json()
def get_fast_packs_menu():
kb = Keyboard(inline=True)
packs = supabase.table("settings").select("*").eq("is_premium", False).order("price").execute().data
for pack in packs:
kb.add(Text(f"📸 {pack['name']} - {int(pack['price'])} р"), color=KeyboardButtonColor.PRIMARY).row()
# Добавляем кнопку бесплатных генераций в самый низ списка пакетов
kb.add(Text("🎁 Бесплатные генерации"), color=KeyboardButtonColor.PRIMARY).row()
kb.add(Text("⬅️ В главное меню"), color=KeyboardButtonColor.PRIMARY)
return kb.get_json()
def get_gender_menu():
return (Keyboard(inline=True)
.add(Text("🙋‍♂️ Мужской"), color=KeyboardButtonColor.PRIMARY)
.add(Text("🙋‍♀️ Женский"), color=KeyboardButtonColor.PRIMARY)
).get_json()
def get_styles_menu(gender, page=1):
kb = Keyboard(inline=True)
g = gender if gender else "female"
if g in prompts.FACE_SWAP_STYLES:
all_styles = [s for s in prompts.FACE_SWAP_STYLES[g].keys() if s != "🎁 Одна фотография в подарок"]
per_page = 6
start = (page - 1) * per_page
end = start + per_page
current_styles = all_styles[start:end]
for i in range(0, len(current_styles), 2):
kb.add(Text(current_styles[i]), color=KeyboardButtonColor.PRIMARY)
if i + 1 < len(current_styles):
kb.add(Text(current_styles[i + 1]), color=KeyboardButtonColor.PRIMARY)
kb.row()
if page == 1 and len(all_styles) > per_page:
kb.add(Text("Ещё стили ➡️"), color=KeyboardButtonColor.PRIMARY).row()
elif page > 1:
kb.add(Text("⬅️ Назад к стилям"), color=KeyboardButtonColor.PRIMARY).row()
kb.add(Text("⬅️ В главное меню"), color=KeyboardButtonColor.PRIMARY)
return kb.get_json()
# --- ОБРАБОТЧИКИ ---
@bot.on.message(text=["Начать", "начать", "Меню", "start", "⬅️ В главное меню"])
async def start_handler(message: Message):
# Ловим ref_id из сообщения (параметр в ссылке vk.me/bot?ref=ID)
ref_id = message.ref
user = get_user(message.from_id, ref_id)
supabase.table("users").update({"status": "main"}).eq("id", message.from_id).execute()
if not user['free_used']:
msg = ("Привет 👋\n"
"Здесь ты можешь создать стильные AI-фотографии с собой в главной роли ✨\n\n"
"🎁 Для начала — бесплатный бонус: одна генерация.\n"
"Нажми «Одна фотография в подарок» и попробуй 👇\n\n")
else:
msg = ("Выбери режим в меню ниже — и начнём создавать твои новые образы ✨")
await message.answer(msg, keyboard=get_main_menu(user['free_used']))
@bot.on.message(text="🆘 Поддержка")
async def support_handler(message: Message):
supabase.table("users").update({"status": "wait_support"}).eq("id", message.from_id).execute()
await message.answer("Опишите вашу ситуацию прямо здесь, в чате. Я передам сообщение специалисту.")
@bot.on.message(text="🔄 Замена фото")
async def photo_change_main_handler(message: Message):
await message.answer("Что вы хотите изменить?", keyboard=get_photo_change_menu())
@bot.on.message(text="🎁 Бесплатные генерации")
async def free_gens_handler(message: Message):
link = f"https://vk.me/club{GROUP_ID}?ref={message.from_id}"
msg = (f"🎁 Получай бесплатные генерации за друзей!\n\n"
f"Твоя личная ссылка: {link}\n\n"
f"Как это работает?\n"
f"1. Отправь эту ссылку другу\n"
f"2. Когда он совершит свою первую оплату, ты мгновенно получишь +{REFERRAL_BONUS} генерации на свой баланс!\n\n"
f"Количество друзей не ограничено. Больше друзей — больше фото! ✨")
await message.answer(msg)
@bot.on.message(text=["🎁 Одна фотография в подарок", "✅ Я подписался"])
async def free_trial_handler(message: Message):
user = get_user(message.from_id)
if user['free_used']:
return await message.answer("Вы уже использовали бесплатную генерацию!", keyboard=get_main_menu(True))
if not user['gender']:
supabase.table("users").update({"status": "wait_gender_free"}).eq("id", message.from_id).execute()
return await message.answer("Выберите ваш пол:", keyboard=get_gender_menu())
if not user['primary_photo']:
supabase.table("users").update({"status": "wait_primary_photo"}).eq("id", message.from_id).execute()
return await message.answer("📸 Отправь одно своё фото — я сделаю стильный образ\n\n"
"Чтобы результат получился классным, важно:\n"
"• на фото только ты\n"
"• лицо хорошо освещено и видно целиком\n"
"• без очков, шапок и эффектов\n"
"• без других людей в кадре\n\n"
"Через пару минут ты получишь стильное нейро-изображение ✨\n\n"
"Жду фотографию ⬇️")
if not await is_subscribed(message.from_id):
subscribe_kb = (
Keyboard(inline=True)
.add(OpenLink(f"https://vk.com/club{GROUP_ID}", "📢 Подписаться"))
.row()
.add(Text("✅ Я подписался"), color=KeyboardButtonColor.POSITIVE)
).get_json()
return await message.answer(
"🎁 Для бесплатной генерации подпишись на наше сообщество.\n\n"
"После подписки нажмите кнопку «✅ Я подписался» ниже 👇",
keyboard=subscribe_kb
)
await message.answer("⏳ Генерирую фотографию... Подожди около 30 секунд")
gender = user['gender']
prompt = prompts.FACE_SWAP_STYLES[gender]["🎁 Одна фотография в подарок"]
paths_to_clean = []
bucket_path, temporary_url = await upload_to_bucket(user['primary_photo'], message.from_id, "src")
if bucket_path: paths_to_clean.append(bucket_path)
res_url = await generator.generate_image(
prompt=prompt,
gender=gender,
photo_url=temporary_url if temporary_url else user['primary_photo'],
lora_model=None
)
if res_url:
res_path, res_url_bucket = await upload_to_bucket(res_url, message.from_id, "res")
if res_path: paths_to_clean.append(res_path)
if res_url_bucket:
try:
file_stream = await download_file(res_url_bucket)
if file_stream:
photo_vk = await uploader.upload(file_stream, file_name="result.png")
await message.answer("Отлично! Твоя фотография готова ✨", attachment=photo_vk)
promo_text = (
"Теперь ты можешь выбрать один из 10 готовых стилей для фото ✨\n"
"или создать собственный образ в Творческом режиме.\n\n"
"Выбирай, с чего начнём 👇"
)
await message.answer(promo_text, keyboard=get_main_menu(True))
supabase.table("users").update({"free_used": True}).eq("id", message.from_id).execute()
else:
await message.answer("Ошибка: не удалось получить файл из хранилища. Попробуйте еще раз позже.")
except Exception as e:
logger.error(f"Ошибка в free_trial VK upload: {e}")
await message.answer("Техническая ошибка при отправке фото.")
else:
await message.answer("Ошибка при обработке результата.")
else:
await message.answer("Произошла ошибка при генерации. Попробуйте другое фото.")
if paths_to_clean:
supabase.storage.from_(BUCKET_NAME).remove(paths_to_clean)
@bot.on.message(text=["💰 Кошелек"])
async def balance_handler(message: Message):
user = get_user(message.from_id)
photo_st = "Есть" if user['primary_photo'] else "Нет"
gender_ru = "Не выбран"
if user['gender'] == "male":
gender_ru = "Мужской"
elif user['gender'] == "female":
gender_ru = "Женский"
kb = Keyboard()
kb.add(Text("🛒 Купить генерации"), color=KeyboardButtonColor.PRIMARY).row()
# Кнопка бесплатные генерации в кошелек
kb.add(Text("🎁 Бесплатные генерации"), color=KeyboardButtonColor.PRIMARY).row()
kb.add(Text("⬅️ В главное меню"), color=KeyboardButtonColor.PRIMARY)
await message.answer(
f"🪙 Доступно: {user['balance']} генераций\n"
f"👤 Пол: {gender_ru}\n"
f"🖼 Закрепленная фотография: {photo_st}",
keyboard=kb.get_json()
)
@bot.on.message(text="🛒 Купить генерации")
async def shop_handler(message: Message):
await message.answer("Выберите количество генераций:", keyboard=get_fast_packs_menu())
@bot.on.message(text="🖤 Стили фото")
async def styles_handler(message: Message):
user = get_user(message.from_id)
if not user['gender']: return await message.answer("Сначала выберите пол:", keyboard=get_gender_menu())
await message.answer("Здесь можно получить готовое фото примерно за 30 секунд.\n\n"
"В каждом стиле есть 10 разных вариантов изображений.\n"
"При каждой генерации они выбираются случайным образом.\n\n"
"Даже если выбрать один и тот же стиль несколько раз,\n"
"фотографии каждый раз будут разными.\n\n"
"Выбери стиль ниже 👇", keyboard=get_styles_menu(user['gender'], page=1))
@bot.on.message(text="Ещё стили ➡️")
async def styles_page_2_handler(message: Message):
user = get_user(message.from_id)
await message.answer("Вторая страница стилей 👇", keyboard=get_styles_menu(user['gender'], page=2))
@bot.on.message(text="⬅️ Назад к стилям")
async def styles_page_1_back_handler(message: Message):
user = get_user(message.from_id)
await message.answer("Первая страница стилей 👇", keyboard=get_styles_menu(user['gender'], page=1))
@bot.on.message(text="✨ Творческий режим")
async def creator_mode(message: Message):
user = get_user(message.from_id)
if user['balance'] <= 0:
return await message.answer("В твоем кошельке 0 генераций. Пополни баланс:",
keyboard=get_fast_packs_menu())
supabase.table("users").update({"status": "wait_custom_prompt"}).eq("id", message.from_id).execute()
kb = Keyboard(inline=True).add(Text("⬅️ В главное меню"), color=KeyboardButtonColor.PRIMARY).get_json()
await message.answer(
"Опиши идею для будущего фото ✨\n"
"Напиши, где ты находишься, во что одет(а) и какое настроение хочешь передать.\n"
"Можно описать даже самую смелую задумку — я превращу её в фото👇",
keyboard=kb)
# --- ГЛАВНЫЕ ОБРАБОТЧИКИ ---
@bot.on.message()
async def main_handler(message: Message):
user = get_user(message.from_id)
text = message.text
if text == "⬅️ В главное меню":
supabase.table("users").update({"status": "main"}).eq("id", message.from_id).execute()
return await start_handler(message)
# Логика поддержки
if user['status'] == "wait_support":
try:
await bot.api.messages.send(peer_id=ADMIN_ID, random_id=0,
message=f"🆘 Новое обращение!\nОт: [id{message.from_id}|пользователя]\nТекст: {text}")
except Exception as e:
logger.error(f"Support notification error: {e}")
supabase.table("users").update({"status": "main"}).eq("id", message.from_id).execute()
return await message.answer("Ваше обращение передано в поддержку, вам скоро ответят, ожидайте.",
keyboard=get_main_menu(user['free_used']))
# 1. Оплата
packs = supabase.table("settings").select("*").execute().data
for pack in packs:
if pack['name'] in text and str(int(pack['price'])) in text:
pay_url = (
f"https://yoomoney.ru/quickpay/confirm.xml?"
f"receiver={YOOMONEY_WALLET}&"
f"quickpay-form=button&"
f"sum={pack['price']}&"
f"label={message.from_id}&"
f"targets=Оплата+пакета+{pack['name'].replace(' ', '+')}"
)
kb = Keyboard(inline=True).add(OpenLink(pay_url, "💳 Перейти к оплате")).get_json()
return await message.answer(f"Вы выбрали: {pack['name']}\nК оплате: {int(pack['price'])} руб.", keyboard=kb)
# 2. Пол
if text in ["🙋‍♂️ Мужской", "🙋‍♀️ Женский"]:
g = "male" if "Мужской" in text else "female"
supabase.table("users").update({"gender": g}).eq("id", message.from_id).execute()
if user['status'] == "wait_gender_change" or user['status'] == "wait_gender_free":
supabase.table("users").update({"status": "main"}).eq("id", message.from_id).execute()
if user['status'] == "wait_gender_free" and not user['primary_photo']:
supabase.table("users").update({"status": "wait_primary_photo"}).eq("id", message.from_id).execute()
return await message.answer("📸 Отправь одно своё фото — я сделаю стильный образ\n\n"
"Чтобы результат получился классным, важно:\n"
"• на фото только ты\n"
"• лицо хорошо освещено и видно целиком\n"
"• без очков, шапок и эффектов\n"
"• без других людей в кадре\n\n"
"Через пару минут ты получишь стильное нейро-изображение ✨\n\n"
"Жду фотографию ⬇️")
return await message.answer("✅ Сохранено", keyboard=get_main_menu(user['free_used']))
supabase.table("users").update({"status": "main"}).eq("id", message.from_id).execute()
return await message.answer(f"✅ Пол сохранен", keyboard=get_main_menu(user['free_used']))
# 3. Кнопки смены
if text == "🔄 Сменить пол":
supabase.table("users").update({"status": "wait_gender_change"}).eq("id", message.from_id).execute()
return await message.answer("Выбери пол:", keyboard=get_gender_menu())
if text == "🔄 Сменить фото":
supabase.table("users").update({"status": "wait_primary_photo"}).eq("id", message.from_id).execute()
return await message.answer("Отправь новое фото лица:")
# 4. Прием фото
if user['status'] == "wait_primary_photo" or (message.attachments and not user['primary_photo']):
if message.attachments:
photo_sizes = message.attachments[0].photo.sizes
photo_url = sorted(photo_sizes, key=lambda x: x.width)[-1].url
supabase.table("users").update({"primary_photo": photo_url, "status": "main"}).eq("id",
message.from_id).execute()
if not user['free_used']:
return await free_trial_handler(message)
return await message.answer("✅ Фото лица принято", keyboard=get_main_menu(user['free_used']))
# 5. ГЕНЕРАЦИЯ
gender = user['gender'] or "female"
styles_dict = prompts.FACE_SWAP_STYLES.get(gender, {})
if text in styles_dict or user['status'] == "wait_custom_prompt":
if user['balance'] <= 0 and user['status'] != "wait_custom_prompt":
if not (text in styles_dict and not user['free_used'] and text == "🎁 Одна фотография в подарок"):
return await message.answer("В твоем кошельке 0 генераций. Пополни баланс:",
keyboard=get_fast_packs_menu())
if user['status'] == "wait_custom_prompt" and user['balance'] <= 0:
return await message.answer("В твоем кошельке 0 генераций. Пополни баланс:", keyboard=get_fast_packs_menu())
if not user['primary_photo']:
supabase.table("users").update({"status": "wait_primary_photo"}).eq("id", message.from_id).execute()
return await message.answer("Сначала отправь фото своего лица")
if text in styles_dict:
options = styles_dict[text]
last_val = user.get('last_style') or ""
last_idx = -1
if "|" in last_val and last_val.split("|")[0] == text:
try:
last_idx = int(last_val.split("|")[1])
except:
last_idx = -1
available_indices = [i for i in range(len(options)) if i != last_idx]
new_idx = random.choice(available_indices)
prompt = options[new_idx]
new_last_style = f"{text}|{new_idx}"
else:
prompt = text
new_last_style = "custom_prompt|0"
await message.answer("⏳ Генерирую фотографию... Подожди около 30 секунд")
paths_to_clean = []
bucket_path, temporary_url = await upload_to_bucket(user['primary_photo'], message.from_id, "src")
if bucket_path: paths_to_clean.append(bucket_path)
res_url = await generator.generate_image(prompt=prompt, gender=gender,
photo_url=temporary_url if temporary_url else user['primary_photo'])
if res_url:
res_path, res_url_bucket = await upload_to_bucket(res_url, message.from_id, "res")
if res_path: paths_to_clean.append(res_path)
if res_url_bucket:
try:
file_stream = await download_file(res_url_bucket)
if file_stream:
photo_vk = await uploader.upload(file_stream, file_name="style.png")
await message.answer("Отлично! Твое фото готово ✨", attachment=photo_vk)
supabase.rpc("increment_balance", {"user_id": message.from_id, "amount": -1}).execute()
supabase.table("users").update({"last_style": new_last_style, "status": "main"}).eq("id",
message.from_id).execute()
else:
await message.answer("Ошибка: не удалось получить файл из хранилища. Баланс сохранен.")
except Exception as e:
logger.error(f"VK Style Error: {e}")
await message.answer("ВК не принял фото. Баланс сохранен.")
else:
await message.answer("Ошибка нейросети. Попробуйте другой запрос.")
if paths_to_clean: supabase.storage.from_(BUCKET_NAME).remove(paths_to_clean)
return
await message.answer("Выбирай действие в меню 👇", keyboard=get_main_menu(user['free_used']))
@app.route('/yoomoney', methods=['POST'])
def yoomoney_webhook():
d = request.form
check = f"{d.get('notification_type')}&{d.get('operation_id')}&{d.get('amount')}&{d.get('currency')}&{d.get('datetime')}&{d.get('sender')}&{d.get('codepro')}&{YOOMONEY_SECRET}&{d.get('label')}"
my_hash = hashlib.sha1(check.encode('utf-8')).hexdigest()
if my_hash == d.get('sha1_hash'):
try:
u_id = int(d.get('label'))
received_amount = float(d.get('withdraw_amount') or d.get('amount'))
packs = supabase.table("settings").select("*").execute().data
for pack in packs:
if abs(received_amount - float(pack['price'])) / float(pack['price']) < 0.05:
# 1. Начисляем баланс самому пользователю
supabase.rpc("increment_balance", {"user_id": u_id, "amount": pack['amount']}).execute()
# 2. Логика реферального бонуса
u_data_res = supabase.table("users").select("invited_by", "ref_bonus_done").eq("id", u_id).execute()
if u_data_res.data:
u_data = u_data_res.data[0]
if u_data['invited_by'] and not u_data['ref_bonus_done']:
inviter_id = u_data['invited_by']
# Начисляем бонус тому, кто пригласил
supabase.rpc("increment_balance",
{"user_id": inviter_id, "amount": REFERRAL_BONUS}).execute()
# Помечаем, что бонус за этого реферала выдан
supabase.table("users").update({"ref_bonus_done": True}).eq("id", u_id).execute()
if loop:
asyncio.run_coroutine_threadsafe(
bot.api.messages.send(peer_id=inviter_id, random_id=0,
message=f"🎁 Твой друг совершил первую оплату! Тебе начислено +{REFERRAL_BONUS} бонусных генераций."),
loop
)
if loop:
asyncio.run_coroutine_threadsafe(bot.api.messages.send(peer_id=u_id, random_id=0,
message=f"✅ Спасибо! Оплата прошла успешно! Теперь выбирай стиль фото в меню"),
loop)
break
except Exception as e:
logger.error(f"Webhook error: {e}")
return "OK", 200
@app.route('/')
def index():
return "<h1>VK Bot is Running</h1>", 200
if __name__ == "__main__":
loop = asyncio.get_event_loop()
threading.Thread(target=lambda: app.run(host="0.0.0.0", port=7860, use_reloader=False), daemon=True).start()
bot.run_forever()