cedpsam/latentdiffusion

1from torchvision.datasets.utils import download_url2from ldm.util import instantiate_from_config3import torch4import os5# todo ?6from google.colab import files7from IPython.display import Image as ipyimg8import ipywidgets as widgets9from PIL import Image10from numpy import asarray11from einops import rearrange, repeat12import torch, torchvision13from ldm.models.diffusion.ddim import DDIMSampler14from ldm.util import ismap15import time16from omegaconf import OmegaConf17 18 19def download_models(mode):20 21    if mode == "superresolution":22        # this is the small bsr light model23        url_conf = 'https://heibox.uni-heidelberg.de/f/31a76b13ea27482981b4/?dl=1'24        url_ckpt = 'https://heibox.uni-heidelberg.de/f/578df07c8fc04ffbadf3/?dl=1'25 26        path_conf = 'logs/diffusion/superresolution_bsr/configs/project.yaml'27        path_ckpt = 'logs/diffusion/superresolution_bsr/checkpoints/last.ckpt'28 29        download_url(url_conf, path_conf)30        download_url(url_ckpt, path_ckpt)31 32        path_conf = path_conf + '/?dl=1' # fix it33        path_ckpt = path_ckpt + '/?dl=1' # fix it34        return path_conf, path_ckpt35 36    else:37        raise NotImplementedError38 39 40def load_model_from_config(config, ckpt):41    print(f"Loading model from {ckpt}")42    pl_sd = torch.load(ckpt, map_location="cpu")43    global_step = pl_sd["global_step"]44    sd = pl_sd["state_dict"]45    model = instantiate_from_config(config.model)46    m, u = model.load_state_dict(sd, strict=False)47    model.cuda()48    model.eval()49    return {"model": model}, global_step50 51 52def get_model(mode):53    path_conf, path_ckpt = download_models(mode)54    config = OmegaConf.load(path_conf)55    model, step = load_model_from_config(config, path_ckpt)56    return model57 58 59def get_custom_cond(mode):60    dest = "data/example_conditioning"61 62    if mode == "superresolution":63        uploaded_img = files.upload()64        filename = next(iter(uploaded_img))65        name, filetype = filename.split(".") # todo assumes just one dot in name !66        os.rename(f"{filename}", f"{dest}/{mode}/custom_{name}.{filetype}")67 68    elif mode == "text_conditional":69        w = widgets.Text(value='A cake with cream!', disabled=True)70        display(w)71 72        with open(f"{dest}/{mode}/custom_{w.value[:20]}.txt", 'w') as f:73            f.write(w.value)74 75    elif mode == "class_conditional":76        w = widgets.IntSlider(min=0, max=1000)77        display(w)78        with open(f"{dest}/{mode}/custom.txt", 'w') as f:79            f.write(w.value)80 81    else:82        raise NotImplementedError(f"cond not implemented for mode{mode}")83 84 85def get_cond_options(mode):86    path = "data/example_conditioning"87    path = os.path.join(path, mode)88    onlyfiles = [f for f in sorted(os.listdir(path))]89    return path, onlyfiles90 91 92def select_cond_path(mode):93    path = "data/example_conditioning"  # todo94    path = os.path.join(path, mode)95    onlyfiles = [f for f in sorted(os.listdir(path))]96 97    selected = widgets.RadioButtons(98        options=onlyfiles,99        description='Select conditioning:',100        disabled=False101    )102    display(selected)103    selected_path = os.path.join(path, selected.value)104    return selected_path105 106 107def get_cond(mode, selected_path):108    example = dict()109    if mode == "superresolution":110        up_f = 4111        visualize_cond_img(selected_path)112 113        c = Image.open(selected_path)114        c = torch.unsqueeze(torchvision.transforms.ToTensor()(c), 0)115        c_up = torchvision.transforms.functional.resize(c, size=[up_f * c.shape[2], up_f * c.shape[3]], antialias=True)116        c_up = rearrange(c_up, '1 c h w -> 1 h w c')117        c = rearrange(c, '1 c h w -> 1 h w c')118        c = 2. * c - 1.119 120        c = c.to(torch.device("cuda"))121        example["LR_image"] = c122        example["image"] = c_up123 124    return example125 126 127def visualize_cond_img(path):128    display(ipyimg(filename=path))129 130 131def run(model, selected_path, task, custom_steps, resize_enabled=False, classifier_ckpt=None, global_step=None):132 133    example = get_cond(task, selected_path)134 135    save_intermediate_vid = False136    n_runs = 1137    masked = False138    guider = None139    ckwargs = None140    mode = 'ddim'141    ddim_use_x0_pred = False142    temperature = 1.143    eta = 1.144    make_progrow = True145    custom_shape = None146 147    height, width = example["image"].shape[1:3]148    split_input = height >= 128 and width >= 128149 150    if split_input:151        ks = 128152        stride = 64153        vqf = 4  #154        model.split_input_params = {"ks": (ks, ks), "stride": (stride, stride),155                                    "vqf": vqf,156                                    "patch_distributed_vq": True,157                                    "tie_braker": False,158                                    "clip_max_weight": 0.5,159                                    "clip_min_weight": 0.01,160                                    "clip_max_tie_weight": 0.5,161                                    "clip_min_tie_weight": 0.01}162    else:163        if hasattr(model, "split_input_params"):164            delattr(model, "split_input_params")165 166    invert_mask = False167 168    x_T = None169    for n in range(n_runs):170        if custom_shape is not None:171            x_T = torch.randn(1, custom_shape[1], custom_shape[2], custom_shape[3]).to(model.device)172            x_T = repeat(x_T, '1 c h w -> b c h w', b=custom_shape[0])173 174        logs = make_convolutional_sample(example, model,175                                         mode=mode, custom_steps=custom_steps,176                                         eta=eta, swap_mode=False , masked=masked,177                                         invert_mask=invert_mask, quantize_x0=False,178                                         custom_schedule=None, decode_interval=10,179                                         resize_enabled=resize_enabled, custom_shape=custom_shape,180                                         temperature=temperature, noise_dropout=0.,181                                         corrector=guider, corrector_kwargs=ckwargs, x_T=x_T, save_intermediate_vid=save_intermediate_vid,182                                         make_progrow=make_progrow,ddim_use_x0_pred=ddim_use_x0_pred183                                         )184    return logs185 186 187@torch.no_grad()188def convsample_ddim(model, cond, steps, shape, eta=1.0, callback=None, normals_sequence=None,189                    mask=None, x0=None, quantize_x0=False, img_callback=None,190                    temperature=1., noise_dropout=0., score_corrector=None,191                    corrector_kwargs=None, x_T=None, log_every_t=None192                    ):193 194    ddim = DDIMSampler(model)195    bs = shape[0]  # dont know where this comes from but wayne196    shape = shape[1:]  # cut batch dim197    print(f"Sampling with eta = {eta}; steps: {steps}")198    samples, intermediates = ddim.sample(steps, batch_size=bs, shape=shape, conditioning=cond, callback=callback,199                                         normals_sequence=normals_sequence, quantize_x0=quantize_x0, eta=eta,200                                         mask=mask, x0=x0, temperature=temperature, verbose=False,201                                         score_corrector=score_corrector,202                                         corrector_kwargs=corrector_kwargs, x_T=x_T)203 204    return samples, intermediates205 206 207@torch.no_grad()208def make_convolutional_sample(batch, model, mode="vanilla", custom_steps=None, eta=1.0, swap_mode=False, masked=False,209                              invert_mask=True, quantize_x0=False, custom_schedule=None, decode_interval=1000,210                              resize_enabled=False, custom_shape=None, temperature=1., noise_dropout=0., corrector=None,211                              corrector_kwargs=None, x_T=None, save_intermediate_vid=False, make_progrow=True,ddim_use_x0_pred=False):212    log = dict()213 214    z, c, x, xrec, xc = model.get_input(batch, model.first_stage_key,215                                        return_first_stage_outputs=True,216                                        force_c_encode=not (hasattr(model, 'split_input_params')217                                                            and model.cond_stage_key == 'coordinates_bbox'),218                                        return_original_cond=True)219 220    log_every_t = 1 if save_intermediate_vid else None221 222    if custom_shape is not None:223        z = torch.randn(custom_shape)224        print(f"Generating {custom_shape[0]} samples of shape {custom_shape[1:]}")225 226    z0 = None227 228    log["input"] = x229    log["reconstruction"] = xrec230 231    if ismap(xc):232        log["original_conditioning"] = model.to_rgb(xc)233        if hasattr(model, 'cond_stage_key'):234            log[model.cond_stage_key] = model.to_rgb(xc)235 236    else:237        log["original_conditioning"] = xc if xc is not None else torch.zeros_like(x)238        if model.cond_stage_model:239            log[model.cond_stage_key] = xc if xc is not None else torch.zeros_like(x)240            if model.cond_stage_key =='class_label':241                log[model.cond_stage_key] = xc[model.cond_stage_key]242 243    with model.ema_scope("Plotting"):244        t0 = time.time()245        img_cb = None246 247        sample, intermediates = convsample_ddim(model, c, steps=custom_steps, shape=z.shape,248                                                eta=eta,249                                                quantize_x0=quantize_x0, img_callback=img_cb, mask=None, x0=z0,250                                                temperature=temperature, noise_dropout=noise_dropout,251                                                score_corrector=corrector, corrector_kwargs=corrector_kwargs,252                                                x_T=x_T, log_every_t=log_every_t)253        t1 = time.time()254 255        if ddim_use_x0_pred:256            sample = intermediates['pred_x0'][-1]257 258    x_sample = model.decode_first_stage(sample)259 260    try:261        x_sample_noquant = model.decode_first_stage(sample, force_not_quantize=True)262        log["sample_noquant"] = x_sample_noquant263        log["sample_diff"] = torch.abs(x_sample_noquant - x_sample)264    except:265        pass266 267    log["sample"] = x_sample268    log["time"] = t1 - t0269 270    return log