added the note for generation taking time

app.py

@@ -19,7 +19,7 @@ from ldm.models.diffusion.plms import PLMSSampler
 
 def load_model_from_config(config, ckpt, verbose=False):
     print(f"Loading model from {ckpt}")
-    pl_sd = torch.load(ckpt, map_location="cpu")
+    pl_sd = torch.load(ckpt, map_location="cpu") # TODO: change for GPU resources
     # pl_sd = torch.load(ckpt)#, map_location="cpu")
     sd = pl_sd["state_dict"]
     model = instantiate_from_config(config.model)
@@ -31,7 +31,7 @@ def load_model_from_config(config, ckpt, verbose=False):
         print("unexpected keys:")
         print(u)
 
-    # model.cuda()
+    # model.cuda() # TODO: change for GPU resources
     model.eval()
     return model
 
@@ -50,8 +50,8 @@ def masking_embed(embedding, levels=1):
 # LOAD MODEL GLOBALLY
 ckpt_path = './model_files/fishes/epoch=000119.ckpt'
 config_path = './model_files/fishes/2024-03-01T23-15-36-project.yaml'
-config = OmegaConf.load(config_path)  # TODO: Optionally download from same location as ckpt and chnage this logic
-model = load_model_from_config(config, ckpt_path)  # TODO: check path
+config = OmegaConf.load(config_path)
+model = load_model_from_config(config, ckpt_path) 
 
 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 model = model.to(device)
@@ -78,12 +78,7 @@ def generate_image(fish_name, masking_level_input,
                 return key
 
 
-    if opt.plms:
-        sampler = PLMSSampler(model)
-    else:
-        sampler = DDIMSampler(model)
-
-
+    sampler = DDIMSampler(model)
 
     prompt = opt.prompt
     all_images = []
@@ -169,12 +164,6 @@ if __name__ == "__main__":
         help="number of ddim sampling steps",
     )
 
-    parser.add_argument(
-        "--plms",
-        action='store_true',
-        help="use plms sampling",
-    )
-
     parser.add_argument(
         "--ddim_eta",
         type=float,
@@ -205,8 +194,6 @@ if __name__ == "__main__":
     opt = parser.parse_args()
 
     title = "🎞️ Phylo Diffusion - Generating Fish Images Tool"
-    description = "Write the Species name to generate an image for.\n For Trait Masking: Specify the Level information as well"
-
 
     def load_example(prompt, level, option, components):
         components['prompt_input'].value = prompt
@@ -214,14 +201,16 @@ if __name__ == "__main__":
     
     def setup_interface():
         with gr.Blocks() as demo:
-            gr.Markdown("# Phylo Diffusion - Generating Fish Images Tool")
+
+            gr.Markdown("# Phylo-Diffusion: Generating Fish Images Tool")
             gr.Markdown("### Write the Species name to generate a fish image")
-            gr.Markdown("### Select one of the experiments: Trait Masking or Trait Swapping")
+            gr.Markdown("### 1. Trait Masking: Specify the Level information to mask")
+            gr.Markdown("### 2. Trait Swapping: Specify the species name to swap trait with and at what level")
     
             with gr.Row():
                 with gr.Column():
-                    gr.Markdown("## Generate Images Based on Prompts")
-                    gr.Markdown("Enter a prompt to generate an image:")
+                    # gr.Markdown("## Generate Images Based on Prompts")
+                    gr.Markdown("**NOTE:** The demo is currently running on free CPU resources provided by Hugging Face, so it may take up to 10 minutes to generate an image. We're working on securing additional resources to speed up the process. Thank you for your patience!")
                     prompt_input = gr.Textbox(label="Species Name")
     
                     # Radio button to select experiment type, with no default selection
@@ -248,7 +237,6 @@ if __name__ == "__main__":
                     gr.Markdown("## Select an example:")
                     examples = [
                         ("Gambusia Affinis", "None", "", "Level 3"),
-                        ("Lepomis Auritus", "None", "", "Level 3"),
                         ("Lepomis Auritus", "Level 3", "", "Level 3"),
                         ("Noturus nocturnus", "None", "Notropis dorsalis", "Level 2")
                     ]