88 lines
3.3 KiB
Python
88 lines
3.3 KiB
Python
import discord
|
|
import logging
|
|
import aiohttp
|
|
import io
|
|
from discord.ext import commands
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class VoiceGeneratorProfile():
|
|
def __init__(self, name, uuid, model="audio-cpp-chatterbox"):
|
|
self.name = name
|
|
self.model = model
|
|
self.uuid = uuid
|
|
|
|
|
|
class VoiceCog(commands.Cog):
|
|
def __init__(self, bot: commands.Bot):
|
|
self._running = False
|
|
self._current_voice = "Pricey"
|
|
self.api_url = "http://10.6.9.4:8080"
|
|
|
|
async def cog_load(self):
|
|
data = await self._get_voices()
|
|
self.voices = [VoiceGeneratorProfile(d['name'], d['id']) for d in data.get('data')]
|
|
self._current_voice = self.voices[0]
|
|
|
|
async def _get_voices(self):
|
|
async with aiohttp.ClientSession() as session:
|
|
data = await session.get(f"{self.api_url}/api/voice-profiles")
|
|
data = await data.json()
|
|
return data
|
|
|
|
@commands.command(name="voice_set", description="Change the active voice by name")
|
|
async def voice_set(self, ctx, voice_name):
|
|
"""Set the voice for the generator to use"""
|
|
for v in self.voices:
|
|
if v.name.lower() == voice_name.lower():
|
|
self._current_voice = v
|
|
await ctx.send(f"Swapped voice to {v.name}")
|
|
return
|
|
await ctx.send("Voice not found asshole")
|
|
|
|
@commands.command(name="voice_list", description="List available voices")
|
|
async def voice_list(self, ctx: commands.Context):
|
|
"""List available voices"""
|
|
await ctx.channel.send(f"{"\n".join([f" - {v.name}" for v in self.voices])}")
|
|
|
|
@commands.command(name="voice_info", description="See the status of the voice generator :tm:")
|
|
async def voice_info(self, ctx: commands.Context):
|
|
"""View status of the voice generator"""
|
|
data = f"""
|
|
- Running: {self._running}
|
|
- Current Voice: {self._current_voice.name}
|
|
"""
|
|
await ctx.channel.send(data)
|
|
|
|
@commands.command(name="voice_gen", description="Generate an audio clip using the currently configured voice")
|
|
async def voice_gen(self, ctx: commands.Context, prompt: str):
|
|
"""Generate some voices"""
|
|
if self._running:
|
|
await ctx.send("All lines are currently busy, please try again later")
|
|
return
|
|
data = {"model": self._current_voice.model,
|
|
"voice": f"localai://voice-profiles/{self._current_voice.uuid}",
|
|
"input": prompt,
|
|
"stream": False}
|
|
await ctx.channel.send("Working on it!")
|
|
self._running = True
|
|
try:
|
|
headers = {"Content-Type": "application/json"}
|
|
async with aiohttp.ClientSession(headers=headers) as session:
|
|
resp = await session.post(f"{self.api_url}/tts", json=data)
|
|
d = await resp.read()
|
|
# logging.debug(d)
|
|
voice_data = io.BytesIO(d)
|
|
self._running = False
|
|
filedata = discord.File(voice_data, filename=f"{self._current_voice.name.replace(" ", "_").lower()}-{prompt[:10].replace(" ", "_").lower()}.wav")
|
|
await ctx.channel.send(file=filedata)
|
|
except Exception as e:
|
|
self._running = False
|
|
logging.error(e)
|
|
|
|
|
|
async def setup(bot: commands.Bot) -> None:
|
|
"""Register the voice_gen cog with the bot."""
|
|
await bot.add_cog(VoiceCog(bot))
|