Upload app.py

app.py

@@ -1,5 +1,6 @@
 """
 CLIP Image Embedding API - Lightweight version for HF Spaces free tier
+Supports both URL and base64 image input
 """
 
 import gradio as gr
@@ -8,6 +9,8 @@ from PIL import Image
 from transformers import CLIPProcessor, CLIPModel
 import requests
 from io import BytesIO
+import base64
+import re
 
 # Use CPU and smaller memory footprint
 model = None
@@ -21,18 +24,34 @@ def load_model():
         model.eval()
     return model, processor
 
-def get_embedding_from_url(image_url: str):
-    """Get CLIP embedding from image URL"""
+def get_embedding(image_input: str):
+    """Get CLIP embedding from image URL or base64 string"""
     try:
-        if not image_url or not image_url.startswith('http'):
-            return {"success": False, "error": "Please provide a valid image URL"}
+        if not image_input:
+            return {"success": False, "error": "Please provide an image URL or base64 string"}
         
         # Load model on first use
         model, processor = load_model()
         
-        # Download image
-        response = requests.get(image_url, timeout=30)
-        image = Image.open(BytesIO(response.content)).convert('RGB')
+        image = None
+        
+        # Check if it's base64 (data:image/... or raw base64)
+        if image_input.startswith('data:image'):
+            # Extract base64 data after the comma
+            base64_data = image_input.split(',')[1] if ',' in image_input else image_input
+            image_bytes = base64.b64decode(base64_data)
+            image = Image.open(BytesIO(image_bytes)).convert('RGB')
+        elif not image_input.startswith('http'):
+            # Try as raw base64
+            try:
+                image_bytes = base64.b64decode(image_input)
+                image = Image.open(BytesIO(image_bytes)).convert('RGB')
+            except:
+                return {"success": False, "error": "Invalid input: provide URL or base64"}
+        else:
+            # It's a URL - download it
+            response = requests.get(image_input, timeout=30)
+            image = Image.open(BytesIO(response.content)).convert('RGB')
         
         # Get embedding
         inputs = processor(images=image, return_tensors="pt")
@@ -52,11 +71,11 @@ def get_embedding_from_url(image_url: str):
 
 # Gradio interface with API enabled
 demo = gr.Interface(
-    fn=get_embedding_from_url,
-    inputs=gr.Textbox(label="Image URL", placeholder="https://example.com/image.jpg"),
+    fn=get_embedding,
+    inputs=gr.Textbox(label="Image (URL or base64)", placeholder="https://example.com/image.jpg or data:image/jpeg;base64,..."),
     outputs=gr.JSON(label="Result"),
     title="CLIP Embedding API",
-    description="Get 512-dim CLIP embeddings for images",
+    description="Get 512-dim CLIP embeddings from image URL or base64",
     api_name="predict"
 )