updated app for making it a radio button to select between the experiment

app.py

@@ -147,7 +147,6 @@ def generate_image(fish_name, masking_level_input,
     # to image
     grid = 255. * rearrange(grid, 'c h w -> h w c').cpu().numpy()
     final_image = Image.fromarray(grid.astype(np.uint8))
-    # final_image.save(os.path.join(sample_path, f'{class_name.replace(" ", "-")}.png'))
 
     return final_image
 
@@ -163,13 +162,6 @@ if __name__ == "__main__":
         help="the prompt to render"
     )
 
-    parser.add_argument(
-        "--outdir",
-        type=str,
-        nargs="?",
-        help="dir to write results to",
-        default="outputs/txt2img-samples"
-    )
     parser.add_argument(
         "--ddim_steps",
         type=int,
@@ -196,20 +188,6 @@ if __name__ == "__main__":
         help="sample this often",
     )
 
-    parser.add_argument(
-        "--H",
-        type=int,
-        default=256,
-        help="image height, in pixel space",
-    )
-
-    parser.add_argument(
-        "--W",
-        type=int,
-        default=256,
-        help="image width, in pixel space",
-    )
-
     parser.add_argument(
         "--n_samples",
         type=int,
@@ -217,20 +195,6 @@ if __name__ == "__main__":
         help="how many samples to produce for the given prompt",
     )
 
-    parser.add_argument(
-        "--output_dir_name",
-        type=str,
-        default='default_file',
-        help="name of folder",
-    )
-
-    parser.add_argument(
-        "--postfix",
-        type=str,
-        default='',
-        help="name of folder",
-    )
-
     parser.add_argument(
         "--scale",
         type=float,
@@ -247,86 +211,97 @@ if __name__ == "__main__":
     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")
-
+            gr.Markdown("### Select one of the experiments: Trait Masking or Trait Swapping")
+    
             with gr.Row():
                 with gr.Column():
                     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")
-                    with gr.Row():
-                        masking_level_input = gr.Dropdown(label="Select Ancestral Level", choices=["None", "Level 3", "Level 2"], value="None")
-                        # masking_node_input = gr.Dropdown(label="Select Internal", choices=["0", "1", "2", "3", "4", "5", "6", "7", "8"], value="0")
-
-                    gr.Markdown("Trait Swapping")
-                    with gr.Row():
-                        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")
+    
+                    # 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=512, height=512)
-
-
+                    output_image = gr.Image(label="Generated Image", width=256, height=256)
+    
                     # 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")]
-
+                        ("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}")
+                            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
                         button.click(
-                            fn=lambda text=text, level=level, swap_text=swap_text, swap_level=swap_level: (text, level, swap_text, swap_level),
+                            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]
+                            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
                         )
-                
-
-            # Display an image of the architecture
-
-
+    
+            # Update visibility of inputs based on experiment selection
+            def update_inputs(experiment_type):
+                if experiment_type == "Trait Masking":
+                    return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
+                elif experiment_type == "Trait Swapping":
+                    return gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
+                else:
+                    # No experiment selected
+                    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+    
+            experiment_choice.change(
+                fn=update_inputs,
+                inputs=[experiment_choice],
+                outputs=[masking_level_input, swap_fish_name, swap_level_input]
+            )
+    
+            # Submit button functionality
             submit_button.click(
                 fn=generate_image,
-                inputs=[prompt_input, masking_level_input,
-                        swap_fish_name, swap_level_input],
+                inputs=[prompt_input, masking_level_input, swap_fish_name, swap_level_input],
                 outputs=output_image
             )
-
-        return demo
-
-    # # Launch the interface
-    # iface = setup_interface()
     
-        # iface = gr.Interface(
-    # fn=generate_image,
-    # inputs=gr.Textbox(label="Prompt"),
-    # outputs=[
-    #     gr.Image(label="Generated Image"),
-    # ] 
-    # )
+        return demo
+        
 
     iface = setup_interface()
-
     iface.launch(share=True)