bot.py

import asyncio
import io

import discord
import scipy.io.wavfile
from discord import app_commands
from discord.ext import commands
from pocket_tts import TTSModel

# For rapid testing
MY_GUILD = discord.Object(id=1117895866255159457)


class WrappedBytesIOSource(discord.FFmpegPCMAudio):
    def __init__(self, stream: io.BytesIO, **kwargs):
        super().__init__(stream, pipe=True, **kwargs)
        self._stream = stream

    def cleanup(self):
        super().cleanup()

        if self._stream:
            self._stream.close()
            self._stream = None


class VoiceRequest:
    def __init__(self, interaction: discord.Interaction, message: str):
        self.interaction = interaction
        self.message = message


class GuildVoicePlayer:
    def __init__(self, interaction: discord.Interaction):
        self.bot = interaction.client
        self.guild = interaction.guild
        self.channel = interaction.channel

        self.queue = asyncio.Queue()
        self.next = asyncio.Event()

        self.voice_worker = self.bot.loop.create_task(self.player_loop())

    async def player_loop(self):
        await self.bot.wait_until_ready()

        try:
            while not self.bot.is_closed():
                self.next.clear()

                request: VoiceRequest = await self.queue.get()
                await self.play_voice(request)

                await self.next.wait()

        except asyncio.CancelledError:
            pass

    async def enqueue(self, interaction: discord.Interaction, message: str):
        request = VoiceRequest(interaction, message)
        await self.queue.put(request)

    def play_next(self, error=None):
        if error:
            print(f"Player error: {error}")

        self.bot.loop.call_soon_threadsafe(self.next.set)

    async def play_voice(self, request: VoiceRequest):
        voice_client = self.guild.voice_client

        if not voice_client:  # Bot is not in a voice channel.
            # Try to join the user's
            if request.interaction.user.voice:
                voice_client = await request.interaction.user.voice.channel.connect()
        else:  # Bot is already in a voice channel.
            # Try to move to the user's VC if they're in one, otherwise stay put.
            if (
                request.interaction.user.voice
                and voice_client.channel != request.interaction.user.voice.channel
            ):
                voice_client = await voice_client.move_to(
                    request.interaction.user.voice.channel
                )

        byte_io = await bot.loop.run_in_executor(None, generate_audio, request.message)

        source = WrappedBytesIOSource(byte_io)

        description = (
            f"**{request.interaction.user.display_name}** said: {request.message}"
        )
        embed = discord.Embed(
            description=description[:4095],
            timestamp=discord.utils.utcnow(),
            color=request.interaction.user.color,
        )
        embed.add_field(name="Voice Channel", value=voice_client.channel.mention)
        embed.set_thumbnail(url=request.interaction.user.avatar.url)
        embed.set_author(
            name=str(request.interaction.guild),
            icon_url=request.interaction.guild.icon.url,
        )

        await request.interaction.followup.send(embed=embed)

        voice_client.play(source, after=self.play_next)


class SpeechBot(commands.Bot):
    # Suppress error on the User attribute being None since it fills up later
    user: discord.ClientUser

    def __init__(self):
        intents = discord.Intents.default()
        intents.message_content = True
        self.tts_model = TTSModel.load_model()
        self.voice_state = self.tts_model.get_state_for_audio_prompt("azelma")

        self.players = {}

        super().__init__(command_prefix="!@#$!", intents=intents)

    async def setup_hook(self):
        self.tree.copy_global_to(guild=MY_GUILD)
        await self.tree.sync(guild=MY_GUILD)
        print("Slash commands synced!")

    async def on_ready(self):
        print(f"Logged in as {self.user} (ID: {self.user.id})")
        print("------")

    async def on_voice_state_update(self, member, before, after):
        # If the bot itself is disconnected
        if member == bot.user and after.channel is None:
            guild_id = before.channel.guild.id
            if guild_id in self.players:
                print("Teardown NOT IMPLEMENTED")
                pass  # Implement clean-up logic


bot = SpeechBot()


def generate_audio(message: str):
    audio_tensor = bot.tts_model.generate_audio(bot.voice_state, message)
    byte_io = io.BytesIO()
    scipy.io.wavfile.write(byte_io, bot.tts_model.sample_rate, audio_tensor.numpy())
    byte_io.seek(0)

    return byte_io


@bot.tree.command(description="Stop playing audio.")
@commands.guild_only()
async def stop(interaction: discord.Interaction):
    return await interaction.response.send_message("Not yet implemented, sorry.")


@bot.tree.command(description="Skip to the next message.")
@commands.guild_only()
async def skip(interaction: discord.Interaction):
    if interaction.guild.voice_client and interaction.guild.voice_client.is_playing():
        # Stopping the client triggers the 'after' callback,
        # which sets self.next and moves the loop forward.
        interaction.guild.voice_client.stop()

        await interaction.response.send_message("Voice message stopped.")


@bot.tree.command(description="Say something in voice chat.")
@app_commands.describe(message="What should I say?")
@commands.guild_only()
async def say(interaction: discord.Interaction, message: str):
    if interaction.guild.id not in bot.players:
        player = GuildVoicePlayer(interaction)
        bot.players[interaction.guild.id] = player

    player = bot.players[interaction.guild.id]

    await player.enqueue(interaction, message)
    await interaction.response.defer()


bot.run("BOT_TOKEN_HERE")