Dynamatrix/DiffBIR-OpenXLab
0
1"""SAMPLING ONLY."""2 3import torch4import numpy as np5from tqdm import tqdm6 7from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like, extract_into_tensor8 9 10class DDIMSampler(object):11 def __init__(self, model, schedule="linear", **kwargs):12 super().__init__()13 self.model = model14 self.ddpm_num_timesteps = model.num_timesteps15 self.schedule = schedule16 17 def register_buffer(self, name, attr):18 if type(attr) == torch.Tensor:19 if attr.device != torch.device("cuda"):20 attr = attr.to(torch.device("cuda"))21 setattr(self, name, attr)22 23 def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):24 self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,25 num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose)26 alphas_cumprod = self.model.alphas_cumprod27 assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'28 to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)29 30 self.register_buffer('betas', to_torch(self.model.betas))31 self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))32 self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))33 34 # calculations for diffusion q(x_t | x_{t-1}) and others35 self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))36 self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))37 self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))38 self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))39 self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))40 41 # ddim sampling parameters42 ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),43 ddim_timesteps=self.ddim_timesteps,44 eta=ddim_eta,verbose=verbose)45 self.register_buffer('ddim_sigmas', ddim_sigmas)46 self.register_buffer('ddim_alphas', ddim_alphas)47 self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)48 self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))49 sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(50 (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (51 1 - self.alphas_cumprod / self.alphas_cumprod_prev))52 self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)53 54 @torch.no_grad()55 def sample(self,56 S,57 batch_size,58 shape,59 conditioning=None,60 callback=None,61 normals_sequence=None,62 img_callback=None,63 quantize_x0=False,64 eta=0.,65 mask=None,66 x0=None,67 temperature=1.,68 noise_dropout=0.,69 score_corrector=None,70 corrector_kwargs=None,71 verbose=True,72 x_T=None,73 log_every_t=100,74 unconditional_guidance_scale=1.,75 unconditional_conditioning=None, # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...76 dynamic_threshold=None,77 ucg_schedule=None,78 **kwargs79 ):80 if conditioning is not None:81 if isinstance(conditioning, dict):82 ctmp = conditioning[list(conditioning.keys())[0]]83 while isinstance(ctmp, list): ctmp = ctmp[0]84 cbs = ctmp.shape[0]85 if cbs != batch_size:86 print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")87 88 elif isinstance(conditioning, list):89 for ctmp in conditioning:90 if ctmp.shape[0] != batch_size:91 print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")92 93 else:94 if conditioning.shape[0] != batch_size:95 print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")96 97 self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)98 # sampling99 C, H, W = shape100 size = (batch_size, C, H, W)101 print(f'Data shape for DDIM sampling is {size}, eta {eta}')102 103 samples, intermediates = self.ddim_sampling(conditioning, size,104 callback=callback,105 img_callback=img_callback,106 quantize_denoised=quantize_x0,107 mask=mask, x0=x0,108 ddim_use_original_steps=False,109 noise_dropout=noise_dropout,110 temperature=temperature,111 score_corrector=score_corrector,112 corrector_kwargs=corrector_kwargs,113 x_T=x_T,114 log_every_t=log_every_t,115 unconditional_guidance_scale=unconditional_guidance_scale,116 unconditional_conditioning=unconditional_conditioning,117 dynamic_threshold=dynamic_threshold,118 ucg_schedule=ucg_schedule119 )120 return samples, intermediates121 122 @torch.no_grad()123 def ddim_sampling(self, cond, shape,124 x_T=None, ddim_use_original_steps=False,125 callback=None, timesteps=None, quantize_denoised=False,126 mask=None, x0=None, img_callback=None, log_every_t=100,127 temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,128 unconditional_guidance_scale=1., unconditional_conditioning=None, dynamic_threshold=None,129 ucg_schedule=None):130 device = self.model.betas.device131 b = shape[0]132 if x_T is None:133 img = torch.randn(shape, device=device)134 else:135 img = x_T136 137 if timesteps is None:138 timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps139 elif timesteps is not None and not ddim_use_original_steps:140 subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1141 timesteps = self.ddim_timesteps[:subset_end]142 143 intermediates = {'x_inter': [img], 'pred_x0': [img]}144 time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else np.flip(timesteps)145 total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]146 print(f"Running DDIM Sampling with {total_steps} timesteps")147 148 iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)149 150 for i, step in enumerate(iterator):151 index = total_steps - i - 1152 ts = torch.full((b,), step, device=device, dtype=torch.long)153 154 if mask is not None:155 assert x0 is not None156 img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass?157 img = img_orig * mask + (1. - mask) * img158 159 if ucg_schedule is not None:160 assert len(ucg_schedule) == len(time_range)161 unconditional_guidance_scale = ucg_schedule[i]162 163 outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,164 quantize_denoised=quantize_denoised, temperature=temperature,165 noise_dropout=noise_dropout, score_corrector=score_corrector,166 corrector_kwargs=corrector_kwargs,167 unconditional_guidance_scale=unconditional_guidance_scale,168 unconditional_conditioning=unconditional_conditioning,169 dynamic_threshold=dynamic_threshold)170 img, pred_x0 = outs171 if callback: callback(i)172 if img_callback: img_callback(pred_x0, i)173 174 if index % log_every_t == 0 or index == total_steps - 1:175 intermediates['x_inter'].append(img)176 intermediates['pred_x0'].append(pred_x0)177 178 return img, intermediates179 180 @torch.no_grad()181 def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,182 temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,183 unconditional_guidance_scale=1., unconditional_conditioning=None,184 dynamic_threshold=None):185 b, *_, device = *x.shape, x.device186 187 if unconditional_conditioning is None or unconditional_guidance_scale == 1.:188 model_output = self.model.apply_model(x, t, c)189 else:190 x_in = torch.cat([x] * 2)191 t_in = torch.cat([t] * 2)192 if isinstance(c, dict):193 assert isinstance(unconditional_conditioning, dict)194 c_in = dict()195 for k in c:196 if isinstance(c[k], list):197 c_in[k] = [torch.cat([198 unconditional_conditioning[k][i],199 c[k][i]]) for i in range(len(c[k]))]200 else:201 c_in[k] = torch.cat([202 unconditional_conditioning[k],203 c[k]])204 elif isinstance(c, list):205 c_in = list()206 assert isinstance(unconditional_conditioning, list)207 for i in range(len(c)):208 c_in.append(torch.cat([unconditional_conditioning[i], c[i]]))209 else:210 c_in = torch.cat([unconditional_conditioning, c])211 model_uncond, model_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)212 model_output = model_uncond + unconditional_guidance_scale * (model_t - model_uncond)213 214 if self.model.parameterization == "v":215 e_t = self.model.predict_eps_from_z_and_v(x, t, model_output)216 else:217 e_t = model_output218 219 if score_corrector is not None:220 assert self.model.parameterization == "eps", 'not implemented'221 e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)222 223 alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas224 alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev225 sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas226 sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas227 # select parameters corresponding to the currently considered timestep228 a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)229 a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)230 sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)231 sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)232 233 # current prediction for x_0234 if self.model.parameterization != "v":235 pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()236 else:237 pred_x0 = self.model.predict_start_from_z_and_v(x, t, model_output)238 239 if quantize_denoised:240 pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)241 242 if dynamic_threshold is not None:243 raise NotImplementedError()244 245 # direction pointing to x_t246 dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t247 noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature248 if noise_dropout > 0.:249 noise = torch.nn.functional.dropout(noise, p=noise_dropout)250 x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise251 return x_prev, pred_x0252 253 @torch.no_grad()254 def encode(self, x0, c, t_enc, use_original_steps=False, return_intermediates=None,255 unconditional_guidance_scale=1.0, unconditional_conditioning=None, callback=None):256 num_reference_steps = self.ddpm_num_timesteps if use_original_steps else self.ddim_timesteps.shape[0]257 258 assert t_enc <= num_reference_steps259 num_steps = t_enc260 261 if use_original_steps:262 alphas_next = self.alphas_cumprod[:num_steps]263 alphas = self.alphas_cumprod_prev[:num_steps]264 else:265 alphas_next = self.ddim_alphas[:num_steps]266 alphas = torch.tensor(self.ddim_alphas_prev[:num_steps])267 268 x_next = x0269 intermediates = []270 inter_steps = []271 for i in tqdm(range(num_steps), desc='Encoding Image'):272 t = torch.full((x0.shape[0],), i, device=self.model.device, dtype=torch.long)273 if unconditional_guidance_scale == 1.:274 noise_pred = self.model.apply_model(x_next, t, c)275 else:276 assert unconditional_conditioning is not None277 e_t_uncond, noise_pred = torch.chunk(278 self.model.apply_model(torch.cat((x_next, x_next)), torch.cat((t, t)),279 torch.cat((unconditional_conditioning, c))), 2)280 noise_pred = e_t_uncond + unconditional_guidance_scale * (noise_pred - e_t_uncond)281 282 xt_weighted = (alphas_next[i] / alphas[i]).sqrt() * x_next283 weighted_noise_pred = alphas_next[i].sqrt() * (284 (1 / alphas_next[i] - 1).sqrt() - (1 / alphas[i] - 1).sqrt()) * noise_pred285 x_next = xt_weighted + weighted_noise_pred286 if return_intermediates and i % (287 num_steps // return_intermediates) == 0 and i < num_steps - 1:288 intermediates.append(x_next)289 inter_steps.append(i)290 elif return_intermediates and i >= num_steps - 2:291 intermediates.append(x_next)292 inter_steps.append(i)293 if callback: callback(i)294 295 out = {'x_encoded': x_next, 'intermediate_steps': inter_steps}296 if return_intermediates:297 out.update({'intermediates': intermediates})298 return x_next, out299 300 @torch.no_grad()301 def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):302 # fast, but does not allow for exact reconstruction303 # t serves as an index to gather the correct alphas304 if use_original_steps:305 sqrt_alphas_cumprod = self.sqrt_alphas_cumprod306 sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod307 else:308 sqrt_alphas_cumprod = torch.sqrt(self.ddim_alphas)309 sqrt_one_minus_alphas_cumprod = self.ddim_sqrt_one_minus_alphas310 311 if noise is None:312 noise = torch.randn_like(x0)313 return (extract_into_tensor(sqrt_alphas_cumprod, t, x0.shape) * x0 +314 extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x0.shape) * noise)315 316 @torch.no_grad()317 def decode(self, x_latent, cond, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None,318 use_original_steps=False, callback=None):319 320 timesteps = np.arange(self.ddpm_num_timesteps) if use_original_steps else self.ddim_timesteps321 timesteps = timesteps[:t_start]322 323 time_range = np.flip(timesteps)324 total_steps = timesteps.shape[0]325 print(f"Running DDIM Sampling with {total_steps} timesteps")326 327 iterator = tqdm(time_range, desc='Decoding image', total=total_steps)328 x_dec = x_latent329 for i, step in enumerate(iterator):330 index = total_steps - i - 1331 ts = torch.full((x_latent.shape[0],), step, device=x_latent.device, dtype=torch.long)332 x_dec, _ = self.p_sample_ddim(x_dec, cond, ts, index=index, use_original_steps=use_original_steps,333 unconditional_guidance_scale=unconditional_guidance_scale,334 unconditional_conditioning=unconditional_conditioning)335 if callback: callback(i)336 return x_dec