import numpy as np import torch from tqdm import tqdm def _extract_into_tensor(arr, timesteps, broadcast_shape): """ Extract values from a 1-D numpy array for a batch of indices. :param arr: the 1-D numpy array. :param timesteps: a tensor of indices into the array to extract. :param broadcast_shape: a larger shape of K dimensions with the batch dimension equal to the length of timesteps. :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims. """ res = torch.from_numpy(arr).to(device=timesteps.device)[timesteps].float() while len(res.shape) < len(broadcast_shape): res = res[..., None] return res.expand(broadcast_shape) class DiffSynthSampler: def __init__(self, timesteps, beta_start=0.0001, beta_end=0.02, device=None, mute=False, height=128, max_batchsize=16, max_width=256, channels=4, train_width=64, noise_strategy="repeat"): if device is None: self.device = "cuda" if torch.cuda.is_available() else "cpu" else: self.device = device self.height = height self.train_width = train_width self.max_batchsize = max_batchsize self.max_width = max_width self.channels = channels self.num_timesteps = timesteps self.timestep_map = list(range(self.num_timesteps)) self.betas = np.array(np.linspace(beta_start, beta_end, self.num_timesteps), dtype=np.float64) self.respaced = False self.define_beta_schedule() self.CFG = 1.0 self.mute = mute self.noise_strategy = noise_strategy def get_deterministic_noise_tensor_non_repeat(self, batchsize, width, reference_noise=None): if reference_noise is None: large_noise_tensor = torch.randn((self.max_batchsize, self.channels, self.height, self.max_width), device=self.device) else: assert reference_noise.shape == (batchsize, self.channels, self.height, self.max_width), "reference_noise shape mismatch" large_noise_tensor = reference_noise return large_noise_tensor[:batchsize, :, :, :width], None def get_deterministic_noise_tensor(self, batchsize, width, reference_noise=None): if self.noise_strategy == "repeat": noise, concat_points = self.get_deterministic_noise_tensor_repeat(batchsize, width, reference_noise=reference_noise) return noise, concat_points else: noise, concat_points = self.get_deterministic_noise_tensor_non_repeat(batchsize, width, reference_noise=reference_noise) return noise, concat_points def get_deterministic_noise_tensor_repeat(self, batchsize, width, reference_noise=None): # 生成与训练数据长度相等的噪音 if reference_noise is None: train_noise_tensor = torch.randn((self.max_batchsize, self.channels, self.height, self.train_width), device=self.device) else: assert reference_noise.shape == (batchsize, self.channels, self.height, self.train_width), "reference_noise shape mismatch" train_noise_tensor = reference_noise release_width = int(self.train_width * 1.0 / 4) first_part_width = self.train_width - release_width first_part = train_noise_tensor[:batchsize, :, :, :first_part_width] release_part = train_noise_tensor[:batchsize, :, :, -release_width:] # 如果所需 length 小于等于 origin length，去掉 first_part 的中间部分 if width <= self.train_width: _first_part_head_width = int((width - release_width) / 2) _first_part_tail_width = width - release_width - _first_part_head_width all_parts = [first_part[:, :, :, :_first_part_head_width], first_part[:, :, :, -_first_part_tail_width:], release_part] # 沿第四维度拼接张量 noise_tensor = torch.cat(all_parts, dim=3) # 记录拼接点的位置 concat_points = [0] for part in all_parts[:-1]: next_point = concat_points[-1] + part.size(3) concat_points.append(next_point) return noise_tensor, concat_points # 如果所需 length 大于 origin length，不断地从中间插入 first_part 的中间部分 else: # 计算需要重复front_width的次数 repeats = (width - release_width) // first_part_width extra = (width - release_width) % first_part_width _repeat_first_part_head_width = int(first_part_width / 2) _repeat_first_part_tail_width = first_part_width - _repeat_first_part_head_width repeated_first_head_parts = [first_part[:, :, :, :_repeat_first_part_head_width] for _ in range(repeats)] repeated_first_tail_parts = [first_part[:, :, :, -_repeat_first_part_tail_width:] for _ in range(repeats)] # 计算起始索引 _middle_part_start_index = (first_part_width - extra) // 2 # 切片张量以获取中间部分 middle_part = first_part[:, :, :, _middle_part_start_index: _middle_part_start_index + extra] all_parts = repeated_first_head_parts + [middle_part] + repeated_first_tail_parts + [release_part] # 沿第四维度拼接张量 noise_tensor = torch.cat(all_parts, dim=3) # 记录拼接点的位置 concat_points = [0] for part in all_parts[:-1]: next_point = concat_points[-1] + part.size(3) concat_points.append(next_point) return noise_tensor, concat_points def define_beta_schedule(self): assert self.respaced == False, "This schedule has already been respaced!" # define alphas self.alphas = 1.0 - self.betas self.alphas_cumprod = np.cumprod(self.alphas, axis=0) self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1]) self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0) # calculations for diffusion q(x_t | x_{t-1}) and others self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod) self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod) self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod) self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod) self.sqrt_recip_alphas = np.sqrt(1.0 / self.alphas) self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1) # calculations for posterior q(x_{t-1} | x_t, x_0) self.posterior_variance = (self.betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)) def activate_classifier_free_guidance(self, CFG, unconditional_condition): assert ( not unconditional_condition is None) or CFG == 1.0, "For CFG != 1.0, unconditional_condition must be available" self.CFG = CFG self.unconditional_condition = unconditional_condition def respace(self, use_timesteps=None): if not use_timesteps is None: last_alpha_cumprod = 1.0 new_betas = [] self.timestep_map = [] for i, _alpha_cumprod in enumerate(self.alphas_cumprod): if i in use_timesteps: new_betas.append(1 - _alpha_cumprod / last_alpha_cumprod) last_alpha_cumprod = _alpha_cumprod self.timestep_map.append(i) self.num_timesteps = len(use_timesteps) self.betas = np.array(new_betas) self.define_beta_schedule() self.respaced = True def generate_linear_noise(self, shape, variance=1.0, first_endpoint=None, second_endpoint=None): assert shape[1] == self.channels, "shape[1] != self.channels" assert shape[2] == self.height, "shape[2] != self.height" noise = torch.empty(*shape, device=self.device) # 第三种情况：两个端点都不是None，进行线性插值 if first_endpoint is not None and second_endpoint is not None: for i in range(shape[0]): alpha = i / (shape[0] - 1) # 插值系数 noise[i] = alpha * second_endpoint + (1 - alpha) * first_endpoint return noise # 返回插值后的结果，不需要进行后续的均值和方差调整 else: # 第一个端点不是None if first_endpoint is not None: noise[0] = first_endpoint if shape[0] > 1: noise[1], _ = self.get_deterministic_noise_tensor(1, shape[3])[0] else: noise[0], _ = self.get_deterministic_noise_tensor(1, shape[3])[0] if shape[0] > 1: noise[1], _ = self.get_deterministic_noise_tensor(1, shape[3])[0] # 生成其他的噪声点 for i in range(2, shape[0]): noise[i] = 2 * noise[i - 1] - noise[i - 2] # 当只有一个端点被指定时 current_var = noise.var() stddev_ratio = torch.sqrt(variance / current_var) noise = noise * stddev_ratio # 如果第一个端点被指定，进行平移调整 if first_endpoint is not None: shift = first_endpoint - noise[0] noise += shift return noise def q_sample(self, x_start, t, noise=None): """ Diffuse the data for a given number of diffusion steps. In other words, sample from q(x_t | x_0). :param x_start: the initial data batch. :param t: the number of diffusion steps (minus 1). Here, 0 means one step. :param noise: if specified, the split-out normal noise. :return: A noisy version of x_start. """ assert x_start.shape[1] == self.channels, "shape[1] != self.channels" assert x_start.shape[2] == self.height, "shape[2] != self.height" if noise is None: # noise = torch.randn_like(x_start) noise, _ = self.get_deterministic_noise_tensor(x_start.shape[0], x_start.shape[3]) assert noise.shape == x_start.shape return ( _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise ) @torch.no_grad() def ddim_sample(self, model, x, t, condition=None, ddim_eta=0.0): map_tensor = torch.tensor(self.timestep_map, device=t.device, dtype=t.dtype) mapped_t = map_tensor[t] # Todo: add CFG if self.CFG == 1.0: pred_noise = model(x, mapped_t, condition) else: unconditional_condition = self.unconditional_condition.unsqueeze(0).repeat( *([x.shape[0]] + [1] * len(self.unconditional_condition.shape))) x_in = torch.cat([x] * 2) t_in = torch.cat([mapped_t] * 2) c_in = torch.cat([unconditional_condition, condition]) noise_uncond, noise = model(x_in, t_in, c_in).chunk(2) pred_noise = noise_uncond + self.CFG * (noise - noise_uncond) # Todo: END alpha_cumprod_t = _extract_into_tensor(self.alphas_cumprod, t, x.shape) alpha_cumprod_t_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape) pred_x0 = (x - torch.sqrt((1. - alpha_cumprod_t)) * pred_noise) / torch.sqrt(alpha_cumprod_t) sigmas_t = ( ddim_eta * torch.sqrt((1 - alpha_cumprod_t_prev) / (1 - alpha_cumprod_t)) * torch.sqrt(1 - alpha_cumprod_t / alpha_cumprod_t_prev) ) pred_dir_xt = torch.sqrt(1 - alpha_cumprod_t_prev - sigmas_t ** 2) * pred_noise step_noise, _ = self.get_deterministic_noise_tensor(x.shape[0], x.shape[3]) x_prev = torch.sqrt(alpha_cumprod_t_prev) * pred_x0 + pred_dir_xt + sigmas_t * step_noise return x_prev def p_sample(self, model, x, t, condition=None, sampler="ddim"): if sampler == "ddim": return self.ddim_sample(model, x, t, condition=condition, ddim_eta=0.0) elif sampler == "ddpm": return self.ddim_sample(model, x, t, condition=condition, ddim_eta=1.0) else: raise NotImplementedError() def get_dynamic_masks(self, n_masks, shape, concat_points, mask_flexivity=0.8): release_length = int(self.train_width / 4) assert shape[3] == (concat_points[-1] + release_length), "shape[3] != (concat_points[-1] + release_length)" fraction_lengths = [concat_points[i + 1] - concat_points[i] for i in range(len(concat_points) - 1)] # Todo: remove hard-coding n_guidance_steps = int(n_masks * mask_flexivity) n_free_steps = n_masks - n_guidance_steps masks = [] # Todo: 在一半的 steps 内收缩 mask。也就是说，在后程对 release 以外的区域不做inpaint，而是 img2img for i in range(n_guidance_steps): # mask = 1, freeze step_i_mask = torch.zeros((shape[0], 1, shape[2], shape[3]), dtype=torch.float32).to(self.device) step_i_mask[:, :, :, -release_length:] = 1.0 for fraction_index in range(len(fraction_lengths)): _fraction_mask_length = int((n_guidance_steps - 1 - i) / (n_guidance_steps - 1) * fraction_lengths[fraction_index]) if fraction_index == 0: step_i_mask[:, :, :, :_fraction_mask_length] = 1.0 elif fraction_index == len(fraction_lengths) - 1: if not _fraction_mask_length == 0: step_i_mask[:, :, :, -_fraction_mask_length - release_length:] = 1.0 else: fraction_mask_start_position = int((fraction_lengths[fraction_index] - _fraction_mask_length) / 2) step_i_mask[:, :, :, concat_points[fraction_index] + fraction_mask_start_position:concat_points[ fraction_index] + fraction_mask_start_position + _fraction_mask_length] = 1.0 masks.append(step_i_mask) for i in range(n_free_steps): step_i_mask = torch.zeros((shape[0], 1, shape[2], shape[3]), dtype=torch.float32).to(self.device) step_i_mask[:, :, :, -release_length:] = 1.0 masks.append(step_i_mask) masks.reverse() return masks @torch.no_grad() def p_sample_loop(self, model, shape, initial_noise=None, start_noise_level_ratio=1.0, end_noise_level_ratio=0.0, return_tensor=False, condition=None, guide_img=None, mask=None, sampler="ddim", inpaint=False, use_dynamic_mask=False, mask_flexivity=0.8): assert shape[1] == self.channels, "shape[1] != self.channels" assert shape[2] == self.height, "shape[2] != self.height" initial_noise, _ = self.get_deterministic_noise_tensor(shape[0], shape[3], reference_noise=initial_noise) assert initial_noise.shape == shape, "initial_noise.shape != shape" start_noise_level_index = int(self.num_timesteps * start_noise_level_ratio) # not included!!! end_noise_level_index = int(self.num_timesteps * end_noise_level_ratio) timesteps = reversed(range(end_noise_level_index, start_noise_level_index)) # configure initial img assert (start_noise_level_ratio == 1.0) or ( not guide_img is None), "A guide_img must be given to sample from a non-pure-noise." if guide_img is None: img = initial_noise else: guide_img, concat_points = self.get_deterministic_noise_tensor_repeat(shape[0], shape[3], reference_noise=guide_img) assert guide_img.shape == shape, "guide_img.shape != shape" if start_noise_level_index > 0: t = torch.full((shape[0],), start_noise_level_index-1, device=self.device).long() # -1 for start_noise_level_index not included img = self.q_sample(guide_img, t, noise=initial_noise) else: print("Zero noise added to the guidance latent representation.") img = guide_img # get masks n_masks = start_noise_level_index - end_noise_level_index if use_dynamic_mask: masks = self.get_dynamic_masks(n_masks, shape, concat_points, mask_flexivity) else: masks = [mask for _ in range(n_masks)] imgs = [img] current_mask = None for i in tqdm(timesteps, total=start_noise_level_index - end_noise_level_index, disable=self.mute): # if i == 3: # return [img], initial_noise # 第1排，第1列 img = self.p_sample(model, img, torch.full((shape[0],), i, device=self.device, dtype=torch.long), condition=condition, sampler=sampler) # if i == 3: # return [img], initial_noise # 第1排，第2列 if inpaint: if i > 0: t = torch.full((shape[0],), int(i-1), device=self.device).long() img_noise_t = self.q_sample(guide_img, t, noise=initial_noise) # if i == 3: # return [img_noise_t], initial_noise # 第2排，第2列 current_mask = masks.pop() img = current_mask * img_noise_t + (1 - current_mask) * img # if i == 3: # return [img], initial_noise # 第1.5排，最后1列 else: img = current_mask * guide_img + (1 - current_mask) * img if return_tensor: imgs.append(img) else: imgs.append(img.cpu().numpy()) return imgs, initial_noise def sample(self, model, shape, return_tensor=False, condition=None, sampler="ddim", initial_noise=None, seed=None): if not seed is None: torch.manual_seed(seed) return self.p_sample_loop(model, shape, initial_noise=initial_noise, start_noise_level_ratio=1.0, end_noise_level_ratio=0.0, return_tensor=return_tensor, condition=condition, sampler=sampler) def interpolate(self, model, shape, variance, first_endpoint=None, second_endpoint=None, return_tensor=False, condition=None, sampler="ddim", seed=None): if not seed is None: torch.manual_seed(seed) linear_noise = self.generate_linear_noise(shape, variance, first_endpoint=first_endpoint, second_endpoint=second_endpoint) return self.p_sample_loop(model, shape, initial_noise=linear_noise, start_noise_level_ratio=1.0, end_noise_level_ratio=0.0, return_tensor=return_tensor, condition=condition, sampler=sampler) def img_guided_sample(self, model, shape, noising_strength, guide_img, return_tensor=False, condition=None, sampler="ddim", initial_noise=None, seed=None): if not seed is None: torch.manual_seed(seed) assert guide_img.shape[-1] == shape[-1], "guide_img.shape[:-1] != shape[:-1]" return self.p_sample_loop(model, shape, start_noise_level_ratio=noising_strength, end_noise_level_ratio=0.0, return_tensor=return_tensor, condition=condition, sampler=sampler, guide_img=guide_img, initial_noise=initial_noise) def inpaint_sample(self, model, shape, noising_strength, guide_img, mask, return_tensor=False, condition=None, sampler="ddim", initial_noise=None, use_dynamic_mask=False, end_noise_level_ratio=0.0, seed=None, mask_flexivity=0.8): if not seed is None: torch.manual_seed(seed) return self.p_sample_loop(model, shape, start_noise_level_ratio=noising_strength, end_noise_level_ratio=end_noise_level_ratio, return_tensor=return_tensor, condition=condition, guide_img=guide_img, mask=mask, sampler=sampler, inpaint=True, initial_noise=initial_noise, use_dynamic_mask=use_dynamic_mask, mask_flexivity=mask_flexivity)