File size: 3,041 Bytes
8a0bc2e
 
58c201b
 
 
 
 
 
 
 
8a0bc2e
f82ba90
8a0bc2e
 
 
 
940edd3
54389af
58c201b
8a0bc2e
54389af
58c201b
 
54389af
8a0bc2e
58c201b
 
54389af
 
58c201b
54389af
 
8a0bc2e
54389af
8a0bc2e
 
54389af
940edd3
58c201b
 
 
 
 
 
 
8a0bc2e
54389af
58c201b
54389af
 
7f192d1
 
8a0bc2e
 
 
 
 
54389af
 
 
458a3b3
58c201b
 
 
 
 
 
 
8a0bc2e
458a3b3
58c201b
 
 
 
 
 
 
 
 
 
 
8a0bc2e
 
eab20c1
 
8a0bc2e
 
 
 
 
 
 
eab20c1
f7d01ae
54389af
8a0bc2e
54389af
 
8a0bc2e
 
 
 
 
 
 
 
ea3dcdc
58c201b
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
import os
import threading

import discord
import gradio as gr
from deepfloydif import deepfloydif_stage_1, deepfloydif_stage_2_react_check
from discord import app_commands
from discord.ext import commands
from falcon import continue_falcon, try_falcon
from musicgen import music_create


# HF GUILD SETTINGS
MY_GUILD_ID = 1077674588122648679 if os.getenv("TEST_ENV", False) else 879548962464493619
MY_GUILD = discord.Object(id=MY_GUILD_ID)
DISCORD_TOKEN = os.environ.get("DISCORD_TOKEN", None)


class Bot(commands.Bot):
    """This structure allows slash commands to work instantly."""

    def __init__(self):
        super().__init__(command_prefix="/", intents=discord.Intents.all())

    async def setup_hook(self):
        await self.tree.sync(guild=discord.Object(MY_GUILD_ID))
        print(f"Synced slash commands for {self.user}.")


client = Bot()


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


@client.hybrid_command(
    name="falcon",
    with_app_command=True,
    description="Enter some text to chat with the bot! Like this: /falcon Hello, how are you?",
)
@app_commands.guilds(MY_GUILD)
async def falcon(ctx, prompt: str):
    """Command that begins a new conversation with Falcon"""
    try:
        await try_falcon(ctx, prompt)
    except Exception as e:
        print(f"Error: {e}")


@client.event
async def on_message(message):
    """Checks channel and continues Falcon conversation if it's the right Discord Thread"""
    try:
        await continue_falcon(message)
    except Exception as e:
        print(f"Error: {e}")


@client.hybrid_command(
    name="deepfloydif",
    with_app_command=True,
    description="Enter a prompt to generate an image! Can generate realistic text, too!",
)
@app_commands.guilds(MY_GUILD)
async def deepfloydif(ctx, prompt: str):
    """DeepfloydIF stage 1 generation"""
    try:
        await deepfloydif_stage_1(ctx, prompt, client)
    except Exception as e:
        print(f"Error: {e}")


@client.hybrid_command(name="musicgen", with_app_command=True, description="Enter a prompt to generate music!")
@app_commands.guilds(MY_GUILD)
async def musicgen(ctx, prompt: str):
    """Generates music based on a prompt"""
    try:
        await music_create(ctx, prompt)
    except Exception as e:
        print(f"Error: {e}")


@client.event
async def on_reaction_add(reaction, user):
    """Checks for a reaction in order to call dfif2"""
    try:
        await deepfloydif_stage_2_react_check(reaction, user)
    except Exception as e:
        print(f"Error: {e} (known error, does not cause issues, low priority)")


def run_bot():
    client.run(DISCORD_TOKEN)


threading.Thread(target=run_bot).start()
"""This allows us to run the Discord bot in a Python thread"""
with gr.Blocks() as demo:
    gr.Markdown("""
    # Huggingbots Server
    This space hosts the huggingbots discord bot.
    Currently supported models are Falcon and DeepfloydIF
    """)
demo.queue(concurrency_count=100)
demo.launch()