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 = './model_files/fishes/epoch=000119.ckpt' config_path = './model_files/fishes/2024-03-01T23-15-36-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) 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 = opt.prompt all_images = [] labels = [] class_to_node = './model_files/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()} 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)) 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( "--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( "--n_samples", type=int, default=1, help="how many samples to produce for the given prompt", ) 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 Fish 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 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("### 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") # 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 = [ ("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, 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 ) # 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], outputs=output_image ) return demo iface = setup_interface() iface.launch(share=True)