Spaces:
Running
on
Zero
Running
on
Zero
Upload 2 files
Browse files
diffusers_helper/k_diffusion/uni_pc_fm.py
ADDED
@@ -0,0 +1,155 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Better Flow Matching UniPC by Lvmin Zhang
|
2 |
+
# (c) 2025
|
3 |
+
# CC BY-SA 4.0
|
4 |
+
# Attribution-ShareAlike 4.0 International Licence
|
5 |
+
|
6 |
+
|
7 |
+
import torch
|
8 |
+
|
9 |
+
from tqdm.auto import trange
|
10 |
+
|
11 |
+
|
12 |
+
def expand_dims(v, dims):
|
13 |
+
return v[(...,) + (None,) * (dims - 1)]
|
14 |
+
|
15 |
+
|
16 |
+
class FlowMatchUniPC:
|
17 |
+
def __init__(self, model, extra_args, variant='bh1'):
|
18 |
+
self.model = model
|
19 |
+
self.variant = variant
|
20 |
+
self.extra_args = extra_args
|
21 |
+
|
22 |
+
def model_fn(self, x, t):
|
23 |
+
return self.model(x, t, **self.extra_args)
|
24 |
+
|
25 |
+
def update_fn(self, x, model_prev_list, t_prev_list, t, order):
|
26 |
+
assert order <= len(model_prev_list)
|
27 |
+
dims = x.dim()
|
28 |
+
|
29 |
+
t_prev_0 = t_prev_list[-1]
|
30 |
+
lambda_prev_0 = - torch.log(t_prev_0)
|
31 |
+
lambda_t = - torch.log(t)
|
32 |
+
model_prev_0 = model_prev_list[-1]
|
33 |
+
|
34 |
+
h = lambda_t - lambda_prev_0
|
35 |
+
|
36 |
+
rks = []
|
37 |
+
D1s = []
|
38 |
+
for i in range(1, order):
|
39 |
+
t_prev_i = t_prev_list[-(i + 1)]
|
40 |
+
model_prev_i = model_prev_list[-(i + 1)]
|
41 |
+
lambda_prev_i = - torch.log(t_prev_i)
|
42 |
+
rk = ((lambda_prev_i - lambda_prev_0) / h)[0]
|
43 |
+
rks.append(rk)
|
44 |
+
D1s.append((model_prev_i - model_prev_0) / rk)
|
45 |
+
|
46 |
+
rks.append(1.)
|
47 |
+
rks = torch.tensor(rks, device=x.device)
|
48 |
+
|
49 |
+
R = []
|
50 |
+
b = []
|
51 |
+
|
52 |
+
hh = -h[0]
|
53 |
+
h_phi_1 = torch.expm1(hh)
|
54 |
+
h_phi_k = h_phi_1 / hh - 1
|
55 |
+
|
56 |
+
factorial_i = 1
|
57 |
+
|
58 |
+
if self.variant == 'bh1':
|
59 |
+
B_h = hh
|
60 |
+
elif self.variant == 'bh2':
|
61 |
+
B_h = torch.expm1(hh)
|
62 |
+
else:
|
63 |
+
raise NotImplementedError('Bad variant!')
|
64 |
+
|
65 |
+
for i in range(1, order + 1):
|
66 |
+
R.append(torch.pow(rks, i - 1))
|
67 |
+
b.append(h_phi_k * factorial_i / B_h)
|
68 |
+
factorial_i *= (i + 1)
|
69 |
+
h_phi_k = h_phi_k / hh - 1 / factorial_i
|
70 |
+
|
71 |
+
R = torch.stack(R)
|
72 |
+
b = torch.tensor(b, device=x.device)
|
73 |
+
|
74 |
+
use_predictor = len(D1s) > 0
|
75 |
+
|
76 |
+
if use_predictor:
|
77 |
+
D1s = torch.stack(D1s, dim=1)
|
78 |
+
if order == 2:
|
79 |
+
rhos_p = torch.tensor([0.5], device=b.device)
|
80 |
+
else:
|
81 |
+
rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
|
82 |
+
else:
|
83 |
+
D1s = None
|
84 |
+
rhos_p = None
|
85 |
+
|
86 |
+
if order == 1:
|
87 |
+
rhos_c = torch.tensor([0.5], device=b.device)
|
88 |
+
else:
|
89 |
+
rhos_c = torch.linalg.solve(R, b)
|
90 |
+
|
91 |
+
x_t_ = expand_dims(t / t_prev_0, dims) * x - expand_dims(h_phi_1, dims) * model_prev_0
|
92 |
+
|
93 |
+
if use_predictor:
|
94 |
+
pred_res = torch.tensordot(D1s, rhos_p, dims=([1], [0]))
|
95 |
+
else:
|
96 |
+
pred_res = 0
|
97 |
+
|
98 |
+
x_t = x_t_ - expand_dims(B_h, dims) * pred_res
|
99 |
+
model_t = self.model_fn(x_t, t)
|
100 |
+
|
101 |
+
if D1s is not None:
|
102 |
+
corr_res = torch.tensordot(D1s, rhos_c[:-1], dims=([1], [0]))
|
103 |
+
else:
|
104 |
+
corr_res = 0
|
105 |
+
|
106 |
+
D1_t = (model_t - model_prev_0)
|
107 |
+
x_t = x_t_ - expand_dims(B_h, dims) * (corr_res + rhos_c[-1] * D1_t)
|
108 |
+
|
109 |
+
return x_t, model_t
|
110 |
+
|
111 |
+
def sample(self, x, sigmas, callback=None, disable_pbar=False):
|
112 |
+
order = min(3, len(sigmas) - 2)
|
113 |
+
model_prev_list, t_prev_list = [], []
|
114 |
+
try:
|
115 |
+
for i in trange(len(sigmas) - 1, disable=disable_pbar):
|
116 |
+
vec_t = sigmas[i].expand(x.shape[0])
|
117 |
+
|
118 |
+
if i == 0:
|
119 |
+
model_prev_list = [self.model_fn(x, vec_t)]
|
120 |
+
t_prev_list = [vec_t]
|
121 |
+
elif i < order:
|
122 |
+
init_order = i
|
123 |
+
x, model_x = self.update_fn(x, model_prev_list, t_prev_list, vec_t, init_order)
|
124 |
+
model_prev_list.append(model_x)
|
125 |
+
t_prev_list.append(vec_t)
|
126 |
+
else:
|
127 |
+
x, model_x = self.update_fn(x, model_prev_list, t_prev_list, vec_t, order)
|
128 |
+
model_prev_list.append(model_x)
|
129 |
+
t_prev_list.append(vec_t)
|
130 |
+
|
131 |
+
model_prev_list = model_prev_list[-order:]
|
132 |
+
t_prev_list = t_prev_list[-order:]
|
133 |
+
|
134 |
+
if callback is not None:
|
135 |
+
try:
|
136 |
+
callback({'x': x, 'i': i, 'denoised': model_prev_list[-1]})
|
137 |
+
except KeyboardInterrupt as e:
|
138 |
+
print(f"User interruption detected: {e}")
|
139 |
+
# Return the last available result
|
140 |
+
return model_prev_list[-1]
|
141 |
+
except KeyboardInterrupt as e:
|
142 |
+
print(f"Process interrupted: {e}")
|
143 |
+
# Return the last available result if we have one
|
144 |
+
if model_prev_list:
|
145 |
+
return model_prev_list[-1]
|
146 |
+
else:
|
147 |
+
# If no results yet, re-raise the exception
|
148 |
+
raise
|
149 |
+
|
150 |
+
return model_prev_list[-1]
|
151 |
+
|
152 |
+
|
153 |
+
def sample_unipc(model, noise, sigmas, extra_args=None, callback=None, disable=False, variant='bh1'):
|
154 |
+
assert variant in ['bh1', 'bh2']
|
155 |
+
return FlowMatchUniPC(model, extra_args=extra_args, variant=variant).sample(noise, sigmas=sigmas, callback=callback, disable_pbar=disable)
|
diffusers_helper/k_diffusion/wrapper.py
ADDED
@@ -0,0 +1,51 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
|
3 |
+
|
4 |
+
def append_dims(x, target_dims):
|
5 |
+
return x[(...,) + (None,) * (target_dims - x.ndim)]
|
6 |
+
|
7 |
+
|
8 |
+
def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0):
|
9 |
+
if guidance_rescale == 0:
|
10 |
+
return noise_cfg
|
11 |
+
|
12 |
+
std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
|
13 |
+
std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
|
14 |
+
noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
|
15 |
+
noise_cfg = guidance_rescale * noise_pred_rescaled + (1.0 - guidance_rescale) * noise_cfg
|
16 |
+
return noise_cfg
|
17 |
+
|
18 |
+
|
19 |
+
def fm_wrapper(transformer, t_scale=1000.0):
|
20 |
+
def k_model(x, sigma, **extra_args):
|
21 |
+
dtype = extra_args['dtype']
|
22 |
+
cfg_scale = extra_args['cfg_scale']
|
23 |
+
cfg_rescale = extra_args['cfg_rescale']
|
24 |
+
concat_latent = extra_args['concat_latent']
|
25 |
+
|
26 |
+
original_dtype = x.dtype
|
27 |
+
sigma = sigma.float()
|
28 |
+
|
29 |
+
x = x.to(dtype)
|
30 |
+
timestep = (sigma * t_scale).to(dtype)
|
31 |
+
|
32 |
+
if concat_latent is None:
|
33 |
+
hidden_states = x
|
34 |
+
else:
|
35 |
+
hidden_states = torch.cat([x, concat_latent.to(x)], dim=1)
|
36 |
+
|
37 |
+
pred_positive = transformer(hidden_states=hidden_states, timestep=timestep, return_dict=False, **extra_args['positive'])[0].float()
|
38 |
+
|
39 |
+
if cfg_scale == 1.0:
|
40 |
+
pred_negative = torch.zeros_like(pred_positive)
|
41 |
+
else:
|
42 |
+
pred_negative = transformer(hidden_states=hidden_states, timestep=timestep, return_dict=False, **extra_args['negative'])[0].float()
|
43 |
+
|
44 |
+
pred_cfg = pred_negative + cfg_scale * (pred_positive - pred_negative)
|
45 |
+
pred = rescale_noise_cfg(pred_cfg, pred_positive, guidance_rescale=cfg_rescale)
|
46 |
+
|
47 |
+
x0 = x.float() - pred.float() * append_dims(sigma, x.ndim)
|
48 |
+
|
49 |
+
return x0.to(dtype=original_dtype)
|
50 |
+
|
51 |
+
return k_model
|