#!/usr/bin/env python3
"""
AIGC Telegram Bot — AI Cover Generation Pipeline for Telegram
Features:
- Upload audio or provide YouTube URL to generate AI voice covers
- Choose from multiple RVC voice models
- Customize pitch, index rate, f0 method, reverb, and more
- Download generated covers directly in Telegram
- Admin controls for managing models and users
Usage:
python -m src.bot
"""
import asyncio
import html
import json
import logging
import os
import signal
import sys
import time
import traceback
from pathlib import Path
from typing import Optional
from telegram import (
Update,
BotCommand,
InlineKeyboardButton,
InlineKeyboardMarkup,
)
from telegram.ext import (
Application,
ApplicationBuilder,
CallbackQueryHandler,
CommandHandler,
ContextTypes,
ConversationHandler,
MessageHandler,
filters,
)
from .config import load_config, load_pipeline_config, get_path_config, BotConfig, PipelineConfig, PathConfig
from .state import state_manager, UserState, InputType, UserSession
from .pipeline import PipelineRunner
# ============================================================
# Setup logging
# ============================================================
logging.basicConfig(
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("aigc_bot")
def setup_logging(config: BotConfig):
"""Configure logging based on bot config."""
root_logger = logging.getLogger()
root_logger.setLevel(getattr(logging, config.log_level.upper(), logging.INFO))
if config.log_file:
file_handler = logging.FileHandler(config.log_file)
file_handler.setFormatter(
logging.Formatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s")
)
root_logger.addHandler(file_handler)
# ============================================================
# Constants
# ============================================================
CONVO_STATES = {
"INPUT_TYPE": 1,
"SONG_INPUT": 2,
"VOICE_MODEL": 3,
"CONFIRM": 4,
"SETTINGS_MENU": 5,
"PITCH_INPUT": 6,
"MODE_SELECT": 7,
}
# Callback data prefixes
CB_PREFIX = {
"INPUT_YT": "input_yt",
"INPUT_FILE": "input_file",
"MODEL_SELECT": "mdl_",
"CONFIRM_START": "confirm_start",
"CANCEL": "cancel",
"SETTINGS": "settings",
"SET_PITCH": "set_pitch",
"SET_MODE": "set_mode",
"SET_FORMAT": "set_fmt_",
"SET_F0": "set_f0_",
"SET_DENOISE": "set_denoise",
"SET_KEEP": "set_keep",
"SET_REVERB": "set_reverb",
"BACK_MENU": "back_menu",
}
# ============================================================
# AIGC Telegram Bot Class
# ============================================================
class AIGCTelegramBot:
"""
Main Telegram Bot application for AI Cover Generation.
"""
def __init__(self):
self.config: BotConfig = load_config()
self.pipeline_config: PipelineConfig = load_pipeline_config()
self.paths: PathConfig = get_path_config()
self.pipeline: Optional[PipelineRunner] = None
self.application: Optional[Application] = None
self._progress_messages: dict = {} # user_id -> last progress message
setup_logging(self.config)
# Validate config
if not self.config.bot_token:
logger.error("BOT_TOKEN is not set! Please configure it in .env file.")
sys.exit(1)
async def initialize(self):
"""Initialize the bot and pipeline."""
logger.info("Initializing AIGC Telegram Bot...")
# Initialize pipeline runner
self.pipeline = PipelineRunner(
models_dir=str(self.paths.models_dir),
output_dir=str(self.paths.output_dir),
temp_dir=str(self.paths.temp_dir),
mdx_models_dir=str(self.paths.mdx_models_dir),
max_workers=self.pipeline_config.max_concurrent_jobs,
)
logger.info(f"Pipeline initialized with {self.pipeline_config.max_concurrent_jobs} workers")
# Build application
self.application = (
ApplicationBuilder()
.token(self.config.bot_token)
.concurrent_updates(True)
.build()
)
# Register handlers
self._register_handlers()
# Register bot commands
await self._register_commands()
logger.info(f"Bot '{self.config.bot_name}' initialized successfully")
def _register_handlers(self):
"""Register all Telegram bot handlers."""
app = self.application
# Start handler
app.add_handler(CommandHandler("start", self.cmd_start))
# Help handler
app.add_handler(CommandHandler("help", self.cmd_help))
# Cancel handler
app.add_handler(CommandHandler("cancel", self.cmd_cancel))
# Status handler
app.add_handler(CommandHandler("status", self.cmd_status))
# Models list handler
app.add_handler(CommandHandler("models", self.cmd_models))
# Settings handler
app.add_handler(CommandHandler("settings", self.cmd_settings))
# Admin commands
app.add_handler(CommandHandler("admin", self.cmd_admin))
# Cover generation conversation handler
cover_handler = ConversationHandler(
entry_points=[
CommandHandler("cover", self.cmd_cover),
CommandHandler("newcover", self.cmd_cover),
],
states={
CONVO_STATES["INPUT_TYPE"]: [
CallbackQueryHandler(self.on_input_type_select, pattern=f"^({'|'.join([CB_PREFIX['INPUT_YT'], CB_PREFIX['INPUT_FILE'], CB_PREFIX['CANCEL']])})$"),
],
CONVO_STATES["SONG_INPUT"]: [
MessageHandler(filters.TEXT & ~filters.COMMAND, self.on_song_url_input),
MessageHandler(filters.AUDIO | filters.VOICE, self.on_song_file_input),
CommandHandler("cancel", self.cmd_cancel),
],
CONVO_STATES["VOICE_MODEL"]: [
CallbackQueryHandler(self.on_model_select, pattern=f"^{CB_PREFIX['MODEL_SELECT']}"),
CallbackQueryHandler(self.on_cancel, pattern=f"^{CB_PREFIX['CANCEL']}$"),
],
CONVO_STATES["CONFIRM"]: [
CallbackQueryHandler(self.on_confirm_start, pattern=f"^{CB_PREFIX['CONFIRM_START']}$"),
CallbackQueryHandler(self.on_settings_from_confirm, pattern=f"^{CB_PREFIX['SETTINGS']}$"),
CallbackQueryHandler(self.on_cancel, pattern=f"^{CB_PREFIX['CANCEL']}$"),
],
},
fallbacks=[
CommandHandler("cancel", self.cmd_cancel),
CommandHandler("start", self.cmd_start),
],
allow_reentry=True,
)
app.add_handler(cover_handler)
# Settings conversation handler
settings_handler = ConversationHandler(
entry_points=[
CallbackQueryHandler(self.on_settings_callback, pattern=f"^{CB_PREFIX['SET_PITCH']}|{CB_PREFIX['SET_MODE']}|{CB_PREFIX['SET_FORMAT']}|{CB_PREFIX['SET_F0']}|{CB_PREFIX['SET_DENOISE']}|{CB_PREFIX['SET_KEEP']}|{CB_PREFIX['SET_REVERB']}|{CB_PREFIX['BACK_MENU']}$"),
],
states={
CONVO_STATES["PITCH_INPUT"]: [
MessageHandler(filters.TEXT & ~filters.COMMAND, self.on_pitch_input),
],
CONVO_STATES["MODE_SELECT"]: [
CallbackQueryHandler(self.on_mode_select),
],
},
fallbacks=[
CommandHandler("cancel", self.cmd_cancel),
],
allow_reentry=True,
)
app.add_handler(settings_handler)
# Error handler
app.add_error_handler(self.error_handler)
# Log all messages for debugging
if self.config.log_level == "DEBUG":
app.add_handler(MessageHandler(filters.ALL, self._log_message), group=-1)
async def _register_commands(self):
"""Register bot commands with Telegram."""
commands = [
BotCommand("start", "Start the bot and see welcome message"),
BotCommand("cover", "Generate a new AI cover"),
BotCommand("newcover", "Start a new cover generation"),
BotCommand("models", "List available voice models"),
BotCommand("settings", "Adjust generation settings"),
BotCommand("status", "Check bot status and queue"),
BotCommand("cancel", "Cancel current operation"),
BotCommand("help", "Show help and usage guide"),
]
await self.application.bot.set_my_commands(commands)
# ============================================================
# Auth & Access Control
# ============================================================
def _is_authorized(self, user_id: int) -> bool:
"""Check if a user is authorized to use the bot."""
if self.config.allowed_user_ids is None:
return True # Public bot
return user_id in self.config.allowed_user_ids
def _is_admin(self, user_id: int) -> bool:
"""Check if a user is an admin."""
return user_id in self.config.admin_ids
# ============================================================
# Command Handlers
# ============================================================
async def cmd_start(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle /start command."""
user = update.effective_user
if not self._is_authorized(user.id):
await update.message.reply_text(
"Sorry, this bot is private and you don't have access.\n"
"Contact the bot administrator for access."
)
return ConversationHandler.END
state_manager.clear_session(user.id)
welcome_text = (
f"Hi {html.escape(user.first_name)}! Welcome to {html.escape(self.config.bot_name)}\n\n"
"I can generate AI voice covers of your favorite songs using RVC voice models.\n\n"
"How to use:\n"
"1. Send /cover to start\n"
"2. Choose to upload audio or paste a YouTube link\n"
"3. Select a voice model\n"
"4. Adjust settings (optional)\n"
"5. Confirm and wait for your cover!\n\n"
"Commands:\n"
"/cover - Start cover generation\n"
"/models - List available voice models\n"
"/settings - Adjust generation settings\n"
"/status - Check bot status\n"
"/help - Show detailed help"
)
keyboard = InlineKeyboardMarkup([
[InlineKeyboardButton("Generate Cover", callback_data="start_cover")],
[InlineKeyboardButton("View Models", callback_data="view_models")],
[InlineKeyboardButton("Settings", callback_data="settings")],
])
await update.message.reply_text(
welcome_text,
parse_mode="HTML",
reply_markup=keyboard,
)
# Handle inline button presses for start menu
# These are captured by the general callback handler below
return ConversationHandler.END
async def cmd_help(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle /help command with detailed usage guide."""
help_text = (
"AIGC Cover Bot - Help Guide\n\n"
"Quick Start:\n"
"1. Send /cover to begin\n"
"2. Choose input method (YouTube URL or audio file upload)\n"
"3. For YouTube: paste the video link\n"
" For file: upload an MP3, WAV, or other audio file\n"
"4. Select a voice model from the list\n"
"5. Review settings and tap 'Generate Cover'\n\n"
"Voice Models:\n"
"Voice models determine how the AI will sound. Each model is trained "
"on a different voice. Use /models to see all available models.\n\n"
"Settings:\n"
"/settings lets you customize:\n"
"- Pitch Change: Shift the AI voice pitch in octaves. "
"Use +1 for male-to-female, -1 for female-to-male.\n"
"- Pitch Change All: Shift pitch of all audio (vocals + instruments).\n"
"- Index Rate: Timbre similarity (0-1). Higher = more like training data.\n"
"- F0 Method: Pitch detection algorithm. rmvpe = best clarity.\n"
"- Reverb: Add room reverb to the AI vocals.\n"
"- Denoise: Apply extra noise reduction.\n"
"- Output Format: MP3 (smaller) or WAV (best quality).\n\n"
"Inference Modes:\n"
"- Full: Separate vocals + RVC voice conversion + mix\n"
"- RVC Only: Voice conversion only (skip vocal separation)\n"
"- MDX Only: Vocal separation only (skip voice conversion)\n\n"
"Tips:\n"
"- YouTube links must be direct video URLs (not playlists)\n"
"- Audio files up to 20MB are supported\n"
"- Processing time depends on song length (usually 1-3 minutes)\n"
"- Use /cancel to abort at any time\n\n"
"Admin Commands:\n"
"/admin stats - View bot usage statistics\n"
"/admin cleanup - Clean up old temp files"
)
await update.message.reply_text(help_text, parse_mode="HTML")
async def cmd_cancel(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle /cancel command."""
user = update.effective_user
state_manager.clear_session(user.id)
state_manager.set_state(user.id, UserState.IDLE)
await update.message.reply_text(
"Operation cancelled. Send /cover to start a new one."
)
return ConversationHandler.END
async def cmd_models(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle /models command - list available voice models."""
if not self.pipeline:
await update.message.reply_text("Pipeline is not initialized yet.")
return
models = self.pipeline.get_available_models()
if not models:
await update.message.reply_text(
"No voice models available. Please add models to the "
f"{self.paths.models_dir} directory.\n\n"
"Each model should be in its own folder containing a .pth file "
"(and optionally a .index file).",
parse_mode="HTML",
)
return
models_text = (
f"Available Voice Models ({len(models)}):\n\n"
)
for i, model_name in enumerate(models, 1):
models_text += f" {i}. {html.escape(model_name)}\n"
models_text += (
"\nSend /cover to use one of these models."
)
await update.message.reply_text(models_text, parse_mode="HTML")
async def cmd_status(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle /status command."""
active_count = state_manager.get_active_count()
models_count = len(self.pipeline.get_available_models()) if self.pipeline else 0
status_text = (
f"{html.escape(self.config.bot_name)} - Status\n\n"
f"Active jobs: {active_count}/{self.pipeline_config.max_concurrent_jobs}\n"
f"Available models: {models_count}\n"
f"Output format: {self.pipeline_config.default_output_format}\n"
f"Inference mode: {self.pipeline_config.default_inference_mode}\n"
)
session = state_manager.get_session(update.effective_user.id)
status_text += f"\nYour state: {session.state.value}"
await update.message.reply_text(status_text, parse_mode="HTML")
async def cmd_settings(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle /settings command."""
session = state_manager.get_session(update.effective_user.id)
await self._send_settings_menu(update.effective_chat.id, session, context)
return ConversationHandler.END
async def cmd_admin(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle /admin command."""
user = update.effective_user
if not self._is_admin(user.id):
await update.message.reply_text("You don't have admin access.")
return
if not context.args:
await update.message.reply_text(
"Admin Commands:\n"
"/admin stats - View statistics\n"
"/admin cleanup - Clean temp files\n"
"/admin broadcast - Send message to all users",
parse_mode="HTML",
)
return
action = context.args[0].lower()
if action == "stats":
models = self.pipeline.get_available_models() if self.pipeline else []
await update.message.reply_text(
f"Bot Stats\n"
f"Models: {len(models)}\n"
f"Active jobs: {state_manager.get_active_count()}\n"
f"Max concurrent: {self.pipeline_config.max_concurrent_jobs}",
parse_mode="HTML",
)
elif action == "cleanup":
if self.pipeline:
self.pipeline.cleanup_temp_files()
await update.message.reply_text("Temp files cleaned up.")
else:
await update.message.reply_text("Pipeline not initialized.")
elif action == "broadcast" and len(context.args) > 1:
msg = " ".join(context.args[1:])
await update.message.reply_text("Broadcast feature not yet implemented.")
else:
await update.message.reply_text("Unknown admin command.")
# ============================================================
# Cover Generation Flow
# ============================================================
async def cmd_cover(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle /cover command - start cover generation flow."""
user = update.effective_user
if not self._is_authorized(user.id):
await update.message.reply_text("You don't have access to this bot.")
return ConversationHandler.END
if state_manager.is_processing(user.id):
await update.message.reply_text(
"You already have a cover being generated. "
"Please wait or use /cancel to abort."
)
return ConversationHandler.END
state_manager.clear_session(user.id)
session = state_manager.get_session(user.id)
# Load defaults from config
session.pitch_change = self.pipeline_config.default_pitch_change
session.index_rate = self.pipeline_config.default_index_rate
session.filter_radius = self.pipeline_config.default_filter_radius
session.volume_envelope = self.pipeline_config.default_volume_envelope
session.f0_method = self.pipeline_config.default_f0_method
session.hop_length = self.pipeline_config.default_hop_length
session.protect = self.pipeline_config.default_protect
session.reverb_size = self.pipeline_config.default_reverb_size
session.reverb_wet = self.pipeline_config.default_reverb_wet
session.reverb_dry = self.pipeline_config.default_reverb_dry
session.reverb_damping = self.pipeline_config.default_reverb_damping
session.main_gain = self.pipeline_config.default_main_gain
session.backup_gain = self.pipeline_config.default_backup_gain
session.inst_gain = self.pipeline_config.default_inst_gain
session.output_format = self.pipeline_config.default_output_format
session.inference_mode = self.pipeline_config.default_inference_mode
# Step 1: Ask for input type
keyboard = InlineKeyboardMarkup([
[
InlineKeyboardButton(
"YouTube URL",
callback_data=CB_PREFIX["INPUT_YT"]
),
InlineKeyboardButton(
"Upload Audio File",
callback_data=CB_PREFIX["INPUT_FILE"]
),
],
[InlineKeyboardButton("Cancel", callback_data=CB_PREFIX["CANCEL"])],
])
await update.message.reply_text(
"Step 1/3: Choose Input Method\n\n"
"How would you like to provide the song?",
parse_mode="HTML",
reply_markup=keyboard,
)
return CONVO_STATES["INPUT_TYPE"]
async def on_input_type_select(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle input type selection (YouTube or File)."""
query = update.callback_query
await query.answer()
user_id = query.from_user.id
session = state_manager.get_session(user_id)
if query.data == CB_PREFIX["INPUT_YT"]:
session.input_type = InputType.YOUTUBE
state_manager.set_state(user_id, UserState.AWAITING_YOUTUBE_URL)
await query.edit_message_text(
"Step 2/3: Provide Song\n\n"
"Send me a YouTube video URL.\n"
"Example: https://www.youtube.com/watch?v=dQw4w9WgXcQ\n\n"
"Note: Only direct video URLs are supported (no playlists).",
parse_mode="HTML",
)
return CONVO_STATES["SONG_INPUT"]
elif query.data == CB_PREFIX["INPUT_FILE"]:
session.input_type = InputType.FILE
state_manager.set_state(user_id, UserState.AWAITING_SONG_INPUT)
await query.edit_message_text(
"Step 2/3: Provide Song\n\n"
"Upload an audio file (MP3, WAV, OGG, FLAC).\n"
f"Maximum file size: {self.pipeline_config.max_file_size_mb}MB\n\n"
"Send the audio file now:",
parse_mode="HTML",
)
return CONVO_STATES["SONG_INPUT"]
elif query.data == CB_PREFIX["CANCEL"]:
state_manager.clear_session(user_id)
await query.edit_message_text("Cancelled. Send /cover to start again.")
return ConversationHandler.END
async def on_song_url_input(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle YouTube URL input from user."""
user_id = update.effective_user.id
session = state_manager.get_session(user_id)
url = update.message.text.strip()
# Validate YouTube URL
if not self._is_valid_youtube_url(url):
await update.message.reply_text(
"That doesn't look like a valid YouTube URL.\n"
"Please send a direct YouTube video link like:\n"
"https://www.youtube.com/watch?v=...",
parse_mode="HTML",
)
return CONVO_STATES["SONG_INPUT"]
session.song_input = url
await self._show_model_selection(update.effective_chat.id, session, context)
return CONVO_STATES["VOICE_MODEL"]
async def on_song_file_input(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle audio file upload from user."""
user_id = update.effective_user.id
session = state_manager.get_session(user_id)
# Get the audio file
audio = update.message.audio or update.message.voice
if not audio:
await update.message.reply_text("Please send an audio file.")
return CONVO_STATES["SONG_INPUT"]
# Check file size
file_size_mb = audio.file_size / (1024 * 1024) if audio.file_size else 0
if file_size_mb > self.pipeline_config.max_file_size_mb:
await update.message.reply_text(
f"File too large ({file_size_mb:.1f}MB). "
f"Maximum is {self.pipeline_config.max_file_size_mb}MB."
)
return CONVO_STATES["SONG_INPUT"]
# Download the file to temp directory
status_msg = await update.message.reply_text("Downloading audio file...")
try:
file = await audio.get_file()
temp_filename = f"input_{user_id}_{int(time.time())}.ogg"
temp_path = str(self.paths.temp_dir / temp_filename)
await file.download_to_drive(temp_path)
session.song_input = temp_path
session.song_file_id = audio.file_id
await status_msg.edit_text("Audio file downloaded. Choose a voice model:")
except Exception as e:
logger.error(f"Failed to download audio: {e}")
await status_msg.edit_text(f"Failed to download audio file: {html.escape(str(e))}")
return ConversationHandler.END
await self._show_model_selection(update.effective_chat.id, session, context)
return CONVO_STATES["VOICE_MODEL"]
async def _show_model_selection(self, chat_id: int, session: UserSession, context: ContextTypes.DEFAULT_TYPE):
"""Display voice model selection keyboard."""
models = self.pipeline.get_available_models() if self.pipeline else []
if not models:
await context.bot.send_message(
chat_id,
"No voice models available. Please add models to the "
f"{self.paths.models_dir} directory and try again.",
parse_mode="HTML",
)
return ConversationHandler.END
# Build inline keyboard with models (paginated)
buttons = []
for model in models:
buttons.append([InlineKeyboardButton(
model,
callback_data=f"{CB_PREFIX['MODEL_SELECT']}{model}"
)])
buttons.append([InlineKeyboardButton("Cancel", callback_data=CB_PREFIX["CANCEL"])])
keyboard = InlineKeyboardMarkup(buttons)
text = (
f"Step 3/3: Select Voice Model\n\n"
f"Choose a voice model for the cover ({len(models)} available):\n"
)
if session.pitch_change != 0:
text += f"\nCurrent pitch change: {session.pitch_change} octaves"
await context.bot.send_message(chat_id, text, parse_mode="HTML", reply_markup=keyboard)
async def on_model_select(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle voice model selection."""
query = update.callback_query
await query.answer()
user_id = query.from_user.id
session = state_manager.get_session(user_id)
model_name = query.data.replace(CB_PREFIX["MODEL_SELECT"], "")
session.voice_model = model_name
# Show confirmation with current settings
await self._show_confirmation(query, session, context)
return CONVO_STATES["CONFIRM"]
async def _show_confirmation(self, query_or_msg, session: UserSession, context: ContextTypes.DEFAULT_TYPE):
"""Show generation confirmation with settings summary."""
pitch_display = f"+{session.pitch_change}" if session.pitch_change > 0 else str(session.pitch_change)
mode_labels = {
"full": "Full (Separate + Convert + Mix)",
"mdx": "MDX Only (Vocal Separation)",
"rvc": "RVC Only (Voice Conversion)",
}
input_display = (
"YouTube URL" if session.input_type == InputType.YOUTUBE
else "Audio File Upload"
)
song_display = session.song_input
if session.input_type == InputType.YOUTUBE:
if len(song_display) > 50:
song_display = song_display[:50] + "..."
text = (
"Confirm Cover Generation\n\n"
f"Song: {html.escape(song_display)}\n"
f"Input: {input_display}\n"
f"Voice Model: {html.escape(session.voice_model or 'None')}\n\n"
"Settings:\n"
f" Mode: {mode_labels.get(session.inference_mode, session.inference_mode)}\n"
f" Pitch: {pitch_display} octaves"
)
if session.pitch_change_all != 0:
text += f" (All: +{session.pitch_change_all})"
text += (
f"\n F0 Method: {session.f0_method}\n"
f" Index Rate: {session.index_rate}\n"
f" Filter Radius: {session.filter_radius}\n"
f" Protect: {session.protect}\n"
f" Output: {session.output_format.upper()}\n"
f" Extra Denoise: {'Yes' if session.extra_denoise else 'No'}\n"
)
keyboard = InlineKeyboardMarkup([
[
InlineKeyboardButton(
"Generate Cover",
callback_data=CB_PREFIX["CONFIRM_START"]
),
],
[
InlineKeyboardButton(
"Adjust Settings",
callback_data=CB_PREFIX["SETTINGS"]
),
InlineKeyboardButton(
"Cancel",
callback_data=CB_PREFIX["CANCEL"]
),
],
])
if hasattr(query_or_msg, "edit_message_text"):
await query_or_msg.edit_message_text(text, parse_mode="HTML", reply_markup=keyboard)
else:
await context.bot.send_message(
query_or_msg.chat_id, text, parse_mode="HTML", reply_markup=keyboard
)
async def on_confirm_start(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle confirmation to start generation."""
query = update.callback_query
await query.answer("Starting generation...")
user_id = query.from_user.id
session = state_manager.get_session(user_id)
# Mark user as processing
state_manager.start_task(user_id, f"cover_{int(time.time())}")
processing_msg = await query.edit_message_text(
"Generating your AI cover...\n\n"
"This may take a few minutes depending on the song length. "
"I'll send you the result when it's ready!"
)
self._progress_messages[user_id] = processing_msg.message_id
# Start generation in background
asyncio.create_task(
self._run_generation(user_id, query.effective_chat.id, session, context)
)
return ConversationHandler.END
async def on_settings_from_confirm(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle settings button from confirmation screen."""
query = update.callback_query
await query.answer()
user_id = query.from_user.id
session = state_manager.get_session(user_id)
await self._send_settings_menu(query.effective_chat.id, session, context)
# Stay in CONFIRM state so user can go back
return CONVO_STATES["CONFIRM"]
async def on_cancel(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle cancel button."""
query = update.callback_query
await query.answer("Cancelled")
user_id = query.from_user.id
state_manager.clear_session(user_id)
await query.edit_message_text(
"Cancelled. Send /cover to start again."
)
return ConversationHandler.END
# ============================================================
# Settings Menus
# ============================================================
async def _send_settings_menu(self, chat_id: int, session: UserSession, context: ContextTypes.DEFAULT_TYPE):
"""Send the settings menu."""
pitch_display = f"+{session.pitch_change}" if session.pitch_change > 0 else str(session.pitch_change)
mode_labels = {
"full": "Full Pipeline",
"mdx": "MDX Only",
"rvc": "RVC Only",
}
text = (
"Generation Settings\n\n"
f"Voice Model: {html.escape(session.voice_model or 'Not selected')}\n"
f"Mode: {mode_labels.get(session.inference_mode, session.inference_mode)}\n"
f"Pitch Change: {pitch_display} octaves\n"
f"Pitch All: {'+' if session.pitch_change_all >= 0 else ''}{session.pitch_change_all}\n"
f"F0 Method: {session.f0_method}\n"
f"Index Rate: {session.index_rate}\n"
f"Filter Radius: {session.filter_radius}\n"
f"Volume Envelope: {session.volume_envelope}\n"
f"Protect: {session.protect}\n"
f"Output Format: {session.output_format.upper()}\n"
f"Reverb Size: {session.reverb_size}\n"
f"Reverb Wetness: {session.reverb_wet}\n"
f"Extra Denoise: {'On' if session.extra_denoise else 'Off'}\n"
)
keyboard = InlineKeyboardMarkup([
[
InlineKeyboardButton("Change Pitch", callback_data=CB_PREFIX["SET_PITCH"]),
InlineKeyboardButton("Change Mode", callback_data=CB_PREFIX["SET_MODE"]),
],
[
InlineKeyboardButton("Format: MP3" if session.output_format == "mp3" else "Format: WAV",
callback_data=f"{CB_PREFIX['SET_FORMAT']}{'wav' if session.output_format == 'mp3' else 'mp3'}"),
InlineKeyboardButton("Denoise: " + ("ON" if session.extra_denoise else "OFF"),
callback_data=CB_PREFIX["SET_DENOISE"]),
],
[
InlineKeyboardButton("F0: rmvpe" if session.f0_method == "rmvpe" else f"F0: {session.f0_method}",
callback_data=f"{CB_PREFIX['SET_F0']}{'mangio-crepe' if session.f0_method == 'rmvpe' else 'rmvpe'}"),
InlineKeyboardButton("Reverb", callback_data=CB_PREFIX["SET_REVERB"]),
],
[InlineKeyboardButton("Back", callback_data=CB_PREFIX["BACK_MENU"])],
])
await context.bot.send_message(chat_id, text, parse_mode="HTML", reply_markup=keyboard)
async def on_settings_callback(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle settings button callbacks."""
query = update.callback_query
await query.answer()
user_id = query.from_user.id
session = state_manager.get_session(user_id)
if query.data == CB_PREFIX["SET_PITCH"]:
await query.edit_message_text(
"Send the pitch change value in octaves.\n\n"
"Examples:\n"
" 1 = one octave up (male to female)\n"
" -1 = one octave down (female to male)\n"
" 0 = no change\n"
" 2 = two octaves up\n\n"
"Current: " + (f"+{session.pitch_change}" if session.pitch_change > 0 else str(session.pitch_change)) + " octaves",
parse_mode="HTML",
)
return CONVO_STATES["PITCH_INPUT"]
elif query.data == CB_PREFIX["SET_MODE"]:
keyboard = InlineKeyboardMarkup([
[
InlineKeyboardButton("Full Pipeline", callback_data="mode_full"),
InlineKeyboardButton("RVC Only", callback_data="mode_rvc"),
],
[InlineKeyboardButton("MDX Only", callback_data="mode_mdx")],
])
await query.edit_message_text(
"Select Inference Mode\n\n"
"Full Pipeline: Separate vocals + RVC conversion + mix\n"
"RVC Only: Voice conversion only\n"
"MDX Only: Vocal separation only",
parse_mode="HTML",
reply_markup=keyboard,
)
return CONVO_STATES["MODE_SELECT"]
elif query.data.startswith(CB_PREFIX["SET_FORMAT"]):
new_format = query.data.replace(CB_PREFIX["SET_FORMAT"], "")
session.output_format = new_format
await self._send_settings_menu(query.effective_chat.id, session, context)
elif query.data.startswith(CB_PREFIX["SET_F0"]):
new_f0 = query.data.replace(CB_PREFIX["SET_F0"], "")
session.f0_method = new_f0
await self._send_settings_menu(query.effective_chat.id, session, context)
elif query.data == CB_PREFIX["SET_DENOISE"]:
session.extra_denoise = not session.extra_denoise
await self._send_settings_menu(query.effective_chat.id, session, context)
elif query.data == CB_PREFIX["SET_REVERB"]:
# Cycle through reverb presets: None -> Light -> Medium -> Heavy -> None
if session.reverb_wet < 0.1:
session.reverb_wet, session.reverb_size = 0.1, 0.1
label = "Light"
elif session.reverb_wet < 0.3:
session.reverb_wet, session.reverb_size = 0.3, 0.2
label = "Medium"
elif session.reverb_wet < 0.5:
session.reverb_wet, session.reverb_size = 0.5, 0.4
label = "Heavy"
else:
session.reverb_wet, session.reverb_size = 0.0, 0.0
label = "None"
await self._send_settings_menu(query.effective_chat.id, session, context)
elif query.data == CB_PREFIX["BACK_MENU"]:
await query.edit_message_text("Settings saved. Send /cover to generate a cover.")
return ConversationHandler.END
async def on_pitch_input(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle pitch value input."""
user_id = update.effective_user.id
session = state_manager.get_session(user_id)
try:
pitch_value = int(update.message.text.strip())
except ValueError:
await update.message.reply_text(
"Invalid value. Please enter an integer (e.g., 1, -1, 0, 2)."
)
return CONVO_STATES["PITCH_INPUT"]
if pitch_value < -12 or pitch_value > 12:
await update.message.reply_text(
"Pitch must be between -12 and +12 octaves."
)
return CONVO_STATES["PITCH_INPUT"]
session.pitch_change = pitch_value
await update.message.reply_text(
f"Pitch set to {pitch_value} octaves.",
parse_mode="HTML",
)
await self._send_settings_menu(update.effective_chat.id, session, context)
return ConversationHandler.END
async def on_mode_select(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle inference mode selection."""
query = update.callback_query
await query.answer()
user_id = query.from_user.id
session = state_manager.get_session(user_id)
if query.data == "mode_full":
session.inference_mode = "full"
elif query.data == "mode_rvc":
session.inference_mode = "rvc"
elif query.data == "mode_mdx":
session.inference_mode = "mdx"
await self._send_settings_menu(query.effective_chat.id, session, context)
return ConversationHandler.END
# ============================================================
# Generation Runner
# ============================================================
async def _run_generation(
self,
user_id: int,
chat_id: int,
session: UserSession,
context: ContextTypes.DEFAULT_TYPE,
):
"""Run the cover generation pipeline and send the result."""
progress_msg_id = self._progress_messages.get(user_id)
async def progress_callback(percent: float, desc: str = ""):
"""Update progress message."""
if progress_msg_id and percent % 10 < 5: # Update every ~10%
try:
bar_len = 20
filled = int(bar_len * percent / 100)
bar = "█" * filled + "░" * (bar_len - filled)
text = (
f"Generating AI Cover...\n\n"
f"[{bar}] {percent:.0f}%\n"
f"{html.escape(desc or 'Processing...')}"
)
await context.bot.edit_message_text(
text,
chat_id=chat_id,
message_id=progress_msg_id,
parse_mode="HTML",
)
except Exception:
pass # Ignore edit failures (rate limits, etc.)
try:
result_path = await self.pipeline.process_cover(
song_input=session.song_input,
voice_model=session.voice_model,
pitch_change=session.pitch_change,
keep_files=session.keep_files,
main_gain=session.main_gain,
backup_gain=session.backup_gain,
inst_gain=session.inst_gain,
index_rate=session.index_rate,
filter_radius=session.filter_radius,
volume_envelope=session.volume_envelope,
f0_method=session.f0_method,
hop_length=session.hop_length,
protect=session.protect,
pitch_change_all=session.pitch_change_all,
reverb_rm_size=session.reverb_size,
reverb_wet=session.reverb_wet,
reverb_dry=session.reverb_dry,
reverb_damping=session.reverb_damping,
output_format=session.output_format,
extra_denoise=session.extra_denoise,
inference_mode=session.inference_mode,
progress_callback=progress_callback,
)
# Send the result
if result_path and os.path.exists(result_path):
file_size = os.path.getsize(result_path) / (1024 * 1024)
# Check if file is within Telegram's 50MB limit
if file_size > 50:
# File too large for Telegram - send note
await context.bot.send_message(
chat_id,
f"Cover generated but file is too large for Telegram ({file_size:.1f}MB).\n"
f"File saved at: {html.escape(result_path)}",
parse_mode="HTML",
)
else:
# Get the base filename
filename = os.path.basename(result_path)
caption = (
f"Your AI cover is ready!\n\n"
f"Voice: {html.escape(session.voice_model or 'Unknown')}\n"
f"Pitch: {session.pitch_change}\n"
f"Format: {session.output_format.upper()}\n"
)
await context.bot.send_audio(
chat_id=chat_id,
audio=open(result_path, "rb"),
caption=caption,
parse_mode="HTML",
)
else:
await context.bot.send_message(
chat_id,
"Generation completed but no output file was produced. "
"This might indicate an issue with the input or model."
)
except Exception as e:
error_text = f"Generation failed: {html.escape(str(e))}"
logger.error(f"Generation failed for user {user_id}: {e}")
logger.debug(traceback.format_exc())
if progress_msg_id:
try:
await context.bot.edit_message_text(
f"Error\n\n{error_text}\n\n"
"Try again with /cover or check your settings.",
chat_id=chat_id,
message_id=progress_msg_id,
parse_mode="HTML",
)
except Exception:
await context.bot.send_message(chat_id, error_text, parse_mode="HTML")
else:
await context.bot.send_message(chat_id, error_text, parse_mode="HTML")
finally:
state_manager.end_task(user_id)
self._progress_messages.pop(user_id, None)
# Cleanup temp files
if session.song_input and os.path.exists(session.song_input) and session.input_type == InputType.FILE:
try:
os.remove(session.song_input)
except Exception:
pass
# ============================================================
# Utility Methods
# ============================================================
@staticmethod
def _is_valid_youtube_url(url: str) -> bool:
"""Validate a YouTube URL."""
from urllib.parse import urlparse
parsed = urlparse(url)
valid_hosts = {
"www.youtube.com",
"youtube.com",
"m.youtube.com",
"youtu.be",
"music.youtube.com",
}
if parsed.hostname in valid_hosts:
if parsed.hostname == "youtu.be" and parsed.path[1:]:
return True
if "v=" in parsed.query:
return True
if parsed.path.startswith("/watch"):
return True
return False
async def _log_message(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Log all incoming messages (debug mode only)."""
logger.debug(f"Message from {update.effective_user.id}: {update.message.text or '[media]'}")
# ============================================================
# Error Handler
# ============================================================
async def error_handler(self, update: object, context: ContextTypes.DEFAULT_TYPE):
"""Handle all uncaught exceptions."""
error = context.error
tb = "".join(traceback.format_exception(type(error), error, error.__traceback__))
logger.error(f"Unhandled exception: {error}\n{tb}")
if update and hasattr(update, "effective_message") and update.effective_message:
try:
await update.effective_message.reply_text(
f"An unexpected error occurred: {html.escape(str(error))}\n"
"Please try again or contact the admin."
)
except Exception:
pass
# ============================================================
# Run
# ============================================================
def run(self):
"""Start the bot (blocking)."""
if not self.application:
raise RuntimeError("Bot not initialized. Call initialize() first.")
logger.info(f"Starting {self.config.bot_name}...")
# Setup signal handlers for graceful shutdown
def signal_handler(sig, frame):
logger.info("Received shutdown signal")
if self.pipeline:
self.pipeline.shutdown()
if self.application:
asyncio.create_task(self.application.shutdown())
sys.exit(0)
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
# Start polling
self.application.run_polling(
drop_pending_updates=True,
allowed_updates=Update.ALL_TYPES,
)
async def run_async(self):
"""Start the bot asynchronously."""
if not self.application:
raise RuntimeError("Bot not initialized. Call initialize() first.")
await self.application.initialize()
await self.application.start()
await self.application.updater.start_polling(
drop_pending_updates=True,
allowed_updates=Update.ALL_TYPES,
)
logger.info(f"{self.config.bot_name} is running")
# ============================================================
# Entry Point
# ============================================================
def main():
"""Main entry point."""
bot = AIGCTelegramBot()
async def startup():
await bot.initialize()
await bot.run_async()
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
loop.run_until_complete(startup())
loop.run_forever()
except (KeyboardInterrupt, SystemExit):
logger.info("Shutting down...")
finally:
if bot.pipeline:
bot.pipeline.shutdown()
loop.close()
if __name__ == "__main__":
main()