File size: 5,435 Bytes
6fb60de 0e8c7ef 6fb60de 0e8c7ef 6fb60de |
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 110 111 112 113 114 115 116 |
import discord
from discord import app_commands
import gradio as gr
from gradio_client import Client
import os
import threading
import json
import random
from PIL import Image
import asyncio
import glob
MY_GUILD_ID = 1077674588122648679 if os.getenv("TEST_ENV", False) else 879548962464493619
MY_GUILD = discord.Object(id=MY_GUILD_ID)
HF_TOKEN = os.getenv('HF_TOKEN')
DISCORD_TOKEN = os.environ.get("DISCORD_TOKEN", Nonedeepfloyd_client = Client("huggingface-projects/IF", HF_TOKEN)
BOT_USER_ID = 1086256910572986469
DEEPFLOYD_CHANNEL_ID = 1121834257959092234
#-------------------------------------------------------------------------------------------------------------------------------------
# deepfloydif stage 1 generation
def inference(prompt):
negative_prompt = ''
seed = random.randint(0, 1000)
#seed = 1
number_of_images = 4
guidance_scale = 7
custom_timesteps_1 = 'smart50'
number_of_inference_steps = 50
stage_1_results, stage_1_param_path, stage_1_result_path = df.predict(
prompt, negative_prompt, seed, number_of_images, guidance_scale, custom_timesteps_1, number_of_inference_steps, api_name='/generate64')
return [stage_1_results, stage_1_param_path, stage_1_result_path]
#-------------------------------------------------------------------------------------------------------------------------------------
async def try_deepfloydif(interaction, prompt):
thread = None
try:
global BOT_USER_ID
global DEEPFLOYD_CHANNEL_ID
if interaction.user.id != BOT_USER_ID:
if interaction.channel.id == DEEPFLOYD_CHANNEL_ID:
await interaction.response.send_message(f"Working on it!")
channel = interaction.channel
message = await channel.send(f"DeepfloydIF Thread")
await message.add_reaction('🔁')
thread = await message.create_thread(name=f'{prompt}', auto_archive_duration=60)
await thread.send(f"[DISCLAIMER: HuggingBot is a **highly experimental** beta feature; Additional information" \
f" on the DeepfloydIF model can be found here: https://huggingface.co/spaces/DeepFloyd/IF")
dfif_command_message_id = message.id # used for updating the 'status' of our generations using reaction emojis
await thread.send(f'{interaction.user.mention}Generating images in thread, can take ~1 minute...')
# generation
loop = asyncio.get_running_loop()
result = await loop.run_in_executor(None, inference, prompt)
stage_1_results = result[0]
stage_1_result_path = result[2]
partialpath = stage_1_result_path[5:]
png_files = list(glob.glob(f"{stage_1_results}/**/*.png"))
img1 = None
img2 = None
img3 = None
img4 = None
if png_files:
png_file_index = 0
images = load_image(png_files, stage_1_results, png_file_index)
img1 = images[0]
img2 = images[1]
img3 = images[2]
img4 = images[3]
combined_image = Image.new('RGB', (img1.width * 2, img1.height * 2))
combined_image.paste(img1, (0, 0))
combined_image.paste(img2, (img1.width, 0))
combined_image.paste(img3, (0, img1.height))
combined_image.paste(img4, (img1.width, img1.height))
combined_image_path = os.path.join(stage_1_results, f'{partialpath}{dfif_command_message_id}.png')
combined_image.save(combined_image_path)
with open(combined_image_path, 'rb') as f:
combined_image_dfif = await thread.send(f'{interaction.user.mention}React with the image number you want to upscale!', file=discord.File(
f, f'{partialpath}{dfif_command_message_id}.png')) # named something like: tmpgtv4qjix1111269940599738479.png
emoji_list = ['↖️', '↗️', '↙️', '↘️']
await react1234(emoji_list, combined_image_dfif)
await message.remove_reaction('🔁', client.user)
await message.add_reaction('✔️')
except Exception as e:
print(f"Error: {e}")
#-------------------------------------------------------------------------------------------------------------------------------------
def load_image(png_files, stage_1_results, png_file_index):
for file in png_files:
png_file = png_files[png_file_index]
png_path = os.path.join(stage_1_results, png_file)
if png_file_index == 0:
img1 = Image.open(png_path)
if png_file_index == 1:
img2 = Image.open(png_path)
if png_file_index == 2:
img3 = Image.open(png_path)
if png_file_index == 3:
img4 = Image.open(png_path)
png_file_index = png_file_index + 1
return [img1, img2, img3, img4]
#-------------------------------------------------------------------------------------------------------------------------------------
|