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():
examples = [[Image.open("metal.png"), "a dragon, in 3d melting gold metal",0.9, 0.5, 0, 5, 28, 28, 0, False,False, 2, False,Image.open("dragon.png") ],
[Image.open("doll.png"), "anime illustration",0.9, 0.5, 0, 6, 28, 28, 0, False, False, 2, False ,Image.open("anime.png")],
[Image.open("doll.png"), "raccoon, made of yarn",0.9, 0.5, 0, 4, 28, 28, 0, False, False, 2, False , Image.open("raccoon.png")],
[Image.open("cat.jpg"), "a parrot", 0.9 ,0.5,2, 8,28, 28,0, False , False, 1, False,Image.open("parrot.png")],
[Image.open("cat.jpg"), "a tiger", 0.9 ,0.5,0, 4,8, 8,789385745, False , False, 1, True,Image.open("tiger.png")],
[Image.open("metal.png"),"a dragon, in 3d melting gold metal",0.9, 0.5, 0, 4, 8, 8, 789385745, False,True, 2, True , Image.open("dragon.png")],
]
return examples
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:
return 0.9, 0.5, 0, 6, 28, 28, False,False
else:
if enable_hyper_flux:
return 0.9, 0.5, 0, 4, 8, 8, False,False
else:
return 0.9, 0.5, 2, 7, 28, 28, False,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(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",
)
with gr.Row():
enable_hyper_flux = gr.Checkbox(label="8-step LoRA", value=False, info="")
stylezation = gr.Checkbox(label="stylzation")
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, prompt,eta,gamma,start_timestep, stop_timestep, num_inversion_steps, num_inference_steps, seed, randomize_seed, eta_decay, decay_power, enable_hyper_flux ]
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_decay, decay_power]
)
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()