import gradio as gr
import torch
from PIL import Image
from run import StableRemix, run_remixing

pipe = StableRemix.from_pretrained(
    "stabilityai/stable-diffusion-2-1-unclip",
    torch_dtype=torch.float16,
    variant="fp16"
)
pipe = pipe.to('cuda')
pipe.enable_attention_slicing()

print('pipe loaded')


def remix(image1, image2, alpha):
    # style_img = Image.open(args.style_img).convert('RGB')

    # images = run_remixing(pipe, image1, image1, [0.6, 0.65, 0.7])
    images = run_remixing(pipe, image1, image2, [alpha])
    return images[0]
    for idx, image in enumerate(images):
        path = args.save_dir / f'remix_{idx}.png'
        print('Saving remix to', path)
        image.save(path)


demo = gr.Interface(
    fn=remix, inputs=[gr.Image(image_mode='RGB', shape=[512, 512]), gr.Image(image_mode='RGB', shape=[512, 512]), gr.Slider(0.0, 1.0, 0.6)], outputs="image")
demo.launch()