app.py

# app.py — veureu/svision (Salamandra Vision 7B · ZeroGPU) — compatible con ENGINE
import os
import json
from typing import Dict, List, Optional, Tuple, Union

import gradio as gr
import spaces
import torch
from facenet_pytorch import MTCNN, InceptionResnetV1
import numpy as np
from PIL import Image
from transformers import AutoProcessor, LlavaOnevisionForConditionalGeneration

MODEL_ID = os.environ.get("MODEL_ID", "BSC-LT/salamandra-7b-vision")
DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

_model = None
_processor = None
_mtcnn = None
_facenet = None


def _load_face_models() -> Tuple[MTCNN, InceptionResnetV1]:
    global _mtcnn, _facenet
    if _mtcnn is None or _facenet is None:
        device = DEVICE if DEVICE == "cuda" and torch.cuda.is_available() else "cpu"
        _mtcnn = MTCNN(image_size=160, margin=0, post_process=True, device=device)
        _facenet = InceptionResnetV1(pretrained="vggface2").eval().to(device)
    return _mtcnn, _facenet


def _lazy_load() -> Tuple[LlavaOnevisionForConditionalGeneration, AutoProcessor]:
    global _model, _processor
    if _model is None or _processor is None:
        _processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
        _model = LlavaOnevisionForConditionalGeneration.from_pretrained(
            MODEL_ID,
            dtype=DTYPE,
            low_cpu_mem_usage=True,
            trust_remote_code=True,
            use_safetensors=True,
            device_map=None,
        )
        _model.to(DEVICE)
    return _model, _processor


def _compose_prompt(user_text: str, context: Optional[Dict] = None) -> List[Dict]:
    """Construye el chat template con imagen + texto + contexto opcional."""
    ctx_txt = ""
    if context:
        try:
            # breve, sin ruido
            ctx_txt = "\n\nContexto adicional:\n" + json.dumps(context, ensure_ascii=False)[:2000]
        except Exception:
            pass
    user_txt = (user_text or "Describe la imagen con detalle.") + ctx_txt
    convo = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": user_txt},
            ],
        }
    ]
    return convo


@spaces.GPU  # en HF Spaces usará GPU cuando haya disponibilidad (ZeroGPU)
def _infer_one(image: Image.Image, text: str, max_new_tokens: int = 256, temperature: float = 0.7,
               context: Optional[Dict] = None) -> str:
    # Reducir el tamaño de la imagen para ahorrar memoria en la GPU
    image.thumbnail((1024, 1024))

    model, processor = _lazy_load()
    prompt = processor.apply_chat_template(_compose_prompt(text, context), add_generation_prompt=True)
    inputs = processor(images=image, text=prompt, return_tensors="pt").to(DEVICE, dtype=DTYPE)
    with torch.inference_mode():
        out = model.generate(**inputs, max_new_tokens=int(max_new_tokens), temperature=float(temperature))
    return processor.decode(out[0], skip_special_tokens=True).strip()


# ----------------------------- API helpers -----------------------------------

def describe_raw(image: Image.Image, text: str = "Describe la imagen con detalle.",
                 max_new_tokens: int = 256, temperature: float = 0.7) -> Dict[str, str]:
    result = _infer_one(image, text, max_new_tokens, temperature, context=None)
    return {"text": result}


def describe_batch(images: List[Image.Image], context_json: str,
                   max_new_tokens: int = 256, temperature: float = 0.7) -> List[str]:
    """Endpoint batch para ENGINE: lista de imágenes + contexto (JSON) → lista de textos."""
    try:
        context = json.loads(context_json) if context_json else None
    except Exception:
        context = None
    outputs: List[str] = []
    for img in images:
        outputs.append(_infer_one(img, text="Describe la imagen con detalle.", max_new_tokens=max_new_tokens,
                                  temperature=temperature, context=context))
    return outputs


@spaces.GPU
def face_image_embedding(image: Image.Image) -> List[float] | None:
    try:
        mtcnn, facenet = _load_face_models()
        # detectar y extraer cara
        face = mtcnn(image)

        if face is None:
            return None

        # FaceNet espera tensor shape (1,3,160,160)
        device = DEVICE if DEVICE == "cuda" and torch.cuda.is_available() else "cpu"
        face = face.unsqueeze(0).to(device)

        # obtener embedding
        with torch.no_grad():
            emb = facenet(face).cpu().numpy()[0]

        # normalizar igual que tu código original
        emb = emb / np.linalg.norm(emb)

        return emb.astype(float).tolist()

    except Exception as e:
        print(f"Fallo embedding cara: {e}")
        return None


# ----------------------------- UI & Endpoints --------------------------------

with gr.Blocks(title="Salamandra Vision 7B · ZeroGPU") as demo:
    gr.Markdown("## Salamandra-Vision 7B · ZeroGPU\nImagen + texto → descripción.")
    with gr.Row():
        with gr.Column():
            in_img = gr.Image(label="Imagen", type="pil")
            in_txt = gr.Textbox(label="Texto/prompt", value="Describe la imagen con detalle (ES/CA).")
            max_new = gr.Slider(16, 1024, value=256, step=16, label="max_new_tokens")
            temp = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="temperature")
            btn = gr.Button("Generar", variant="primary")
        with gr.Column():
            out = gr.Textbox(label="Descripción", lines=18)

    # UI
    btn.click(_infer_one, [in_img, in_txt, max_new, temp], out, api_name="describe", concurrency_limit=1)

    # API simple (multipart) compatible con tu versión anterior
    # demo.load(
    #     None,
    #     [gr.Image(label="image", type="pil"),
    #      gr.Textbox(value="Describe la imagen con detalle."),
    #      gr.Slider(16, 1024, value=256),
    #      gr.Slider(0.0, 1.5, value=0.7)],
    #     describe_raw,
    #     api_name="describe_raw"
    # )

    # API BATCH para ENGINE (Gradio Client): images + context_json → list[str]
    # Firma que espera el VisionClient del engine (api_name="/predict")
    batch_in_images = gr.Gallery(label="Imágenes (batch)", show_label=False, columns=4, height="auto")
    batch_context = gr.Textbox(label="context_json", value="{}", lines=4)
    batch_max = gr.Slider(16, 1024, value=256, step=16, label="max_new_tokens")
    batch_temp = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="temperature")
    batch_btn = gr.Button("Describir lote")
    batch_out = gr.JSON(label="Descripciones (lista)")

    # Nota: Gradio Gallery entrega rutas/obj; nos apoyamos en el cliente para cargar archivos
    batch_btn.click(describe_batch, [batch_in_images, batch_context, batch_max, batch_temp], batch_out,
                    api_name="predict", concurrency_limit=1)

    with gr.Row():
        face_img = gr.Image(label="Imagen para embedding facial", type="pil")
        face_btn = gr.Button("Obtener embedding facial")
        face_out = gr.JSON(label="Embedding facial (vector)")
    face_btn.click(face_image_embedding, [face_img], face_out, api_name="face_image_embedding", concurrency_limit=1)


demo.queue(max_size=16).launch()