fallenshock commited on
Commit
336dbcf
1 Parent(s): 2969ec0

added files

Browse files
Files changed (3) hide show
  1. FlowEdit_utils.py +404 -0
  2. app.py +272 -0
  3. 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>&nbsp;|&nbsp;
175
+ <a href="https://matankleiner.github.io/flowedit/">[Project Page]</a>&nbsp;|&nbsp;
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