import gradio as gr
import os
import torchvision.transforms as T
from upstash_vector import Index
from datasets import load_dataset
from transformers import AutoFeatureExtractor, AutoModel

index = Index.from_env()


model_ckpt = "google/vit-base-patch16-224-in21k"
extractor = AutoFeatureExtractor.from_pretrained(model_ckpt)
model = AutoModel.from_pretrained(model_ckpt)
hidden_dim = model.config.hidden_size
dataset = load_dataset("BounharAbdelaziz/Face-Aging-Dataset")

# Data transformation chain.
transformation_chain = T.Compose(
    [
        T.Resize(extractor.size["height"]),
        T.CenterCrop(extractor.size["height"]),
        T.ToTensor(),
        T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
)

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # Find Your Twins

        Upload your face and find the most similar people from the X dataset. Powered by [Upstash Vector](https://upstash.com) 🚀
        """
    )

    with gr.Tab("Basic"):
        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(type="pil")
            with gr.Column(scale=3):
                output_image = gr.Gallery(height=800)


        @input_image.upload(inputs=input_image, outputs=output_image)
        def find_similar_faces(image):
            t_image = transformation_chain(image)
            inputs = extractor(images=t_image, return_tensors="pt")
            outputs = model(**inputs)
            embed = outputs.last_hidden_state[0][0]
            result = index.query(vector=embed, top_k=4)
            return [dataset["train"][int(vector.id)]["image"] for vector in result]

    with gr.Tab("Advanced"):
        with gr.Row():
            with gr.Column(scale=1):
                adv_input_image = gr.Image(type="pil")
                adv_image_count = gr.Number(9, label="Image Count")

            with gr.Column(scale=3):
                adv_output_image = gr.Gallery(height=1000)


        @adv_input_image.upload(inputs=[adv_input_image, adv_image_count], outputs=[adv_output_image])
        def find_similar_faces(image, count):
            t_image = transformation_chain(image)
            inputs = extractor(images=t_image, return_tensors="pt")
            outputs = model(**inputs)
            embed = outputs.last_hidden_state[0][0]
            result = index.query(vector=embed, top_k=max(1, min(19, count)))
            return [dataset["train"][int(vector.id)]["image"] for vector in result]

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