Update app.py

app.py

@@ -2,8 +2,8 @@ import gradio as gr
 from datasets import load_dataset
 import numpy as np
 from PIL import Image
-from sentence_transformers import SentenceTransformer
 import torch
+from transformers import CLIPModel, CLIPProcessor
 
 def flux_to_gray(flux_array):
     a = np.array(flux_array, dtype=np.float32)
@@ -12,37 +12,32 @@ def flux_to_gray(flux_array):
         axis = int(np.argmin(a.shape))
         a = np.nanmean(a, axis=axis)
     a = np.nan_to_num(a, nan=0.0, posinf=0.0, neginf=0.0)
-    lo = np.nanpercentile(a, 1)
-    hi = np.nanpercentile(a, 99)
+    lo = np.nanpercentile(a, 1); hi = np.nanpercentile(a, 99)
     if not np.isfinite(lo) or not np.isfinite(hi) or hi <= lo:
         lo, hi = float(np.nanmin(a)), float(np.nanmax(a))
     norm = np.clip((a - lo) / (hi - lo + 1e-9), 0, 1)
     arr = (norm * 255).astype(np.uint8)
     return Image.fromarray(arr, mode="L")
 
-# well known model from sentence-transformers
-model = SentenceTransformer("clip-ViT-B-32")  # alias for sentence-transformers/clip-ViT-B-32
+model_id = "openai/clip-vit-base-patch32"
+model = CLIPModel.from_pretrained(model_id)
+processor = CLIPProcessor.from_pretrained(model_id)
 
-def test_single_embedding():
+def test_clip():
     ds = load_dataset("MultimodalUniverse/jwst", split="train", streaming=True)
     rec = next(iter(ds))
     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
+        image_inputs = processor(images=pil, return_tensors="pt")
+        image_feats = model.get_image_features(**image_inputs)  # [1, 512]
+    return pil, f"image_features shape: {tuple(image_feats.shape)}", f"object_id: {rec.get('object_id')}"
 
 demo = gr.Interface(
-    fn=test_single_embedding,
+    fn=test_clip,
     inputs=None,
-    outputs=[
-        gr.Image(type="pil", label="Preview"),
-        gr.Textbox(label="Info", lines=1),
-        gr.Textbox(label="Embedding", lines=1),
-    ],
-    title="JWST → CLIP embedding check",
-    description="Embeds one JWST image with CLIP to confirm the pipeline."
+    outputs=[gr.Image(type="pil", label="Preview"),
+             gr.Textbox(label="Shape", lines=1),
+             gr.Textbox(label="Info", lines=1)],
+    title="JWST → CLIP embedding check (transformers)"
 )
-
 demo.launch()