import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel, PeftConfig
from PIL import Image
import requests
from io import BytesIO
import torchvision.datasets as datasets
import numpy as np

# Load SigLIP for image embeddings
from model.siglip import SigLIPModel

def get_cifar_examples():
    # Load CIFAR10 test set
    cifar10_test = datasets.CIFAR10(root='./data', train=False, download=True)
    
    # CIFAR10 classes
    classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 
              'dog', 'frog', 'horse', 'ship', 'truck']
    
    # Get one example from each class
    examples = []
    used_classes = set()
    
    for idx in range(len(cifar10_test)):
        img, label = cifar10_test[idx]
        if classes[label] not in used_classes:
            # Save the image temporarily
            img_path = f"examples/{classes[label]}_example.jpg"
            img.save(img_path)
            examples.append(img_path)
            used_classes.add(classes[label])
            
        if len(used_classes) == 10:  # We have one example from each class
            break
    
    return examples

def load_models():
    # Load SigLIP model
    siglip = SigLIPModel()
    
    # Load base Phi model
    base_model = AutoModelForCausalLM.from_pretrained(
        "microsoft/Phi-3-mini-4k-instruct",
        trust_remote_code=True,
        device_map="auto",
        torch_dtype=torch.float32
    )
    
    # Load our fine-tuned LoRA adapter
    model = PeftModel.from_pretrained(
        base_model,
        "jatingocodeo/phi-vlm",  # Your uploaded model
        device_map="auto"
    )
    
    tokenizer = AutoTokenizer.from_pretrained("jatingocodeo/phi-vlm")
    
    return siglip, model, tokenizer

def generate_description(image, siglip, model, tokenizer):
    # Convert image to RGB if needed
    if image.mode != "RGB":
        image = image.convert("RGB")
    
    # Resize image to match SigLIP's expected size
    image = image.resize((32, 32))
    
    # Get image embedding from SigLIP
    image_embedding = siglip.encode_image(image)
    
    # Prepare prompt
    prompt = """Below is an image. Please describe it in detail.

Image: <image>
Description: """
    
    # Tokenize input
    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=128
    ).to(model.device)
    
    # Generate description
    with torch.no_grad():
        outputs = model(
            **inputs,
            image_embeddings=image_embedding.unsqueeze(0),
            max_new_tokens=100,
            temperature=0.7,
            do_sample=True,
            top_p=0.9
        )
    
    # Decode and return the generated text
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_text.split("Description: ")[-1].strip()

# Load models
print("Loading models...")
siglip, model, tokenizer = load_models()

# Create Gradio interface
def process_image(image):
    description = generate_description(image, siglip, model, tokenizer)
    return description

# Get CIFAR10 examples
examples = get_cifar_examples()

# Define interface
iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Textbox(label="Generated Description"),
    title="Image Description Generator",
    description="""Upload an image and get a detailed description generated by our fine-tuned VLM model.
                  Below are sample images from CIFAR10 dataset that you can try.""",
    examples=[[ex] for ex in examples]  # Format examples for Gradio
)

# Launch the interface
if __name__ == "__main__":
    iface.launch()