import gradio as gr
from transformers import AutoImageProcessor, AutoModel
import torch
from PIL import Image
import base64
import io

processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
model = AutoModel.from_pretrained("facebook/dinov2-base")

def get_embedding(base64_str):
    header, encoded = base64_str.split(",", 1)
    image_data = base64.b64decode(encoded)
    image = Image.open(io.BytesIO(image_data)).convert("RGB")

    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        embeddings = model(**inputs).last_hidden_state[:, 0]
    return embeddings.squeeze().tolist()

iface = gr.Interface(
    fn=get_embedding,
    inputs="text",  # ahora recibimos un string base64
    outputs="json",
    description="Microservicio para extraer embeddings desde base64."
)

iface.queue()  # 👈 Esta línea activa el sistema de event_id y polling
iface.launch()