from diffusers import DiffusionPipeline
import gradio as gr
import torch
import math
import PIL

if torch.cuda.is_available():
    device = "cuda"
    dtype = torch.float16
else:
    device = "cpu"
    dtype = torch.bfloat16

pipe = DiffusionPipeline.from_pretrained("kakaobrain/karlo-v1-alpha-image-variations", torch_dtype=dtype, custom_pipeline='unclip_image_interpolation')
pipe.to(device)

def unclip_image_interpolation(
  start_image,
  end_image,
  steps,
  seed
):
    generator = torch.Generator()
    generator.manual_seed(seed)

    images = [start_image, end_image]
    output = pipe(image=images, steps=steps, generator=generator)
    return output.images

inputs = [
  gr.Image(type="pil"),
  gr.Image(type="pil"),
  gr.Slider(minimum=2, maximum=12, default=5, step=1, label="Steps"),
  gr.Number(0, label="Seed", precision=0)
]

output = gr.Gallery(
            label="Generated images", show_label=False, elem_id="gallery"
        ).style(grid=[2], height="auto")

examples = [
  ["starry_night.jpg","dogs.jpg", 5, 20],
  ["flowers.jpg", "dogs.jpg", 5, 42],
  ["starry_night.jpg","flowers.jpg", 6, 9011]
]

title = "UnClip Image Interpolation Pipeline"

demo_app = gr.Interface(
    fn=unclip_image_interpolation,
    inputs=inputs,
    outputs=output,
    title=title,
    theme='huggingface',
    examples=examples,
    cache_examples=False
)
demo_app.launch(debug=True, enable_queue=True)