import gradio as gr
import torch
from diffusers import KandinskyPriorPipeline, KandinskyPipeline
import qrcode

IMAGE_WIDTH = 512
IMAGE_HEIGHT = 512

device = 0
if not torch.cuda.is_available():
    print("GPU isn't available. Running on CPU.")
    device = -1

pipe_prior = KandinskyPriorPipeline.from_pretrained(
    "kandinsky-community/kandinsky-2-1-prior", torch_dtype=torch.float16
)
pipe_prior.to("cuda")

pipe = KandinskyPipeline.from_pretrained(
    "kandinsky-community/kandinsky-2-1", torch_dtype=torch.float16)
pipe.to("cuda")


def interpolate_image(qr_code_value, image, qr_image_weight, images_weight):
    qr_code = qrcode.make(qr_code_value).resize((IMAGE_WIDTH, IMAGE_HEIGHT))
    image_texts = [image, qr_code]
    weights = [images_weight, qr_image_weight]
    prompt = ""
    prior_out = pipe_prior.interpolate(image_texts, weights)

    result = pipe(prompt, **prior_out, height=IMAGE_HEIGHT, width=IMAGE_WIDTH)

    return result.images


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            qr_input = gr.Text(label="QR Code Value")
            image = gr.Image(label="Images", type="pil")

            with gr.Row():
                qr_image_weight = gr.Slider(
                    label="QR Image Weight", value=0.5, maximum=1)
                images_weight = gr.Slider(
                    label="Input Images Weight", value=0.3, maximum=1)

            submit_button = gr.Button("Submit")

        with gr.Column():
            output = gr.Gallery(label="Output")
            output.style(grid=4)

    submit_button.click(
        fn=interpolate_image, inputs=[
            qr_input, image, qr_image_weight, images_weight], outputs=output, api_name="interpolate"
    )

    gr.Examples(
        examples=[["http://huggingface.co", "Starry_Night.jpg"]],
        inputs=[qr_input, image],
        outputs=output
    )
    

demo.theme = gr.themes.Base()
demo.launch()