Spaces:
Running
on
Zero
Running
on
Zero
fallenshock
commited on
Commit
•
336dbcf
1
Parent(s):
2969ec0
added files
Browse files- FlowEdit_utils.py +404 -0
- app.py +272 -0
- requirements.txt +81 -0
FlowEdit_utils.py
ADDED
@@ -0,0 +1,404 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from typing import Optional, Tuple, Union
|
2 |
+
import torch
|
3 |
+
from diffusers import FlowMatchEulerDiscreteScheduler
|
4 |
+
from tqdm import tqdm
|
5 |
+
import numpy as np
|
6 |
+
|
7 |
+
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
|
8 |
+
|
9 |
+
|
10 |
+
|
11 |
+
def scale_noise(
|
12 |
+
scheduler,
|
13 |
+
sample: torch.FloatTensor,
|
14 |
+
timestep: Union[float, torch.FloatTensor],
|
15 |
+
noise: Optional[torch.FloatTensor] = None,
|
16 |
+
) -> torch.FloatTensor:
|
17 |
+
"""
|
18 |
+
Foward process in flow-matching
|
19 |
+
|
20 |
+
Args:
|
21 |
+
sample (`torch.FloatTensor`):
|
22 |
+
The input sample.
|
23 |
+
timestep (`int`, *optional*):
|
24 |
+
The current timestep in the diffusion chain.
|
25 |
+
|
26 |
+
Returns:
|
27 |
+
`torch.FloatTensor`:
|
28 |
+
A scaled input sample.
|
29 |
+
"""
|
30 |
+
# if scheduler.step_index is None:
|
31 |
+
scheduler._init_step_index(timestep)
|
32 |
+
|
33 |
+
sigma = scheduler.sigmas[scheduler.step_index]
|
34 |
+
sample = sigma * noise + (1.0 - sigma) * sample
|
35 |
+
|
36 |
+
return sample
|
37 |
+
|
38 |
+
|
39 |
+
# for flux
|
40 |
+
def calculate_shift(
|
41 |
+
image_seq_len,
|
42 |
+
base_seq_len: int = 256,
|
43 |
+
max_seq_len: int = 4096,
|
44 |
+
base_shift: float = 0.5,
|
45 |
+
max_shift: float = 1.16,
|
46 |
+
):
|
47 |
+
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
|
48 |
+
b = base_shift - m * base_seq_len
|
49 |
+
mu = image_seq_len * m + b
|
50 |
+
return mu
|
51 |
+
|
52 |
+
|
53 |
+
|
54 |
+
def calc_v_sd3(pipe, src_tar_latent_model_input, src_tar_prompt_embeds, src_tar_pooled_prompt_embeds, src_guidance_scale, tar_guidance_scale, t):
|
55 |
+
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
56 |
+
timestep = t.expand(src_tar_latent_model_input.shape[0])
|
57 |
+
# joint_attention_kwargs = {}
|
58 |
+
# # add timestep to joint_attention_kwargs
|
59 |
+
# joint_attention_kwargs["timestep"] = timestep[0]
|
60 |
+
# joint_attention_kwargs["timestep_idx"] = i
|
61 |
+
|
62 |
+
|
63 |
+
with torch.no_grad():
|
64 |
+
# # predict the noise for the source prompt
|
65 |
+
noise_pred_src_tar = pipe.transformer(
|
66 |
+
hidden_states=src_tar_latent_model_input,
|
67 |
+
timestep=timestep,
|
68 |
+
encoder_hidden_states=src_tar_prompt_embeds,
|
69 |
+
pooled_projections=src_tar_pooled_prompt_embeds,
|
70 |
+
joint_attention_kwargs=None,
|
71 |
+
return_dict=False,
|
72 |
+
)[0]
|
73 |
+
|
74 |
+
# perform guidance source
|
75 |
+
if pipe.do_classifier_free_guidance:
|
76 |
+
src_noise_pred_uncond, src_noise_pred_text, tar_noise_pred_uncond, tar_noise_pred_text = noise_pred_src_tar.chunk(4)
|
77 |
+
noise_pred_src = src_noise_pred_uncond + src_guidance_scale * (src_noise_pred_text - src_noise_pred_uncond)
|
78 |
+
noise_pred_tar = tar_noise_pred_uncond + tar_guidance_scale * (tar_noise_pred_text - tar_noise_pred_uncond)
|
79 |
+
|
80 |
+
return noise_pred_src, noise_pred_tar
|
81 |
+
|
82 |
+
|
83 |
+
|
84 |
+
def calc_v_flux(pipe, latents, prompt_embeds, pooled_prompt_embeds, guidance, text_ids, latent_image_ids, t):
|
85 |
+
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
86 |
+
timestep = t.expand(latents.shape[0])
|
87 |
+
# joint_attention_kwargs = {}
|
88 |
+
# # add timestep to joint_attention_kwargs
|
89 |
+
# joint_attention_kwargs["timestep"] = timestep[0]
|
90 |
+
# joint_attention_kwargs["timestep_idx"] = i
|
91 |
+
|
92 |
+
|
93 |
+
with torch.no_grad():
|
94 |
+
# # predict the noise for the source prompt
|
95 |
+
noise_pred = pipe.transformer(
|
96 |
+
hidden_states=latents,
|
97 |
+
timestep=timestep / 1000,
|
98 |
+
guidance=guidance,
|
99 |
+
encoder_hidden_states=prompt_embeds,
|
100 |
+
txt_ids=text_ids,
|
101 |
+
img_ids=latent_image_ids,
|
102 |
+
pooled_projections=pooled_prompt_embeds,
|
103 |
+
joint_attention_kwargs=None,
|
104 |
+
return_dict=False,
|
105 |
+
)[0]
|
106 |
+
|
107 |
+
return noise_pred
|
108 |
+
|
109 |
+
|
110 |
+
|
111 |
+
@torch.no_grad()
|
112 |
+
def FlowEditSD3(pipe,
|
113 |
+
scheduler,
|
114 |
+
x_src,
|
115 |
+
src_prompt,
|
116 |
+
tar_prompt,
|
117 |
+
negative_prompt,
|
118 |
+
T_steps: int = 50,
|
119 |
+
n_avg: int = 1,
|
120 |
+
src_guidance_scale: float = 3.5,
|
121 |
+
tar_guidance_scale: float = 13.5,
|
122 |
+
n_min: int = 0,
|
123 |
+
n_max: int = 15,):
|
124 |
+
|
125 |
+
device = x_src.device
|
126 |
+
|
127 |
+
timesteps, T_steps = retrieve_timesteps(scheduler, T_steps, device, timesteps=None)
|
128 |
+
|
129 |
+
num_warmup_steps = max(len(timesteps) - T_steps * scheduler.order, 0)
|
130 |
+
pipe._num_timesteps = len(timesteps)
|
131 |
+
pipe._guidance_scale = src_guidance_scale
|
132 |
+
|
133 |
+
# src prompts
|
134 |
+
(
|
135 |
+
src_prompt_embeds,
|
136 |
+
src_negative_prompt_embeds,
|
137 |
+
src_pooled_prompt_embeds,
|
138 |
+
src_negative_pooled_prompt_embeds,
|
139 |
+
) = pipe.encode_prompt(
|
140 |
+
prompt=src_prompt,
|
141 |
+
prompt_2=None,
|
142 |
+
prompt_3=None,
|
143 |
+
negative_prompt=negative_prompt,
|
144 |
+
do_classifier_free_guidance=pipe.do_classifier_free_guidance,
|
145 |
+
device=device,
|
146 |
+
)
|
147 |
+
|
148 |
+
# tar prompts
|
149 |
+
pipe._guidance_scale = tar_guidance_scale
|
150 |
+
(
|
151 |
+
tar_prompt_embeds,
|
152 |
+
tar_negative_prompt_embeds,
|
153 |
+
tar_pooled_prompt_embeds,
|
154 |
+
tar_negative_pooled_prompt_embeds,
|
155 |
+
) = pipe.encode_prompt(
|
156 |
+
prompt=tar_prompt,
|
157 |
+
prompt_2=None,
|
158 |
+
prompt_3=None,
|
159 |
+
negative_prompt=negative_prompt,
|
160 |
+
do_classifier_free_guidance=pipe.do_classifier_free_guidance,
|
161 |
+
device=device,
|
162 |
+
)
|
163 |
+
|
164 |
+
# CFG prep
|
165 |
+
src_tar_prompt_embeds = torch.cat([src_negative_prompt_embeds, src_prompt_embeds, tar_negative_prompt_embeds, tar_prompt_embeds], dim=0)
|
166 |
+
src_tar_pooled_prompt_embeds = torch.cat([src_negative_pooled_prompt_embeds, src_pooled_prompt_embeds, tar_negative_pooled_prompt_embeds, tar_pooled_prompt_embeds], dim=0)
|
167 |
+
|
168 |
+
# initialize our ODE Zt_edit_1=x_src
|
169 |
+
zt_edit = x_src.clone()
|
170 |
+
|
171 |
+
for i, t in tqdm(enumerate(timesteps)):
|
172 |
+
|
173 |
+
if T_steps - i > n_max:
|
174 |
+
continue
|
175 |
+
|
176 |
+
t_i = t/1000
|
177 |
+
if i+1 < len(timesteps):
|
178 |
+
t_im1 = (timesteps[i+1])/1000
|
179 |
+
else:
|
180 |
+
t_im1 = torch.zeros_like(t_i).to(t_i.device)
|
181 |
+
|
182 |
+
if T_steps - i > n_min:
|
183 |
+
|
184 |
+
# Calculate the average of the V predictions
|
185 |
+
V_delta_avg = torch.zeros_like(x_src)
|
186 |
+
for k in range(n_avg):
|
187 |
+
|
188 |
+
fwd_noise = torch.randn_like(x_src).to(x_src.device)
|
189 |
+
|
190 |
+
zt_src = (1-t_i)*x_src + (t_i)*fwd_noise
|
191 |
+
|
192 |
+
zt_tar = zt_edit + zt_src - x_src
|
193 |
+
|
194 |
+
src_tar_latent_model_input = torch.cat([zt_src, zt_src, zt_tar, zt_tar]) if pipe.do_classifier_free_guidance else (zt_src, zt_tar)
|
195 |
+
|
196 |
+
Vt_src, Vt_tar = calc_v_sd3(pipe, src_tar_latent_model_input,src_tar_prompt_embeds, src_tar_pooled_prompt_embeds, src_guidance_scale, tar_guidance_scale, t)
|
197 |
+
|
198 |
+
V_delta_avg += (1/n_avg) * (Vt_tar - Vt_src) # - (hfg-1)*( x_src))
|
199 |
+
|
200 |
+
# propagate direct ODE
|
201 |
+
zt_edit = zt_edit.to(torch.float32)
|
202 |
+
|
203 |
+
zt_edit = zt_edit + (t_im1 - t_i) * V_delta_avg
|
204 |
+
|
205 |
+
zt_edit = zt_edit.to(V_delta_avg.dtype)
|
206 |
+
|
207 |
+
else: # i >= T_steps-n_min # regular sampling for last n_min steps
|
208 |
+
|
209 |
+
if i == T_steps-n_min:
|
210 |
+
# initialize SDEDIT-style generation phase
|
211 |
+
fwd_noise = torch.randn_like(x_src).to(x_src.device)
|
212 |
+
xt_src = scale_noise(scheduler, x_src, t, noise=fwd_noise)
|
213 |
+
xt_tar = zt_edit + xt_src - x_src
|
214 |
+
|
215 |
+
src_tar_latent_model_input = torch.cat([xt_tar, xt_tar, xt_tar, xt_tar]) if pipe.do_classifier_free_guidance else (xt_src, xt_tar)
|
216 |
+
|
217 |
+
_, noise_pred_tar = calc_v_sd3(pipe, src_tar_latent_model_input,src_tar_prompt_embeds, src_tar_pooled_prompt_embeds, src_guidance_scale, tar_guidance_scale, t)
|
218 |
+
|
219 |
+
xt_tar = xt_tar.to(torch.float32)
|
220 |
+
|
221 |
+
prev_sample = xt_tar + (t_im1 - t_im1) * (noise_pred_tar)
|
222 |
+
|
223 |
+
prev_sample = prev_sample.to(noise_pred_tar.dtype)
|
224 |
+
|
225 |
+
xt_tar = prev_sample
|
226 |
+
|
227 |
+
return zt_edit if n_min == 0 else xt_tar
|
228 |
+
|
229 |
+
|
230 |
+
|
231 |
+
@torch.no_grad()
|
232 |
+
def FlowEditFLUX(pipe,
|
233 |
+
scheduler,
|
234 |
+
x_src,
|
235 |
+
src_prompt,
|
236 |
+
tar_prompt,
|
237 |
+
negative_prompt,
|
238 |
+
T_steps: int = 28,
|
239 |
+
n_avg: int = 1,
|
240 |
+
src_guidance_scale: float = 1.5,
|
241 |
+
tar_guidance_scale: float = 5.5,
|
242 |
+
n_min: int = 0,
|
243 |
+
n_max: int = 24,):
|
244 |
+
|
245 |
+
device = x_src.device
|
246 |
+
orig_height, orig_width = x_src.shape[2]*pipe.vae_scale_factor//2, x_src.shape[3]*pipe.vae_scale_factor//2
|
247 |
+
num_channels_latents = pipe.transformer.config.in_channels // 4
|
248 |
+
|
249 |
+
pipe.check_inputs(
|
250 |
+
prompt=src_prompt,
|
251 |
+
prompt_2=None,
|
252 |
+
height=orig_height,
|
253 |
+
width=orig_width,
|
254 |
+
callback_on_step_end_tensor_inputs=None,
|
255 |
+
max_sequence_length=512,
|
256 |
+
)
|
257 |
+
|
258 |
+
x_src, latent_src_image_ids = pipe.prepare_latents(batch_size= x_src.shape[0], num_channels_latents=num_channels_latents, height=orig_height, width=orig_width, dtype=x_src.dtype, device=x_src.device, generator=None,latents=x_src)
|
259 |
+
x_src_packed = pipe._pack_latents(x_src, x_src.shape[0], num_channels_latents, x_src.shape[2], x_src.shape[3])
|
260 |
+
latent_tar_image_ids = latent_src_image_ids
|
261 |
+
|
262 |
+
# 5. Prepare timesteps
|
263 |
+
sigmas = np.linspace(1.0, 1 / T_steps, T_steps)
|
264 |
+
image_seq_len = x_src_packed.shape[1]
|
265 |
+
mu = calculate_shift(
|
266 |
+
image_seq_len,
|
267 |
+
scheduler.config.base_image_seq_len,
|
268 |
+
scheduler.config.max_image_seq_len,
|
269 |
+
scheduler.config.base_shift,
|
270 |
+
scheduler.config.max_shift,
|
271 |
+
)
|
272 |
+
timesteps, T_steps = retrieve_timesteps(
|
273 |
+
scheduler,
|
274 |
+
T_steps,
|
275 |
+
device,
|
276 |
+
timesteps=None,
|
277 |
+
sigmas=sigmas,
|
278 |
+
mu=mu,
|
279 |
+
)
|
280 |
+
|
281 |
+
num_warmup_steps = max(len(timesteps) - T_steps * pipe.scheduler.order, 0)
|
282 |
+
pipe._num_timesteps = len(timesteps)
|
283 |
+
|
284 |
+
|
285 |
+
# src prompts
|
286 |
+
(
|
287 |
+
src_prompt_embeds,
|
288 |
+
src_pooled_prompt_embeds,
|
289 |
+
src_text_ids,
|
290 |
+
|
291 |
+
) = pipe.encode_prompt(
|
292 |
+
prompt=src_prompt,
|
293 |
+
prompt_2=None,
|
294 |
+
device=device,
|
295 |
+
)
|
296 |
+
|
297 |
+
# tar prompts
|
298 |
+
pipe._guidance_scale = tar_guidance_scale
|
299 |
+
(
|
300 |
+
tar_prompt_embeds,
|
301 |
+
tar_pooled_prompt_embeds,
|
302 |
+
tar_text_ids,
|
303 |
+
) = pipe.encode_prompt(
|
304 |
+
prompt=tar_prompt,
|
305 |
+
prompt_2=None,
|
306 |
+
device=device,
|
307 |
+
)
|
308 |
+
|
309 |
+
# handle guidance
|
310 |
+
if pipe.transformer.config.guidance_embeds:
|
311 |
+
src_guidance = torch.tensor([src_guidance_scale], device=device)
|
312 |
+
src_guidance = src_guidance.expand(x_src_packed.shape[0])
|
313 |
+
tar_guidance = torch.tensor([tar_guidance_scale], device=device)
|
314 |
+
tar_guidance = tar_guidance.expand(x_src_packed.shape[0])
|
315 |
+
else:
|
316 |
+
src_guidance = None
|
317 |
+
tar_guidance = None
|
318 |
+
|
319 |
+
# initialize our ODE Zt_edit_1=x_src
|
320 |
+
zt_edit = x_src_packed.clone()
|
321 |
+
|
322 |
+
for i, t in tqdm(enumerate(timesteps)):
|
323 |
+
|
324 |
+
if T_steps - i > n_max:
|
325 |
+
continue
|
326 |
+
|
327 |
+
scheduler._init_step_index(t)
|
328 |
+
t_i = scheduler.sigmas[scheduler.step_index]
|
329 |
+
if i < len(timesteps):
|
330 |
+
t_im1 = scheduler.sigmas[scheduler.step_index + 1]
|
331 |
+
else:
|
332 |
+
t_im1 = t_i
|
333 |
+
|
334 |
+
if T_steps - i > n_min:
|
335 |
+
|
336 |
+
# Calculate the average of the V predictions
|
337 |
+
V_delta_avg = torch.zeros_like(x_src_packed)
|
338 |
+
|
339 |
+
for k in range(n_avg):
|
340 |
+
|
341 |
+
|
342 |
+
fwd_noise = torch.randn_like(x_src_packed).to(x_src_packed.device)
|
343 |
+
|
344 |
+
zt_src = (1-t_i)*x_src_packed + (t_i)*fwd_noise
|
345 |
+
|
346 |
+
zt_tar = zt_edit + zt_src - x_src_packed
|
347 |
+
|
348 |
+
# Merge in the future to avoid double computation
|
349 |
+
Vt_src = calc_v_flux(pipe,
|
350 |
+
latents=zt_src,
|
351 |
+
prompt_embeds=src_prompt_embeds,
|
352 |
+
pooled_prompt_embeds=src_pooled_prompt_embeds,
|
353 |
+
guidance=src_guidance,
|
354 |
+
text_ids=src_text_ids,
|
355 |
+
latent_image_ids=latent_src_image_ids,
|
356 |
+
t=t)
|
357 |
+
|
358 |
+
Vt_tar = calc_v_flux(pipe,
|
359 |
+
latents=zt_tar,
|
360 |
+
prompt_embeds=tar_prompt_embeds,
|
361 |
+
pooled_prompt_embeds=tar_pooled_prompt_embeds,
|
362 |
+
guidance=tar_guidance,
|
363 |
+
text_ids=tar_text_ids,
|
364 |
+
latent_image_ids=latent_tar_image_ids,
|
365 |
+
t=t)
|
366 |
+
|
367 |
+
V_delta_avg += (1/n_avg) * (Vt_tar - Vt_src) # - (hfg-1)*( x_src))
|
368 |
+
|
369 |
+
# propagate direct ODE
|
370 |
+
zt_edit = zt_edit.to(torch.float32)
|
371 |
+
|
372 |
+
zt_edit = zt_edit + (t_im1 - t_i) * V_delta_avg
|
373 |
+
|
374 |
+
zt_edit = zt_edit.to(V_delta_avg.dtype)
|
375 |
+
|
376 |
+
else: # i >= T_steps-n_min # regular sampling last n_min steps
|
377 |
+
|
378 |
+
if i == T_steps-n_min:
|
379 |
+
# initialize SDEDIT-style generation phase
|
380 |
+
fwd_noise = torch.randn_like(x_src_packed).to(x_src_packed.device)
|
381 |
+
xt_src = scale_noise(scheduler, x_src_packed, t, noise=fwd_noise)
|
382 |
+
xt_tar = zt_edit + xt_src - x_src_packed
|
383 |
+
|
384 |
+
Vt_tar = calc_v_flux(pipe,
|
385 |
+
latents=xt_tar,
|
386 |
+
prompt_embeds=tar_prompt_embeds,
|
387 |
+
pooled_prompt_embeds=tar_pooled_prompt_embeds,
|
388 |
+
guidance=tar_guidance,
|
389 |
+
text_ids=tar_text_ids,
|
390 |
+
latent_image_ids=latent_tar_image_ids,
|
391 |
+
t=t)
|
392 |
+
|
393 |
+
|
394 |
+
xt_tar = xt_tar.to(torch.float32)
|
395 |
+
|
396 |
+
prev_sample = xt_tar + (t_im1 - t_i) * (Vt_tar)
|
397 |
+
|
398 |
+
prev_sample = prev_sample.to(Vt_tar.dtype)
|
399 |
+
xt_tar = prev_sample
|
400 |
+
out = zt_edit if n_min == 0 else xt_tar
|
401 |
+
unpacked_out = pipe._unpack_latents(out, orig_height, orig_width, pipe.vae_scale_factor)
|
402 |
+
return unpacked_out
|
403 |
+
|
404 |
+
|
app.py
ADDED
@@ -0,0 +1,272 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
import torch
|
3 |
+
from diffusers import FluxPipeline, StableDiffusion3Pipeline
|
4 |
+
from PIL import Image
|
5 |
+
|
6 |
+
import random
|
7 |
+
import numpy as np
|
8 |
+
import spaces
|
9 |
+
|
10 |
+
from FlowEdit_utils import FlowEditSD3, FlowEditFLUX
|
11 |
+
|
12 |
+
|
13 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
14 |
+
# device = "cpu"
|
15 |
+
# model_type = 'SD3'
|
16 |
+
|
17 |
+
pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
|
18 |
+
scheduler = pipe.scheduler
|
19 |
+
pipe = pipe.to(device)
|
20 |
+
loaded_model = 'SD3'
|
21 |
+
|
22 |
+
|
23 |
+
def on_model_change(model_type):
|
24 |
+
if model_type == 'SD3':
|
25 |
+
|
26 |
+
T_steps_value = 50
|
27 |
+
|
28 |
+
src_guidance_scale_value = 3.5
|
29 |
+
|
30 |
+
tar_guidance_scale_value = 13.5
|
31 |
+
|
32 |
+
n_max_value = 33
|
33 |
+
|
34 |
+
elif model_type == 'FLUX':
|
35 |
+
|
36 |
+
T_steps_value = 28
|
37 |
+
|
38 |
+
src_guidance_scale_value = 1.5
|
39 |
+
|
40 |
+
tar_guidance_scale_value = 5.5
|
41 |
+
|
42 |
+
n_max_value = 24
|
43 |
+
|
44 |
+
else:
|
45 |
+
raise NotImplementedError(f"Model type {model_type} not implemented")
|
46 |
+
|
47 |
+
return T_steps_value, src_guidance_scale_value, tar_guidance_scale_value, n_max_value
|
48 |
+
|
49 |
+
def get_examples():
|
50 |
+
case = [
|
51 |
+
["inputs/cat.png", "SD3", 50, 3.5, 13.5, 33, "a cat sitting in the grass", "a puppy sitting in the grass", 0, 1, 42],
|
52 |
+
["inputs/gas_station.png", "SD3", 50, 3.5, 13.5, 33, "cars are parked in front of a gas station with a sign that says \"CAFE\"", "cars are parked in front of a gas station with a sign that says \"CVPR\"", 0, 1, 42],
|
53 |
+
["inputs/iguana.png", "SD3", 50, 3.5, 13.5, 31, "A large orange lizard sitting on a rock near the ocean. The lizard is positioned in the center of the scene, with the ocean waves visible in the background. The rock is located close to the water, providing a picturesque setting for the lizard''s resting spot.", "A large dragon sitting on a rock near the ocean. The dragon is positioned in the center of the scene, with the ocean waves visible in the background. The rock is located close to the water, providing a picturesque setting for the dragon''s resting spot.", 0, 1, 42],
|
54 |
+
["inputs/cat.png", "FLUX", 28, 1.5, 5.5, 24, "a cat sitting in the grass", "a puppy sitting in the grass", 0, 1, 42],
|
55 |
+
["inputs/gas_station.png", "FLUX", 28, 1.5, 5.5, 24, "cars are parked in front of a gas station with a sign that says \"CAFE\"", "cars are parked in front of a gas station with a sign that says \"CVPR\"", 0, 1, 23],
|
56 |
+
["inputs/steak.png", "FLUX", 28, 1.5, 5.5, 24, "A steak accompanied by a side of leaf salad.", "A bread roll accompanied by a side of leaf salad.", 0, 1, 42],
|
57 |
+
]
|
58 |
+
return case
|
59 |
+
|
60 |
+
|
61 |
+
@spaces.GPU()
|
62 |
+
def FlowEditRun(
|
63 |
+
image_src: str,
|
64 |
+
model_type: str,
|
65 |
+
T_steps: int,
|
66 |
+
src_guidance_scale: float,
|
67 |
+
tar_guidance_scale: float,
|
68 |
+
n_max: int,
|
69 |
+
src_prompt: str,
|
70 |
+
tar_prompt: str,
|
71 |
+
n_min: int,
|
72 |
+
n_avg: int,
|
73 |
+
seed: int,
|
74 |
+
|
75 |
+
):
|
76 |
+
|
77 |
+
if not len(src_prompt):
|
78 |
+
raise gr.Error("source prompt cannot be empty")
|
79 |
+
if not len(tar_prompt):
|
80 |
+
raise gr.Error("target prompt cannot be empty")
|
81 |
+
|
82 |
+
global pipe
|
83 |
+
global scheduler
|
84 |
+
global loaded_model
|
85 |
+
|
86 |
+
# reload model only if different from the loaded model
|
87 |
+
if loaded_model != model_type:
|
88 |
+
|
89 |
+
if model_type == 'FLUX':
|
90 |
+
# pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.float16)
|
91 |
+
pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.float16)
|
92 |
+
loaded_model = 'FLUX'
|
93 |
+
elif model_type == 'SD3':
|
94 |
+
pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
|
95 |
+
loaded_model = 'SD3'
|
96 |
+
else:
|
97 |
+
raise NotImplementedError(f"Model type {model_type} not implemented")
|
98 |
+
|
99 |
+
scheduler = pipe.scheduler
|
100 |
+
pipe = pipe.to(device)
|
101 |
+
|
102 |
+
|
103 |
+
|
104 |
+
|
105 |
+
# set seed
|
106 |
+
random.seed(seed)
|
107 |
+
np.random.seed(seed)
|
108 |
+
torch.manual_seed(seed)
|
109 |
+
torch.cuda.manual_seed_all(seed)
|
110 |
+
# load image
|
111 |
+
image = Image.open(image_src)
|
112 |
+
# crop image to have both dimensions divisibe by 16 - avoids issues with resizing
|
113 |
+
image = image.crop((0, 0, image.width - image.width % 16, image.height - image.height % 16))
|
114 |
+
image_src = pipe.image_processor.preprocess(image)
|
115 |
+
# image_tar = pipe.image_processor.postprocess(image_src)
|
116 |
+
# return image_tar[0]
|
117 |
+
|
118 |
+
# cast image to half precision
|
119 |
+
image_src = image_src.to(device).half()
|
120 |
+
|
121 |
+
with torch.autocast("cuda"), torch.inference_mode():
|
122 |
+
x0_src_denorm = pipe.vae.encode(image_src).latent_dist.mode()
|
123 |
+
x0_src = (x0_src_denorm - pipe.vae.config.shift_factor) * pipe.vae.config.scaling_factor
|
124 |
+
# send to cuda
|
125 |
+
x0_src = x0_src.to(device)
|
126 |
+
|
127 |
+
negative_prompt = "" # optionally add support for negative prompts (SD3)
|
128 |
+
|
129 |
+
if model_type == 'SD3':
|
130 |
+
x0_tar = FlowEditSD3(pipe,
|
131 |
+
scheduler,
|
132 |
+
x0_src,
|
133 |
+
src_prompt,
|
134 |
+
tar_prompt,
|
135 |
+
negative_prompt,
|
136 |
+
T_steps,
|
137 |
+
n_avg,
|
138 |
+
src_guidance_scale,
|
139 |
+
tar_guidance_scale,
|
140 |
+
n_min,
|
141 |
+
n_max,)
|
142 |
+
|
143 |
+
elif model_type == 'FLUX':
|
144 |
+
x0_tar = FlowEditFLUX(pipe,
|
145 |
+
scheduler,
|
146 |
+
x0_src,
|
147 |
+
src_prompt,
|
148 |
+
tar_prompt,
|
149 |
+
negative_prompt,
|
150 |
+
T_steps,
|
151 |
+
n_avg,
|
152 |
+
src_guidance_scale,
|
153 |
+
tar_guidance_scale,
|
154 |
+
n_min,
|
155 |
+
n_max,)
|
156 |
+
else:
|
157 |
+
raise NotImplementedError(f"Sampler type {model_type} not implemented")
|
158 |
+
|
159 |
+
|
160 |
+
x0_tar_denorm = (x0_tar / pipe.vae.config.scaling_factor) + pipe.vae.config.shift_factor
|
161 |
+
with torch.autocast("cuda"), torch.inference_mode():
|
162 |
+
image_tar = pipe.vae.decode(x0_tar_denorm, return_dict=False)[0]
|
163 |
+
image_tar = pipe.image_processor.postprocess(image_tar)
|
164 |
+
|
165 |
+
|
166 |
+
return image_tar[0]
|
167 |
+
|
168 |
+
|
169 |
+
# title = "FlowEdit: Inversion-Free Text-Based Editing Using Pre-Trained Flow Models"
|
170 |
+
|
171 |
+
intro = """
|
172 |
+
<h1 style="font-weight: 1000; text-align: center; margin: 0px;">FlowEdit: Inversion-Free Text-Based Editing Using Pre-Trained Flow Models</h1>
|
173 |
+
<h3 style="margin-bottom: 10px; text-align: center;">
|
174 |
+
<a href="https://arxiv.org/">[Paper]</a> |
|
175 |
+
<a href="https://matankleiner.github.io/flowedit/">[Project Page]</a> |
|
176 |
+
<a href="https://github.com/fallenshock/FlowEdit">[Code]</a>
|
177 |
+
</h3>
|
178 |
+
Gradio demo for FlowEdit: Inversion-Free Text-Based Editing Using Pre-Trained Flow Models. See our project page for more details.
|
179 |
+
|
180 |
+
<br>
|
181 |
+
<br>Edit your image using Flow models! upload an image, add a description of it, and specify the edits you want to make.
|
182 |
+
<h3>Notes:</h3>
|
183 |
+
|
184 |
+
<ol>
|
185 |
+
<li>We use FLUX.1 dev and SD3 for the demo. The models are large and may take a while to load.</li>
|
186 |
+
<li>We recommend 1024x1024 images for the best results. If the input images are too large, there may be out-of-memory errors.</li>
|
187 |
+
<li>Default hyperparameters for each model used in the paper are provided as examples. Feel free to experiment with them as well.</li>
|
188 |
+
</ol>
|
189 |
+
|
190 |
+
"""
|
191 |
+
|
192 |
+
# article = """
|
193 |
+
# 📝 **Citation**
|
194 |
+
# ```bibtex
|
195 |
+
# @article{aaa,
|
196 |
+
# author = {},
|
197 |
+
# title = {},
|
198 |
+
# journal = {},
|
199 |
+
# year = {2024},
|
200 |
+
# url = {}
|
201 |
+
# }
|
202 |
+
# ```
|
203 |
+
# """
|
204 |
+
|
205 |
+
|
206 |
+
with gr.Blocks() as demo:
|
207 |
+
|
208 |
+
|
209 |
+
gr.HTML(intro)
|
210 |
+
|
211 |
+
with gr.Row(equal_height=True):
|
212 |
+
image_src = gr.Image(type="filepath", label="Source Image", value="inputs/cat.png",)
|
213 |
+
image_tar = gr.Image(label="Output", type="pil", show_label=True, format="png",),
|
214 |
+
|
215 |
+
with gr.Row():
|
216 |
+
src_prompt = gr.Textbox(lines=2, label="Source Prompt", value="a cat sitting in the grass")
|
217 |
+
|
218 |
+
with gr.Row():
|
219 |
+
tar_prompt = gr.Textbox(lines=2, label="Target Prompt", value="a puppy sitting in the grass")
|
220 |
+
|
221 |
+
with gr.Row():
|
222 |
+
model_type = gr.Dropdown(["SD3", "FLUX"], label="Model Type", value="SD3")
|
223 |
+
T_steps = gr.Number(value=50, label="Total Steps", minimum=1, maximum=50)
|
224 |
+
n_max = gr.Number(value=33, label="n_max (control the strength of the edit)")
|
225 |
+
|
226 |
+
with gr.Row():
|
227 |
+
src_guidance_scale = gr.Slider(minimum=1.0, maximum=30.0, value=3.5, label="src_guidance_scale")
|
228 |
+
tar_guidance_scale = gr.Slider(minimum=1.0, maximum=30.0, value=13.5, label="tar_guidance_scale")
|
229 |
+
|
230 |
+
with gr.Row():
|
231 |
+
submit_button = gr.Button("Run FlowEdit", variant="primary")
|
232 |
+
|
233 |
+
with gr.Accordion(label="Advanced Settings", open=False):
|
234 |
+
# additional inputs
|
235 |
+
n_min = gr.Number(value=0, label="n_min (for improved style edits)")
|
236 |
+
n_avg = gr.Number(value=1, label="n_avg (improve structure at the cost of runtime)", minimum=1)
|
237 |
+
seed = gr.Number(value=42, label="seed")
|
238 |
+
|
239 |
+
|
240 |
+
|
241 |
+
submit_button.click(
|
242 |
+
fn=FlowEditRun,
|
243 |
+
inputs=[
|
244 |
+
image_src,
|
245 |
+
model_type,
|
246 |
+
T_steps,
|
247 |
+
src_guidance_scale,
|
248 |
+
tar_guidance_scale,
|
249 |
+
n_max,
|
250 |
+
src_prompt,
|
251 |
+
tar_prompt,
|
252 |
+
n_min,
|
253 |
+
n_avg,
|
254 |
+
seed,
|
255 |
+
],
|
256 |
+
outputs=[
|
257 |
+
image_tar[0],
|
258 |
+
],
|
259 |
+
)
|
260 |
+
|
261 |
+
gr.Examples(
|
262 |
+
label="Examples",
|
263 |
+
examples=get_examples(),
|
264 |
+
inputs=[image_src, model_type, T_steps, src_guidance_scale, tar_guidance_scale, n_max, src_prompt, tar_prompt, n_min, n_avg, seed],
|
265 |
+
)
|
266 |
+
|
267 |
+
model_type.input(fn=on_model_change, inputs=[model_type], outputs=[T_steps, src_guidance_scale, tar_guidance_scale, n_max])
|
268 |
+
|
269 |
+
|
270 |
+
# gr.HTML(article)
|
271 |
+
demo.queue()
|
272 |
+
demo.launch( )
|
requirements.txt
ADDED
@@ -0,0 +1,81 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
accelerate==1.2.0
|
2 |
+
aiofiles==23.2.1
|
3 |
+
annotated-types==0.7.0
|
4 |
+
anyio==4.7.0
|
5 |
+
certifi==2024.8.30
|
6 |
+
charset-normalizer==3.4.0
|
7 |
+
click==8.1.7
|
8 |
+
diffusers==0.31.0
|
9 |
+
fastapi==0.115.6
|
10 |
+
ffmpy==0.4.0
|
11 |
+
filelock==3.16.1
|
12 |
+
fsspec==2024.10.0
|
13 |
+
gradio==5.8.0
|
14 |
+
gradio_client==1.5.1
|
15 |
+
h11==0.14.0
|
16 |
+
httpcore==1.0.7
|
17 |
+
httpx==0.28.1
|
18 |
+
huggingface-hub==0.26.5
|
19 |
+
idna==3.10
|
20 |
+
importlib_metadata==8.5.0
|
21 |
+
Jinja2==3.1.4
|
22 |
+
markdown-it-py==3.0.0
|
23 |
+
MarkupSafe==2.1.5
|
24 |
+
mdurl==0.1.2
|
25 |
+
mpmath==1.3.0
|
26 |
+
networkx==3.4.2
|
27 |
+
numpy==2.2.0
|
28 |
+
nvidia-cublas-cu12==12.1.3.1
|
29 |
+
nvidia-cuda-cupti-cu12==12.1.105
|
30 |
+
nvidia-cuda-nvrtc-cu12==12.1.105
|
31 |
+
nvidia-cuda-runtime-cu12==12.1.105
|
32 |
+
nvidia-cudnn-cu12==9.1.0.70
|
33 |
+
nvidia-cufft-cu12==11.0.2.54
|
34 |
+
nvidia-curand-cu12==10.3.2.106
|
35 |
+
nvidia-cusolver-cu12==11.4.5.107
|
36 |
+
nvidia-cusparse-cu12==12.1.0.106
|
37 |
+
nvidia-nccl-cu12==2.20.5
|
38 |
+
nvidia-nvjitlink-cu12==12.6.85
|
39 |
+
nvidia-nvtx-cu12==12.1.105
|
40 |
+
orjson==3.10.12
|
41 |
+
packaging==24.2
|
42 |
+
pandas==2.2.3
|
43 |
+
pillow==11.0.0
|
44 |
+
protobuf==5.29.1
|
45 |
+
psutil==5.9.8
|
46 |
+
pydantic==2.10.3
|
47 |
+
pydantic_core==2.27.1
|
48 |
+
pydub==0.25.1
|
49 |
+
Pygments==2.18.0
|
50 |
+
python-dateutil==2.9.0.post0
|
51 |
+
python-multipart==0.0.19
|
52 |
+
pytz==2024.2
|
53 |
+
PyYAML==6.0.2
|
54 |
+
regex==2024.11.6
|
55 |
+
requests==2.32.3
|
56 |
+
rich==13.9.4
|
57 |
+
ruff==0.8.2
|
58 |
+
safehttpx==0.1.6
|
59 |
+
safetensors==0.4.5
|
60 |
+
semantic-version==2.10.0
|
61 |
+
sentencepiece==0.2.0
|
62 |
+
setuptools==75.6.0
|
63 |
+
shellingham==1.5.4
|
64 |
+
six==1.17.0
|
65 |
+
sniffio==1.3.1
|
66 |
+
spaces==0.31.0
|
67 |
+
starlette==0.41.3
|
68 |
+
sympy==1.13.3
|
69 |
+
tokenizers==0.21.0
|
70 |
+
tomlkit==0.13.2
|
71 |
+
torch==2.4.1
|
72 |
+
tqdm==4.67.1
|
73 |
+
transformers==4.47.0
|
74 |
+
triton==3.0.0
|
75 |
+
typer==0.15.1
|
76 |
+
typing_extensions==4.12.2
|
77 |
+
tzdata==2024.2
|
78 |
+
urllib3==2.2.3
|
79 |
+
uvicorn==0.32.1
|
80 |
+
websockets==14.1
|
81 |
+
zipp==3.21.0
|