revert to friday version

app.py

@@ -7,8 +7,6 @@ import json
 import random
 import time
 import re
-
-breaking bot for a moment
 from discord import Embed, Color
 from discord.ext import commands
 # test
@@ -29,78 +27,156 @@ def convert_to_timezone(dt, tz):
     return dt.astimezone(tz).strftime("%Y-%m-%d %H:%M:%S %Z")
 
 DFIF_TOKEN = os.getenv('HF_TOKEN')
-dfif = Client("huggingface-projects/IF", DFIF_TOKEN)
-#sdlu = Client("huggingface-projects/stable-diffusion-latent-upscaler", DFIF_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)
 
 
-#---------------------------------------------------------------------------------------------------------------------------------------------- 
+#----------------------------------------------------------------------------------------------------------------------------------------------
 
-# Set up discord bot 
-class MyClient(discord.Client):
-    async def on_ready(self):
-        print('Logged on as', self.user)
+@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
 
-	#------------------------------------------
-    async def on_message(self, message):
-        prompt = str(message.content)
-        channel = message.channel
-        await channel.send(prompt)
+#---------------------------------------------------------------------------------------------------------------------------------------------- 
+# Stage 1
+@bot.command()
+async def deepfloydif(ctx, *, prompt: str):
+    try:
+        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️⃣']   
+
+        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
+        await thread.send(f'{ctx.author.mention}Generating images in thread, can take ~1 minute...')
         
-        if message.content.startswith('!dfif'):
-              
-
-            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 = dfif.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)
+        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]
 
-            sent_message = await message.reply(f'Here is the combined image. React with the image number you want to upscale!', file=discord.File(f,'combined_image.png'))
+            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)
 
 
+        #1
+        #sent_message = None
+
+        #2
+        # create thread -> send new message inside thread + combined_image -> add reactions -> dfif2
+        #sent_message = await ctx.reply('Creating thread...')
+
+        #3
+        #ctx.message
+        # could replace sent_message with initial user command message? try later
+        #thread = await ctx.message.create_thread(name='Image Upscaling Thread')
+
+        
+        #thread = await sent_message.create_thread(name='Image Upscaling Thread')
+
+        #4
+        with open(combined_image_path, 'rb') as f:
+            threadmsg = await thread.send(f'{ctx.author.mention}React with the image number you want to upscale!', 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 threadmsg.add_reaction(emoji)   
+
+        reaction, user = await bot.wait_for('reaction_add', check=check_reaction)
+        if str(reaction.emoji) == '1️⃣':
+            await thread.send(f"{ctx.author.mention}Upscaling the first image...")
+            index = 0
+            await dfif2(ctx, index, stage_1_result_path, thread)
+        elif str(reaction.emoji) == '2️⃣':
+            await thread.send(f"{ctx.author.mention}Upscaling the second image...")
+            index = 1
+            await dfif2(ctx, index, stage_1_result_path, thread)            
+        elif str(reaction.emoji) == '3️⃣':
+            await thread.send(f"{ctx.author.mention}Upscaling the third image...")
+            index = 2
+            await dfif2(ctx, index, stage_1_result_path, thread)            
+        elif str(reaction.emoji) == '4️⃣':
+            await thread.send(f"{ctx.author.mention}Upscaling the fourth image...")
+            index = 3
+            await dfif2(ctx, index, stage_1_result_path, thread)            
+
+    #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.')
+        await ctx.message.add_reaction('❌')
+
+#----------------------------------------------------------------------------------------------------------------------------
+# Stage 2
+async def dfif2(ctx, 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')
     
-DISCORD_TOKEN = os.environ.get("GRADIOTEST_TOKEN", None)
-intents = discord.Intents.default()
-intents.message_content = True
-client = MyClient(intents=intents)
+        
+        with open(result_path, 'rb') as f:
+            await thread.send(f'{ctx.author.mention}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('✔️')
+            
+
+        
+
+    except Exception as e:
+        print(f"Error: {e}")
+        await ctx.reply('An error occurred while processing stage 2 upscaling. Please try again later.')
+        await ctx.message.add_reaction('❌')
+#----------------------------------------------------------------------------------------------------------------------------
+
+
+        
 
 def run_bot():
-    client.run(DISCORD_TOKEN)
+    bot.run(DISCORD_TOKEN)
 
 threading.Thread(target=run_bot).start()