import spaces
import gradio as gr
import torch
import os

from diffusers import (
    DDPMScheduler,
    StableDiffusionXLImg2ImgPipeline,
    AutoencoderKL,
)

from diffusers.utils import load_image

os.system("pip install torch_tensorrt==2.4.0")

BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"--------->Device: {device}")

vae = AutoencoderKL.from_pretrained(
    "madebyollin/sdxl-vae-fp16-fix", 
    torch_dtype=torch.float16,
)

base_pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
    BASE_MODEL, 
    vae=vae,
    torch_dtype=torch.float16, 
    variant="fp16", 
    use_safetensors=True,
)
base_pipe = base_pipe.to(device, silence_dtype_warnings=True)
base_pipe.scheduler = DDPMScheduler.from_pretrained(
    BASE_MODEL,
    subfolder="scheduler",
)

backend = "torch_tensorrt"
import torch_tensorrt
print('Compiling model...')
compiledModel = torch.compile(
    base_pipe.unet,
    backend=backend,
    options={
        "truncate_long_and_double": True,
        "enabled_precisions": {torch.float32, torch.float16},
    },
    dynamic=False,
)

base_pipe.unet = compiledModel

init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img.png")
generated_image = base_pipe(
    init_image,
    num_inference_steps=5,
).images[0]

generated_image.save("/tmp/gradio/generated_image.png")


def create_demo() -> gr.Blocks:

    @spaces.GPU(duration=30)
    def text_to_image(
        prompt:str,
        steps:int,
    ):
        import torch_tensorrt
        print('Compiling model...')
        compiledModel = torch.compile(
            base_pipe.unet,
            backend=backend,
            options={
                "truncate_long_and_double": True,
                "enabled_precisions": {torch.float32, torch.float16},
            },
            dynamic=False,
        )
        print('Model compiled!')

        print('Saving compiled model...')
        torch_tensorrt.save(compiledModel, "compiled_pipe.ep")
        print('Compiled model saved!')

    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(label="Prompt", placeholder="Write a prompt here", lines=2, value="A beautiful sunset over the city")
            with gr.Column():
                steps = gr.Slider(minimum=1, maximum=100, value=5, step=1, label="Num Steps")
                g_btn = gr.Button("Generate")
                
        with gr.Row():
            with gr.Column():
                generated_image = gr.Image(label="Generated Image", type="pil", interactive=False)
            with gr.Column():
                time_cost = gr.Textbox(label="Time Cost", lines=1, interactive=False)
        
        g_btn.click(
            fn=text_to_image,
            inputs=[prompt, steps],
            # outputs=[generated_image, time_cost],
            outputs=[],
        )

    return demo