app.py

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("cuda")

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()