import discord import os import threading import gradio as gr import requests import json import random import time import re from discord import Embed, Color from discord.ext import commands # restart #unstick #unstick from gradio_client import Client from PIL import Image #from ratelimiter import RateLimiter # from datetime import datetime from pytz import timezone # import asyncio semaphore = asyncio.Semaphore(1) # for concurrent tasks zurich_tz = timezone("Europe/Zurich") def convert_to_timezone(dt, tz): return dt.astimezone(tz).strftime("%Y-%m-%d %H:%M:%S %Z") DFIF_TOKEN = os.getenv('HF_TOKEN') df = Client("huggingface-projects/IF", DFIF_TOKEN) sdlu = Client("huggingface-projects/stable-diffusion-latent-upscaler", DFIF_TOKEN) DISCORD_TOKEN = os.environ.get("GRADIOTEST_TOKEN", None) intents = discord.Intents.default() intents.message_content = True bot = commands.Bot(command_prefix='!', intents=intents) #------------------------------------------------------------------------------------------------------------------------------------------- @bot.event async def on_ready(): print('Logged on as', bot.user) bot.log_channel = bot.get_channel(1100458786826747945) # 1100458786826747945 = bot-test, 1107006391547342910 = lunarbot server #------------------------------------------------------------------------------------------------------------------------------------------- # Stage 1 @bot.command() async def deepfloydif(ctx, *, prompt: str): try: await ctx.message.add_reaction('👍') prompt = prompt.strip()[:100] # Limit the prompt length to 100 characters prompt = re.sub(r'[^\w\s]', '', prompt) # Remove special characters def check_reaction(reaction, user): return user == ctx.author and str(reaction.emoji) in ['1️⃣', '2️⃣', '3️⃣', '4️⃣'] number_of_images = 4 current_time = int(time.time()) random.seed(current_time) seed = random.randint(0, 2**32 - 1) stage_1_results, stage_1_param_path, stage_1_result_path = df.predict (prompt, "blur", seed, number_of_images, 7.0, 'smart100', 50, api_name="/generate64") png_files = [f for f in os.listdir(stage_1_results) if f.endswith('.png')] if png_files: first_png = png_files[0] second_png = png_files[1] third_png = png_files[2] fourth_png = png_files[3] first_png_path = os.path.join(stage_1_results, first_png) second_png_path = os.path.join(stage_1_results, second_png) third_png_path = os.path.join(stage_1_results, third_png) fourth_png_path = os.path.join(stage_1_results, fourth_png) img1 = Image.open(first_png_path) img2 = Image.open(second_png_path) img3 = Image.open(third_png_path) img4 = Image.open(fourth_png_path) 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, 'combined_image.png') combined_image.save(combined_image_path) # Trigger the second stage prediction #await dfif2(ctx, stage_1_result_path) message = await ctx.reply('Here is the combined image. React with the image you want to upscale!') with open(combined_image_path, 'rb') as f: sent_message = await ctx.send(file=discord.File(f, 'combined_image.png')) # bot reacts with appropriate emojis to the post, both showing the user what options they have, # as well as showing which post to react to. for emoji in ['1️⃣', '2️⃣', '3️⃣', '4️⃣']: await sent_message.add_reaction(emoji) reaction, user = await bot.wait_for('reaction_add', check=check_reaction) if str(reaction.emoji) == '1️⃣': await ctx.send("You chose the first image!") index = 0 await dfif2(ctx, index, stage_1_result_path) elif str(reaction.emoji) == '2️⃣': await ctx.send("You chose the second image!") index = 1 await dfif2(ctx, index, stage_1_result_path) elif str(reaction.emoji) == '3️⃣': await ctx.send("You chose the third image!") index = 2 await dfif2(ctx, index, stage_1_result_path) elif str(reaction.emoji) == '4️⃣': await ctx.send("You chose the fourth image!") index = 3 await dfif2(ctx, index, stage_1_result_path) #deepfloydif try/except except Exception as e: print(f"Error: {e}") await ctx.reply('An error occurred while processing your request. Please wait 5 seconds before retrying.') #---------------------------------------------------------------------------------------------------------------------------- # Stage 2 async def dfif2(ctx, index: int, stage_1_result_path): try: selected_index_for_stage_2 = index seed_2 = 0 guidance_scale_2 = 4 custom_timesteps_2 = 'smart50' number_of_inference_steps_2 = 50 result_path = df.predict(stage_1_result_path, selected_index_for_stage_2, seed_2, guidance_scale_2, custom_timesteps_2, number_of_inference_steps_2, api_name='/upscale256') with open(result_path, 'rb') as f: await ctx.reply('Here is the result of the second stage', file=discord.File(f, 'result.png')) except Exception as e: print(f"Error: {e}") await ctx.reply('An error occurred while processing stage 2 upscaling. Please try again later.') #---------------------------------------------------------------------------------------------------------------------------- async def process_dfif2(ctx, index, stage_1_result_path): async with semaphore: await dfif2(ctx, index, stage_1_result_path) #---------------------------------------------------------------------------------------------------------------------------- def run_bot(): bot.run(DISCORD_TOKEN) threading.Thread(target=run_bot).start() def greet(name): return "Hello " + name + "!" demo = gr.Interface(fn=greet, inputs="text", outputs="text") demo.launch()