Update cover.py

cover.py

@@ -1,22 +1,15 @@
 import os
 from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
-from telegram.ext import (
-    Application,
-    CommandHandler,
-    CallbackQueryHandler,
-    MessageHandler,
-    filters,
-    ContextTypes,
-)
-from main import song_cover_pipeline  # Your song generation pipeline
-from webui import download_online_model  # Your model download function
+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
 
 # Define paths
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 output_dir = os.path.join(BASE_DIR, 'song_output')
 os.makedirs(output_dir, exist_ok=True)
 
-# Start command shows the main menu
+# /start command handler: show main menu
 async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
     keyboard = [
         [InlineKeyboardButton("Generate Song", callback_data='generate')],
@@ -24,215 +17,175 @@ async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
         [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)
+    await update.message.reply_text('Welcome to AICoverGen! Choose an option below:', reply_markup=reply_markup)
 
-# Callback query handler for all button presses
-async def handle_callback(update: Update, context: ContextTypes.DEFAULT_TYPE):
+# CallbackQuery handler for inline buttons
+async def button(update: Update, context: ContextTypes.DEFAULT_TYPE):
     query = update.callback_query
     await query.answer()
-    data = query.data
 
     # Main menu selections
-    if data == 'generate':
-        context.user_data.clear()
+    if query.data == 'generate':
         context.user_data['mode'] = 'generate'
-        context.user_data['keep_files'] = False  # default value
-        # Ask which input method to use
+        # Prepare an inline keyboard for song generation options
         keyboard = [
-            [InlineKeyboardButton("Input YouTube Link", callback_data='gen_input_link')],
-            [InlineKeyboardButton("Upload Audio", callback_data='gen_input_audio')]
+            [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)
-        await query.edit_message_text("Choose input method for your song:", reply_markup=reply_markup)
-        return
+        # 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)
 
-    if data == 'download_model':
-        context.user_data.clear()
+    elif query.data == 'download_model':
         context.user_data['mode'] = 'download_model'
+        # Prepare an inline keyboard for model download options
         keyboard = [
-            [InlineKeyboardButton("Set URL", callback_data='dl_set_url')],
-            [InlineKeyboardButton("Set Model Name", callback_data='dl_set_model_name')],
-            [InlineKeyboardButton("Submit Download", callback_data='dl_submit')]
+            [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)
-        await query.edit_message_text("Download Model - set options using the buttons below:", reply_markup=reply_markup)
-        return
+        # 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)
 
-    if data == 'help':
+    elif query.data == 'help':
         help_text = (
             "To generate a song:\n"
             "1. Click 'Generate Song'.\n"
-            "2. Choose input method: YouTube Link or Upload Audio.\n"
-            "3. Provide the song input (link or audio file).\n"
-            "4. Enter the model name.\n"
-            "5. Choose pitch: Female (1) or Male (-1).\n"
-            "6. Toggle Keep Files if desired.\n"
-            "7. Submit to generate the song.\n\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. Set the URL and Model Name using the buttons.\n"
-            "3. Submit to download the model."
+            "2. Use the buttons to set Model URL and Model Name.\n"
+            "3. Press 'Submit' to download the model."
         )
-        await query.edit_message_text(help_text)
-        return
+        await query.edit_message_text(text=help_text)
 
-    # --- GENERATE MODE CALLBACKS ---
-    if context.user_data.get('mode') == 'generate':
-        if data == 'gen_input_link':
-            context.user_data['song_input_method'] = 'link'
-            context.user_data['generate_step'] = 'song_input'
-            await query.edit_message_text("Please send the YouTube link for the song.")
-            return
-        if data == 'gen_input_audio':
-            context.user_data['song_input_method'] = 'audio'
-            context.user_data['generate_step'] = 'audio_input'
-            await query.edit_message_text("Please upload the audio file for the song.")
-            return
-        if data.startswith('pitch_'):
-            # Expecting callback data like 'pitch_1' or 'pitch_-1'
-            pitch_val = int(data.split('_')[1])
-            context.user_data['pitch'] = pitch_val
-            # After setting pitch, offer toggle for keep_files and a submit button
-            keyboard = [
-                [InlineKeyboardButton(f"Keep Files: {'On' if context.user_data.get('keep_files') else 'Off'}", callback_data='toggle_keep_files')],
-                [InlineKeyboardButton("Submit", callback_data='submit_generate')]
-            ]
-            reply_markup = InlineKeyboardMarkup(keyboard)
-            await query.edit_message_text("Pitch set. You can toggle Keep Files or submit.", reply_markup=reply_markup)
-            return
-        if data == 'toggle_keep_files':
-            # Toggle the boolean value
-            current = context.user_data.get('keep_files', False)
-            context.user_data['keep_files'] = not current
-            keyboard = [
-                [InlineKeyboardButton(f"Keep Files: {'On' if context.user_data['keep_files'] else 'Off'}", callback_data='toggle_keep_files')],
-                [InlineKeyboardButton("Submit", callback_data='submit_generate')]
-            ]
-            reply_markup = InlineKeyboardMarkup(keyboard)
-            await query.edit_message_text("Keep Files toggled. You can toggle again or submit.", reply_markup=reply_markup)
-            return
-        if data == 'submit_generate':
-            # Make sure all required parameters are provided
-            if context.user_data.get('song_input_method') == 'link' and not context.user_data.get('song_link'):
-                await query.edit_message_text("You haven't provided the YouTube link yet.")
-                return
-            if context.user_data.get('song_input_method') == 'audio' and not context.user_data.get('audio_file'):
-                await query.edit_message_text("You haven't uploaded the audio file yet.")
-                return
-            if not context.user_data.get('model_name'):
-                await query.edit_message_text("You haven't provided the model name yet.")
-                return
-            if 'pitch' not in context.user_data:
-                await query.edit_message_text("You haven't selected the pitch yet.")
-                return
+    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)
 
-            await query.edit_message_text("Generating your song, please wait...")
-            song_input = context.user_data.get('song_link') if context.user_data.get('song_input_method') == 'link' else context.user_data.get('audio_file')
-            model_name = context.user_data.get('model_name')
-            pitch = context.user_data.get('pitch')
-            keep_files = context.user_data.get('keep_files')
-            is_webui = False
-            # Call the song generation pipeline (adjust as needed for your implementation)
-            song_output = song_cover_pipeline(song_input, model_name, pitch, keep_files, is_webui)
-            if os.path.exists(song_output):
-                await context.bot.send_audio(chat_id=update.effective_chat.id, audio=open(song_output, 'rb'))
-                os.remove(song_output)
-                await context.bot.send_message(chat_id=update.effective_chat.id, text="Song generated successfully!")
-            else:
-                await context.bot.send_message(chat_id=update.effective_chat.id, text="An error occurred during song generation.")
-            context.user_data.clear()
+    # 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.")
 
-    # --- DOWNLOAD MODE CALLBACKS ---
-    if context.user_data.get('mode') == 'download_model':
-        if data == 'dl_set_url':
-            context.user_data['download_step'] = 'url'
-            await query.edit_message_text("Please send the download URL for the model.")
-            return
-        if data == 'dl_set_model_name':
-            context.user_data['download_step'] = 'model_name'
-            await query.edit_message_text("Please enter the custom model name.")
-            return
-        if data == 'dl_submit':
-            if not context.user_data.get('download_url') or not context.user_data.get('download_model_name'):
-                await query.edit_message_text("Both URL and model name must be set before submitting.")
-                return
-            model_url = context.user_data.get('download_url')
-            model_name = context.user_data.get('download_model_name')
-            try:
-                download_online_model(model_url, model_name)
-                await query.edit_message_text(f"Model '{model_name}' downloaded successfully from {model_url}!")
-            except Exception as e:
-                await query.edit_message_text(f"Failed to download model. Error: {str(e)}")
-            context.user_data.clear()
+    # 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)
 
-# Message handler for receiving text or file inputs
+    # 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
 async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
     mode = context.user_data.get('mode')
-    if mode == 'generate':
-        step = context.user_data.get('generate_step')
-        if step == 'song_input':
-            # Expecting a YouTube link
-            context.user_data['song_link'] = update.message.text.strip()
-            context.user_data['generate_step'] = 'model_name'
-            await update.message.reply_text("YouTube link set. Please enter the model name.")
-            return
-        if step == 'audio_input':
-            # Check for an audio file upload
-            if update.message.audio:
-                audio_file = await update.message.audio.get_file()
-                file_path = os.path.join(output_dir, f"{update.message.audio.file_id}.mp3")
-                await audio_file.download_to_drive(file_path)
-                context.user_data['audio_file'] = file_path
-                context.user_data['generate_step'] = 'model_name'
-                await update.message.reply_text("Audio file received. Please enter the model name.")
-            else:
-                await update.message.reply_text("Please upload a valid audio file.")
-            return
-        if step == 'model_name':
-            context.user_data['model_name'] = update.message.text.strip()
-            # Now ask for pitch via inline keyboard
-            keyboard = [
-                [InlineKeyboardButton("Female (1)", callback_data='pitch_1')],
-                [InlineKeyboardButton("Male (-1)", callback_data='pitch_-1')]
-            ]
-            reply_markup = InlineKeyboardMarkup(keyboard)
-            await update.message.reply_text("Please choose the pitch:", reply_markup=reply_markup)
-            context.user_data['generate_step'] = None  # move on to waiting for pitch
-            return
-        await update.message.reply_text("Please use the provided buttons to navigate the process.")
-    
-    elif mode == 'download_model':
-        download_step = context.user_data.get('download_step')
-        if download_step == 'url':
-            text = update.message.text.strip()
-            if not text.startswith("http"):
-                await update.message.reply_text("Please provide a valid URL.")
-                return
-            context.user_data['download_url'] = text
-            context.user_data['download_step'] = None
-            await update.message.reply_text("Download URL set.")
-            return
-        if download_step == 'model_name':
-            context.user_data['download_model_name'] = update.message.text.strip()
-            context.user_data['download_step'] = None
-            await update.message.reply_text("Model name set.")
-            return
-        await update.message.reply_text("Please use the provided buttons for setting options.")
+    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 main buttons.")
+        await update.message.reply_text("Please choose an option first by clicking one of the buttons.")
 
-# Main function to run the bot
 def main():
-    bot_token = "7722898432"+":AAEfj9"+"s6u"+"bY10"+"7S"+"Wi"+"F6Uy"+"1IKJFF"+"msiqY_BA"
+    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(handle_callback))
-    application.add_handler(MessageHandler(filters.ALL, handle_message))
+    application.add_handler(CallbackQueryHandler(button))
+    application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))
     application.run_polling()
 
 if __name__ == '__main__':