Update cover.py

cover.py

@@ -1,192 +1,161 @@
+Change the bot to multi button option bot, like:
+Inference Options:
+
+1. song_link or audio input
+
+
+2. model_name,
+
+
+3. pitch,
+
+
+4. keep_files
+Download options:
+
+
+5. Url
+
+
+6. model name
+
+
+
 import os
 from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
-from telegram.ext import Application, CommandHandler, CallbackQueryHandler, MessageHandler, filters, ContextTypes
-from main import song_cover_pipeline  # Import your song generation pipeline
-from webui import download_online_model  # Import the model download function
+from telegram.ext import Application, CommandHandler, CallbackQueryHandler, MessageHandler, filters
+from telegram.ext import ContextTypes
+from main import song_cover_pipeline  # Keeping this import from your original main.py
+from webui import download_online_model  # Import the download function
+
+Define paths
 
-# Define paths
-BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+BASE_DIR = os.path.dirname(os.path.abspath(file))
 output_dir = os.path.join(BASE_DIR, 'song_output')
+
+Ensure the output directory exists
+
 os.makedirs(output_dir, exist_ok=True)
 
-# /start command handler: show main menu
+Start command
+
 async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
-    keyboard = [
-        [InlineKeyboardButton("Generate Song", callback_data='generate')],
-        [InlineKeyboardButton("Download Model", callback_data='download_model')],
-        [InlineKeyboardButton("Help", callback_data='help')]
-    ]
-    reply_markup = InlineKeyboardMarkup(keyboard)
-    await update.message.reply_text('Welcome to AICoverGen! Choose an option below:', reply_markup=reply_markup)
-
-# CallbackQuery handler for inline buttons
+keyboard = [
+[InlineKeyboardButton("Generate Song", callback_data='generate')],
+[InlineKeyboardButton("Download Model", callback_data='download_model')],
+[InlineKeyboardButton("Help", callback_data='help')]
+]
+reply_markup = InlineKeyboardMarkup(keyboard)
+await update.message.reply_text('Welcome to AICoverGen! Choose an option below:', reply_markup=reply_markup)
+
+Button handler
+
 async def button(update: Update, context: ContextTypes.DEFAULT_TYPE):
-    query = update.callback_query
-    await query.answer()
-
-    # Main menu selections
-    if query.data == 'generate':
-        context.user_data['mode'] = 'generate'
-        # Prepare an inline keyboard for song generation options
-        keyboard = [
-            [InlineKeyboardButton("Set Song Link/Audio", callback_data='set_input')],
-            [InlineKeyboardButton("Set Model Name", callback_data='set_model')],
-            [InlineKeyboardButton("Set Pitch", callback_data='set_pitch')],
-            [InlineKeyboardButton("Toggle Keep Files", callback_data='toggle_keep')],
-            [InlineKeyboardButton("Submit", callback_data='submit_generation')],
-            [InlineKeyboardButton("Back", callback_data='back')]
-        ]
-        reply_markup = InlineKeyboardMarkup(keyboard)
-        # Initialize storage for generation parameters
-        context.user_data['song_data'] = {
-            'input': None,
-            'model': None,
-            'pitch': None,
-            'keep_files': False
-        }
-        await query.edit_message_text(text="Please set your song generation options:", reply_markup=reply_markup)
-
-    elif query.data == 'download_model':
-        context.user_data['mode'] = 'download_model'
-        # Prepare an inline keyboard for model download options
-        keyboard = [
-            [InlineKeyboardButton("Set Model URL", callback_data='set_model_url')],
-            [InlineKeyboardButton("Set Model Name", callback_data='set_download_model')],
-            [InlineKeyboardButton("Submit", callback_data='submit_download')],
-            [InlineKeyboardButton("Back", callback_data='back')]
-        ]
-        reply_markup = InlineKeyboardMarkup(keyboard)
-        # Initialize storage for download parameters
-        context.user_data['download_data'] = {
-            'url': None,
-            'model': None
-        }
-        await query.edit_message_text(text="Please set your model download options:", reply_markup=reply_markup)
-
-    elif query.data == 'help':
-        help_text = (
-            "To generate a song:\n"
-            "1. Click 'Generate Song'.\n"
-            "2. Use the buttons to set each option:\n"
-            "   • Song Link/Audio input\n"
-            "   • Model Name\n"
-            "   • Pitch (e.g., 1 for female, -1 for male)\n"
-            "   • Toggle Keep Files if needed\n"
-            "3. Press 'Submit' to generate the song.\n\n"
-            "To download a model:\n"
-            "1. Click 'Download Model'.\n"
-            "2. Use the buttons to set Model URL and Model Name.\n"
-            "3. Press 'Submit' to download the model."
-        )
-        await query.edit_message_text(text=help_text)
-
-    elif query.data == 'back':
-        # Return to the main menu
-        keyboard = [
-            [InlineKeyboardButton("Generate Song", callback_data='generate')],
-            [InlineKeyboardButton("Download Model", callback_data='download_model')],
-            [InlineKeyboardButton("Help", callback_data='help')]
-        ]
-        reply_markup = InlineKeyboardMarkup(keyboard)
-        await query.edit_message_text(text="Welcome back to AICoverGen! Choose an option:", reply_markup=reply_markup)
-
-    # Submission for song generation
-    elif query.data == 'submit_generation':
-        song_data = context.user_data.get('song_data', {})
-        input_val = song_data.get('input')
-        model_name = song_data.get('model')
-        pitch = song_data.get('pitch')
-        keep_files = song_data.get('keep_files', False)
-        if not input_val or not model_name or pitch is None:
-            await query.edit_message_text(text="Missing parameters! Ensure you set Song Link/Audio, Model Name, and Pitch.")
-            return
-        song_output = song_cover_pipeline(input_val, model_name, pitch, keep_files, is_webui=False)
-        if os.path.exists(song_output):
-            await query.edit_message_text(text="Song generated successfully. Sending audio...")
-            await context.bot.send_audio(chat_id=update.effective_chat.id, audio=open(song_output, 'rb'))
-            os.remove(song_output)
-        else:
-            await query.edit_message_text(text="An error occurred while generating the song.")
-
-    # Submission for model download
-    elif query.data == 'submit_download':
-        download_data = context.user_data.get('download_data', {})
-        model_url = download_data.get('url')
-        model_name = download_data.get('model')
-        if not model_url or not model_name:
-            await query.edit_message_text(text="Missing parameters! Ensure you set both Model URL and Model Name.")
-            return
-        try:
-            download_online_model(model_url, model_name)
-            await query.edit_message_text(text=f"Model '{model_name}' downloaded successfully from {model_url}!")
-        except Exception as e:
-            await query.edit_message_text(text=f"Failed to download the model. Error: {str(e)}")
-
-    # Handling parameter setting buttons for generation
-    elif query.data == 'set_input':
-        await query.edit_message_text(text="Please send the Song Link or upload an audio file.")
-        context.user_data['awaiting_input'] = 'input'
-    elif query.data == 'set_model':
-        await query.edit_message_text(text="Please send the Model Name.")
-        context.user_data['awaiting_input'] = 'model'
-    elif query.data == 'set_pitch':
-        await query.edit_message_text(text="Please send the Pitch (e.g., 1 for female, -1 for male).")
-        context.user_data['awaiting_input'] = 'pitch'
-    elif query.data == 'toggle_keep':
-        song_data = context.user_data.get('song_data', {})
-        current = song_data.get('keep_files', False)
-        song_data['keep_files'] = not current
-        context.user_data['song_data'] = song_data
-        await query.answer(text=f"Keep Files set to {song_data['keep_files']}", show_alert=True)
-
-    # Handling parameter setting buttons for download
-    elif query.data == 'set_model_url':
-        await query.edit_message_text(text="Please send the Model URL.")
-        context.user_data['awaiting_input'] = 'model_url'
-    elif query.data == 'set_download_model':
-        await query.edit_message_text(text="Please send the Model Name for download.")
-        context.user_data['awaiting_input'] = 'download_model'
-
-# Message handler to capture text input after a parameter-setting prompt
+query = update.callback_query
+await query.answer()
+
+# Determine which option was selected  
+if query.data == 'generate':  
+    # Store the user state as 'generate'  
+    context.user_data['mode'] = 'generate'  
+    await query.edit_message_text(text="Please send the model name, YouTube link, and pitch (e.g., '<model_name> <link> <pitch>')\nNote: pitch 1 for female and pitch -1 for male.")  
+  
+elif query.data == 'download_model':  
+    # Store the user state as 'download_model'  
+    context.user_data['mode'] = 'download_model'  
+    await query.edit_message_text(text="Please send the model name and URL in the format '<model_name> <url>'.")  
+
+elif query.data == 'help':  
+    help_text = (  
+        "To generate a song, follow these steps:\n"  
+        "1. Click 'Generate Song'.\n"  
+        "2. Send a message in the format '<model_name> <link> <pitch>' (e.g., 'model1 https://youtube.com/abc 2').\n"  
+        "3. Wait for the bot to process and return the generated song.\n"  
+        "Pitch: Use 1 for female voice, -1 for male voice.\n\n"  
+        "To download a model:\n"  
+        "1. Click 'Download Model'.\n"  
+        "2. Send a message in the format '<model_name> <url>' to specify the model name and the download URL."  
+    )  
+    await query.edit_message_text(text=help_text)
+
+Message handler
+
 async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
-    mode = context.user_data.get('mode')
-    if 'awaiting_input' in context.user_data:
-        key = context.user_data.pop('awaiting_input')
-        if mode == 'generate':
-            song_data = context.user_data.get('song_data', {})
-            if key == 'input':
-                # Here you could also add logic to handle audio files if needed
-                song_data['input'] = update.message.text
-            elif key == 'model':
-                song_data['model'] = update.message.text
-            elif key == 'pitch':
-                try:
-                    song_data['pitch'] = int(update.message.text)
-                except ValueError:
-                    await update.message.reply_text("Invalid pitch. Please send a valid integer (e.g., 1 or -1).")
-                    return
-            context.user_data['song_data'] = song_data
-            await update.message.reply_text(f"Set {key} to {update.message.text}.")
-        elif mode == 'download_model':
-            download_data = context.user_data.get('download_data', {})
-            if key == 'model_url':
-                download_data['url'] = update.message.text
-            elif key == 'download_model':
-                download_data['model'] = update.message.text
-            context.user_data['download_data'] = download_data
-            await update.message.reply_text(f"Set {key} to {update.message.text}.")
-    else:
-        await update.message.reply_text("Please choose an option first by clicking one of the buttons.")
+# Check which mode the user is in based on previous button choice
+mode = context.user_data.get('mode')
+
+if mode == 'generate':  
+    # Process song generation input  
+    await generate_song(update, context)  
+elif mode == 'download_model':  
+    # Process model download input  
+    await download_model(update, context)  
+else:  
+    # If no mode is selected, ask the user to choose an option  
+    await update.message.reply_text("Please choose an option first by clicking 'Generate Song' or 'Download Model'.")
+
+Generate song handler
+
+async def generate_song(update: Update, context: ContextTypes.DEFAULT_TYPE):
+song_input = update.message.text
+try:
+model_name, song_link, pitch_str = song_input.split()
+pitch = int(pitch_str)
+except ValueError:
+await update.message.reply_text(f"Please send a valid input in the format '<model_name> <link> <pitch>' (e.g., 'model1 https://youtube.com/abc 2').")
+return
+
+keep_files = False  
+is_webui = False  
+
+song_output = song_cover_pipeline(song_link, model_name, pitch, keep_files, is_webui)  
+  
+if os.path.exists(song_output):  
+    await update.message.reply_audio(audio=open(song_output, 'rb'))  
+    os.remove(song_output)  
+else:  
+    await update.message.reply_text(f"An error occurred while generating the song.")
+
+Download model handler with custom name
+
+async def download_model(update: Update, context: ContextTypes.DEFAULT_TYPE):
+model_input = update.message.text
+try:
+# Split the input into model name and URL
+model_name, model_url = model_input.split()
+except ValueError:
+await update.message.reply_text(f"Please send a valid input in the format '<model_name> <url>' (e.g., 'model1 https://model.com/abc').")
+return
+
+if not model_url.startswith("http"):  
+    await update.message.reply_text("Please send a valid URL.")  
+    return  
+
+try:  
+    # Call the function to download the model with a custom name  
+    download_online_model(model_url, model_name)  
+    await update.message.reply_text(f"Model '{model_name}' downloaded successfully from {model_url}!")  
+except Exception as e:  
+    await update.message.reply_text(f"Failed to download the model. Error: {str(e)}")
+
+Main function to run the bot
 
 def main():
-    bot_token = "7722898432:"+"AAEfj9s6ubY107SWiF6Uy1IKJFFmsiqY_BA"
-    if not bot_token:
-        raise ValueError("Bot token not found. Set the TELEGRAM_BOT_TOKEN environment variable.")
-    application = Application.builder().token(bot_token).build()
-    application.add_handler(CommandHandler("start", start))
-    application.add_handler(CallbackQueryHandler(button))
-    application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))
-    application.run_polling()
-
-if __name__ == '__main__':
-    main()
+bot_token = "7722898432:AAEfj9s6ubY107SWiF6Uy1IKJFFmsiqY_BA"
+
+if not bot_token:  
+    raise ValueError("Bot token not found. Set the TELEGRAM_BOT_TOKEN environment variable.")  
+
+application = Application.builder().token(bot_token).build()  
+
+# Handlers  
+application.add_handler(CommandHandler("start", start))  
+application.add_handler(CallbackQueryHandler(button))  
+application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))  # Unified message handler  
+
+# Run the bot  
+application.run_polling()
+
+if name == 'main':
+main()
+