simplifying, testing where errors come from

app.py

@@ -27,156 +27,78 @@ 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)
+dfif = 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:
-        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...')
-        
-        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)
+# Set up discord bot 
+class MyClient(discord.Client):
+    async def on_ready(self):
+        print('Logged on as', self.user)
 
-            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))
+	#------------------------------------------
+    async def on_message(self, message):
+		ctx = await self.get_context(message)
+        
+        if message.content.startswith('!dfif'): #ex: !dfif sleepy cat
+			prompt = string[5:]                 #ex: sleepy cat
+			await dfif (ctx, prompt)            # pass prompt to dfif (for generation) and ctx (for keeping track of command msg)
+			
+
+    		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)
 
-            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)
+            sent_message = await message.reply('Here is the combined image. React with the image number you want to upscale!', file=discord.File(f, 'combined_image.png'))
 
-            # 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')
     
-        
-        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('❌')
-#----------------------------------------------------------------------------------------------------------------------------
-
-
-        
+DISCORD_TOKEN = os.environ.get("GRADIOTEST_TOKEN", None)
+intents = discord.Intents.default()
+intents.message_content = True
+client = MyClient(intents=intents)
 
 def run_bot():
-    bot.run(DISCORD_TOKEN)
+    client.run(DISCORD_TOKEN)
 
 threading.Thread(target=run_bot).start()