import spaces
import base64
from io import BytesIO

import gradio as gr
import PIL.Image
import torch

from diffusers import StableDiffusionPipeline, AutoencoderKL, AutoencoderTiny
from peft import PeftModel

device = "cuda"
weight_type = torch.float16

pipe = StableDiffusionPipeline.from_pretrained("IDKiro/sdxs-512-dreamshaper")
pipe.unet = PeftModel.from_pretrained(pipe.unet, "IDKiro/sdxs-512-dreamshaper-anime")
pipe.unet.merge_and_unload()
pipe.to(device, dtype=weight_type)

vae_tiny = AutoencoderTiny.from_pretrained(
    "IDKiro/sdxs-512-dreamshaper", subfolder="vae"
)
vae_tiny.to(device, dtype=weight_type)

vae_large = AutoencoderKL.from_pretrained(
    "IDKiro/sdxs-512-dreamshaper", subfolder="vae_large"
)
vae_tiny.to(device, dtype=weight_type)


def pil_image_to_data_url(img, format="PNG"):
    buffered = BytesIO()
    img.save(buffered, format=format)
    img_str = base64.b64encode(buffered.getvalue()).decode()
    return f"data:image/{format.lower()};base64,{img_str}"


@spaces.GPU
def run(
    prompt: str,
    device_type="GPU",
    vae_type=None,
    param_dtype="torch.float16",
) -> PIL.Image.Image:
    if vae_type == "tiny vae":
        pipe.vae = vae_tiny
    elif vae_type == "large vae":
        pipe.vae = vae_large

    if device_type == "CPU":
        device = "cpu"
        param_dtype = "torch.float32"
    else:
        device = "cuda"

    pipe.to(
        torch_device=device,
        torch_dtype=torch.float16 if param_dtype == "torch.float16" else torch.float32,
    )

    result = pipe(
        prompt=prompt,
        guidance_scale=0.0,
        num_inference_steps=1,
        output_type="pil",
    ).images[0]

    result_url = pil_image_to_data_url(result)

    return (result, result_url)


examples = [
    "A photo of beautiful mountain with realistic sunset and blue lake, highly detailed, masterpiece",
]

with gr.Blocks(css="style.css") as demo:
    gr.Markdown("# SDXS-512-DreamShaper-Anime")
    gr.Markdown("[SDXS: Real-Time One-Step Latent Diffusion Models with Image Conditions](https://arxiv.org/abs/2403.16627) | [GitHub](https://github.com/IDKiro/sdxs)")
    with gr.Group():
        with gr.Row():
            with gr.Column(min_width=685):
                with gr.Row():
                    prompt = gr.Text(
                        label="Prompt",
                        show_label=False,
                        max_lines=1,
                        placeholder="Enter your prompt",
                        container=False,
                    )
                    run_button = gr.Button("Run", scale=0)

                device_choices = ["GPU", "CPU"]
                device_type = gr.Radio(
                    device_choices,
                    label="Device",
                    value=device_choices[0],
                    interactive=True,
                    info="Thanks to the community for the GPU!",
                )

                vae_choices = ["tiny vae", "large vae"]
                vae_type = gr.Radio(
                    vae_choices,
                    label="Image Decoder Type",
                    value=vae_choices[0],
                    interactive=True,
                    info="To save GPU memory, use tiny vae. For better quality, use large vae.",
                )

                dtype_choices = ["torch.float16", "torch.float32"]
                param_dtype = gr.Radio(
                    dtype_choices,
                    label="torch.weight_type",
                    value=dtype_choices[0],
                    interactive=True,
                    info="To save GPU memory, use torch.float16. For better quality, use torch.float32.",
                )

                download_output = gr.Button(
                    "Download output", elem_id="download_output"
                )

            with gr.Column(min_width=512):
                result = gr.Image(
                    label="Result",
                    height=512,
                    width=512,
                    elem_id="output_image",
                    show_label=False,
                    show_download_button=True,
                )

    gr.Examples(examples=examples, inputs=prompt, outputs=result, fn=run)

    demo.load(None, None, None)

    inputs = [prompt, device_type, vae_type, param_dtype]
    outputs = [result, download_output]
    prompt.submit(fn=run, inputs=inputs, outputs=outputs)
    run_button.click(fn=run, inputs=inputs, outputs=outputs)

if __name__ == "__main__":
    demo.queue().launch(debug=True)