456 lines
17 KiB
Python
456 lines
17 KiB
Python
import asyncio
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import tempfile
|
|
import uuid
|
|
from datetime import datetime
|
|
from PIL import Image
|
|
import subprocess
|
|
import requests
|
|
import shutil
|
|
import glob
|
|
from telegram import Update, InlineKeyboardMarkup, InlineKeyboardButton
|
|
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, ConversationHandler, CallbackQueryHandler
|
|
|
|
RULES_CALLBACK = "show_rules"
|
|
BACK_CALLBACK = "back_to_main"
|
|
|
|
COMFYUI_WORKFLOW = "video_generate.json"
|
|
TELEGRAM_TOKEN = "TELEGRAM_TOKEN"
|
|
COMFYUI_URL = "http://127.0.0.1:8188"
|
|
COMFYUI_OUTPUT_DIR = "output/"
|
|
COMFYUI_INPUT_DIR = "input/"
|
|
COMFYUI_COMMON_PROMT = ["woman take off her clothes, full nudity, full naked, pyssy, boobs, nipples. camera zooms out",""]
|
|
COMFYUI_MAX_SIZE = 580
|
|
COMFYUI_COST = 20
|
|
ADMIN_IDS = [1145483994, 498845556]
|
|
|
|
WAITING_CONFIRM = 0
|
|
|
|
def init_db():
|
|
conn = sqlite3.connect('bot.db')
|
|
c = conn.cursor()
|
|
c.execute('''CREATE TABLE IF NOT EXISTS users (
|
|
user_id INTEGER PRIMARY KEY,
|
|
username TEXT,
|
|
balance INTEGER DEFAULT 0,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)''')
|
|
c.execute('''CREATE TABLE IF NOT EXISTS admin_logs (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
admin_id INTEGER NOT NULL,
|
|
admin_username TEXT,
|
|
target_user_id INTEGER NOT NULL,
|
|
tokens_added INTEGER NOT NULL,
|
|
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)''')
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def log_admin_action(admin_id, admin_username, target_user_id, tokens):
|
|
"""Записывает действие админа в файл"""
|
|
timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
|
log_entry = (
|
|
f"[{timestamp}] Admin {admin_id} (@{admin_username}) "
|
|
f"added {tokens} tokens to user {target_user_id}\n"
|
|
)
|
|
with open('admin_actions.log', 'a', encoding='utf-8') as f:
|
|
f.write(log_entry)
|
|
|
|
def get_user_balance(user_id):
|
|
conn = sqlite3.connect('bot.db')
|
|
c = conn.cursor()
|
|
c.execute('SELECT balance FROM users WHERE user_id = ?', (user_id,))
|
|
result = c.fetchone()
|
|
conn.close()
|
|
return result[0] if result else 0
|
|
|
|
def add_user(user_id, username):
|
|
conn = sqlite3.connect('bot.db')
|
|
c = conn.cursor()
|
|
c.execute('INSERT OR IGNORE INTO users (user_id, username) VALUES (?, ?)', (user_id, username))
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def update_balance(user_id, amount):
|
|
conn = sqlite3.connect('bot.db')
|
|
c = conn.cursor()
|
|
c.execute('UPDATE users SET balance = balance + ? WHERE user_id = ?', (amount, user_id))
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def deduct_balance(user_id, amount):
|
|
conn = sqlite3.connect('bot.db')
|
|
c = conn.cursor()
|
|
c.execute('UPDATE users SET balance = balance - ? WHERE user_id = ?', (amount, user_id))
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def resize_image(image_path, max_size=COMFYUI_MAX_SIZE):
|
|
with Image.open(image_path) as img:
|
|
img = img.convert('RGBA')
|
|
width, height = img.size
|
|
if width > height:
|
|
new_width = min(width, max_size)
|
|
new_height = int(height * new_width / width)
|
|
else:
|
|
new_height = min(height, max_size)
|
|
new_width = int(width * new_height / height)
|
|
resized = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
|
return new_width, new_height, resized
|
|
|
|
def find_video_output_node(workflow_dict):
|
|
for node_id, node in workflow_dict.items():
|
|
inputs = node.get("inputs", {})
|
|
if isinstance(inputs, dict) and "filename_prefix" in inputs:
|
|
return node_id
|
|
raise ValueError("No video output node with 'filename_prefix' found in workflow!")
|
|
|
|
def find_latest_video(output_dir, extensions=('.mp4', '.webm', '.mkv')):
|
|
"""Ищет самый свежий видеофайл в output_dir"""
|
|
all_videos = []
|
|
for root, _, files in os.walk(output_dir):
|
|
for f in files:
|
|
if f.lower().endswith(extensions):
|
|
full_path = os.path.join(root, f)
|
|
all_videos.append(full_path)
|
|
if not all_videos:
|
|
return None
|
|
return max(all_videos, key=os.path.getmtime)
|
|
|
|
def modify_workflow(workflow_str, prompt, image_filename, width, height, video_prefix):
|
|
wf_raw = json.loads(workflow_str)
|
|
wf = {}
|
|
for key, value in wf_raw.items():
|
|
clean_key = key.strip()
|
|
if isinstance(value, dict) and 'inputs' in value:
|
|
inputs_raw = value['inputs']
|
|
if isinstance(inputs_raw, dict):
|
|
clean_inputs = {}
|
|
for in_key, in_val in inputs_raw.items():
|
|
clean_inputs[in_key.strip()] = in_val
|
|
value = {k: v for k, v in value.items() if k != 'inputs'}
|
|
value['inputs'] = clean_inputs
|
|
wf[clean_key] = value
|
|
|
|
if '93' in wf and 'inputs' in wf['93']:
|
|
wf['93']['inputs']['text'] = f"{COMFYUI_COMMON_PROMT}"
|
|
|
|
height = (height // 16) * 16
|
|
width = (width // 16) * 16
|
|
|
|
if '183' in wf and 'inputs' in wf['183']:
|
|
wf['183']['inputs']['value'] = height
|
|
if '184' in wf and 'inputs' in wf['184']:
|
|
wf['184']['inputs']['value'] = width
|
|
|
|
if '193' in wf and 'inputs' in wf['193']:
|
|
wf['193']['inputs']['image'] = image_filename
|
|
|
|
if '214' in wf and 'inputs' in wf['214']:
|
|
wf['214']['inputs']['filename_prefix'] = video_prefix
|
|
|
|
return json.dumps(wf)
|
|
|
|
def find_video_by_prefix(output_dir, prefix, extensions=('.mp4', '.webm', '.mkv')):
|
|
for ext in extensions:
|
|
pattern = os.path.join(output_dir, '**', f"{prefix}*{ext}")
|
|
matches = glob.glob(pattern, recursive=True)
|
|
if matches:
|
|
return max(matches, key=os.path.getmtime)
|
|
return None
|
|
|
|
async def run_comfyui_workflow(workflow_json_str, input_image_path, video_prefix, update: Update):
|
|
workflow_dict = json.loads(workflow_json_str)
|
|
client_id = str(uuid.uuid4())
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.post(f"{COMFYUI_URL}/prompt", json={
|
|
"prompt": workflow_dict,
|
|
"client_id": client_id
|
|
}) as resp:
|
|
if resp.status != 200:
|
|
raise Exception(f"ComfyUI API error {resp.status}")
|
|
prompt_id = (await resp.json())["prompt_id"]
|
|
|
|
progress_msg = await update.message.reply_text("🚀 Запуск генерации...")
|
|
uri = f"ws://127.0.0.1:8188/ws?clientId={client_id}"
|
|
|
|
try:
|
|
async with websockets.connect(uri) as websocket:
|
|
while True:
|
|
try:
|
|
msg = await asyncio.wait_for(websocket.recv(), timeout=300.0)
|
|
data = json.loads(msg)
|
|
|
|
if data["type"] == "executing":
|
|
node = data["data"].get("node")
|
|
if node is None:
|
|
break
|
|
try:
|
|
await progress_msg.edit_text("🔄 Обрабатываю...")
|
|
except BadRequest as e:
|
|
if "Message is not modified" not in str(e):
|
|
raise
|
|
|
|
elif data["type"] == "progress":
|
|
value = data["data"].get("value", 0)
|
|
max_val = data["data"].get("max", 1)
|
|
pct = int(100 * value / max_val) if max_val else 0
|
|
await progress_msg.edit_text(f"⏳ Прогресс: {pct}%")
|
|
|
|
elif data["type"] == "execution_error":
|
|
err = data["data"].get("exception_message", "Ошибка")
|
|
raise Exception(err)
|
|
|
|
except asyncio.TimeoutError:
|
|
break
|
|
|
|
except Exception as e:
|
|
print(f"WebSocket error: {e}")
|
|
|
|
await progress_msg.edit_text("🎥 Собираю видео...")
|
|
for _ in range(60):
|
|
await asyncio.sleep(2)
|
|
video_path = find_latest_video(COMFYUI_OUTPUT_DIR)
|
|
if video_path and os.path.getsize(video_path) > 1024:
|
|
await progress_msg.edit_text("✅ Готово!")
|
|
return video_path
|
|
|
|
raise Exception("Видео не найдено")
|
|
|
|
def clean_metadata(video_path):
|
|
temp_dir = os.path.dirname(video_path)
|
|
with tempfile.NamedTemporaryFile(suffix='.mp4', dir=temp_dir, delete=False) as tmp:
|
|
temp_path = tmp.name
|
|
|
|
try:
|
|
subprocess.run([
|
|
'ffmpeg', '-y', '-i', video_path,
|
|
'-map_metadata', '-1',
|
|
'-c:v', 'libx264',
|
|
'-c:a', 'aac',
|
|
'-pix_fmt', 'yuv420p',
|
|
temp_path
|
|
], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
|
|
shutil.copy2(temp_path, video_path)
|
|
os.unlink(temp_path)
|
|
return video_path
|
|
except Exception as e:
|
|
if os.path.exists(temp_path):
|
|
os.unlink(temp_path)
|
|
raise Exception(f"Metadata cleanup failed: {str(e)}")
|
|
|
|
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
|
|
user = update.effective_user
|
|
add_user(user.id, user.username)
|
|
balance = get_user_balance(user.id)
|
|
keyboard = [[InlineKeyboardButton("ПРАВИЛА", callback_data=RULES_CALLBACK)]]
|
|
await update.message.reply_text(
|
|
f"👋🏻 Привет, Творец!\n\n"
|
|
"Рада видеть тебя в SecretRoom\n\n"
|
|
"✅ Любая фантазия о которой ты мечтал оживёт в этом боте\n\n"
|
|
"Жми «🔮 Оживить фото» и наслаждайся\n\n"
|
|
"Стоимость: 20 токен за запрос.\n\n"
|
|
f"Ваш баланс: {balance} токенов.",
|
|
reply_markup=InlineKeyboardMarkup(keyboard),
|
|
)
|
|
return ConversationHandler.END
|
|
|
|
async def handle_admin_add(update: Update, context: ContextTypes.DEFAULT_TYPE):
|
|
if update.effective_user.id not in ADMIN_IDS:
|
|
await update.message.reply_text("У вас нет прав администратора.")
|
|
return
|
|
try:
|
|
tokens = int(context.args[0])
|
|
target_id = int(context.args[1])
|
|
update_balance(target_id, tokens)
|
|
log_admin_action(
|
|
update.effective_user.id,
|
|
update.effective_user.username or 'unknown',
|
|
target_id,
|
|
tokens
|
|
)
|
|
await update.message.reply_text(f"Добавлено {tokens} токенов пользователю {target_id}.")
|
|
except (IndexError, ValueError):
|
|
await update.message.reply_text("Использование: /add <токены> <user_id>")
|
|
|
|
async def handle_photo(update: Update, context: ContextTypes.DEFAULT_TYPE):
|
|
photo_file = await update.message.photo[-1].get_file()
|
|
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
|
|
temp_image = tmp.name
|
|
await photo_file.download_to_drive(temp_image)
|
|
|
|
width, height, resized_img = resize_image(temp_image)
|
|
resized_path = temp_image.replace('.png', '_resized.png')
|
|
resized_img.save(resized_path, format='PNG')
|
|
os.unlink(temp_image)
|
|
|
|
context.user_data['image_path'] = resized_path
|
|
context.user_data['width'] = width
|
|
context.user_data['height'] = height
|
|
|
|
video_prefix = f"video_{uuid.uuid4().hex[:12]}"
|
|
context.user_data['video_prefix'] = video_prefix
|
|
|
|
cost = COMFYUI_COST
|
|
balance = get_user_balance(update.effective_user.id)
|
|
|
|
keyboard = [
|
|
[InlineKeyboardButton("✅ Начать (Пресет 1)", callback_data="preset_1")],
|
|
[InlineKeyboardButton("🎭 Начать (Пресет 2)", callback_data="preset_2")],
|
|
[InlineKeyboardButton("❌ Отмена", callback_data="cancel_gen")]
|
|
]
|
|
reply_markup = InlineKeyboardMarkup(keyboard)
|
|
|
|
await update.message.reply_text(
|
|
f"Стоимость генерации: {cost} токенов.\n"
|
|
f"Ваш баланс: {balance}\n"
|
|
"Выберите пресет для генерации:",
|
|
reply_markup=reply_markup
|
|
)
|
|
|
|
return ConversationHandler.END
|
|
|
|
async def handle_photo_anywhere(update: Update, context: ContextTypes.DEFAULT_TYPE):
|
|
return await handle_photo(update, context)
|
|
|
|
async def handle_callback(update: Update, context: ContextTypes.DEFAULT_TYPE):
|
|
query = update.callback_query
|
|
await query.answer()
|
|
if query.data == "preset_1":COMFYUI_COMMON_PROMT[0]
|
|
if query.data == "preset_2":COMFYUI_COMMON_PROMT[1]
|
|
if query.data == RULES_CALLBACK:
|
|
keyboard = [[InlineKeyboardButton("← Назад", callback_data=BACK_CALLBACK)]]
|
|
await query.message.reply_text(
|
|
"🖼 Творец, отправь фото ниже!:\n\n"
|
|
"Выбери фото человека или из аниме🙈:\n\n"
|
|
"Для лучшего результата следуй этим правилам:\n\n"
|
|
"✅ Фотография в полный рост\n"
|
|
"✅ Девушка смотрит прямо в камеру\n"
|
|
"✅ Хорошее освещение\n"
|
|
"❌ Нет солнцезащитных очков\n"
|
|
"❌ Не закрывать лицо волосами",
|
|
reply_markup=InlineKeyboardMarkup(keyboard)
|
|
)
|
|
await query.message.delete()
|
|
elif query.data == "preset_1":
|
|
await query.answer()
|
|
await start_generation(update, context, preset_index=0)
|
|
|
|
elif query.data == "preset_2":
|
|
await query.answer()
|
|
await start_generation(update, context, preset_index=1)
|
|
|
|
elif query.data == "cancel_gen":
|
|
await query.answer()
|
|
image_path = context.user_data.get('image_path')
|
|
if image_path and os.path.exists(image_path):
|
|
os.unlink(image_path)
|
|
context.user_data.clear()
|
|
await query.message.edit_text("Операция отменена.")
|
|
|
|
elif query.data == BACK_CALLBACK:
|
|
context.user_data.clear()
|
|
|
|
user = update.effective_user
|
|
add_user(user.id, user.username)
|
|
balance = get_user_balance(user.id)
|
|
keyboard = [[InlineKeyboardButton("ПРАВИЛА", callback_data=RULES_CALLBACK)]]
|
|
await query.message.reply_text(
|
|
f"👋🏻 Привет, Творец!\n\n"
|
|
"Рада видеть тебя в SecretRoom\n\n"
|
|
"✅ Любая фантазия о которой ты мечтал оживёт в этом боте\n\n"
|
|
"Жми «🔮 Оживить фото» и наслаждайся\n\n"
|
|
"Стоимость: 20 токен за запрос.\n\n"
|
|
f"Ваш баланс: {balance} токенов.",
|
|
reply_markup=InlineKeyboardMarkup(keyboard)
|
|
)
|
|
await query.message.delete()
|
|
|
|
async def cancel(update: Update, context: ContextTypes.DEFAULT_TYPE):
|
|
await update.message.reply_text("Отменено.")
|
|
image_path = context.user_data.get('image_path')
|
|
if image_path and os.path.exists(image_path):
|
|
os.unlink(image_path)
|
|
context.user_data.clear()
|
|
return ConversationHandler.END
|
|
|
|
async def start_generation(update: Update, context: ContextTypes.DEFAULT_TYPE, preset_index: int):
|
|
query = update.callback_query
|
|
user_id = update.effective_user.id
|
|
|
|
cost = COMFYUI_COST
|
|
balance = get_user_balance(user_id)
|
|
if balance < cost:
|
|
await query.edit_message_text("❌ Недостаточно токенов.")
|
|
context.user_data.clear()
|
|
return
|
|
|
|
deduct_balance(user_id, cost)
|
|
await query.edit_message_text("🚀 Начинаю генерацию...")
|
|
|
|
try:
|
|
image_path = context.user_data['image_path']
|
|
video_prefix = context.user_data['video_prefix']
|
|
selected_prompt = COMFYUI_COMMON_PROMT[preset_index]
|
|
|
|
os.makedirs(COMFYUI_INPUT_DIR, exist_ok=True)
|
|
image_filename = f"user_image_{uuid.uuid4().hex[:12]}.png"
|
|
input_image_path = os.path.join(COMFYUI_INPUT_DIR, image_filename)
|
|
shutil.copy(image_path, input_image_path)
|
|
|
|
with open(COMFYUI_WORKFLOW, 'r', encoding='utf-8') as f:
|
|
workflow_base = f.read()
|
|
|
|
modified_workflow = modify_workflow(
|
|
workflow_base,
|
|
selected_prompt,
|
|
image_filename,
|
|
context.user_data['width'],
|
|
context.user_data['height'],
|
|
video_prefix
|
|
)
|
|
|
|
video_path = await run_comfyui_workflow(modified_workflow, input_image_path, video_prefix)
|
|
video_path = clean_metadata(video_path)
|
|
|
|
with open(video_path, 'rb') as video_file:
|
|
await query.message.reply_video(video_file, supports_streaming=True)
|
|
|
|
await query.message.reply_text("✅ Генерация завершена!")
|
|
|
|
for p in [image_path, input_image_path, video_path]:
|
|
if os.path.exists(p):
|
|
os.unlink(p)
|
|
|
|
except Exception as e:
|
|
await query.message.reply_text(f"❌ Ошибка генерации: {str(e)}")
|
|
update_balance(user_id, cost)
|
|
|
|
new_balance = get_user_balance(user_id)
|
|
await query.message.reply_text(f"Ваш баланс: {new_balance} токенов.")
|
|
context.user_data.clear()
|
|
|
|
def main():
|
|
init_db()
|
|
app = Application.builder().token(TELEGRAM_TOKEN).build()
|
|
|
|
|
|
conv_handler = ConversationHandler(
|
|
entry_points=[MessageHandler(filters.PHOTO & ~filters.COMMAND, handle_photo)],
|
|
states={},
|
|
fallbacks=[],
|
|
)
|
|
|
|
app.add_handler(CommandHandler("start", start))
|
|
app.add_handler(CommandHandler("add", handle_admin_add))
|
|
app.add_handler(conv_handler)
|
|
app.add_handler(CallbackQueryHandler(handle_callback))
|
|
print(f"Бот запущен. Output dir: {COMFYUI_OUTPUT_DIR}")
|
|
app.run_polling()
|
|
|
|
if __name__ == '__main__':
|
|
main() |