import gradio as gr
import spaces
from PIL import Image
from typing import List
import glob
import os

from pipeline import prepare_white_image, MultiViewGenerator
from util import download_file, unzip_file

download_file("https://huggingface.co/aki-0421/character-360/resolve/main/v2.ckpt", "v2.ckpt")
download_file("https://huggingface.co/hbyang/Hi3D/resolve/main/ckpts.zip", "ckpts.zip")

unzip_file("ckpts.zip", ".")

multi_view_generator = MultiViewGenerator(checkpoint_path="v2.ckpt")

@spaces.GPU(duration=120)
def generate_images(input_image: Image.Image) -> List[Image.Image]:
    white_image = prepare_white_image(input_image=input_image)

    return multi_view_generator.infer(white_image=white_image)

def prepare_example() -> List[Image.Image]:
    image_files = sorted(glob.glob(os.path.join("assets/", '*.png')))

    images = []
    for image_file in image_files:
        image = Image.open(image_file)
        images.append(image)

    return images

with gr.Blocks() as demo:
    gr.Markdown("# Character-360")
    gr.Markdown("Do you want to check out your character in 360°?")
    with gr.Row():
        input_image = gr.Image(label="Input Image", type="pil")
        output_gallery = gr.Gallery(
        label="Generated images", show_label=False, elem_id="gallery"
    , columns=[3], object_fit="contain", height="auto")

    btn = gr.Button("Generate")
    btn.click(generate_images, input_image, output_gallery)

    example_images = prepare_example()
    gr.Markdown("### Example Output")
    gr.Gallery(
        value=example_images,
        label="Example Images",
        columns=[5], object_fit="contain", height="auto"
    )

demo.launch()