audioldm2.py

import asyncio
import os

import discord
from gradio_client import Client
from gradio_client.utils import QueueError

BOT_USER_ID = 1102236653545861151  # real
MUSIC_CHANNEL_ID = 1143183148881035365  # real


HF_TOKEN = os.getenv("HF_TOKEN")
audioldm2 = Client("huggingface-projects/audioldm2-text2audio-text2music", HF_TOKEN)


def audioldm2_create_job(prompt):
    """Generates a sound or music based on a given prompt"""
    try:
        job = audioldm2.submit(prompt, api_name="/text2audio")
        while not job.done():
            pass
        return job

    except Exception as e:
        print(f"audioldm2_create_job Error: {e}")


async def audioldm2_create(ctx, prompt):
    """Runs audioldm2_create_job in executor"""
    try:
        if ctx.author.id != BOT_USER_ID:
            if ctx.channel.id == MUSIC_CHANNEL_ID:
                if os.environ.get("TEST_ENV") == "True":
                    print("Safetychecks passed for audioldm2_create")

                message = await ctx.send(f"**{prompt}** - {ctx.author.mention}")
                if len(prompt) > 99:
                    small_prompt = prompt[:99]
                else:
                    small_prompt = prompt
                thread = await message.create_thread(name=small_prompt, auto_archive_duration=60)

                if os.environ.get("TEST_ENV") == "True":
                    print("Running audioldm2_create_job...")

                loop = asyncio.get_running_loop()
                job = await loop.run_in_executor(None, audioldm2_create_job, prompt)

                try:
                    job.result()
                    video = job.outputs()[0]
                except QueueError:
                    await thread.send("The gradio space powering this bot is really busy! Please try again later!")

                short_filename = prompt[:20]
                video_filename = f"{short_filename}.mp4"

                with open(video, "rb") as file:
                    discord_video_file = discord.File(file, filename=video_filename)
                await thread.send(file=discord_video_file)

    except Exception as e:
        print(f"audioldm2_create Error: {e}")