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 | |
| 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"): | |
| 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() | |
| 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) |