Update deepfloydif.py

deepfloydif.py

@@ -37,6 +37,26 @@ def deepfloydif_stage_2_inference(index, stage_1_result_path):
     result_path = deepfloydif_client.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 react_1234(reaction_emojis, combined_image_dfif):
+    """Sets up 4 reaction emojis so the user can choose an image to upscale for deepfloydif"""
+    for emoji in reaction_emojis:
+        await combined_image_dfif.add_reaction(emoji) 
+
+def load_image(png_files, stage_1_results, png_file_index):
+    """Opens images as variables so we can combine them later"""
+    for file in png_files:
+        png_file = png_files[png_file_index]
+        png_path = os.path.join(stage_1_results, png_file)
+        if png_file_index == 0:
+            img1 = Image.open(png_path)
+        if png_file_index == 1:
+            img2 = Image.open(png_path)
+        if png_file_index == 2:
+            img3 = Image.open(png_path)
+        if png_file_index == 3:
+            img4 = Image.open(png_path)
+        png_file_index = png_file_index + 1
+    return [img1, img2, img3, img4]
     
 async def deepfloydif_stage_1(interaction, prompt, client):
     """DeepfloydIF command (generate images with realistic text using slash commands)"""
@@ -89,55 +109,6 @@ async def deepfloydif_stage_1(interaction, prompt, client):
 
     except Exception as e:
         print(f"Error: {e}")
-  
-def load_image(png_files, stage_1_results, png_file_index):
-    """Opens images as variables so we can combine them later"""
-    for file in png_files:
-        png_file = png_files[png_file_index]
-        png_path = os.path.join(stage_1_results, png_file)
-        if png_file_index == 0:
-            img1 = Image.open(png_path)
-        if png_file_index == 1:
-            img2 = Image.open(png_path)
-        if png_file_index == 2:
-            img3 = Image.open(png_path)
-        if png_file_index == 3:
-            img4 = Image.open(png_path)
-        png_file_index = png_file_index + 1
-    return [img1, img2, img3, img4]
-
-async def react_1234(reaction_emojis, combined_image_dfif):
-    """Sets up 4 reaction emojis so the user can choose an image to upscale for deepfloydif"""
-    for emoji in reaction_emojis:
-        await combined_image_dfif.add_reaction(emoji) 
-
-async def deepfloydif_stage_2(index: int, stage_1_result_path, thread, dfif_command_message_id): 
-    """upscaling function for images generated using /deepfloydif"""
-    try:
-        if os.environ.get('TEST_ENV') == 'True':
-            print("Running deepfloydif_stage_2")        
-        parent_channel = thread.parent
-        dfif_command_message = await parent_channel.fetch_message(dfif_command_message_id)
-        if index == 0:
-            position = "top left"
-        elif index == 1:
-            position = "top right"
-        elif index == 2:
-            position = "bottom left"
-        elif index == 3:
-            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, deepfloydif_stage_2_inference, index, stage_1_result_path) 
-        
-        with open(result_path, 'rb') as f:
-            await thread.send('Here is the upscaled image!', file=discord.File(f, 'result.png'))
-        await thread.edit(archived=True)
-
-    except Exception as e:
-        print(f"Error: {e}")
 
 async def deepfloydif_stage_2_react_check(reaction, user): 
     """Checks for a reaction in order to call dfif2"""
@@ -173,4 +144,31 @@ async def deepfloydif_stage_2_react_check(reaction, user):
 
     except Exception as e:
         print(f"Error: {e} (known error, does not cause issues, low priority)")
-                               
+  
+async def deepfloydif_stage_2(index: int, stage_1_result_path, thread, dfif_command_message_id): 
+    """upscaling function for images generated using /deepfloydif"""
+    try:
+        if os.environ.get('TEST_ENV') == 'True':
+            print("Running deepfloydif_stage_2")        
+        parent_channel = thread.parent
+        dfif_command_message = await parent_channel.fetch_message(dfif_command_message_id)
+        if index == 0:
+            position = "top left"
+        elif index == 1:
+            position = "top right"
+        elif index == 2:
+            position = "bottom left"
+        elif index == 3:
+            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, deepfloydif_stage_2_inference, index, stage_1_result_path) 
+        
+        with open(result_path, 'rb') as f:
+            await thread.send('Here is the upscaled image!', file=discord.File(f, 'result.png'))
+        await thread.edit(archived=True)
+
+    except Exception as e:
+        print(f"Error: {e}")