updated the app

app.py

@@ -46,13 +46,23 @@ def masking_embed(embedding, levels=1):
     embedding[:, :, -replace_size:] = random_noise
     return embedding
 
+
+# LOAD MODEL GLOBALLY
+ckpt_path = '/globalscratch/mridul/ldm/final_runs_eccv/fishes/2024-03-01T23-15-36_HLE_days3/checkpoints/epoch=000119.ckpt'
+config_path = '/globalscratch/mridul/ldm/final_runs_eccv/fishes/2024-03-01T23-15-36_HLE_days3/configs/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
+
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+model = model.to(device)
+
 def generate_image(fish_name, masking_level_input,
                         swap_fish_name, swap_level_input):
     
     fish_name = fish_name.lower()
 
-    ckpt_path = '/globalscratch/mridul/ldm/final_runs_eccv/fishes/2024-03-01T23-15-36_HLE_days3/checkpoints/epoch=000119.ckpt'
-    config_path = '/globalscratch/mridul/ldm/final_runs_eccv/fishes/2024-03-01T23-15-36_HLE_days3/configs/2024-03-01T23-15-36-project.yaml'
+    # ckpt_path = '/globalscratch/mridul/ldm/final_runs_eccv/fishes/2024-03-01T23-15-36_HLE_days3/checkpoints/epoch=000119.ckpt'
+    # config_path = '/globalscratch/mridul/ldm/final_runs_eccv/fishes/2024-03-01T23-15-36_HLE_days3/configs/2024-03-01T23-15-36-project.yaml'
 
 
 
@@ -72,19 +82,19 @@ def generate_image(fish_name, masking_level_input,
             if value == class_name:
                 return key
 
-    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)  # TODO: Optionally download from same location as ckpt and chnage this logic
+    # model = load_model_from_config(config, ckpt_path)  # TODO: check path
 
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-    model = model.to(device)
+    # device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    # model = model.to(device)
 
     if opt.plms:
         sampler = PLMSSampler(model)
     else:
         sampler = DDIMSampler(model)
 
-    os.makedirs(opt.outdir, exist_ok=True)
-    outpath = opt.outdir
+    # os.makedirs(opt.outdir, exist_ok=True)
+    # outpath = opt.outdir
 
     prompt = opt.prompt
     all_images = []
@@ -96,12 +106,13 @@ def generate_image(fish_name, masking_level_input,
 
     class_to_node_dict =  {key.lower(): value for key, value in class_to_node_dict.items()}
 
-    sample_path = os.path.join(outpath, opt.output_dir_name)
-    os.makedirs(sample_path, exist_ok=True)
-    base_count = len(os.listdir(sample_path))
+    # sample_path = os.path.join(outpath, opt.output_dir_name)
+    # os.makedirs(sample_path, exist_ok=True)
+    # base_count = len(os.listdir(sample_path))
 
     prompt = class_to_node_dict[fish_name]
     if swap_fish_name:
+        swap_fish_name = swap_fish_name.lower()
         swap_level = int(swap_level_input.split(" ")[-1]) - 1
         swap_fish = class_to_node_dict[swap_fish_name]
 
@@ -243,12 +254,26 @@ 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
+        components['masking_level_input'].value = level
+        # components['option'].value = option
 
     def setup_interface():
         with gr.Blocks() as demo:
+
+            gr.Markdown("# Phylo Diffusion - Generating Fish Images Tool")
+            gr.Markdown("### Write the Species name to generate a fish image")
+            gr.Markdown("### 1. Trait Masking: Specify the Level information as well")
+            gr.Markdown("### 2. Trait Swapping: Specify the species name to swap trait with at also at what level")
+
             with gr.Row():
                 with gr.Column():
-                    gr.Markdown("### Generate Images Based on Prompts")
+                    gr.Markdown("## Generate Images Based on Prompts")
                     gr.Markdown("Enter a prompt to generate an image:")
                     prompt_input = gr.Textbox(label="Species Name")
                     gr.Markdown("Trait Masking")
@@ -261,13 +286,37 @@ if __name__ == "__main__":
                         swap_fish_name = gr.Textbox(label="Species Name to swap trait with:")
                         swap_level_input = gr.Dropdown(label="Level of swapping", choices=["Level 3", "Level 2"], value="Level 3")
                     submit_button = gr.Button("Generate")
-                    gr.Markdown("### Phylogeny Tree")
+                    gr.Markdown("## Phylogeny Tree")
                     architecture_image = "phylogeny_tree.jpg"  # Update this with the actual path
                     gr.Image(value=architecture_image, label="Phylogeny Tree")
 
                 with gr.Column():
-                    gr.Markdown("### Generated Image")
-                    output_image = gr.Image(label="Generated Image")
+
+                    gr.Markdown("## Generated Image")
+                    output_image = gr.Image(label="Generated Image",  width=512, height=512)
+
+
+                    # Place to put example buttons
+                    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")]
+
+                    for text, level, swap_text, swap_level in examples:
+                        if level == "None" and swap_text == "":
+                            button = gr.Button(f"Species: {text}")
+                        elif level != "None":
+                            button = gr.Button(f"Species: {text} | Masking: {level}")
+                        elif swap_text != "":
+                            button = gr.Button(f"Species: {text} | Swapping with {swap_text} at {swap_level} ")
+                        button.click(
+                            fn=lambda text=text, level=level, swap_text=swap_text, swap_level=swap_level: (text, level, swap_text, swap_level),
+                            inputs=[],
+                            outputs=[prompt_input, masking_level_input, swap_fish_name, swap_level_input]
+                        )
+                
 
             # Display an image of the architecture