import gradio as gr
from diffusers import KandinskyPriorPipeline, KandinskyPipeline
from diffusers.utils import load_image
import torch


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 blend(img1, img2, slider):
    # add all the conditions we want to interpolate, can be either text or image
    images_texts = [img1, img2]
    
    # specify the weights for each condition in images_texts
    weights = [1-slider, slider]
    
    prior_out = pipe_prior.interpolate(images_texts, weights)

    image = pipe(prompt='', **prior_out, height=1024, width=1024).images[0]
    
    return image

with gr.Blocks() as demo:
    gr.Markdown("""
    # Image Blender
    by [Tony Assi](https://www.tonyassi.com/)
    """)
    
    with gr.Row():
        with gr.Column():
            img1 = gr.Image(label='Image 0', type='pil')
            img2 = gr.Image(label='Image 1',type='pil')
            slider = gr.Slider(label='Weight', maximum=1.0, value=0.5)
            btn = gr.Button("Blend")
        with gr.Column():
            output = gr.Image(label='Result')

    gr.Examples(
        [['./cat.png', './starry_night.jpg', 0.5]],
        [img1, img2, slider],
        output,
        blend,
        cache_examples=True,
    )
    
    btn.click(fn=blend, inputs=[img1, img2, slider], outputs=output)

demo.launch()