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

orig_start_prompt = "a photograph of an adult lion"
orig_end_prompt = "a photograph of a lion cub"

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

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

def unclip_text_interpolation(
  start_prompt,
  end_prompt,
  steps,
  seed
):
    generator = torch.Generator()
    generator.manual_seed(seed)

    output = pipe(start_prompt, end_prompt, steps, enable_sequential_cpu_offload=False, generator=generator)
    return output.images

inputs = [
  gr.Textbox(lines=2, default=orig_start_prompt, label="Start Prompt"),
  gr.Textbox(lines=2, default=orig_end_prompt, label="End Prompt"),
  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 = [
  [orig_start_prompt, orig_end_prompt, 5, 42],
  ["a photo of a landscape in winter","a photo of a landscape in fall", 5, 20],
  ["a photo of a victorian house", "a photo of a modern house", 5, 20]
]

title = "UnClip Text Interpolation Pipeline"

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