added butterflies app

butterflies_app.py

@@ -0,0 +1,330 @@
+import torch
+import gradio as gr
+
+
+import argparse, os, sys, glob
+import torch
+import pickle
+import numpy as np
+from omegaconf import OmegaConf
+from PIL import Image
+from tqdm import tqdm, trange
+from einops import rearrange
+from torchvision.utils import make_grid
+
+from ldm.util import instantiate_from_config
+from ldm.models.diffusion.ddim import DDIMSampler
+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")
+    sd = pl_sd["state_dict"]
+    model = instantiate_from_config(config.model)
+    m, u = model.load_state_dict(sd, strict=False)
+    if len(m) > 0 and verbose:
+        print("missing keys:")
+        print(m)
+    if len(u) > 0 and verbose:
+        print("unexpected keys:")
+        print(u)
+
+    model.cuda()
+    model.eval()
+    return model
+
+
+def masking_embed(embedding, levels=1):
+    """
+    size of embedding - nx1xd, n: number of samples, d - 512
+    replacing the last 128*levels from the embedding
+    """
+    replace_size = 128*levels
+    random_noise = torch.randn(embedding.shape[0], embedding.shape[1], replace_size)
+    embedding[:, :, -replace_size:] = random_noise
+    return embedding
+
+
+# LOAD MODEL GLOBALLY
+ckpt_path = '/globalscratch/mridul/ldm/butterflies/model_runs/2024-06-18T21-37-12_HLE_lr1e-6_custom_NEW/checkpoints/epoch=000233.ckpt'
+config_path = '/globalscratch/mridul/ldm/butterflies/model_runs/2024-06-18T21-37-12_HLE_lr1e-6_custom_NEW/configs/2024-06-18T21-37-12-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)
+
+class_to_node = '/projects/ml4science/mridul/data/cambridge_butterfly/level_encodings/butterflies_hle_4levels_custom_NEW.pkl'
+with open(class_to_node, 'rb') as pickle_file:
+    class_to_node_dict = pickle.load(pickle_file)
+
+class_to_node_dict =  {key.lower(): value for key, value in class_to_node_dict.items()}
+species_name_to_class = {'_'.join(x.split('_')[2:]):x for x in class_to_node_dict.keys()}
+
+species_names = list(species_name_to_class.keys())
+
+def generate_image(fish_name, masking_level_input,
+                        swap_fish_name, swap_level_input):
+    
+    # fish_name = fish_name.lower()
+
+
+    # label_to_class_mapping = {0: 'Alosa-chrysochloris', 1: 'Carassius-auratus', 2: 'Cyprinus-carpio', 3: 'Esox-americanus', 
+    # 4: 'Gambusia-affinis', 5: 'Lepisosteus-osseus', 6: 'Lepisosteus-platostomus', 7: 'Lepomis-auritus', 8: 'Lepomis-cyanellus', 
+    # 9: 'Lepomis-gibbosus', 10: 'Lepomis-gulosus', 11: 'Lepomis-humilis', 12: 'Lepomis-macrochirus', 13: 'Lepomis-megalotis', 
+    # 14: 'Lepomis-microlophus', 15: 'Morone-chrysops', 16: 'Morone-mississippiensis', 17: 'Notropis-atherinoides', 
+    # 18: 'Notropis-blennius', 19: 'Notropis-boops', 20: 'Notropis-buccatus', 21: 'Notropis-buchanani', 22: 'Notropis-dorsalis', 
+    # 23: 'Notropis-hudsonius', 24: 'Notropis-leuciodus', 25: 'Notropis-nubilus', 26: 'Notropis-percobromus', 
+    # 27: 'Notropis-stramineus', 28: 'Notropis-telescopus', 29: 'Notropis-texanus', 30: 'Notropis-volucellus', 
+    # 31: 'Notropis-wickliffi', 32: 'Noturus-exilis', 33: 'Noturus-flavus', 34: 'Noturus-gyrinus', 35: 'Noturus-miurus', 
+    # 36: 'Noturus-nocturnus', 37: 'Phenacobius-mirabilis'}
+
+    # def get_label_from_class(class_name):
+    #     for key, value in label_to_class_mapping.items():
+    #         if value == class_name:
+    #             return key
+
+
+    if opt.plms:
+        sampler = PLMSSampler(model)
+    else:
+        sampler = DDIMSampler(model)
+
+
+    prompt = class_to_node_dict[species_name_to_class[fish_name]]
+
+    ### Trait Swapping
+    if swap_fish_name!='None':
+        # swap_fish_name = swap_fish_name.lower()
+        swap_level = int(swap_level_input.split(" ")[-1]) - 1
+        swap_fish = class_to_node_dict[species_name_to_class[swap_fish_name]]
+
+        swap_fish_split = swap_fish[0].split(',')
+        fish_name_split = prompt[0].split(',')
+        fish_name_split[swap_level] = swap_fish_split[swap_level]
+
+        prompt = [','.join(fish_name_split)]
+
+    all_samples=list()
+    with torch.no_grad():
+        with model.ema_scope():
+            uc = None
+            for n in trange(opt.n_iter, desc="Sampling"):
+
+                all_prompts = opt.n_samples * (prompt)
+                all_prompts = [tuple(all_prompts)]
+                c = model.get_learned_conditioning({'class_to_node': all_prompts})
+                if masking_level_input != "None":
+                    masked_level = int(masking_level_input.split(" ")[-1])
+                    masked_level = 4-masked_level
+                    c = masking_embed(c, levels=masked_level)
+                shape = [3, 64, 64]
+                samples_ddim, _ = sampler.sample(S=opt.ddim_steps,
+                                                conditioning=c,
+                                                batch_size=opt.n_samples,
+                                                shape=shape,
+                                                verbose=False,
+                                                unconditional_guidance_scale=opt.scale,
+                                                unconditional_conditioning=uc,
+                                                eta=opt.ddim_eta)
+
+                x_samples_ddim = model.decode_first_stage(samples_ddim)
+                x_samples_ddim = torch.clamp((x_samples_ddim+1.0)/2.0, min=0.0, max=1.0)
+
+                all_samples.append(x_samples_ddim)
+
+    ###### to make grid
+    # additionally, save as grid
+    grid = torch.stack(all_samples, 0)
+    grid = rearrange(grid, 'n b c h w -> (n b) c h w')
+    grid = make_grid(grid, nrow=opt.n_samples)
+
+    # 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
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    # parser.add_argument(
+    #     "--prompt",
+    #     type=str,
+    #     nargs="?",
+    #     default="a painting of a virus monster playing guitar",
+    #     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,
+        default=200,
+        help="number of ddim sampling steps",
+    )
+
+    parser.add_argument(
+        "--plms",
+        action='store_true',
+        help="use plms sampling",
+    )
+
+    parser.add_argument(
+        "--ddim_eta",
+        type=float,
+        default=1.0,
+        help="ddim eta (eta=0.0 corresponds to deterministic sampling",
+    )
+    parser.add_argument(
+        "--n_iter",
+        type=int,
+        default=1,
+        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,
+        default=3,
+        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,
+        # default=5.0,
+        default=1.0,
+        help="unconditional guidance scale: eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))",
+    )
+    opt = parser.parse_args()
+
+    title = "🎞️ Phylo Diffusion - Generating Butterfly 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 Butterfly Images Tool")
+            gr.Markdown("### Write the Species name to generate a butterfly 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("Select a species to generate an image:")
+                    # prompt_input = gr.Textbox(label="Species Name")
+                    prompt_input = gr.Dropdown(label="Select Butterfly", choices=species_names, value="None")
+                    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.Dropdown(label="Select species Name to swap trait with:", choices=species_names, value="None")
+                        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")
+                    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=768, 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")]
+
+                    # 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
+
+
+            submit_button.click(
+                fn=generate_image,
+                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"),
+    # ] 
+    # )
+
+    iface = setup_interface()
+
+    iface.launch(share=True)