-
Notifications
You must be signed in to change notification settings - Fork 3
/
main.py
165 lines (136 loc) · 7.1 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
# © [2024] Malith-Rukshan. All rights reserved.
# Repository: https://github.com/Malith-Rukshan/Suno-AI-BOT
import asyncio
import logging
import os
from telegram.constants import ParseMode
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import Application, CommandHandler, ContextTypes, MessageHandler, CallbackQueryHandler, filters
import suno
# Configure logging
logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO)
logger = logging.getLogger(__name__)
logging.getLogger("httpx").setLevel(logging.WARNING)
# Disable SunoAI Logs
# logging.getLogger("SunoAI").setLevel(logging.WARNING)
# Initialize Suno AI Library
SUNO_COOKIE = os.getenv("SUNO_COOKIE")
client = suno.Suno(cookie=SUNO_COOKIE)
# Store user session data
chat_states = {}
# Keyboard options for user selection
def get_base_keyboard():
return InlineKeyboardMarkup([
[InlineKeyboardButton("🎨 Custom", callback_data="custom")],
[InlineKeyboardButton("🏞️ Default", callback_data="default")]
])
# Welcome message with Markdown
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
welcome_message = (
"👋 Hello! Welcome to the *Suno AI Music Generator Bot*! 🎶\n\n"
"👉 Use /generate to start creating your unique music track. 🚀\n"
"👉 Use /credits to check your credits balance.\n\n"
"📥 This bot utilizes the [SunoAI API](https://github.com/Malith-Rukshan/Suno-API). You can also deploy your own version of this bot! For more details, visit our [GitHub repo](https://github.com/Malith-Rukshan/Suno-AI-BOT)."
)
await update.message.reply_markdown(welcome_message,disable_web_page_preview=False)
# Handler for the get credits
async def credits_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
credit_info_message = (
"*💰Credits Stat*\n\n"
"ᗚ Available : {}\n"
"ᗚ Usage : {}"
)
try:
credits = await asyncio.to_thread(client.get_credits)
except Exception as e:
return await update.message.reply_text(f"⁉️ Failed to get credits info: {e}")
await update.message.reply_text(credit_info_message.format(credits.credits_left,credits.monthly_usage),parse_mode=ParseMode.MARKDOWN)
# Handler for the generate command
async def generate(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
await update.message.reply_text('Select mode: custom or not. 🤔', reply_markup=get_base_keyboard())
chat_states[update.effective_chat.id] = {}
# Command to cancel and clear state
async def cancel(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
chat_id = update.effective_chat.id
if chat_id in chat_states:
chat_states.pop(chat_id, None)
await update.message.reply_text('Generation canceled. 🚫 You can start again with /generate.')
# Handler for button presses
async def button_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
query = update.callback_query
await query.answer()
chat_id = int(update.effective_chat.id)
chat_states[chat_id]['mode'] = query.data
if query.data == "custom":
await query.message.reply_text("🎤 Send lyrics first.")
else:
await query.message.reply_text("🎤 Send song description.")
return await context.application.bot.delete_message(chat_id=query.message.chat.id,message_id=query.message.message_id)
async def onMessage(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
chat_id = int(update.effective_chat.id)
# Collects lyrics from the user
if chat_id in chat_states and 'mode' in chat_states[chat_id]:
if not 'lyrics' in chat_states[chat_id]:
chat_states[chat_id]['lyrics'] = update.message.text
if chat_states[chat_id].get('mode') == 'custom':
if not (chat_id in chat_states and 'tags' in chat_states[chat_id] and "Wait-for-tags" == chat_states[chat_id]['tags']):
chat_states[chat_id]['tags'] = "Wait-for-tags"
return await update.message.reply_text("🏷️ Now send tags.\n\nExample : Classical")
# Collects tags (if custom) / generates music
if chat_id in chat_states and 'lyrics' in chat_states[chat_id]:
if chat_states[chat_id].get('mode') == 'custom':
# Custom music generation logic
chat_states[chat_id]['tags'] = update.message.text
await update.message.reply_text("Generating your music... please wait. ⏳")
try:
prompt = f"{chat_states[chat_id]['lyrics']}"
tags = f"{chat_states[chat_id]['tags']}"
# Generate Custom Music
songs = await asyncio.to_thread(
client.generate,
prompt=prompt,
tags=tags,
is_custom=True,
wait_audio=True)
for song in songs:
file_path = await asyncio.to_thread(client.download,song=song)
await context.bot.send_audio(chat_id=chat_id, audio=open(file_path, 'rb'), thumbnail=open("thumb.jpg", 'rb'))
os.remove(file_path)
if chat_id in chat_states:
chat_states.pop(chat_id, None)
except Exception as e:
await update.message.reply_text(f"⁉️ Failed to generate music: {e}")
else:
# Default music generation logic
await update.message.reply_text("Generating your music... please wait. ⏳")
try:
prompt = f"{chat_states[chat_id]['lyrics']}"
# Generate Music by Description
songs = await asyncio.to_thread(
client.generate,
prompt=prompt,
is_custom=False,
wait_audio=True)
for song in songs:
file_path = await asyncio.to_thread(client.download,song=song)
await context.bot.send_audio(chat_id=chat_id, audio=open(file_path, 'rb'), thumbnail=open("thumb.jpg", 'rb'))
os.remove(file_path)
if chat_id in chat_states:
chat_states.pop(chat_id, None)
except Exception as e:
await update.message.reply_text(f"⁉️ Failed to generate music: {e}")
if chat_id in chat_states:
chat_states.pop(chat_id, None)
def main():
BOT_TOKEN = os.getenv("BOT_TOKEN")
application = Application.builder().token(BOT_TOKEN).build()
application.add_handler(CommandHandler("start", start))
application.add_handler(CommandHandler("generate", generate))
application.add_handler(CommandHandler("cancel", cancel))
application.add_handler(CallbackQueryHandler(button_handler))
application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, onMessage))
application.add_handler(CommandHandler("credits", credits_command))
# Run the bot until the user presses Ctrl-C
application.run_polling(allowed_updates=Update.ALL_TYPES)
if __name__ == "__main__":
main()