import gradio as gr
from transformers import AutoImageProcessor, AutoModel
import torch
from PIL import Image
import numpy as np

# Cargar el modelo solo una vez
processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
model = AutoModel.from_pretrained("facebook/dinov2-base")
model.eval()

def get_embedding(image_np):
    image = Image.fromarray(image_np).convert("RGB")
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        embeddings = model(**inputs).last_hidden_state[:, 0]  # CLS token
    return embeddings.squeeze().tolist()

iface = gr.Interface(
    fn=get_embedding,
    inputs=gr.Image(type="numpy"),  # CAMBIO CLAVE AQUÍ
    outputs="json",
    description="Microservicio para extraer embeddings de imágenes usando DINOv2."
)

iface.queue()
iface.launch()