| | |
| | |
| | |
| | |
| |
|
| |
|
| | 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) |
| |
|
