updated monkeys app and fish app

app.py

@@ -61,11 +61,6 @@ def generate_image(fish_name, masking_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'
-
-
-
 
     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', 
@@ -82,19 +77,13 @@ 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
-
-    # 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
+
 
     prompt = opt.prompt
     all_images = []
@@ -105,12 +94,11 @@ def generate_image(fish_name, masking_level_input,
         class_to_node_dict = pickle.load(pickle_file)
 
     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))
+    
 
     prompt = class_to_node_dict[fish_name]
+
+    ### Trait Swapping
     if swap_fish_name:
         swap_fish_name = swap_fish_name.lower()
         swap_level = int(swap_level_input.split(" ")[-1]) - 1
@@ -118,8 +106,6 @@ def generate_image(fish_name, masking_level_input,
 
         swap_fish_split = swap_fish[0].split(',')
         fish_name_split = prompt[0].split(',')
-        # print(swap_fish_split, fish_name_split)
-        # print(swap_level)
         fish_name_split[swap_level] = swap_fish_split[swap_level]
 
         prompt = [','.join(fish_name_split)]

monkeys_app.py

@@ -0,0 +1,345 @@
+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
+config_path = '/globalscratch/mridul/ldm/monkeys/2024-04-03T19-49-17_HLE_days1_lr1e-6_6levels/configs/2024-04-03T19-49-17-project.yaml'
+ckpt_path = '/globalscratch/mridul/ldm/monkeys/2024-04-03T19-49-17_HLE_days1_lr1e-6_6levels/checkpoints/epoch=000335.ckpt'
+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):
+    
+def generate_image(fish_name,
+                        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 = opt.prompt
+    all_images = []
+    labels = []
+
+    # class_to_node = '/fastscratch/mridul/fishes/class_to_ancestral_label.pkl'
+    class_to_node = '/projects/ml4science/mridul/data/monkeys_dataset/monkey_HLE_labels_6levels.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()}
+    
+
+    prompt = class_to_node_dict[fish_name]
+
+    ### Trait Swapping
+    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]
+
+        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=1,
+        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 Monkey 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("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_level_input = "None"
+
+                    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 5","Level 4","Level 3", "Level 2"], value="Level 5")
+                    submit_button = gr.Button("Generate")
+                    gr.Markdown("## Phylogeny Tree")
+                    architecture_image = "monkeys_masking_6levels.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 = [
+                    # ("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, 
+            #             swap_fish_name, swap_level_input],
+            #     outputs=output_image
+            # )
+
+
+
+            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)