Update

app.py

@@ -1,24 +1,34 @@
 import torch
 import torch.nn.functional as F
 import gradio as gr
-from transformers import CLIPProcessor, CLIPModel
+from transformers import CLIPProcessor, CLIPModel, AutoProcessor, AutoModel
 import spaces
 
-# Dictionary of available CLIP models with their image sizes
-CLIP_MODELS = {
-    "ViT-B/32": ("openai/clip-vit-base-patch32", 224),
-    "ViT-B/16": ("openai/clip-vit-base-patch16", 224),
-    "ViT-L/14": ("openai/clip-vit-large-patch14", 224),
-    "ViT-L/14@336px": ("openai/clip-vit-large-patch14-336", 336),
+# Dictionary of available models with their image sizes
+MODELS = {
+    "CLIP ViT-B/32": ("openai/clip-vit-base-patch32", 224, "clip"),
+    "CLIP ViT-B/16": ("openai/clip-vit-base-patch16", 224, "clip"),
+    "CLIP ViT-L/14": ("openai/clip-vit-large-patch14", 224, "clip"),
+    "CLIP ViT-L/14@336px": ("openai/clip-vit-large-patch14-336", 336, "clip"),
+    "SigLIP SO400M/14-384": ("google/siglip-so400m-patch14-384", 384, "siglip"),
+    "SigLIP Large/16-256": ("google/siglip-large-patch16-256", 256, "siglip"),
+    "SigLIP SO400M/14-224": ("google/siglip-so400m-patch14-224", 224, "siglip"),
+    "SigLIP Base/16-384": ("google/siglip-base-patch16-384", 384, "siglip"),
+    "SigLIP Large/16-384": ("google/siglip-large-patch16-384", 384, "siglip"),
 }
 
 # Initialize models and processors
 models = {}
 processors = {}
 
-for model_name, (model_path, _) in CLIP_MODELS.items():
-    models[model_name] = CLIPModel.from_pretrained(model_path).to("cuda")
-    processors[model_name] = CLIPProcessor.from_pretrained(model_path)
+for model_name, (model_path, _, model_type) in MODELS.items():
+    if model_type == "clip":
+        models[model_name] = CLIPModel.from_pretrained(model_path).to("cuda")
+        processors[model_name] = CLIPProcessor.from_pretrained(model_path)
+    elif model_type == "siglip":
+        models[model_name] = AutoModel.from_pretrained(model_path).to("cuda")
+        processors[model_name] = AutoProcessor.from_pretrained(model_path)
+
 
 @spaces.GPU
 def calculate_score(image, text, model_name):
@@ -27,63 +37,76 @@ def calculate_score(image, text, model_name):
     labels = list(filter(None, labels))
     if len(labels) == 0:
         return dict()
-    
+
     model = models[model_name]
     processor = processors[model_name]
-    
+    model_type = MODELS[model_name][2]
+
     # Preprocess the image and text
     inputs = processor(text=labels, images=[image], return_tensors="pt", padding=True)
     inputs = {k: v.to("cuda") for k, v in inputs.items()}
-    
+
     # Calculate embeddings
     with torch.no_grad():
         outputs = model(**inputs)
-        image_embeds = outputs.image_embeds
-        text_embeds = outputs.text_embeds
-    
+        if model_type == "clip":
+            image_embeds = outputs.image_embeds
+            text_embeds = outputs.text_embeds
+        elif model_type == "siglip":
+            image_embeds = outputs.image_embeds
+            text_embeds = outputs.text_embeds
+
     # Normalize embeddings
     image_embeds = F.normalize(image_embeds, p=2, dim=1)
     text_embeds = F.normalize(text_embeds, p=2, dim=1)
-    
+
     # Calculate cosine similarity
     cosine_similarities = torch.mm(text_embeds, image_embeds.t()).squeeze(1)
-    
-    # Convert to percentages
-    percentages = ((cosine_similarities + 1) / 2).cpu().numpy()
-    
-    results_dict = {label: float(score) for label, score in zip(labels, percentages)}
+
+    # Ensure values are between 0 and 1
+    cosine_similarities = torch.clamp(cosine_similarities, min=0, max=1)
+
+    # Convert to numpy array
+    similarities = cosine_similarities.cpu().numpy()
+
+    results_dict = {label: float(score) for label, score in zip(labels, similarities)}
     return results_dict
 
+
 with gr.Blocks() as demo:
-    gr.Markdown("# Multi-Model CLIP Score")
-    gr.Markdown("Calculate the CLIP score (cosine similarity) between the given image and text descriptions using different CLIP model variants")
-    
+    gr.Markdown("# Multi-Model CLIP and SigLIP Score")
+    gr.Markdown(
+        "Calculate the score (cosine similarity) between the given image and text descriptions using different CLIP and SigLIP model variants"
+    )
+
     with gr.Row():
         image_input = gr.Image(type="pil")
         output_label = gr.Label()
-    
+
     with gr.Row():
         text_input = gr.Textbox(label="Descriptions (separated by semicolons)")
-        model_dropdown = gr.Dropdown(choices=list(CLIP_MODELS.keys()), label="CLIP Model", value="ViT-B/16")
-    
+        model_dropdown = gr.Dropdown(
+            choices=list(MODELS.keys()), label="Model", value="CLIP ViT-B/16"
+        )
+
     def process_inputs(image, text, model_name):
         if image is None or text.strip() == "":
             return None
         return calculate_score(image, text, model_name)
-    
+
     inputs = [image_input, text_input, model_dropdown]
     outputs = output_label
-    
+
     image_input.change(fn=process_inputs, inputs=inputs, outputs=outputs)
     text_input.submit(fn=process_inputs, inputs=inputs, outputs=outputs)
     model_dropdown.change(fn=process_inputs, inputs=inputs, outputs=outputs)
-    
+
     gr.Examples(
         examples=[
             [
                 "cat.jpg",
                 "a cat stuck in a door; a cat in the air; a cat sitting; a cat standing; a cat is entering the matrix; a cat is entering the void",
-                "ViT-B/16"
+                "CLIP ViT-B/16",
             ]
         ],
         fn=process_inputs,
@@ -91,4 +114,4 @@ with gr.Blocks() as demo:
         outputs=outputs,
     )
 
-demo.launch()
+demo.launch()