import os

# Redirect all Hugging Face / Torch caches to ephemeral storage
os.environ["HF_HOME"] = "/tmp/hf"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf"
os.environ["TORCH_HOME"] = "/tmp/torch"
os.environ["XDG_CACHE_HOME"] = "/tmp"

import gradio as gr
import torch
import open_clip
import faiss
import numpy as np
import json
from PIL import Image
from transformers import BlipProcessor, BlipForConditionalGeneration

with open("image_paths.json", "r") as f:
    image_paths = json.load(f)

image_embeddings = np.load("image_embeddings.npy")

model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
tokenizer = open_clip.get_tokenizer('ViT-B-32')

d = image_embeddings.shape[1]
index = faiss.IndexFlatIP(d)
index.add(image_embeddings)

processor_cap = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model_cap = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

def search_text(query, top_k=3):
    with torch.no_grad():
        tokenized = tokenizer([query])
        text_embed = model.encode_text(tokenized)
        text_embed = text_embed / text_embed.norm(dim=-1, keepdim=True)
        text_np = text_embed.cpu().numpy()
    _, I = index.search(text_np, top_k)
    return [image_paths[i] for i in I[0]]

def search_image(image, top_k=3):
    image_tensor = preprocess(image).unsqueeze(0)
    with torch.no_grad():
        image_embed = model.encode_image(image_tensor)
        image_embed = image_embed / image_embed.norm(dim=-1, keepdim=True)
        image_np = image_embed.cpu().numpy()
    _, I = index.search(image_np, top_k)
    return [image_paths[i] for i in I[0]]

def generate_captions(paths):
    imgs = [Image.open(p).convert("RGB") for p in paths]
    inputs = processor_cap(images=imgs, return_tensors="pt")
    out = model_cap.generate(**inputs)
    captions = [processor_cap.decode(out[i], skip_special_tokens=True) for i in range(len(imgs))]
    return captions

def predict(input_text, input_image):
    if input_text:
        paths = search_text(input_text)
    elif input_image:
        paths = search_image(input_image)
    else:
        return None, None

    captions = generate_captions(paths)
    return [(Image.open(p), c) for p, c in zip(paths, captions)]

with gr.Blocks() as demo:
    gr.Markdown("## Animal Image Search 🐈 🏞️ 🔎 ")
    gr.Markdown("Search database of 5400 animal images using text or image queries.")
    gr.Markdown("Results come with automatic captions.")
    gr.Markdown("For list of animal categories (90 animals x 60 images = 5400 images) visit [GitHub](https://github.com/TensorCruncher/animal-image-search/blob/main/animals.txt).")
    gr.Markdown("Note: Make sure only one input field (text or image) is populated before searching")

    with gr.Row():
        text_input = gr.Textbox(label="Type your query")
        image_input = gr.Image(type="pil", label="Or upload an image")
    
    gr.Examples(
    examples=[
        "examples/sample1.jpg",
        "examples/sample2.jpg",
        "examples/sample3.jpg"
    ],
    inputs=[image_input],
    label="Choose a sample image"
    )

    output = gr.Gallery(label="Top Matches with Captions")

    submit = gr.Button("Search", variant="primary")
    submit.click(predict, inputs=[text_input, image_input], outputs=output)

demo.launch()