Update app.py

app.py

@@ -3,8 +3,8 @@ from datasets import load_dataset
 import numpy as np
 from PIL import Image
 from sentence_transformers import SentenceTransformer
+import torch
 
-# reuse the same grayscale conversion
 def flux_to_gray(flux_array):
     a = np.array(flux_array, dtype=np.float32)
     a = np.squeeze(a)
@@ -20,14 +20,15 @@ def flux_to_gray(flux_array):
     arr = (norm * 255).astype(np.uint8)
     return Image.fromarray(arr, mode="L")
 
-# load a well-known CLIP model
-model = SentenceTransformer("clip-ViT-B-32")
+# well known model from sentence-transformers
+model = SentenceTransformer("clip-ViT-B-32")  # alias for sentence-transformers/clip-ViT-B-32
 
 def test_single_embedding():
     ds = load_dataset("MultimodalUniverse/jwst", split="train", streaming=True)
     rec = next(iter(ds))
-    pil = flux_to_gray(rec["image"]["flux"]).convert("RGB")  # CLIP expects RGB
-    emb = model.encode(pil, convert_to_numpy=True)
+    pil = flux_to_gray(rec["image"]["flux"]).convert("RGB")
+    with torch.no_grad():
+        emb = model.encode([pil], convert_to_numpy=True, normalize_embeddings=True)  # list input
     info = f"OK. Image embedding shape: {emb.shape}"
     caption = f"object_id: {rec.get('object_id')}"
     return pil, caption, info