import gradio as gr
from gradio_image_annotation import image_annotator
from diffusers import StableDiffusionPipeline
import os
import torch

import spaces #[uncomment to use ZeroGPU]
# Load model
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
).to("cuda" if torch.cuda.is_available() else "cpu")
pipe.safety_checker = None

example_annotation = {
    "image": os.path.join(os.path.dirname(__file__), "background.png"),
    "boxes": [],
}

@spaces.GPU
def get_boxes_json(annotations):
    print(annotations)
    image = annotations["image"]
    width = image.shape[1]
    height = image.shape[0]
    boxes = annotations["boxes"]
    prompt_final = [[]]
    for box in boxes:
        box["xmin"] = box["xmin"] / width
        box["xmax"] = box["xmax"] / width
        box["ymin"] = box["ymin"] / height
        box["ymax"] = box["ymax"] / height
        prompt_final[0].append(box["label"])
    # import pdb; pdb.set_trace()
    prompt = ", ".join(prompt_final[0])
    image = pipe(prompt).images[0]
    return image
    # return annotations["boxes"]

with gr.Blocks() as demo:
   with gr.Tab("DreamRenderer", id="DreamRenderer"):
        with gr.Row():
            with gr.Column(scale=1):
                annotator = image_annotator(
                    example_annotation,
                    height=512,
                    width=512
                )
            with gr.Column(scale=1):
                generated_image = gr.Image(label="Generated Image", height=512, width=512)

        button_get = gr.Button("Generation")
        button_get.click(get_boxes_json, inputs=annotator, outputs=generated_image)


if __name__ == "__main__":
    demo.launch()