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):
    return img1

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='filepath')
            img2 = gr.Image(label='Image 1',type='filepath')
            slider = gr.Slider(label='Weight', maximum=1.0, value=0.5)
            btn = gr.Button("Blend")
        with gr.Column():
            output = gr.Image(label='Result')
    
    
    btn.click(fn=blend, inputs=[img1, img2, slider], outputs=output)

demo.launch()