import spaces
import gradio as gr
import time
import torch

from diffusers import (
    DDPMScheduler,
    AutoPipelineForText2Image,
    AutoencoderTiny,
)

import oneflow as flow
from onediff.infer_compiler import oneflow_compile

BASE_MODEL = "stabilityai/sdxl-turbo"
device = "cuda"

vae = AutoencoderTiny.from_pretrained(
    'madebyollin/taesdxl',
    use_safetensors=True,
    torch_dtype=torch.float16,
).to('cuda')
base_pipe = AutoPipelineForText2Image.from_pretrained(
    BASE_MODEL, 
    vae=vae,
    torch_dtype=torch.float16, 
    variant="fp16", 
    use_safetensors=True,
)
base_pipe.to(device)

base_pipe = base_pipe.to(device, silence_dtype_warnings=True)
base_pipe.scheduler = DDPMScheduler.from_pretrained(
    BASE_MODEL,
    subfolder="scheduler",
)
base_pipe.unet = oneflow_compile(base_pipe.unet)
# base_pipe.vae.decoder = oneflow_compile(base_pipe.vae.decoder)

def create_demo() -> gr.Blocks:

    @spaces.GPU(duration=10)
    def text_to_image(
        prompt:str,
        steps:int,
    ):
        run_task_time = 0
        time_cost_str = ''
        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
        generated_image = base_pipe(
            prompt=prompt,
            num_inference_steps=steps,
        ).images[0]
        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
        return generated_image

    def get_time_cost(run_task_time, time_cost_str):
        now_time = int(time.time()*1000)
        if run_task_time == 0:
            time_cost_str = 'start'
        else:
            if time_cost_str != '': 
                time_cost_str += f'-->'
            time_cost_str += f'{now_time - run_task_time}'
        run_task_time = now_time
        return run_task_time, time_cost_str

    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(label="Inference Steps", min=1, max=30, step=1, value=5)
                g_btn = gr.Button("Generate")
                
        with gr.Row():
            generated_image = gr.Image(label="Generated Image", type="pil", interactive=False)
        
        g_btn.click(
            fn=text_to_image,
            inputs=[prompt, steps],
            outputs=[generated_image],
        )

    return demo