run on cpu

app.py

@@ -1,63 +1,33 @@
 """This space is taken and modified from https://huggingface.co/spaces/merve/compare_clip_siglip"""
-import torch
-from transformers import (
-    AutoModel,
-    AutoProcessor
-)
+from transformers import pipeline
 import gradio as gr
 
 ################################################################################
 # Load the models
 ################################################################################
 sg1_ckpt = "google/siglip-so400m-patch14-384"
-siglip1_model = AutoModel.from_pretrained(sg1_ckpt, device_map="auto").eval()
-siglip1_processor = AutoProcessor.from_pretrained(sg1_ckpt)
+sg1_pipe = pipeline(task="zero-shot-image-classification", model=sg1_ckpt, device="cpu")
 
 sg2_ckpt = "google/siglip2-so400m-patch14-384"
-siglip2_model = AutoModel.from_pretrained(sg2_ckpt, device_map="auto").eval()
-siglip2_processor = AutoProcessor.from_pretrained(sg2_ckpt)
+sg2_pipe = pipeline(task="zero-shot-image-classification", model=sg2_ckpt, device="cpu")
 
 ################################################################################
-# Utilities
+# Run inference
 ################################################################################
-def postprocess(output):
-    return {out["label"]: float(out["score"]) for out in output}
-
-
-def postprocess_siglip(sg1_probs, sg2_probs, labels):
-    sg1_output = {labels[i]: float(sg1_probs[0].cpu().numpy()[i]) for i in range(len(labels))} 
-    sg2_output = {labels[i]: float(sg2_probs[0].cpu().numpy()[i]) for i in range(len(labels))}
-    return sg1_output, sg2_output
-
-def siglip_detector(image, texts):
-    sg1_inputs = siglip1_processor(
-        text=texts, images=image, return_tensors="pt", padding="max_length", max_length=64
-    ).to(siglip1_model.device)
-
-    sg2_inputs = siglip2_processor(
-        text=texts, images=image, return_tensors="pt", padding="max_length", max_length=64
-    ).to(siglip2_model.device)
-
-    with torch.no_grad():
-        sg1_outputs = siglip1_model(**sg1_inputs)
-        sg2_outputs = siglip2_model(**sg2_inputs)
-
-        sg1_logits_per_image = sg1_outputs.logits_per_image
-        sg2_logits_per_image = sg2_outputs.logits_per_image
-
-        sg1_probs = torch.sigmoid(sg1_logits_per_image)
-        sg2_probs = torch.sigmoid(sg2_logits_per_image)
-    return sg1_probs, sg2_probs
-
-
 def infer(image, candidate_labels):
     candidate_labels = [label.lstrip(" ") for label in candidate_labels.split(",")]
-    sg1_probs, sg2_probs = siglip_detector(image, candidate_labels)
-    return postprocess_siglip(
-        sg1_probs, sg2_probs, labels=candidate_labels
-    )
 
+    sg1_socres = sg1_pipe(image, candidate_labels=candidate_labels)
+    sg2_socres = sg2_pipe(image, candidate_labels=candidate_labels)
+    
+    sg1_outputs = {element["label"]:element["score"] for element in sg1_socres}
+    sg2_outputs = {element["label"]:element["score"] for element in sg2_socres}
 
+    return sg1_outputs, sg2_outputs
+
+################################################################################
+# Gradio App
+################################################################################
 with gr.Blocks() as demo:
     gr.Markdown("# Compare SigLIP 1 and SigLIP 2")
     gr.Markdown(