Spaces:
Sleeping
Sleeping
| 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) |