import gradio as gr
import torch
import numpy as np
from PIL import Image
import os

from vslerp import UnCLIPImageInterpolationPipeline  # your pipeline + vSLERP

device = "cuda" if torch.cuda.is_available() else "cpu"

# Load pipeline once
pipe = UnCLIPImageInterpolationPipeline.from_pretrained(
    "kakaobrain/karlo-v1-alpha-image-variations",
    torch_dtype=torch.float16
).to(device)

# Put your own images in a local "bank" folder
IMAGE_BANK = {
    "Example 1": "lj.png",
    "Example 2": "kd.png",
    "Example 3": "vase.png",
    "Example 4": "lamp.jpeg"
}

def run_vslerp(img0, img1, bank0, bank1, slerp_num_steps, vslerp_start_idx, vslerp_end_idx, vslerp_num_steps):
    # Decide input images: uploaded takes precedence, else from bank
    if img0 is None and bank0 != "None":
        img0 = Image.open(IMAGE_BANK[bank0])
    if img1 is None and bank1 != "None":
        img1 = Image.open(IMAGE_BANK[bank1])

    if img0 is None or img1 is None:
        raise ValueError("Please provide two images (either upload or select from bank).")

    images = [img0, img1]
    generator = torch.Generator(device=device).manual_seed(42)

    # Prepare a 2D list for the gallery
    gallery_matrix = []

    vslerp_values = np.linspace(vslerp_start_idx, vslerp_end_idx, vslerp_num_steps)
    for m_val in vslerp_values:
        row = []
        for step in range(slerp_num_steps):
            out = pipe(
                image=images,
                generator=generator,
                steps=slerp_num_steps,
                decoder_guidance_scale=1,
                mean_val=m_val
            )
            row.append(out.images[0])  # assuming pipe returns a list with one image per call
        gallery_matrix.append(row)

    return gallery_matrix


with gr.Blocks() as demo:
    gr.Markdown("## vSLERP Demo")
    gr.Markdown("Note: The run may take a while, please be patient 🙏")

    with gr.Row():
        with gr.Column():
            img0 = gr.Image(label="Upload Image 0", type="pil")
            bank0 = gr.Dropdown(choices=["None"] + list(IMAGE_BANK.keys()), value="None", label="Or choose from bank")
        with gr.Column():
            img1 = gr.Image(label="Upload Image 1", type="pil")
            bank1 = gr.Dropdown(choices=["None"] + list(IMAGE_BANK.keys()), value="None", label="Or choose from bank")

    with gr.Row():
        slerp_num_steps = gr.Slider(3, 6, value=6, step=1, label="slerp_num_steps")
        vslerp_start_idx = gr.Slider(-2, 0, value=-1, step=1, label="vslerp_start_idx")
        vslerp_end_idx = gr.Slider(1, 3, value=3, step=1, label="vslerp_end_idx")
        vslerp_num_steps = gr.Slider(3, 6, value=6, step=1, label="vslerp_num_steps")

    run_btn = gr.Button("Run vSLERP")
    gallery = gr.Gallery(label="Generated Interpolations").style(grid=[4], height="auto")

    run_btn.click(
        run_vslerp,
        inputs=[img0, img1, bank0, bank1, slerp_num_steps, vslerp_start_idx, vslerp_end_idx, vslerp_num_steps],
        outputs=[gallery]
    )

if __name__ == "__main__":
    demo.launch()