Update deepfloydif.py

deepfloydif.py

@@ -10,12 +10,8 @@ import pathlib
 HF_TOKEN = os.getenv("HF_TOKEN")
 deepfloydif_client = Client("huggingface-projects/IF", HF_TOKEN)
 
-BOT_USER_ID = (
-    1086256910572986469 if os.getenv("TEST_ENV", False) else 1102236653545861151
-)
-DEEPFLOYDIF_CHANNEL_ID = (
-    1121834257959092234 if os.getenv("TEST_ENV", False) else 1119313215675973714
-)
+BOT_USER_ID = 1086256910572986469 if os.getenv("TEST_ENV", False) else 1102236653545861151
+DEEPFLOYDIF_CHANNEL_ID = 1121834257959092234 if os.getenv("TEST_ENV", False) else 1119313215675973714
 
 
 def deepfloydif_stage_1_inference(prompt):
@@ -27,9 +23,9 @@ def deepfloydif_stage_1_inference(prompt):
     custom_timesteps_1 = "smart50"
     number_of_inference_steps = 50
     (
-        stage_1_results,
+        stage_1_images,
         stage_1_param_path,
-        stage_1_result_path,
+        path_for_stage_2_upscaling,
     ) = deepfloydif_client.predict(
         prompt,
         negative_prompt,
@@ -40,10 +36,10 @@ def deepfloydif_stage_1_inference(prompt):
         number_of_inference_steps,
         api_name="/generate64",
     )
-    return [stage_1_results, stage_1_param_path, stage_1_result_path]
+    return [stage_1_images, stage_1_param_path, path_for_stage_2_upscaling]
 
 
-def deepfloydif_stage_2_inference(index, stage_1_result_path):
+def deepfloydif_stage_2_inference(index, path_for_stage_2_upscaling):
     """Upscales one of the images from deepfloydif_stage_1_inference based on the chosen index"""
     selected_index_for_stage_2 = index
     seed_2 = 0
@@ -51,7 +47,7 @@ def deepfloydif_stage_2_inference(index, stage_1_result_path):
     custom_timesteps_2 = "smart50"
     number_of_inference_steps_2 = 50
     result_path = deepfloydif_client.predict(
-        stage_1_result_path,
+        path_for_stage_2_upscaling,
         selected_index_for_stage_2,
         seed_2,
         guidance_scale_2,
@@ -68,64 +64,56 @@ async def react_1234(reaction_emojis, combined_image_dfif):
         await combined_image_dfif.add_reaction(emoji)
 
 
-def load_image(png_files, stage_1_results):
+def load_image(png_files, stage_1_images):
     """Opens images as variables so we can combine them later"""
     results = []
     for file in png_files:
-        png_path = os.path.join(stage_1_results, file)
+        png_path = os.path.join(stage_1_images, file)
         results.append(Image.open(png_path))
     return results
 
 
+def combine_images(png_files, stage_1_images, partial_path):
+    if os.environ.get("TEST_ENV") == "True":
+        print("Combining images for deepfloydif_stage_1")
+    images = load_image(png_files, stage_1_images)
+    combined_image = Image.new("RGB", (images[0].width * 2, images[0].height * 2))
+    combined_image.paste(images[0], (0, 0))
+    combined_image.paste(images[1], (images[0].width, 0))
+    combined_image.paste(images[2], (0, images[0].height))
+    combined_image.paste(images[3], (images[0].width, images[0].height))
+    combined_image_path = os.path.join(stage_1_images, f"{partial_path}.png")
+    combined_image.save(combined_image_path)
+
+
 async def deepfloydif_stage_1(interaction, prompt, client):
     """DeepfloydIF command (generate images with realistic text using slash commands)"""
     try:
-        # global BOT_USER_ID
-        # global DEEPFLOYDIF_CHANNEL_ID
         if interaction.user.id != BOT_USER_ID:
             if interaction.channel.id == DEEPFLOYDIF_CHANNEL_ID:
                 if os.environ.get("TEST_ENV") == "True":
-                    print("Safetychecks passed for deepfloydif_stage_1")
+                    print("Safety checks passed for deepfloydif_stage_1")
                 await interaction.response.send_message("Working on it!")
                 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")
-                thread = await message.create_thread(
-                    name=f"{prompt}", auto_archive_duration=60
-                )
+                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"
-                )
-                await thread.send(
-                    f"{interaction.user.mention} Generating images in thread, can take ~1 minute..."
+                    "[DISCLAIMER: HuggingBot is a **highly experimental** beta feature; Additional information on the"
+                    " DeepfloydIF model can be found here: https://huggingface.co/spaces/DeepFloyd/IF"
                 )
+                await thread.send(f"{interaction.user.mention} Generating images in thread, can take ~1 minute...")
 
                 loop = asyncio.get_running_loop()
-                result = await loop.run_in_executor(
-                    None, deepfloydif_stage_1_inference, prompt
-                )
-                stage_1_results = result[0]
-                stage_1_result_path = result[2]
+                result = await loop.run_in_executor(None, deepfloydif_stage_1_inference, prompt)
+                stage_1_images = result[0]
+                path_for_stage_2_upscaling = result[2]
 
-                partial_path = pathlib.Path(stage_1_result_path).name
-                png_files = list(glob.glob(f"{stage_1_results}/**/*.png"))
+                partial_path = pathlib.Path(path_for_stage_2_upscaling).name
+                png_files = list(glob.glob(f"{stage_1_images}/**/*.png"))
 
                 if png_files:
-                    # take all 4 images and combine them into one large 2x2 image (similar to Midjourney)
-                    if os.environ.get("TEST_ENV") == "True":
-                        print("Combining images for deepfloydif_stage_1")
-                    images = load_image(png_files, stage_1_results)
-                    combined_image = Image.new(
-                        "RGB", (images[0].width * 2, images[0].height * 2)
-                    )
-                    combined_image.paste(images[0], (0, 0))
-                    combined_image.paste(images[1], (images[0].width, 0))
-                    combined_image.paste(images[2], (0, images[0].height))
-                    combined_image.paste(images[3], (images[0].width, images[0].height))
-                    combined_image_path = os.path.join(
-                        stage_1_results, f"{partial_path}.png"
-                    )
-                    combined_image.save(combined_image_path)
+                    combine_images(png_files, stage_1_images, partial_path)
                     if os.environ.get("TEST_ENV") == "True":
                         print("Images combined for deepfloydif_stage_1")
                     with open(combined_image_path, "rb") as f:
@@ -136,9 +124,7 @@ async def deepfloydif_stage_1(interaction, prompt, client):
                     emoji_list = ["↖️", "↗️", "↙️", "↘️"]
                     await react_1234(emoji_list, combined_image_dfif)
                 else:
-                    await thread.send(
-                        f"{interaction.user.mention} No PNG files were found, cannot post them!"
-                    )
+                    await thread.send(f"{interaction.user.mention} No PNG files were found, cannot post them!")
     except Exception as e:
         print(f"Error: {e}")
 
@@ -169,18 +155,18 @@ async def deepfloydif_stage_2_react_check(reaction, user):
                             index = 2
                         elif emoji == "↘️":
                             index = 3
-                        stage_1_result_path = full_path
+                        path_for_stage_2_upscaling = full_path
                         thread = reaction.message.channel
                         await deepfloydif_stage_2(
                             index,
-                            stage_1_result_path,
+                            path_for_stage_2_upscaling,
                             thread,
                         )
     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):
+async def deepfloydif_stage_2(index: int, path_for_stage_2_upscaling, thread):
     """upscaling function for images generated using /deepfloydif"""
     try:
         if os.environ.get("TEST_ENV") == "True":
@@ -198,13 +184,11 @@ async def deepfloydif_stage_2(index: int, stage_1_result_path, thread):
         # 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
+            None, deepfloydif_stage_2_inference, index, path_for_stage_2_upscaling
         )
 
         with open(result_path, "rb") as f:
-            await thread.send(
-                "Here is the upscaled image!", file=discord.File(f, "result.png")
-            )
+            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}")