cymic/Waifu_Diffusion_Webui

1import gc2import time3import warnings4 5import numpy as np6import torch7import torchvision8from PIL import Image9from einops import rearrange, repeat10from omegaconf import OmegaConf11 12from ldm.models.diffusion.ddim import DDIMSampler13from ldm.util import instantiate_from_config, ismap14 15warnings.filterwarnings("ignore", category=UserWarning)16 17 18# Create LDSR Class19class LDSR:20    def load_model_from_config(self, half_attention):21        print(f"Loading model from {self.modelPath}")22        pl_sd = torch.load(self.modelPath, map_location="cpu")23        sd = pl_sd["state_dict"]24        config = OmegaConf.load(self.yamlPath)25        model = instantiate_from_config(config.model)26        model.load_state_dict(sd, strict=False)27        model.cuda()28        if half_attention:29            model = model.half()30 31        model.eval()32        return {"model": model}33 34    def __init__(self, model_path, yaml_path):35        self.modelPath = model_path36        self.yamlPath = yaml_path37 38    @staticmethod39    def run(model, selected_path, custom_steps, eta):40        example = get_cond(selected_path)41 42        n_runs = 143        guider = None44        ckwargs = None45        ddim_use_x0_pred = False46        temperature = 1.47        eta = eta48        custom_shape = None49 50        height, width = example["image"].shape[1:3]51        split_input = height >= 128 and width >= 12852 53        if split_input:54            ks = 12855            stride = 6456            vqf = 4  #57            model.split_input_params = {"ks": (ks, ks), "stride": (stride, stride),58                                        "vqf": vqf,59                                        "patch_distributed_vq": True,60                                        "tie_braker": False,61                                        "clip_max_weight": 0.5,62                                        "clip_min_weight": 0.01,63                                        "clip_max_tie_weight": 0.5,64                                        "clip_min_tie_weight": 0.01}65        else:66            if hasattr(model, "split_input_params"):67                delattr(model, "split_input_params")68 69        x_t = None70        logs = None71        for n in range(n_runs):72            if custom_shape is not None:73                x_t = torch.randn(1, custom_shape[1], custom_shape[2], custom_shape[3]).to(model.device)74                x_t = repeat(x_t, '1 c h w -> b c h w', b=custom_shape[0])75 76            logs = make_convolutional_sample(example, model,77                                             custom_steps=custom_steps,78                                             eta=eta, quantize_x0=False,79                                             custom_shape=custom_shape,80                                             temperature=temperature, noise_dropout=0.,81                                             corrector=guider, corrector_kwargs=ckwargs, x_T=x_t,82                                             ddim_use_x0_pred=ddim_use_x0_pred83                                             )84        return logs85 86    def super_resolution(self, image, steps=100, target_scale=2, half_attention=False):87        model = self.load_model_from_config(half_attention)88 89        # Run settings90        diffusion_steps = int(steps)91        eta = 1.092 93        down_sample_method = 'Lanczos'94 95        gc.collect()96        torch.cuda.empty_cache()97 98        im_og = image99        width_og, height_og = im_og.size100        # If we can adjust the max upscale size, then the 4 below should be our variable101        down_sample_rate = target_scale / 4102        wd = width_og * down_sample_rate103        hd = height_og * down_sample_rate104        width_downsampled_pre = int(wd)105        height_downsampled_pre = int(hd)106 107        if down_sample_rate != 1:108            print(109                f'Downsampling from [{width_og}, {height_og}] to [{width_downsampled_pre}, {height_downsampled_pre}]')110            im_og = im_og.resize((width_downsampled_pre, height_downsampled_pre), Image.LANCZOS)111        else:112            print(f"Down sample rate is 1 from {target_scale} / 4 (Not downsampling)")113        logs = self.run(model["model"], im_og, diffusion_steps, eta)114 115        sample = logs["sample"]116        sample = sample.detach().cpu()117        sample = torch.clamp(sample, -1., 1.)118        sample = (sample + 1.) / 2. * 255119        sample = sample.numpy().astype(np.uint8)120        sample = np.transpose(sample, (0, 2, 3, 1))121        a = Image.fromarray(sample[0])122 123        del model124        gc.collect()125        torch.cuda.empty_cache()126        return a127 128 129def get_cond(selected_path):130    example = dict()131    up_f = 4132    c = selected_path.convert('RGB')133    c = torch.unsqueeze(torchvision.transforms.ToTensor()(c), 0)134    c_up = torchvision.transforms.functional.resize(c, size=[up_f * c.shape[2], up_f * c.shape[3]],135                                                    antialias=True)136    c_up = rearrange(c_up, '1 c h w -> 1 h w c')137    c = rearrange(c, '1 c h w -> 1 h w c')138    c = 2. * c - 1.139 140    c = c.to(torch.device("cuda"))141    example["LR_image"] = c142    example["image"] = c_up143 144    return example145 146 147@torch.no_grad()148def convsample_ddim(model, cond, steps, shape, eta=1.0, callback=None, normals_sequence=None,149                    mask=None, x0=None, quantize_x0=False, temperature=1., score_corrector=None,150                    corrector_kwargs=None, x_t=None151                    ):152    ddim = DDIMSampler(model)153    bs = shape[0]154    shape = shape[1:]155    print(f"Sampling with eta = {eta}; steps: {steps}")156    samples, intermediates = ddim.sample(steps, batch_size=bs, shape=shape, conditioning=cond, callback=callback,157                                         normals_sequence=normals_sequence, quantize_x0=quantize_x0, eta=eta,158                                         mask=mask, x0=x0, temperature=temperature, verbose=False,159                                         score_corrector=score_corrector,160                                         corrector_kwargs=corrector_kwargs, x_t=x_t)161 162    return samples, intermediates163 164 165@torch.no_grad()166def make_convolutional_sample(batch, model, custom_steps=None, eta=1.0, quantize_x0=False, custom_shape=None, temperature=1., noise_dropout=0., corrector=None,167                              corrector_kwargs=None, x_T=None, ddim_use_x0_pred=False):168    log = dict()169 170    z, c, x, xrec, xc = model.get_input(batch, model.first_stage_key,171                                        return_first_stage_outputs=True,172                                        force_c_encode=not (hasattr(model, 'split_input_params')173                                                            and model.cond_stage_key == 'coordinates_bbox'),174                                        return_original_cond=True)175 176    if custom_shape is not None:177        z = torch.randn(custom_shape)178        print(f"Generating {custom_shape[0]} samples of shape {custom_shape[1:]}")179 180    z0 = None181 182    log["input"] = x183    log["reconstruction"] = xrec184 185    if ismap(xc):186        log["original_conditioning"] = model.to_rgb(xc)187        if hasattr(model, 'cond_stage_key'):188            log[model.cond_stage_key] = model.to_rgb(xc)189 190    else:191        log["original_conditioning"] = xc if xc is not None else torch.zeros_like(x)192        if model.cond_stage_model:193            log[model.cond_stage_key] = xc if xc is not None else torch.zeros_like(x)194            if model.cond_stage_key == 'class_label':195                log[model.cond_stage_key] = xc[model.cond_stage_key]196 197    with model.ema_scope("Plotting"):198        t0 = time.time()199 200        sample, intermediates = convsample_ddim(model, c, steps=custom_steps, shape=z.shape,201                                                eta=eta,202                                                quantize_x0=quantize_x0, mask=None, x0=z0,203                                                temperature=temperature, score_corrector=corrector, corrector_kwargs=corrector_kwargs,204                                                x_t=x_T)205        t1 = time.time()206 207        if ddim_use_x0_pred:208            sample = intermediates['pred_x0'][-1]209 210    x_sample = model.decode_first_stage(sample)211 212    try:213        x_sample_noquant = model.decode_first_stage(sample, force_not_quantize=True)214        log["sample_noquant"] = x_sample_noquant215        log["sample_diff"] = torch.abs(x_sample_noquant - x_sample)216    except:217        pass218 219    log["sample"] = x_sample220    log["time"] = t1 - t0221 222    return log223