[deepfloydif.py] cleanup

deepfloydif.py

@@ -3,7 +3,6 @@ from discord import app_commands
 import gradio as gr
 from gradio_client import Client
 import os
-import threading 
 import json
 import random
 from PIL import Image
@@ -12,12 +11,10 @@ import glob
 import pathlib
 
 HF_TOKEN = os.getenv('HF_TOKEN')
-DISCORD_TOKEN = os.environ.get("DISCORD_TOKEN", None)
 deepfloyd_client = Client("huggingface-projects/IF", HF_TOKEN)
 
-BOT_USER_ID = 1086256910572986469
-DEEPFLOYD_CHANNEL_ID = 1121834257959092234
-
+BOT_USER_ID = 1086256910572986469 if os.getenv("TEST_ENV", False) else 1102236653545861151
+DEEPFLOYD_CHANNEL_ID = 1121834257959092234 if os.getenv("TEST_ENV", False) else 1119313215675973714
 
 def deepfloyd_stage_1_inference(prompt):
     """Generates an image based on a prompt"""
@@ -27,9 +24,7 @@ def deepfloyd_stage_1_inference(prompt):
     guidance_scale = 7
     custom_timesteps_1 = 'smart50'
     number_of_inference_steps = 50
-    
     stage_1_results, stage_1_param_path, stage_1_result_path = deepfloyd_client.predict(prompt, negative_prompt, seed, number_of_images, guidance_scale, custom_timesteps_1, number_of_inference_steps, api_name='/generate64')
-    
     return [stage_1_results, stage_1_param_path, stage_1_result_path]         
 
 def deepfloyd_stage_2_inference(index, stage_1_result_path):
@@ -44,7 +39,6 @@ def deepfloyd_stage_2_inference(index, stage_1_result_path):
     
 async def deepfloydif_stage_1(interaction, prompt, client):
     """DeepfloydIF command (generate images with realistic text using slash commands)"""
-    thread = None
     try:
         global BOT_USER_ID
         global DEEPFLOYD_CHANNEL_ID
@@ -54,12 +48,9 @@ async def deepfloydif_stage_1(interaction, prompt, client):
                 channel = interaction.channel
                 # interaction.response message can't be used to create a thread, so we create another message
                 message = await channel.send("DeepfloydIF Thread")
-                #await message.add_reaction('<a:loading:1114111677990981692>')
                 thread = await message.create_thread(name=f'{prompt}', auto_archive_duration=60) 
                 await thread.send("[DISCLAIMER: HuggingBot is a **highly experimental** beta feature; Additional information on the DeepfloydIF model can be found here: https://huggingface.co/spaces/DeepFloyd/IF")
-
                 dfif_command_message_id = message.id # used for updating the 'status' of our generations using reaction emojis
-                
                 await thread.send(f'{interaction.user.mention} Generating images in thread, can take ~1 minute...')
 
                 loop = asyncio.get_running_loop()
@@ -86,8 +77,6 @@ async def deepfloydif_stage_1(interaction, prompt, client):
 
                     emoji_list = ['↖️', '↗️', '↙️', '↘️']
                     await react_1234(emoji_list, combined_image_dfif)
-                    #await message.remove_reaction('<a:loading:1114111677990981692>', client.user)
-                    #await message.add_reaction('<:agree:1098629085955113011>')
                 else:
                     await thread.send(f'{interaction.user.mention} No PNG files were found, cannot post them!')
 
@@ -120,8 +109,6 @@ async def deepfloydif_stage_2(index: int, stage_1_result_path, thread, dfif_comm
     try:
         parent_channel = thread.parent
         dfif_command_message = await parent_channel.fetch_message(dfif_command_message_id)
-        #await dfif_command_message.remove_reaction('<:agree:1098629085955113011>', client.user)
-        #await dfif_command_message.add_reaction('<a:loading:1114111677990981692>')
         if index == 0:
             position = "top left"
         elif index == 1:
@@ -138,34 +125,17 @@ async def deepfloydif_stage_2(index: int, stage_1_result_path, thread, dfif_comm
         
         with open(result_path, 'rb') as f:
             await thread.send('Here is the upscaled image!', file=discord.File(f, 'result.png'))
-            
-        #await dfif_command_message.remove_reaction('<a:loading:1114111677990981692>', client.user)
-        #await dfif_command_message.add_reaction('<:agree:1098629085955113011>')
         await thread.edit(archived=True)
 
     except Exception as e:
         print(f"Error: {e}")
-        parent_channel = thread.parent
-        dfif_command_message = await parent_channel.fetch_message(dfif_command_message_id)
-        #await dfif_command_message.remove_reaction('<a:loading:1114111677990981692>', client.user)
-        #await dfif_command_message.add_reaction('<:disagree:1098628957521313892>')  
-        await thread.send(f"Error during stage 2 upscaling, {e}") 
-        await thread.edit(archived=True)
-   
-
-
-
-
-
-
-
 
 async def deepfloydif_stage_2_react_check(reaction, user): 
     """Checks for a reaction in order to call dfif2"""
     try:
         global BOT_USER_ID
         global DEEPFLOYD_CHANNEL_ID
-        if user.id != BOT_USER_ID: # 
+        if user.id != BOT_USER_ID:
             thread = reaction.message.channel
             thread_parent_id = thread.parent.id
             if thread_parent_id == DEEPFLOYD_CHANNEL_ID: