[decorators] revert

app.py

@@ -1,167 +1,160 @@
 import discord
-import gradio_client
-from gradio_client import Client
-import gradio as gr
 import os
 import threading
-
-#for deepfloydif
+import gradio as gr
 import requests
 import json
 import random
-from PIL import Image
-import matplotlib.pyplot as plt
-import matplotlib.image as mpimg
 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
 
-# random + small llama #
 
+DFIF_TOKEN = os.getenv('HF_TOKEN')
+df = Client("huggingface-projects/IF", DFIF_TOKEN)
 
-#todos
-#alert
-#fix error on first command on bot startup
-#stable diffusion upscale
-#buttons for deepfloydIF (1,2,3,4)
-#application commands instead of message content checks (more user-friendly)
 
+DISCORD_TOKEN = os.environ.get("GRADIOTEST_TOKEN", None)
+intents = discord.Intents.default()
+intents.message_content = True
 
+bot = commands.Bot(command_prefix='!', intents=intents)
 
 
-HF_TOKEN = os.getenv('HF_TOKEN')
+#----------------------------------------------------------------------------------------------------------------------------------------------
 
-#deepfloydIF
-#df = Client("DeepFloyd/IF", HF_TOKEN) #not reliable at the moment
-df = Client("huggingface-projects/IF", HF_TOKEN)
+@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
 
-#stable diffusion upscaler
-sdlu = Client("huggingface-projects/stable-diffusion-latent-upscaler", HF_TOKEN)
+#---------------------------------------------------------------------------------------------------------------------------------------------- 
+# 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
 
-# Set up discord bot 
-class MyClient(discord.Client):
-    async def on_ready(self):
-        print('Logged on as', self.user)
+        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...')
+        
+        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')]
 
-    async def on_message(self, message):
+        if png_files:
+            first_png = png_files[0]
+            second_png = png_files[1]
+            third_png = png_files[2]
+            fourth_png = png_files[3]
 
-        #safety checks----------------------------------------------------------------------------------------------------
+            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)
 
-        #   tldr, bot should run if
-        #1) it does not have @offline role
-        #2) user has @verified role
-        #3) bot is in #bot-test channel
-        
-        # bot won't respond to itself, prevents feedback loop + API spam
-        if message.author == self.user:
-            return       
-
-        # if the bot has this role, it won't run
-        OFFLINE_ROLE_ID = 1103676632667017266  # 1103676632667017266 = @offline / under maintenance
-        guild = message.guild
-        bot_member = guild.get_member(self.user.id)
-        if any(role.id == OFFLINE_ROLE_ID for role in bot_member.roles):
-            return
-        
-        # the message author needs this role in order to use the bot
-        REQUIRED_ROLE_ID = 897376942817419265 # 900063512829755413 = @verified, 897376942817419265 = @huggingfolks
-        if not any(role.id == REQUIRED_ROLE_ID for role in message.author.roles):
-            return            
-
-        # channels where bot will accept commands
-        ALLOWED_CHANNEL_IDS = [1100458786826747945] # 1100458786826747945 = #bot-test
-        if message.channel.id not in ALLOWED_CHANNEL_IDS:
-            return
-     
-
-        #deepfloydif----------------------------------------------------------------------------------------------------
-            
-        if message.content.startswith('!deepfloydif'): # change to application commands, more intuitive
+            img1 = Image.open(first_png_path)
+            img2 = Image.open(second_png_path)
+            img3 = Image.open(third_png_path)
+            img4 = Image.open(fourth_png_path)
 
-            #(prompt, negative_prompt, seed, number_of_images, guidance_scale,custom_timesteps_1, number_of_inference_steps, api_name="/generate64") 
-            #-> (stage_1_results, stage_1_param_path, stage_1_result_path)
+            combined_image = Image.new('RGB', (img1.width * 2, img1.height * 2))
 
-            # input prompt
-            prompt = message.content[12:].strip()  
-            
-            negative_prompt = ''
-            seed = random.randint(0, 2**32 - 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')
-
-            #stage_1_results, stage_1_param_path, stage_1_result_path = df.predict("gradio written on a wall", "blur", 1,1,7.0, 'smart100',50, api_name="/generate64")
-            
-            # stage_1_results -> path to directory with png files, so we isolate those
-            png_files = [f for f in os.listdir(stage_1_results) if f.endswith('.png')]
-
-            # merge images into larger, 2x2 image the way midjourney does it
-            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)
-            
-                # create a new blank image with the size of the combined images (2x2)
-                combined_image = Image.new('RGB', (img1.width * 2, img1.height * 2))
-            
-                # paste the individual images into the combined image
-                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.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)
+
+
+        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'))          
             
-                # save the combined image
-                combined_image_path = os.path.join(stage_1_results, 'combined_image.png')
-                combined_image.save(combined_image_path)
+        # 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')
+    
+        
+        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('✔️')
             
-            # send the combined image file as a discord attachment
-            with open(combined_image_path, 'rb') as f:
-                sent_message = await message.reply('Here is the combined image', file=discord.File(f, 'combined_image.png'))
-               
-    #-------------------------------------------------------------------------------------------------------------
 
-            
-       
+        
 
+    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('❌')
+#----------------------------------------------------------------------------------------------------------------------------
 
-    # deepfloydif -> bot_react -> 
-    # await bot_react
-    # on_react -> if we react to a message, then we call dfif2            
 
-DISCORD_TOKEN = os.environ.get("GRADIOTEST_TOKEN", None)
-intents = discord.Intents.default()
-intents.message_content = True
-client = MyClient(intents=intents)
+        
 
 def run_bot():
-  client.run(DISCORD_TOKEN)
+    bot.run(DISCORD_TOKEN)
 
 threading.Thread(target=run_bot).start()
 
@@ -169,4 +162,4 @@ def greet(name):
     return "Hello " + name + "!"
 
 demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()   
+demo.launch()