Spaces:
Sleeping
Sleeping
File size: 2,435 Bytes
6843d1b 4f1b9ee 6843d1b 4f1b9ee 6843d1b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | """
SigLIP 2 Text & Image Encoder -- HuggingFace Space
Encodes text or image queries to 768-dim vectors for the Epstein photo search.
Model: google/siglip2-base-patch16-224
"""
import gradio as gr
import torch
import torch.nn.functional as F
from PIL import Image
from transformers import AutoModel, AutoTokenizer, AutoProcessor
MODEL_NAME = "google/siglip2-base-patch16-224"
print(f"Loading {MODEL_NAME}...")
model = AutoModel.from_pretrained(MODEL_NAME).eval()
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
processor = AutoProcessor.from_pretrained(MODEL_NAME)
print(f"Model loaded. Text hidden size: {model.config.text_config.hidden_size}")
def encode(text: str) -> list:
inputs = tokenizer([text], return_tensors="pt", padding="max_length", max_length=64, truncation=True)
with torch.no_grad():
feats = model.text_model(**inputs).pooler_output
feats = F.normalize(feats, dim=-1)
return feats[0].tolist()
def encode_image(image) -> list:
if image is None:
raise gr.Error("No image provided")
# Gradio 6.x base64 shortcut returns RGBA — SigLIP needs RGB
if isinstance(image, Image.Image):
image = image.convert("RGB")
elif isinstance(image, str):
image = Image.open(image).convert("RGB")
else:
raise gr.Error(f"Unexpected image type: {type(image)}")
inputs = processor(images=[image], return_tensors="pt")
with torch.no_grad():
feats = model.get_image_features(pixel_values=inputs["pixel_values"])
if not isinstance(feats, torch.Tensor):
feats = feats.pooler_output
feats = F.normalize(feats, dim=-1)
return feats[0].tolist()
with gr.Blocks(title="SigLIP 2 Encoder") as demo:
gr.Markdown("# SigLIP 2 Encoder\nEncodes text or images to 768-dim normalized vectors using google/siglip2-base-patch16-224")
with gr.Tab("Text"):
text_input = gr.Textbox(label="Text")
text_output = gr.JSON(label="Embedding (768-dim)")
text_btn = gr.Button("Encode Text")
text_btn.click(fn=encode, inputs=text_input, outputs=text_output, api_name="encode")
with gr.Tab("Image"):
image_input = gr.Image(type="pil", label="Image")
image_output = gr.JSON(label="Embedding (768-dim)")
image_btn = gr.Button("Encode Image")
image_btn.click(fn=encode_image, inputs=image_input, outputs=image_output, api_name="encode_image")
demo.launch()
|