import requests
from PIL import Image
from io import BytesIO
from diffusers import StableDiffusionUpscalePipeline
import torch
import gradio as gr

device = "cuda" if torch.cuda.is_available() else "cpu"

print(f'{device} is available')

model_id = "stabilityai/stable-diffusion-x4-upscaler"
upscale_pipe = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
upscale_pipe = upscale_pipe.to(device)

DEFAULT_SRC_PROMPT = "a person with pefect face"


def create_demo() -> gr.Blocks:
    from inversion_run_base import run as base_run

    @spaces.GPU(duration=15)
    def upscale_image(
        input_image: Image,
        prompt: str,
    ):
        upscaled_image = upscale_pipe(prompt=prompt, image=input_image).images[0]
        extension = 'png'

        path = f"output/{uuid.uuid4()}.{extension}"
        upscaled_image.save(path, quality=100)
        
        return upscaled_image, path, time_cost_str

    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:
        croper = gr.State()
        with gr.Row():
            with gr.Column():
                input_image_prompt = gr.Textbox(lines=1, label="Input Image Prompt", value=DEFAULT_SRC_PROMPT)
            with gr.Column():
                g_btn = gr.Button("Upscale Image")
                
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Input Image", type="pil")
            with gr.Column():
                upscaled_image = gr.Image(label="Upscaled Image", format="png", type="pil", interactive=False)
                download_path = gr.File(label="Download the output image", interactive=False)
                generated_cost = gr.Textbox(label="Time cost by step (ms):", visible=True, interactive=False)
        
        g_btn.click(
            fn=upscale_image,
            inputs=[input_image, input_image_prompt],
            outputs=[upscaled_image, download_path, generated_cost],
        )

    return demo