Update app.py

app.py

@@ -1,10 +1,9 @@
 import gradio as gr
 from colpali_engine.models import ColQwen2, ColQwen2Processor
 import torch
-import base64
 from PIL import Image
-import io
 import logging
+import numpy as np
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -13,25 +12,37 @@ logger = logging.getLogger("colqwen-api")
 # Initialize model
 logger.info("Loading ColQwen2 model...")
 model = ColQwen2.from_pretrained(
-    "vidore/colqwen2-v0.1",
+    "vidore/colqwen2-v1.0",  # Updated to v1.0
     torch_dtype=torch.bfloat16,
     device_map="auto",
 )
-processor = ColQwen2Processor.from_pretrained("vidore/colqwen2-v0.1")
+processor = ColQwen2Processor.from_pretrained("vidore/colqwen2-v1.0")
 model = model.eval()
 logger.info("Model loaded successfully")
 
-def process_image(image_data):
+def process_image(image):
     try:
         logger.info("Processing image")
-        processed = processor.process_images([image_data])
-        logger.info("Image processed")
+        # Convert to PIL Image if needed
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
         
-        with torch.no_grad():
-            embeddings = model(**processed)
-            logger.info(f"Embeddings generated: {embeddings.shape}")
+        # Process image
+        inputs = processor(
+            images=image,
+            return_tensors="pt"
+        ).to(model.device)
         
-        return {"embeddings": embeddings.tolist()}
+        logger.info("Generating embeddings")
+        with torch.no_grad():
+            outputs = model(**inputs)
+            embeddings = outputs.last_hidden_state.mean(dim=1).cpu().numpy()
+            
+        logger.info(f"Embeddings shape: {embeddings.shape}")
+        return {
+            "embeddings": embeddings.tolist(),
+            "shape": embeddings.shape
+        }
     except Exception as e:
         logger.error(f"Error: {str(e)}", exc_info=True)
         raise
@@ -40,8 +51,9 @@ interface = gr.Interface(
     fn=process_image,
     inputs=gr.Image(),
     outputs="json",
-    title="ColQwen2 Embedding API"
+    title="ColQwen2 Embedding API",
+    description="Generate embeddings from images using ColQwen2"
 )
 
-# Add share=True to create public URL
-interface.launch()
+# Launch with API
+interface.launch(server_name="0.0.0.0", server_port=7861)