working on arc

app.py

@@ -2,17 +2,296 @@ import torch
 import gradio as gr
 
 
-def doo(text_prompt):
-    return 
+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
+
+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'
+
+
+
+
+    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
+
+    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 = []
+    labels = []
+
+    class_to_node = '/fastscratch/mridul/fishes/class_to_ancestral_label.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()}
+
+    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_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(',')
+        # 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)]
+
+    all_samples=list()
+    with torch.no_grad():
+        with model.ema_scope():
+            uc = None
+            for n in trange(opt.n_iter, desc="Sampling"):
 
-iface = gr.Interface(
-    fn=doo,
-    inputs=gr.Textbox(label="Prompt"),
-    outputs=[
-        gr.Image(label="Generated Image"),
-    ] 
-)
+                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__":
-    iface.launch()
+    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()
+
+
+    def setup_interface():
+        with gr.Blocks() as demo:
+            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")
+                    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")
+
+            # 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)

sample_level_encoding.py

@@ -32,6 +32,16 @@ def load_model_from_config(config, ckpt, verbose=False):
     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
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -206,6 +216,9 @@ if __name__ == "__main__":
     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()}
+
+
     sample_path = os.path.join(outpath, opt.output_dir_name)
     os.makedirs(sample_path, exist_ok=True)
     base_count = len(os.listdir(sample_path))
@@ -223,6 +236,7 @@ if __name__ == "__main__":
                     all_prompts = opt.n_samples * (prompt)
                     all_prompts = [tuple(all_prompts)]
                     print(class_name, prompt)
+                    breakpoint()
                     c = model.get_learned_conditioning({'class_to_node': all_prompts})
                     shape = [3, 64, 64]
                     samples_ddim, _ = sampler.sample(S=opt.ddim_steps,