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"]]
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)

@spaces.GPU
def invert_and_edit(image, 
                    prompt, 
                    eta, 
                    gamma, 
                    start_timestep, 
                    stop_timestep, 
                    num_inversion_steps,
                    width,
                    height,
                    inverted_latents,
                    image_latents,
                    latent_image_ids,
                    do_inversion,
                    seed,
                    randomize_seed,
                   ):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    if do_inversion:
        inverted_latents_tensor, image_latents_tensor, latent_image_ids_tensor = pipe.invert(image, num_inversion_steps=num_inversion_steps, gamma=gamma)
        inverted_latents = gr.State(value=inverted_latents_tensor)
        image_latents = gr.State(value=image_latents_tensor)
        latent_image_ids = gr.State(value=latent_image_ids_tensor)
        do_inversion = False
    
    
    output = pipe(prompt, 
    inverted_latents=inverted_latents.value,
    image_latents=image_latents.value,
    latent_image_ids=latent_image_ids.value,
    start_timestep=start_timestep, 
    stop_timestep=stop_timestep,
    num_inference_steps=num_inversion_steps,
    eta=eta, 
    ).images[0]
        
    return output, inverted_latents.value, image_latents.value, latent_image_ids.value, 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]
based on the algorithm proposed in [*Semantic Image Inversion and Editing using
Stochastic Rectified Differential Equations*](https://rf-inversion.github.io/data/rf-inversion.pdf)

[[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():          
                    start_timestep = gr.Slider(
                    label="start timestep",
                    info = "lower gamma to ehnace the edits",
                    minimum=0.0,
                    maximum=1.0,
                    step=0.1,
                    value=0.9,
                )
                    stop_timestep = gr.Slider(
                        label="stop timestep",
                        info = "lower gamma to ehnace the edits",
                        minimum=0.0,
                        maximum=1.0,
                        step=0.1,
                        value=0.9,
                    )
                    eta = gr.Slider(
                        label="eta",
                        info = "lower eta to ehnace the edits",
                        minimum=0.0,
                        maximum=1.0,
                        step=0.1,
                        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():
                gamma = gr.Slider(
                    label="gamma",
                    info = "lower gamma to ehnace the edits",
                    minimum=0.0,
                    maximum=1.0,
                    step=0.1,
                    value=0.9,
                )
                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,
            width,
            height,
            inverted_latents,
            image_latents,
            latent_image_ids,
            do_inversion,
            seed,
            randomize_seed
        ],
        outputs=[result, inverted_latents, image_latents, latent_image_ids, do_inversion, seed],
    )

    gr.Examples(
                examples=examples,
                inputs=[input_image, prompt,],
                outputs=[result, inverted_latents, image_latents, latent_image_ids, do_inversion, seed],
                fn=infer,
            )

    input_image.change(
        fn=reset_do_inversion,
        outputs=[do_inversion]
    )

if __name__ == "__main__":
    demo.launch()