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 # test from gradio_client import Client from PIL import Image #from ratelimiter import RateLimiter # from datetime import datetime from pytz import timezone # import asyncio import shutil # for doing image movement magic DFIF_TOKEN = os.getenv('HF_TOKEN') df = Client("huggingface-projects/IF", 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: try: await ctx.message.add_reaction('👍') thread = await ctx.message.create_thread(name=f'{ctx.author} Image Upscaling Thread ') # create thread -> send new message inside thread + combined_image -> add reactions -> dfif2 #current_time = int(time.time()) #random.seed(current_time) 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 await thread.send(f'{ctx.author.mention}Generating images in thread, can take ~1 minute...') except Exception as e: print(f"Error: {e}") await ctx.reply('stage 1 error -> pre generation') await ctx.message.add_reaction('❌') try: 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') partialpath = stage_1_result_path[5:] #magic for later except Exception as e: print(f"Error: {e}") await ctx.reply('stage 1 error -> during generation') await ctx.message.add_reaction('❌') try: png_files = [f for f in os.listdir(stage_1_results) if f.endswith('.png')] if png_files: for i, png_file in enumerate(png_files): png_file_path = os.path.join(stage_1_results, png_file) img = Image.open(png_file_path) image_path = os.path.join(stage_1_results, f'{i+1}{partialpath}.png') img.save(image_path) with open(image_path, 'rb') as f: await thread.send(f'{ctx.author.mention}Image {i+1}', file=discord.File(f, f'{i+1}{partialpath}.png')) await asyncio.sleep(1) await thread.send(f'{ctx.author.mention}React with 👍 to the image you want to upscale!') except Exception as e: print(f"Error: {e}") await ctx.reply('stage 1 error -> posting images in thread') await ctx.message.add_reaction('❌') #deepfloydif try/except except Exception as e: print(f"Error: {e}") await ctx.reply('An error occurred in stage 1') await ctx.message.add_reaction('❌') #---------------------------------------------------------------------------------------------------------------------------- # Stage 2 async def dfif2(index: int, stage_1_result_path, thread): 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 thread.send(f'Here is the upscaled image! :) ', file=discord.File(f, 'result.png')) #await ctx.reply('Here is the result of the second stage', file=discord.File(f, 'result.png')) #await ctx.message.add_reaction('✅') need to fix this except Exception as e: print(f"Error: {e}") #await ctx.reply('An error occured in stage 2') need to fix #await ctx.message.add_reaction('❌') #---------------------------------------------------------------------------------------------------------------------------- @bot.event async def on_reaction_add(reaction, user): # safety checks first if user.bot: return # do they have the required (verified) role? guild = reaction.message.guild # Check if the user has the desired role required_role_id = 897376942817419265 # huggingfolks for now required_role = guild.get_role(required_role_id) if required_role not in user.roles: return emoji = reaction.emoji thread = reaction.message.channel tbd = reaction.message.content #index = 0 decided by index in front of partialpath # magic begins attachment = reaction.message.attachments[0] image_name = attachment.filename # we know image_name will be something like 1tmpgtv4qjix.png # remove .png first indexpartialpath = image_name[:-4] # should be 1tmpgtv4qjix # extract index as an integer (dfif2 needs integer) index = int(indexpartialpath[0]) # should be 1 # extract partialpath partialpath = indexpartialpath[1:] # should be tmpgtv4qjix # add /tmp/ to partialpath, save as new variable fullpath = "/tmp/" + partialpath # should be /tmp/tmpgtv4qjix if emoji == '👍': if reaction.message.attachments: if user.id == reaction.message.mentions[0].id: # all we care about is upscaling whatever image this is stage_1_result_path = fullpath index = index await dfif2(index, stage_1_result_path, thread) 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()