File size: 3,358 Bytes
6df0c1f 8032726 6df0c1f 33b0ce2 6df0c1f 6ee42a9 6df0c1f 6ee42a9 6df0c1f 0b35ab2 |
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 |
import asyncio
import os
import threading
import random
from threading import Event
from typing import Optional
import discord
import gradio as gr
from discord import Permissions
from discord.ext import commands
from discord.utils import oauth_url
import gradio_client as grc
from gradio_client.utils import QueueError
event = Event()
DISCORD_TOKEN = os.getenv("DISCORD_TOKEN")
async def wait(job):
while not job.done():
await asyncio.sleep(0.2)
def get_client(session: Optional[str] = None) -> grc.Client:
client = grc.Client("huggingface-projects/transformers-musicgen", hf_token=os.getenv("HF_TOKEN"))
if session:
client.session_hash = session
return client
intents = discord.Intents.default()
intents.message_content = True
bot = commands.Bot(command_prefix="/", intents=intents)
@bot.event
async def on_ready():
print(f"Logged in as {bot.user} (ID: {bot.user.id})")
synced = await bot.tree.sync()
print(f"Synced commands: {', '.join([s.name for s in synced])}.")
event.set()
print("------")
#-----------------------------------------------------------------------
@bot.hybrid_command(
name="testm",
description="Enter a prompt to generate music!",
)
async def musicgen_command(ctx, prompt: str, seed: int = None):
"""Generates music based on a prompt"""
if ctx.author.id == bot.user.id:
return
if seed is None:
seed = random.randint(1, 10000)
try:
await music_create(ctx, prompt, seed)
except Exception as e:
print(f"Error: {e}")
#-----------------------------------------------------------------------
async def music_create(ctx, prompt, seed):
"""Runs music_create_job in executor"""
try:
message = await ctx.send(f"**{prompt}** - {ctx.author.mention} Generating...")
loop = asyncio.get_running_loop()
client = await loop.run_in_executor(None, get_client, None)
job = client.submit(prompt, seed, api_name="/predict")
await wait(job)
try:
job.result()
files = job.outputs()
media_files = files[0]
audio = media_files[0]
video = media_files[1]
short_filename = prompt[:20]
audio_filename = f"{short_filename}.mp3"
video_filename = f"{short_filename}.mp4"
with open(video, "rb") as file:
discord_video_file = discord.File(file, filename=video_filename)
await ctx.send(file=discord_video_file)
with open(audio, "rb") as file:
discord_audio_file = discord.File(file, filename=audio_filename)
await ctx.send(file=discord_audio_file)
except QueueError:
await ctx.send("The gradio space powering this bot is really busy! Please try again later!")
except Exception as e:
print(f"music_create Error: {e}")
#---------------------------------------------------------------------
def run_bot():
if not DISCORD_TOKEN:
print("DISCORD_TOKEN NOT SET")
event.set()
else:
bot.run(DISCORD_TOKEN)
threading.Thread(target=run_bot).start()
event.wait()
with gr.Blocks() as demo:
gr.Markdown(
"""
# Discord bot of https://huggingface.co/spaces/facebook/MusicGen
"""
)
demo.launch() |