| 25 | |
| 26 | |
| 27 | def test_flux(enable_sde: bool = False): |
| 28 | set_seed(999) |
| 29 | if enable_sde: |
| 30 | orig_scheduler = pipe.scheduler |
| 31 | scheduler=FlowMatchEulerDiscreteSDEScheduler.from_pretrained( |
| 32 | "black-forest-labs/FLUX.1-dev", |
| 33 | subfolder="scheduler", |
| 34 | ) |
| 35 | pipe.scheduler = scheduler |
| 36 | |
| 37 | edited_image = pipe(prompt="a tomato", num_inference_steps=28).images[0] |
| 38 | |
| 39 | save_dir = "./results/" |
| 40 | if not os.path.exists(save_dir): |
| 41 | os.makedirs(save_dir) |
| 42 | image_save_path = os.path.join(save_dir, f"flux_{'sde' if enable_sde else 'ode'}_sampling.png") |
| 43 | edited_image.save(image_save_path) |
| 44 | print('Results saved here: ', image_save_path) |
| 45 | |
| 46 | if enable_sde: |
| 47 | # restore the original scheduler |
| 48 | pipe.scheduler = orig_scheduler |
| 49 | |
| 50 | |
| 51 | def test_rf_inversion_sde_sampling(enable_sde: bool = False): |