Update app.py

app.py

@@ -1,16 +1,46 @@
-# app.py
 import gradio as gr
 from colpali_engine.models import ColQwen2, ColQwen2Processor
 import torch
 import base64
 from PIL import Image
 import io
+import logging
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("colqwen-api")
 
 # Initialize model
+logger.info("Loading ColQwen2 model...")
 model = ColQwen2.from_pretrained(
     "vidore/colqwen2-v0.1",
     torch_dtype=torch.bfloat16,
     device_map="auto",
 )
 processor = ColQwen2Processor.from_pretrained("vidore/colqwen2-v0.1")
-model = model.eval()
+model = model.eval()
+logger.info("Model loaded successfully")
+
+def process_image(image_data):
+    try:
+        logger.info("Processing image")
+        processed = processor.process_images([image_data])
+        logger.info("Image processed")
+        
+        with torch.no_grad():
+            embeddings = model(**processed)
+            logger.info(f"Embeddings generated: {embeddings.shape}")
+        
+        return {"embeddings": embeddings.tolist()}
+    except Exception as e:
+        logger.error(f"Error: {str(e)}", exc_info=True)
+        raise
+
+interface = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(),
+    outputs="json",
+    title="ColQwen2 Embedding API"
+)
+
+interface.launch()