Update app.py

app.py

@@ -23,9 +23,7 @@ import concurrent.futures
 import multiprocessing
 
 import glob
-#-------------------------------------------------------------------------------------------------------------------------------------------------------------
-
-
+#-------------------------------------------------------------------------------------------------------------------------------------
 MY_GUILD = discord.Object(id=1077674588122648679)  # replace with your guild id
 HF_TOKEN = os.getenv('HF_TOKEN')
 DISCORD_TOKEN = os.environ.get("DISCORD_TOKEN", None)
@@ -48,7 +46,7 @@ instructions = "The following is a conversation between a highly knowledgeable "
 "with. It knows a lot, and always tells the truth. The conversation begins."
 
 
-#-------------------------------------------------------------------------------------------------------------------------------------------------------------
+#-------------------------------------------------------------------------------------------------------------------------------------
 # myclient class so we can run this in thread
 # sets up commands and the tree stuff for slash commands
 # Allows commands to work instantly
@@ -66,122 +64,44 @@ class MyClient(discord.Client):
 intents = discord.Intents.all() # could be default maybe
 client = MyClient(intents=intents)
 
-#-------------------------------------------------------------------------------------------------------------------------------------------------------------
+#-------------------------------------------------------------------------------------------------------------------------------------
 @client.event
 async def on_ready():
     print(f'Logged in as {client.user} (ID: {client.user.id})')
     print('------')
-#-------------------------------------------------------------------------------------------------------------------------------------------------------------
-
-
-
-# /falcon -> create thread, store it
-# on_message -> if channelid in falcon_userid_threadid_dictionary
-# -> if userid in falcon_userid_threadid_dictionary
-# -> if falcon_userid_threadid_dictionary[userid] == thread_id
-
-
-@client.tree.command()
-@app_commands.describe(
-    prompt='Enter some text to chat with the bot! Like this: /falcon Hello, how are you?')
-async def falcon(interaction: discord.Interaction, prompt: str):
-    try:  
-        global chathistory # ?
-        global falcon_userid_threadid_dictionary # tracks userid-thread existence
-        global instructions
-        global threadid_conversation
-        global FALCON_CHANNEL_ID
-
-        if not interaction.user.bot:
-            if interaction.channel.id == FALCON_CHANNEL_ID: # 1079459939405279232= test
-                await interaction.response.send_message(f"Working on it!")
-                channel = interaction.channel
-                # 1
-                message = await channel.send(f"Creating thread...")
-                thread = await message.create_thread(name=f'{interaction.user}') # interaction.user
-                await thread.send(f"[DISCLAIMER: HuggingBot is a **highly experimental** beta feature; The Falcon " \
-                f"model and system prompt can be found here: https://huggingface.co/spaces/HuggingFaceH4/falcon-chat]")
-                # generation
-                chathistory = falconclient.predict( 
-                        fn_index=5
-                ) # []    
-                job = falconclient.submit(prompt, chathistory, instructions, 0.8, 0.9, fn_index=1)  # This is not blocking, similar to run_in_executor (but better)
-                while job.done() == False: 
-                    status = job.status() 
-                    #print(status)
-                else:
-                    file_paths = job.outputs()
-                    full_generation = file_paths[-1] # tmp12345678.json
-                with open(full_generation, 'r') as file:
-                    data = json.load(file)
-                    output_text = data[-1][-1] # we output this as the bot
-
-                threadid_conversation[thread.id] = full_generation
-                
-                falcon_userid_threadid_dictionary[thread.id] = interaction.user.id
-                
-                print(output_text)
-                await thread.send(f"{output_text}")  
-
-    except Exception as e:
-        print(f"Error: {e}")
-        #await thread.send(f"{e} cc <@811235357663297546> (falconprivate error)") 
-#--------------------------------------------------------------------------------
-@client.event
-async def on_message(message):
-    try:
-        if not message.author.bot:
-            global falcon_userid_threadid_dictionary # tracks userid-thread existence
-            if message.channel.id in falcon_userid_threadid_dictionary: # is this a valid thread?
-                if falcon_userid_threadid_dictionary[message.channel.id] == message.author.id: # more than that - is this specifically the right user for this thread?
-                    # call job for falcon
-                    #await message.reply("checks succeeded, calling continue_falcon")
-                    await continue_falcon(message)
-
-    except Exception as e:
-        print(f"Error: {e}")
-#----------------------------------------------------------------------------------------------------------------------------
-async def continue_falcon(message):
-    try:
-        global chathistory # ?
-        global instructions
-        global threadid_conversation
-        await message.add_reaction('<a:loading:1121820108189339738>') # test=<a:loading:1121820108189339738>  hf=<a:loading:1114111677990981692>
-        chathistory = threadid_conversation[message.channel.id]
-        prompt = message.content
-        # generation
-        job = falconclient.submit(prompt, chathistory, instructions, 0.8, 0.9, fn_index=1)  # This is not blocking, similar to run_in_executor (but better)
-        while job.done() == False: 
-            status = job.status() 
-            #print(status)
-        else:
-            file_paths = job.outputs()
-            full_generation = file_paths[-1] # tmp12345678.json
-        with open(full_generation, 'r') as file:
-            data = json.load(file)
-            output_text = data[-1][-1] # we output this as the bot
 
-        threadid_conversation[message.channel.id] = full_generation # overwrite the old file
-        falcon_userid_threadid_dictionary[message.channel.id] = message.author.id
-        
-        print(output_text)
-        await message.reply(f"{output_text}")  
-        await message.remove_reaction('<a:loading:1121820108189339738>', client.user) # test=<a:loading:1121820108189339738>  hf=<a:loading:1114111677990981692>  
-
-    except Exception as e:
-        print(f"Error: {e}")
-        await message.reply(f"Error: {e} <@811235357663297546> (continue_falcon error)")
-#----------------------------------------------------------------------------------------------------------------------------
-'''
-@bot.command()
-async def falconclear(ctx):
-    if not ctx.author.bot:
-        if ctx.channel.id == 1079459939405279232: # 1079459939405279232 = test in lunar server
-            if ctx.author.id in falcon_userid_threadid_dictionary:
-                    del falcon_userid_threadid_dictionary[ctx.author.id]
-                    await ctx.reply(f"{ctx.author.mention}'s conversation has been cleared. Feel free to start a new one!")
-'''
-#---------------------------------------------------------------------------------------------------------------------------- 
+#-------------------------------------------------------------------------------------------------------------------------------------  
+# deepfloydif stage 1 generation
+def inference(prompt):
+    negative_prompt = ''
+    seed = random.randint(0, 1000)
+    #seed = 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')
+    
+    return [stage_1_results, stage_1_param_path, stage_1_result_path]        
+#------------------------------------------------------------------------------------------------------------------------------------- 
+# deepfloydif stage 2 upscaling
+def inference2(index, stage_1_result_path):
+    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')
+    
+    return result_path    
+#------------------------------------------------------------------------------------------------------------------------------------- 
+async def react1234(reaction_emojis, combined_image_dfif):
+    for emoji in reaction_emojis:
+        await combined_image_dfif.add_reaction(emoji) 
+#------------------------------------------------------------------------------------------------------------------------------------- 
 @client.tree.command()
 @app_commands.describe(
     prompt='Enter a prompt to generate an image! Can generate realistic text, too!')
@@ -284,29 +204,180 @@ async def dfif(interaction: discord.Interaction, prompt: str):
     except Exception as e:
         print(f"Error: {e}")
         #await thread.send(f"Error: {e} <@811235357663297546> (continue_falcon error)")
-#----------------------------------------------------------------------------------------------------------------------------
-async def react1234(reaction_emojis, combined_image_dfif):
-    for emoji in reaction_emojis:
-        await combined_image_dfif.add_reaction(emoji) 
+#-------------------------------------------------------------------------------------------------------------------------------------   
+async def dfif2(index: int, stage_1_result_path, thread, dfif_command_message_id): # add safetychecks
+    try:
+        parent_channel = thread.parent
+        dfif_command_message = await parent_channel.fetch_message(dfif_command_message_id)
+        await dfif_command_message.remove_reaction('✅', bot.user)
+        await dfif_command_message.add_reaction('<a:loading:1121820108189339738>')
+            
+        number = index + 1
+        if number == 1:
+            position = "top left"
+        elif number == 2:
+            position = "top right"
+        elif number == 3:
+            position = "bottom left"
+        elif number == 4:
+            position = "bottom right" 
+        await thread.send(f"Upscaling the {position} image...")  
+        
+        # run blocking function in executor
+        loop = asyncio.get_running_loop()
+        result_path = await loop.run_in_executor(None, inference2, index, stage_1_result_path)
+
+        #await thread.send(f"✅upscale done")          
+        with open(result_path, 'rb') as f:
+            await thread.send(f'Here is the upscaled image! :) ', file=discord.File(f, 'result.png'))
+            
+        await dfif_command_message.remove_reaction('<a:loading:1121820108189339738>', bot.user)
+        await dfif_command_message.add_reaction('✅')
+        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:1121820108189339738>', bot.user)
+        await dfif_command_message.add_reaction('❌')  
+        await thread.send(f"Error during stage 2 upscaling, {e}") 
+        await fullqueue(e, thread) 
+        await thread.edit(archived=True)
+#-------------------------------------------------------------------------------------------------------------------------------------
+@bot.event
+async def on_reaction_add(reaction, user):    # ctx = await bot.get_context(reaction.message)? could try later, might simplify
+    try:
+        global DEEPFLOYD_CHANNEL_ID
+        if not user.bot: 
+            channel_id = reaction.message.channel.id
+            if channel_id == DEEPFLOYD_CHANNEL_ID: # should be whatever the deepfloydif channel is
+                if reaction.message.attachments:
+                    if user.id == reaction.message.mentions[0].id:  #  if user.id == reaction.message.mentions[0].id:           
+                        #await reaction.message.channel.send("✅reaction detected")
+                        attachment = reaction.message.attachments[0]
+                        image_name = attachment.filename # named something like: tmpgtv4qjix1111269940599738479.png
+                        # remove .png first
+                        partialpathmessageid = image_name[:-4] # should be tmpgtv4qjix1111269940599738479 
+                        # extract partialpath, messageid
+                        partialpath = partialpathmessageid[:11] # tmpgtv4qjix
+                        messageid = partialpathmessageid[11:] # 1111269940599738479
+                        # add /tmp/ to partialpath, save as new variable
+                        fullpath = "/tmp/" + partialpath # should be /tmp/tmpgtv4qjix
+                        #await reaction.message.channel.send(f"✅fullpath extracted, {fullpath}")        
+                        emoji = reaction.emoji
+                        
+                        if emoji == "↖️":
+                            index = 0
+                        elif emoji == "↗️":
+                            index = 1
+                        elif emoji == "↙️":
+                            index = 2
+                        elif emoji == "↘️":
+                            index = 3 
+                            
+                        #await reaction.message.channel.send(f"✅index extracted, {index}")         
+                        index = index
+                        stage_1_result_path = fullpath
+                        thread = reaction.message.channel
+                        dfif_command_message_id = messageid
+                        #await reaction.message.channel.send(f"✅calling dfif2")  
+                        await dfif2(index, stage_1_result_path, thread, dfif_command_message_id)
 
-#----------------------------------------------------------------------------------------------------------------------------
-#----------------------------------------------------------------------------------------------------------------------------------------------  
-# deepfloydif stage 1 generation ✅
-def inference(prompt):
-    negative_prompt = ''
-    seed = random.randint(0, 1000)
-    #seed = 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')
-    
-    return [stage_1_results, stage_1_param_path, stage_1_result_path]        
-#----------------------------------------------------------------------------------------------------------------------------        
+    except Exception as e:
+        print(f"Error: {e} (known error, does not cause issues, fix later)")
+#------------------------------------------------------------------------------------------------------------------------------------- 
+@client.tree.command()
+@app_commands.describe(
+    prompt='Enter some text to chat with the bot! Like this: /falcon Hello, how are you?')
+async def falcon(interaction: discord.Interaction, prompt: str):
+    try:  
+        global chathistory # ?
+        global falcon_userid_threadid_dictionary # tracks userid-thread existence
+        global instructions
+        global threadid_conversation
+        global FALCON_CHANNEL_ID
+
+        if not interaction.user.bot:
+            if interaction.channel.id == FALCON_CHANNEL_ID: # 1079459939405279232= test
+                await interaction.response.send_message(f"Working on it!")
+                channel = interaction.channel
+                # 1
+                message = await channel.send(f"Creating thread...")
+                thread = await message.create_thread(name=f'{interaction.user}') # interaction.user
+                await thread.send(f"[DISCLAIMER: HuggingBot is a **highly experimental** beta feature; The Falcon " \
+                f"model and system prompt can be found here: https://huggingface.co/spaces/HuggingFaceH4/falcon-chat]")
+                # generation
+                chathistory = falconclient.predict( 
+                        fn_index=5
+                ) # []    
+                job = falconclient.submit(prompt, chathistory, instructions, 0.8, 0.9, fn_index=1)  # This is not blocking, similar to run_in_executor (but better)
+                while job.done() == False: 
+                    status = job.status() 
+                    #print(status)
+                else:
+                    file_paths = job.outputs()
+                    full_generation = file_paths[-1] # tmp12345678.json
+                with open(full_generation, 'r') as file:
+                    data = json.load(file)
+                    output_text = data[-1][-1] # we output this as the bot
+
+                threadid_conversation[thread.id] = full_generation
+                
+                falcon_userid_threadid_dictionary[thread.id] = interaction.user.id
+                
+                print(output_text)
+                await thread.send(f"{output_text}")  
+
+    except Exception as e:
+        print(f"Error: {e}")
+        #await thread.send(f"{e} cc <@811235357663297546> (falconprivate error)") 
+#------------------------------------------------------------------------------------------------------------------------------------- 
+async def continue_falcon(message):
+    try:
+        global chathistory # ?
+        global instructions
+        global threadid_conversation
+        await message.add_reaction('<a:loading:1121820108189339738>') # test=<a:loading:1121820108189339738>  hf=<a:loading:1114111677990981692>
+        chathistory = threadid_conversation[message.channel.id]
+        prompt = message.content
+        # generation
+        job = falconclient.submit(prompt, chathistory, instructions, 0.8, 0.9, fn_index=1)  # This is not blocking, similar to run_in_executor (but better)
+        while job.done() == False: 
+            status = job.status() 
+            #print(status)
+        else:
+            file_paths = job.outputs()
+            full_generation = file_paths[-1] # tmp12345678.json
+        with open(full_generation, 'r') as file:
+            data = json.load(file)
+            output_text = data[-1][-1] # we output this as the bot
+
+        threadid_conversation[message.channel.id] = full_generation # overwrite the old file
+        falcon_userid_threadid_dictionary[message.channel.id] = message.author.id
+        
+        print(output_text)
+        await message.reply(f"{output_text}")  
+        await message.remove_reaction('<a:loading:1121820108189339738>', client.user) # test=<a:loading:1121820108189339738>  hf=<a:loading:1114111677990981692>  
+
+    except Exception as e:
+        print(f"Error: {e}")
+        await message.reply(f"Error: {e} <@811235357663297546> (continue_falcon error)")
+#-------------------------------------------------------------------------------------------------------------------------------------
+@client.event
+async def on_message(message):
+    try:
+        if not message.author.bot:
+            global falcon_userid_threadid_dictionary # tracks userid-thread existence
+            if message.channel.id in falcon_userid_threadid_dictionary: # is this a valid thread?
+                if falcon_userid_threadid_dictionary[message.channel.id] == message.author.id: # more than that - is this specifically the right user for this thread?
+                    # call job for falcon
+                    #await message.reply("checks succeeded, calling continue_falcon")
+                    await continue_falcon(message)
+
+    except Exception as e:
+        print(f"Error: {e}")
+#-------------------------------------------------------------------------------------------------------------------------------------
 # running in thread
 DISCORD_TOKEN = os.environ.get("DISCORD_TOKEN", None)