from torchvision.datasets.utils import download_url from ldm.util import instantiate_from_config import torch import os # todo ? from google.colab import files from IPython.display import Image as ipyimg import ipywidgets as widgets from PIL import Image from numpy import asarray from einops import rearrange, repeat import torch, torchvision from ldm.models.diffusion.ddim import DDIMSampler from ldm.util import ismap import time from omegaconf import OmegaConf def download_models(mode): if mode == "superresolution": # this is the small bsr light model url_conf = 'https://heibox.uni-heidelberg.de/f/31a76b13ea27482981b4/?dl=1' url_ckpt = 'https://heibox.uni-heidelberg.de/f/578df07c8fc04ffbadf3/?dl=1' path_conf = 'logs/diffusion/superresolution_bsr/configs/project.yaml' path_ckpt = 'logs/diffusion/superresolution_bsr/checkpoints/last.ckpt' download_url(url_conf, path_conf) download_url(url_ckpt, path_ckpt) path_conf = path_conf + '/?dl=1' # fix it path_ckpt = path_ckpt + '/?dl=1' # fix it return path_conf, path_ckpt else: raise NotImplementedError def load_model_from_config(config, ckpt): print(f"Loading model from {ckpt}") pl_sd = torch.load(ckpt, map_location="cpu") global_step = pl_sd["global_step"] sd = pl_sd["state_dict"] model = instantiate_from_config(config.model) m, u = model.load_state_dict(sd, strict=False) model.cuda() model.eval() return {"model": model}, global_step def get_model(mode): path_conf, path_ckpt = download_models(mode) config = OmegaConf.load(path_conf) model, step = load_model_from_config(config, path_ckpt) return model def get_custom_cond(mode): dest = "data/example_conditioning" if mode == "superresolution": uploaded_img = files.upload() filename = next(iter(uploaded_img)) name, filetype = filename.split(".") # todo assumes just one dot in name ! os.rename(f"{filename}", f"{dest}/{mode}/custom_{name}.{filetype}") elif mode == "text_conditional": w = widgets.Text(value='A cake with cream!', disabled=True) display(w) with open(f"{dest}/{mode}/custom_{w.value[:20]}.txt", 'w') as f: f.write(w.value) elif mode == "class_conditional": w = widgets.IntSlider(min=0, max=1000) display(w) with open(f"{dest}/{mode}/custom.txt", 'w') as f: f.write(w.value) else: raise NotImplementedError(f"cond not implemented for mode{mode}") def get_cond_options(mode): path = "data/example_conditioning" path = os.path.join(path, mode) onlyfiles = [f for f in sorted(os.listdir(path))] return path, onlyfiles def select_cond_path(mode): path = "data/example_conditioning" # todo path = os.path.join(path, mode) onlyfiles = [f for f in sorted(os.listdir(path))] selected = widgets.RadioButtons( options=onlyfiles, description='Select conditioning:', disabled=False ) display(selected) selected_path = os.path.join(path, selected.value) return selected_path def get_cond(mode, selected_path): example = dict() if mode == "superresolution": up_f = 4 visualize_cond_img(selected_path) c = Image.open(selected_path) c = torch.unsqueeze(torchvision.transforms.ToTensor()(c), 0) c_up = torchvision.transforms.functional.resize(c, size=[up_f * c.shape[2], up_f * c.shape[3]], antialias=True) c_up = rearrange(c_up, '1 c h w -> 1 h w c') c = rearrange(c, '1 c h w -> 1 h w c') c = 2. * c - 1. c = c.to(torch.device("cuda")) example["LR_image"] = c example["image"] = c_up return example def visualize_cond_img(path): display(ipyimg(filename=path)) def run(model, selected_path, task, custom_steps, resize_enabled=False, classifier_ckpt=None, global_step=None): example = get_cond(task, selected_path) save_intermediate_vid = False n_runs = 1 masked = False guider = None ckwargs = None mode = 'ddim' ddim_use_x0_pred = False temperature = 1. eta = 1. make_progrow = True custom_shape = None height, width = example["image"].shape[1:3] split_input = height >= 128 and width >= 128 if split_input: ks = 128 stride = 64 vqf = 4 # model.split_input_params = {"ks": (ks, ks), "stride": (stride, stride), "vqf": vqf, "patch_distributed_vq": True, "tie_braker": False, "clip_max_weight": 0.5, "clip_min_weight": 0.01, "clip_max_tie_weight": 0.5, "clip_min_tie_weight": 0.01} else: if hasattr(model, "split_input_params"): delattr(model, "split_input_params") invert_mask = False x_T = None for n in range(n_runs): if custom_shape is not None: x_T = torch.randn(1, custom_shape[1], custom_shape[2], custom_shape[3]).to(model.device) x_T = repeat(x_T, '1 c h w -> b c h w', b=custom_shape[0]) logs = make_convolutional_sample(example, model, mode=mode, custom_steps=custom_steps, eta=eta, swap_mode=False , masked=masked, invert_mask=invert_mask, quantize_x0=False, custom_schedule=None, decode_interval=10, resize_enabled=resize_enabled, custom_shape=custom_shape, temperature=temperature, noise_dropout=0., corrector=guider, corrector_kwargs=ckwargs, x_T=x_T, save_intermediate_vid=save_intermediate_vid, make_progrow=make_progrow,ddim_use_x0_pred=ddim_use_x0_pred ) return logs @torch.no_grad() def convsample_ddim(model, cond, steps, shape, eta=1.0, callback=None, normals_sequence=None, mask=None, x0=None, quantize_x0=False, img_callback=None, temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None, x_T=None, log_every_t=None ): ddim = DDIMSampler(model) bs = shape[0] # dont know where this comes from but wayne shape = shape[1:] # cut batch dim print(f"Sampling with eta = {eta}; steps: {steps}") samples, intermediates = ddim.sample(steps, batch_size=bs, shape=shape, conditioning=cond, callback=callback, normals_sequence=normals_sequence, quantize_x0=quantize_x0, eta=eta, mask=mask, x0=x0, temperature=temperature, verbose=False, score_corrector=score_corrector, corrector_kwargs=corrector_kwargs, x_T=x_T) return samples, intermediates @torch.no_grad() def make_convolutional_sample(batch, model, mode="vanilla", custom_steps=None, eta=1.0, swap_mode=False, masked=False, invert_mask=True, quantize_x0=False, custom_schedule=None, decode_interval=1000, resize_enabled=False, custom_shape=None, temperature=1., noise_dropout=0., corrector=None, corrector_kwargs=None, x_T=None, save_intermediate_vid=False, make_progrow=True,ddim_use_x0_pred=False): log = dict() z, c, x, xrec, xc = model.get_input(batch, model.first_stage_key, return_first_stage_outputs=True, force_c_encode=not (hasattr(model, 'split_input_params') and model.cond_stage_key == 'coordinates_bbox'), return_original_cond=True) log_every_t = 1 if save_intermediate_vid else None if custom_shape is not None: z = torch.randn(custom_shape) print(f"Generating {custom_shape[0]} samples of shape {custom_shape[1:]}") z0 = None