import gradio as gr
import spaces
import torch
from gradio_imageslider import ImageSlider
from diffusers import DiffusionPipeline, AutoencoderTiny

from controlnet_union import ControlNetModel_Union
from custom_pipeline import FluxWithCFGPipeline

# Device and model setup
dtype = torch.float16
pipe = FluxWithCFGPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", torch_dtype=dtype
)
pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
# pipe.load_lora_weights("ostris/OpenFLUX.1", weight_name="openflux1-v0.1.0-fast-lora.safetensors", adapter_name="fast")
# pipe.set_adapters("fast")
# pipe.fuse_lora(adapter_names=["fast"], lora_scale=1.0)
pipe.to("cuda")
# pipe.transformer.to(memory_format=torch.channels_last)
# pipe.transformer = torch.compile(
#     pipe.transformer, mode="max-autotune", fullgraph=True
# )
torch.cuda.empty_cache()

@spaces.GPU(duration=25)
def fill_image(prompt, image, paste_back):
    (
        prompt_embeds,
        negative_prompt_embeds,
        pooled_prompt_embeds,
        negative_pooled_prompt_embeds,
    ) = pipe.encode_prompt(prompt, "cuda", True)

    source = image["background"]
    mask = image["layers"][0]

    alpha_channel = mask.split()[3]
    binary_mask = alpha_channel.point(lambda p: p > 0 and 255)
    cnet_image = source.copy()
    cnet_image.paste(0, (0, 0), binary_mask)

    for image in pipe(
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        pooled_prompt_embeds=pooled_prompt_embeds,
        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
        image=cnet_image,
    ):
        yield image, cnet_image

    print(f"{paste_back=}")

    if paste_back:
        image = image.convert("RGBA")
        cnet_image.paste(image, (0, 0), binary_mask)
    else:
        cnet_image = image

    yield source, cnet_image


def clear_result():
    return gr.update(value=None)


title = """<h1 align="center">FLUX Fast Inpaint</h1>
<div align="center">Draw the mask over the subject you want to erase or change and write what you want to inpaint it with.</div>
"""

with gr.Blocks() as demo:
    gr.HTML(title)
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                info="Describe what to inpaint the mask with",
                lines=3,
            )
        with gr.Column():
            with gr.Row():
                with gr.Column():
                    run_button = gr.Button("Generate")

                with gr.Column():
                    paste_back = gr.Checkbox(True, label="Paste back original")

    with gr.Row():
        input_image = gr.ImageMask(
            type="pil", label="Input Image", crop_size=(1024, 1024), layers=False
        )

        result = ImageSlider(
            interactive=False,
            label="Generated Image",
        )

    use_as_input_button = gr.Button("Use as Input Image", visible=False)

    def use_output_as_input(output_image):
        return gr.update(value=output_image[1])

    use_as_input_button.click(
        fn=use_output_as_input, inputs=[result], outputs=[input_image]
    )

    run_button.click(
        fn=clear_result,
        inputs=None,
        outputs=result,
    ).then(
        fn=lambda: gr.update(visible=False),
        inputs=None,
        outputs=use_as_input_button,
    ).then(
        fn=fill_image,
        inputs=[prompt, input_image, paste_back],
        outputs=result,
    ).then(
        fn=lambda: gr.update(visible=True),
        inputs=None,
        outputs=use_as_input_button,
    )

    prompt.submit(
        fn=clear_result,
        inputs=None,
        outputs=result,
    ).then(
        fn=lambda: gr.update(visible=False),
        inputs=None,
        outputs=use_as_input_button,
    ).then(
        fn=fill_image,
        inputs=[prompt, input_image, paste_back],
        outputs=result,
    ).then(
        fn=lambda: gr.update(visible=True),
        inputs=None,
        outputs=use_as_input_button,
    )


demo.launch(share=False)