[MyClient] dictionary test, try/except blocks

app.py

@@ -38,128 +38,153 @@ sdlu = Client("huggingface-projects/stable-diffusion-latent-upscaler", HF_TOKEN)
 
 # Set up discord bot 
 class MyClient(discord.Client):
+    desired_messages = {}
+    
     async def on_ready(self):
         print('Logged on as', self.user)
 
 
     async def on_message(self, message):
-
-        #safety checks----------------------------------------------------------------------------------------------------
-
-        #   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
-
-        # function for reacting
-        async def bot_react(message):
-            emojis = ['1️⃣', '2️⃣', '3️⃣', '4️⃣']
-            for emoji in emojis:
-                await message.add_reaction(emoji)        
-
-        #deepfloydif----------------------------------------------------------------------------------------------------
-            
-        if message.content.startswith('!deepfloydif'): # change to application commands, more intuitive
-
-            #(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)
-
-            # input prompt
-            prompt = message.content[12:].strip()  
+        try:
+            #safety checks----------------------------------------------------------------------------------------------------
+    
+            #   tldr, bot should run if
+            #1) it does not have @offline role
+            #2) user has @verified role
+            #3) bot is in #bot-test channel
             
-            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
+            # 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
             
-            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]
-
-                '''
-                [],[],[],[] -> [][]
-                               [][]
+            # 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
+    
+            # function for reacting
+            async def bot_react(message):
+                emojis = ['1️⃣', '2️⃣', '3️⃣', '4️⃣']
+                for emoji in emojis:
+                    await message.add_reaction(emoji)        
+    
+            #deepfloydif----------------------------------------------------------------------------------------------------      
                 
-                '''
-            
-                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))
-            
-                # save the combined image
-                combined_image_path = os.path.join(stage_1_results, 'combined_image.png')
-                combined_image.save(combined_image_path)
-            
-            # 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'))
+            if message.content.startswith('!deepfloydif'): # change to application commands, more intuitive
+    
+                #(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)
+    
+                # 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))
+                
+                    # save the combined image
+                    combined_image_path = os.path.join(stage_1_results, 'combined_image.png')
+                    combined_image.save(combined_image_path)
+                
+                # 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'))
+    
+                self.desired_messages[message.author.id] = sent_message.id               
+    
+        except Exception as e:
+            print(f'An error occurred while processing the message: {e}')
+            await message.reply('An error occurred during stage 1.')        
+        
 
             await bot_react(sent_message)                
     #-------------------------------------------------------------------------------------------------------------
-
-            
-       
-
-
-    # deepfloydif -> bot_react -> 
-    # await bot_react
-    # on_react -> if we react to a message, then we call dfif2            
+    async def on_reaction_add(self, reaction, user):
+        try:  
+            if user != self.user:
+                # Check if the reaction is from one of the desired users and to the correct message
+                for user_id, message_id in self.desired_messages.items():
+                    if user.id == user_id and reaction.message.id == message_id:
+                        # Remove the user ID from the dictionary after processing the reaction
+                        del self.desired_messages[user_id]
+    
+                        # Get the index corresponding to the reacted emoji
+                        emoji_index = get_index_from_emoji(reaction.emoji)
+    
+                        # Use the index for the stage 2 upscaling
+                        result_path = df.predict(stage_1_result_path, emoji_index, seed_2, 
+                                                 guidance_scale_2, custom_timesteps_2, number_of_inference_steps_2, api_name='/upscale256')
+    
+                        with open(result_path, 'rb') as f:
+                            await reaction.message.reply('Here is the upscaled image!', file=discord.File(f, 'result.png'))
+                        break  # Exit the loop once a matching user and message pair is found
+      
+        except Exception as e:
+            print(f'An error occurred while processing the reaction: {e}')
+            await reaction.message.reply('An error occurred during stage 2.')
+          
 
 DISCORD_TOKEN = os.environ.get("GRADIOTEST_TOKEN", None)
 intents = discord.Intents.default()