Spaces:
Running
on
Zero
Running
on
Zero
| import gradio as gr | |
| import numpy as np | |
| import random | |
| import torch | |
| import spaces | |
| from PIL import Image | |
| import os | |
| from huggingface_hub import hf_hub_download | |
| import torch | |
| from diffusers import DiffusionPipeline | |
| from huggingface_hub import hf_hub_download | |
| # Constants | |
| MAX_SEED = np.iinfo(np.int32).max | |
| MAX_IMAGE_SIZE = 1024 | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", | |
| custom_pipeline="pipeline_flux_rf_inversion", | |
| torch_dtype=torch.bfloat16) | |
| pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), lora_scale=0.125) | |
| pipe.fuse_lora(lora_scale=0.125) | |
| pipe.to(DEVICE) | |
| examples = [[Image.open("cat.jpg"), "a tiger", 0, 0.7, 0.5, 0.9, 8, 8, 789385745, False]] | |
| def reset_do_inversion(): | |
| return True | |
| def resize_img(image, max_size=1024): | |
| width, height = image.size | |
| scaling_factor = min(max_size / width, max_size / height) | |
| new_width = int(width * scaling_factor) | |
| new_height = int(height * scaling_factor) | |
| return image.resize((new_width, new_height), Image.LANCZOS) | |
| def check_style(stylezation, eta, eta_decay, decay_power, start_timestep, stop_timestep): | |
| return eta, eta_decay, decay_power, start_timestep, stop_timestep | |
| def invert_and_edit(image, | |
| prompt, | |
| eta, | |
| gamma, | |
| start_timestep, | |
| stop_timestep, | |
| num_inversion_steps, | |
| num_inference_steps, | |
| width, | |
| height, | |
| inverted_latents, | |
| image_latents, | |
| latent_image_ids, | |
| do_inversion, | |
| seed, | |
| randomize_seed, | |
| eta_decay, | |
| decay_power, | |
| ): | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| if do_inversion: | |
| inverted_latents, image_latents, latent_image_ids = pipe.invert(image, num_inversion_steps=num_inversion_steps, gamma=gamma) | |
| do_inversion = False | |
| output = pipe(prompt, | |
| inverted_latents = inverted_latents.to(DEVICE), | |
| image_latents = image_latents.to(DEVICE), | |
| latent_image_ids = latent_image_ids.to(DEVICE), | |
| start_timestep = start_timestep/num_inference_steps, | |
| stop_timestep = stop_timestep/num_inference_steps, | |
| num_inference_steps = num_inference_steps, | |
| eta=eta, | |
| decay_eta = eta_decay, | |
| eta_decay_power = decay_power, | |
| ).images[0] | |
| return output, inverted_latents.cpu(), image_latents.cpu(), latent_image_ids.cpu(), do_inversion, seed | |
| # UI CSS | |
| css = """ | |
| #col-container { | |
| margin: 0 auto; | |
| max-width: 960px; | |
| } | |
| """ | |
| # Create the Gradio interface | |
| with gr.Blocks(css=css) as demo: | |
| inverted_latents = gr.State() | |
| image_latents = gr.State() | |
| latent_image_ids = gr.State() | |
| do_inversion = gr.State(False) | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown(f"""# RF inversion 🖌️🏞️ | |
| ### Edit real images with FLUX.1 [dev] | |
| following the algorithm proposed in [*Semantic Image Inversion and Editing using | |
| Stochastic Rectified Differential Equations* by Rout et al.](https://rf-inversion.github.io/data/rf-inversion.pdf) | |
| based on the implementations of [@raven38](https://github.com/raven38) & [@DarkMnDragon](https://github.com/DarkMnDragon) 🙌🏻 | |
| [[non-commercial license](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)] [[project page](https://rf-inversion.github.io/) [[arxiv](https://arxiv.org/pdf/2410.10792)] | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_image = gr.Image( | |
| label="Input Image", | |
| type="pil" | |
| ) | |
| prompt = gr.Text( | |
| label="Edit Prompt", | |
| max_lines=1, | |
| placeholder="describe the edited output", | |
| ) | |
| stylezation = gr.Checkbox(label="stylzation", value=True) | |
| with gr.Row(): | |
| start_timestep = gr.Slider( | |
| label="start timestep", | |
| info = "decrease to enhace fidelity to original image", | |
| minimum=0, | |
| maximum=28, | |
| step=1, | |
| value=0, | |
| ) | |
| stop_timestep = gr.Slider( | |
| label="stop timestep", | |
| info = "increase to enhace fidelity to original image", | |
| minimum=0, | |
| maximum=28, | |
| step=1, | |
| value=4, | |
| ) | |
| eta = gr.Slider( | |
| label="eta", | |
| info = "lower eta to ehnace the edits", | |
| minimum=0.0, | |
| maximum=1.0, | |
| step=0.01, | |
| value=0.9, | |
| ) | |
| run_button = gr.Button("Edit", variant="primary") | |
| with gr.Column(): | |
| result = gr.Image(label="Result") | |
| with gr.Accordion("Advanced Settings", open=False): | |
| seed = gr.Slider( | |
| label="Seed", | |
| minimum=0, | |
| maximum=MAX_SEED, | |
| step=1, | |
| value=42, | |
| ) | |
| randomize_seed = gr.Checkbox(label="Randomize seed", value=True) | |
| with gr.Row(): | |
| num_inference_steps = gr.Slider( | |
| label="num inference steps", | |
| minimum=1, | |
| maximum=50, | |
| step=1, | |
| value=8, | |
| ) | |
| eta_decay = gr.Checkbox(label="eta decay", value=False) | |
| decay_power = gr.Slider( | |
| label="eta decay power", | |
| minimum=0, | |
| maximum=5, | |
| step=1, | |
| value=1, | |
| ) | |
| with gr.Row(): | |
| gamma = gr.Slider( | |
| label="gamma", | |
| info = "increase gamma to enhance realism", | |
| minimum=0.0, | |
| maximum=1.0, | |
| step=0.01, | |
| value=0.5, | |
| ) | |
| num_inversion_steps = gr.Slider( | |
| label="num inversion steps", | |
| minimum=1, | |
| maximum=50, | |
| step=1, | |
| value=8, | |
| ) | |
| with gr.Row(): | |
| width = gr.Slider( | |
| label="Width", | |
| minimum=256, | |
| maximum=MAX_IMAGE_SIZE, | |
| step=32, | |
| value=1024, | |
| ) | |
| height = gr.Slider( | |
| label="Height", | |
| minimum=256, | |
| maximum=MAX_IMAGE_SIZE, | |
| step=32, | |
| value=1024, | |
| ) | |
| run_button.click( | |
| fn=invert_and_edit, | |
| inputs=[ | |
| input_image, | |
| prompt, | |
| eta, | |
| gamma, | |
| start_timestep, | |
| stop_timestep, | |
| num_inversion_steps, | |
| num_inference_steps, | |
| width, | |
| height, | |
| inverted_latents, | |
| image_latents, | |
| latent_image_ids, | |
| do_inversion, | |
| seed, | |
| randomize_seed, | |
| eta_decay, | |
| decay_power, | |
| ], | |
| outputs=[result, inverted_latents, image_latents, latent_image_ids, do_inversion, seed], | |
| ) | |
| gr.Examples( | |
| examples=examples, | |
| inputs=[input_image, prompt,start_timestep, stop_timestep, gamma, eta, num_inversion_steps, num_inference_steps, seed, randomize_seed ], | |
| outputs=[result, inverted_latents, image_latents, latent_image_ids, do_inversion, seed], | |
| fn=invert_and_edit, | |
| ) | |
| input_image.change( | |
| fn=reset_do_inversion, | |
| outputs=[do_inversion] | |
| ) | |
| num_inversion_steps.change( | |
| fn=reset_do_inversion, | |
| outputs=[do_inversion] | |
| ) | |
| stylezation.change( | |
| fn=check_style, | |
| inputs=[stylezation, eta, eta_decay, decay_power, start_timestep, stop_timestep], | |
| outputs=[eta, eta_decay, decay_power, start_timestep, stop_timestep] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |