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.enable_lora() pipe.to(DEVICE) def get_examples(): case = [ [Image.open("metal.png"),"dragon.png", "a dragon, in 3d melting gold metal",0.9, 0.5, 0, 5, 28, 28, 0, False,False, 2, False, "text/image guided stylzation" ], [Image.open("doll.png"),"anime.png", "anime illustration",0.9, 0.5, 0, 6, 28, 28, 0, False, False, 2, False,"text/image guided stylzation" ], [Image.open("doll.png"), "raccoon.png", "raccoon, made of yarn",0.9, 0.5, 0, 4, 28, 28, 0, False, False, 2, False, "local subject edits" ], [Image.open("cat.jpg"),"parrot.png", "a parrot", 0.9 ,0.5,2, 8,28, 28,0, False , False, 1, False, "local subject edits"], [Image.open("cat.jpg"),"tiger.png", "a tiger", 0.9 ,0.5,0, 4,8, 8,789385745, False , False, 1, True, "local subject edits"], [Image.open("metal.png"), "dragon.png","a dragon, in 3d melting gold metal",0.9, 0.5, 0, 4, 8, 8, 789385745, False,True, 2, True , "text/image guided stylzation"], ] return case 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, enable_hyper_flux): if stylezation == "text/image guided stylzation": return 0.9, 0.5, 0, 6, 28, 28, False else: if enable_hyper_flux: return 0.9, 0.5, 0, 4, 8, 8, False else: return 0.9, 0.5, 2, 7, 28, 28, False def check_hyper_flux_lora(enable_hyper_flux): if enable_hyper_flux: 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) return 8, 8, 4 else: pipe.unfuse_lora() return 28, 28, 6 @spaces.GPU(duration=85) def invert_and_edit(image, prompt, eta, gamma, start_timestep, stop_timestep, num_inversion_steps, num_inference_steps, seed, randomize_seed, eta_decay, decay_power, width = 1024, height = 1024, inverted_latents = None, image_latents = None, latent_image_ids = None, do_inversion = True, ): 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(True) 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", ) with gr.Row(): enable_hyper_flux = gr.Checkbox(label="8-step LoRA", value=False, info="may reduce edit quality", visible=False) stylezation = gr.Radio(["local subject edits", "text/image guided stylzation"], label="edit type", info="") with gr.Row(): start_timestep = gr.Slider( label="start timestep", info = "increase to enhance fidelity, decrease to enhance realism", 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=6, ) 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=28, ) 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=28, ) 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, seed, randomize_seed, eta_decay, decay_power, width, height, inverted_latents, image_latents, latent_image_ids, do_inversion ], outputs=[result, inverted_latents, image_latents, latent_image_ids, do_inversion, seed], ) gr.Examples( examples=get_examples(), inputs=[input_image,result, prompt,eta,gamma,start_timestep, stop_timestep, num_inversion_steps, num_inference_steps, seed, randomize_seed, eta_decay, decay_power, enable_hyper_flux,stylezation ], outputs=[result], ) input_image.change( fn=reset_do_inversion, outputs=[do_inversion] ) num_inversion_steps.change( fn=reset_do_inversion, outputs=[do_inversion] ) seed.change( fn=reset_do_inversion, outputs=[do_inversion] ) stylezation.change( fn=check_style, inputs=[stylezation], outputs=[eta, gamma, start_timestep, stop_timestep, num_inversion_steps, num_inference_steps, eta_decay] ) enable_hyper_flux.change( fn=check_hyper_flux_lora, inputs=[enable_hyper_flux], outputs=[num_inversion_steps, num_inference_steps, stop_timestep] ) if __name__ == "__main__": demo.launch()