File size: 1,807 Bytes
a181839 |
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 |
import asyncio
import glob
import os
import random
import discord
from gradio_client import Client
HF_TOKEN = os.getenv("HF_TOKEN")
wuerstchen_client = Client("warp-ai/Wuerstchen", HF_TOKEN)
BOT_USER_ID = 1102236653545861151
WUERSTCHEN_CHANNEL_ID = 1151792944676864041
def wuerstchen_inference(prompt, client):
negative_prompt = ""
seed = random.randint(0, 1000)
width = 1024
height = 1024
prior_num_inference_steps = 60
prior_guidance_scale = 4
decoder_num_inference_steps = 12
decoder_guidance_scale = 0
num_images_per_prompt = 1
job = wuerstchen_client.submit(
prompt,
negative_prompt,
seed,
width,
height,
prior_num_inference_steps,
prior_guidance_scale,
decoder_num_inference_steps,
decoder_guidance_scale,
num_images_per_prompt,
api_name="/run",
)
while not job.done():
pass
else:
return job
async def run_wuerstchen(ctx, prompt, client):
"""wuerstchen"""
try:
if ctx.author.id != BOT_USER_ID:
if ctx.channel.id == WUERSTCHEN_CHANNEL_ID:
channel = client.get_channel(WUERSTCHEN_CHANNEL_ID)
message = await ctx.send(f"**{prompt}** - {ctx.author.mention} <a:loading:1114111677990981692>")
loop = asyncio.get_running_loop()
job = await loop.run_in_executor(None, wuerstchen_inference, prompt, client)
png_files = list(glob.glob(f"{job.outputs()[-1]}/**/*.png"))
await message.delete()
with open(png_files[0], "rb") as f:
await channel.send(f"**{prompt}** - {ctx.author.mention}", file=discord.File(f, "wuerstchen.png"))
except Exception as e:
print(f"Error: {e}")
|