Use cached generated images for example buttons

app.py

@@ -1,7 +1,6 @@
 import torch
 import gradio as gr
 
-
 import argparse, os, sys, glob
 import torch
 import pickle
@@ -16,6 +15,7 @@ from ldm.util import instantiate_from_config
 from ldm.models.diffusion.ddim import DDIMSampler
 from ldm.models.diffusion.plms import PLMSSampler
 
+CACHE_DIR = "examples"
 
 def load_model_from_config(config, ckpt, verbose=False):
     print(f"Loading model from {ckpt}")
@@ -206,33 +206,33 @@ if __name__ == "__main__":
             gr.Markdown("### Write the Species name to generate a fish image")
             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("## 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
                     experiment_choice = gr.Radio(label="Select Experiment", choices=["Trait Masking", "Trait Swapping"], value=None)
-    
+
                     # Trait Masking Inputs (hidden initially)
                     masking_level_input = gr.Dropdown(label="Select Ancestral Level", choices=["None", "Level 3", "Level 2"], value="None", visible=False)
-    
+
                     # Trait Swapping Inputs (hidden initially)
                     swap_fish_name = gr.Textbox(label="Species Name to swap trait with:", visible=False)
                     swap_level_input = gr.Dropdown(label="Level of swapping", choices=["Level 3", "Level 2"], value="Level 3", visible=False)
-    
+
                     submit_button = gr.Button("Generate")
-    
+
                     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", width=256, height=256)
-    
+
                     # Place to put example buttons
                     gr.Markdown("## Select an example:")
                     examples = [
@@ -240,29 +240,30 @@ if __name__ == "__main__":
                         ("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}")
-                            experiment_type = "None"
-                        elif level != "None":
-                            button = gr.Button(f"Species: {text} | Masking: {level}")
-                            experiment_type = "Trait Masking"
-                        elif swap_text != "":
-                            button = gr.Button(f"Species: {text} | Swapping with {swap_text} at {swap_level} ")
-                            experiment_type = "Trait Swapping"
-    
-                        # Update radio button, fields and auto-trigger the "Generate" action
+
+                    example_images = [
+                        "eg1.webp",
+                        "eg2.webp",
+                        "eg3.webp",
+                    ]
+
+                    for idx, (species, masking, swap_species, swap_level) in enumerate(examples):
+                        # Descriptive button text
+                        if masking != "None" and swap_species == "":
+                            button_text = f"Species: {species} | Masking: {masking}"
+                        elif swap_species:
+                            button_text = f"Species: {species} | Swapping with {swap_species} at {swap_level}"
+                        else:
+                            button_text = f"Species: {species}"
+
+                        # Create button
+                        button = gr.Button(button_text)
+
+                        # Attach the function to load cached images
                         button.click(
-                            fn=lambda text=text, level=level, swap_text=swap_text, swap_level=swap_level, experiment_type=experiment_type: (
-                                text, level, swap_text, swap_level, experiment_type
-                            ),
-                            inputs=[],
-                            outputs=[prompt_input, masking_level_input, swap_fish_name, swap_level_input, experiment_choice]
-                        ).then(
-                            fn=generate_image, 
-                            inputs=[prompt_input, masking_level_input, swap_fish_name, swap_level_input], 
-                            outputs=output_image
+                            fn=lambda index=idx: os.path.join(CACHE_DIR, example_images[index]),
+                            inputs=[],  # No inputs required
+                            outputs=output_image  # Display the cached image
                         )
     
             # Update visibility of inputs based on experiment selection