| 49 | |
| 50 | |
| 51 | def test_rf_inversion_sde_sampling(enable_sde: bool = False): |
| 52 | set_seed(999) |
| 53 | pipe_rf_inversion = RFInversionFluxPipelineSDE.from_pipe(pipe) |
| 54 | |
| 55 | inverted_latents, image_latents, latent_image_ids = pipe_rf_inversion.invert( |
| 56 | image=example_image, |
| 57 | num_inversion_steps=28, |
| 58 | gamma=0.5 |
| 59 | ) |
| 60 | |
| 61 | edited_image = pipe_rf_inversion( |
| 62 | prompt="a tomato", |
| 63 | inverted_latents=inverted_latents, |
| 64 | image_latents=image_latents, |
| 65 | latent_image_ids=latent_image_ids, |
| 66 | start_timestep=0, |
| 67 | stop_timestep=7/28, |
| 68 | num_inference_steps=28, |
| 69 | eta=0.9, |
| 70 | enable_sde=enable_sde, |
| 71 | ).images[0] |
| 72 | |
| 73 | save_dir = "./results/" |
| 74 | if not os.path.exists(save_dir): |
| 75 | os.makedirs(save_dir) |
| 76 | image_save_path = os.path.join(save_dir, f"rf_inversion_{'sde' if enable_sde else 'ode'}_sampling.png") |
| 77 | edited_image.save(image_save_path) |
| 78 | print('Results saved here: ', image_save_path) |
| 79 | |
| 80 | |
| 81 | if __name__ == "__main__": |