|
|
|
|
|
|
|
|
|
|
|
|
| import torch
|
|
|
| from tqdm.auto import trange
|
|
|
|
|
| def expand_dims(v, dims):
|
| return v[(...,) + (None,) * (dims - 1)]
|
|
|
|
|
| class FlowMatchUniPC:
|
| def __init__(self, model, extra_args, variant='bh1'):
|
| self.model = model
|
| self.variant = variant
|
| self.extra_args = extra_args
|
|
|
| def model_fn(self, x, t):
|
| return self.model(x, t, **self.extra_args)
|
|
|
| def update_fn(self, x, model_prev_list, t_prev_list, t, order):
|
| assert order <= len(model_prev_list)
|
| dims = x.dim()
|
|
|
| t_prev_0 = t_prev_list[-1]
|
| lambda_prev_0 = - torch.log(t_prev_0)
|
| lambda_t = - torch.log(t)
|
| model_prev_0 = model_prev_list[-1]
|
|
|
| h = lambda_t - lambda_prev_0
|
|
|
| rks = []
|
| D1s = []
|
| for i in range(1, order):
|
| t_prev_i = t_prev_list[-(i + 1)]
|
| model_prev_i = model_prev_list[-(i + 1)]
|
| lambda_prev_i = - torch.log(t_prev_i)
|
| rk = ((lambda_prev_i - lambda_prev_0) / h)[0]
|
| rks.append(rk)
|
| D1s.append((model_prev_i - model_prev_0) / rk)
|
|
|
| rks.append(1.)
|
| rks = torch.tensor(rks, device=x.device)
|
|
|
| R = []
|
| b = []
|
|
|
| hh = -h[0]
|
| h_phi_1 = torch.expm1(hh)
|
| h_phi_k = h_phi_1 / hh - 1
|
|
|
| factorial_i = 1
|
|
|
| if self.variant == 'bh1':
|
| B_h = hh
|
| elif self.variant == 'bh2':
|
| B_h = torch.expm1(hh)
|
| else:
|
| raise NotImplementedError('Bad variant!')
|
|
|
| for i in range(1, order + 1):
|
| R.append(torch.pow(rks, i - 1))
|
| b.append(h_phi_k * factorial_i / B_h)
|
| factorial_i *= (i + 1)
|
| h_phi_k = h_phi_k / hh - 1 / factorial_i
|
|
|
| R = torch.stack(R)
|
| b = torch.tensor(b, device=x.device)
|
|
|
| use_predictor = len(D1s) > 0
|
|
|
| if use_predictor:
|
| D1s = torch.stack(D1s, dim=1)
|
| if order == 2:
|
| rhos_p = torch.tensor([0.5], device=b.device)
|
| else:
|
| rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
|
| else:
|
| D1s = None
|
| rhos_p = None
|
|
|
| if order == 1:
|
| rhos_c = torch.tensor([0.5], device=b.device)
|
| else:
|
| rhos_c = torch.linalg.solve(R, b)
|
|
|
| x_t_ = expand_dims(t / t_prev_0, dims) * x - expand_dims(h_phi_1, dims) * model_prev_0
|
|
|
| if use_predictor:
|
| pred_res = torch.tensordot(D1s, rhos_p, dims=([1], [0]))
|
| else:
|
| pred_res = 0
|
|
|
| x_t = x_t_ - expand_dims(B_h, dims) * pred_res
|
| model_t = self.model_fn(x_t, t)
|
|
|
| if D1s is not None:
|
| corr_res = torch.tensordot(D1s, rhos_c[:-1], dims=([1], [0]))
|
| else:
|
| corr_res = 0
|
|
|
| D1_t = (model_t - model_prev_0)
|
| x_t = x_t_ - expand_dims(B_h, dims) * (corr_res + rhos_c[-1] * D1_t)
|
|
|
| return x_t, model_t
|
|
|
| def sample(self, x, sigmas, callback=None, disable_pbar=False):
|
| order = min(3, len(sigmas) - 2)
|
| model_prev_list, t_prev_list = [], []
|
| for i in trange(len(sigmas) - 1, disable=disable_pbar):
|
| vec_t = sigmas[i].expand(x.shape[0])
|
|
|
| if i == 0:
|
| model_prev_list = [self.model_fn(x, vec_t)]
|
| t_prev_list = [vec_t]
|
| elif i < order:
|
| init_order = i
|
| x, model_x = self.update_fn(x, model_prev_list, t_prev_list, vec_t, init_order)
|
| model_prev_list.append(model_x)
|
| t_prev_list.append(vec_t)
|
| else:
|
| x, model_x = self.update_fn(x, model_prev_list, t_prev_list, vec_t, order)
|
| model_prev_list.append(model_x)
|
| t_prev_list.append(vec_t)
|
|
|
| model_prev_list = model_prev_list[-order:]
|
| t_prev_list = t_prev_list[-order:]
|
|
|
| if callback is not None:
|
| callback({'x': x, 'i': i, 'denoised': model_prev_list[-1]})
|
|
|
| return model_prev_list[-1]
|
|
|
|
|
| def sample_unipc(model, noise, sigmas, extra_args=None, callback=None, disable=False, variant='bh1'):
|
| assert variant in ['bh1', 'bh2']
|
| return FlowMatchUniPC(model, extra_args=extra_args, variant=variant).sample(noise, sigmas=sigmas, callback=callback, disable_pbar=disable)
|
|
|
