app

app.py

@@ -1,7 +1,122 @@
 import gradio as gr
+import torch
+from datasets import load_dataset
+from qdrant_client import QdrantClient
+from qdrant_client.http import models
+from colpali_engine.models import ColQwen2, ColQwen2Processor
+from PIL import Image
+import requests
+from io import BytesIO
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Initialize the model, processor, and Qdrant client
+model_name = "vidore/colqwen2-v0.1"
+colpali_model = ColQwen2.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="cuda:0")
+colpali_processor = ColQwen2Processor.from_pretrained(model_name)
+qdrant_client = QdrantClient(":memory:")
+collection_name = "image_collection"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+# Load the dataset (this should be done only once when setting up the app)
+dataset = load_dataset("davanstrien/loc-nineteenth-century-song-sheets", split="train")
+
+def setup_qdrant():
+    # Create a collection in Qdrant
+    qdrant_client.recreate_collection(
+        collection_name=collection_name,
+        vectors_config=models.VectorParams(
+            size=colpali_model.config.hidden_size,
+            distance=models.Distance.COSINE,
+            multivector_config=models.MultiVectorConfig(
+                comparator=models.MultiVectorComparator.MAX_SIM
+            ),
+        ),
+    )
+
+    # Index the dataset (this should be done only once when setting up the app)
+    batch_size = 32
+    for i in range(0, len(dataset), batch_size):
+        batch = dataset[i:i+batch_size]
+        images = batch['image']
+        with torch.no_grad():
+            batch_images = colpali_processor.process_images(images).to(colpali_model.device)
+            image_embeddings = colpali_model(**batch_images)
+
+        points = []
+        for j, embedding in enumerate(image_embeddings):
+            multivector = embedding.cpu().float().numpy().tolist()
+            points.append(models.PointStruct(
+                id=i+j,
+                vector=multivector,
+                payload={
+                    "item_id": batch['item_id'][j],
+                    "item_url": batch['item_url'][j]
+                }
+            ))
+
+        qdrant_client.upsert(
+            collection_name=collection_name,
+            points=points
+        )
+
+    print("Indexing complete!")
+
+def search_similar_images(query, top_k=5, mode="text"):
+    with torch.no_grad():
+        if mode == "text":
+            batch_query = colpali_processor.process_queries([query]).to(colpali_model.device)
+        else:  # Image mode
+            batch_query = colpali_processor.process_images([query]).to(colpali_model.device)
+        query_embedding = colpali_model(**batch_query)
+
+    multivector_query = query_embedding[0].cpu().float().numpy().tolist()
+
+    search_result = qdrant_client.search(
+        collection_name=collection_name,
+        query_vector=multivector_query,
+        limit=top_k
+    )
+
+    return search_result
+
+def process_results(results):
+    output = []
+    for result in results:
+        item_url = result.payload['item_url']
+        score = result.score
+        output.append((item_url, f"Score: {score:.4f}"))
+    return output
+
+def text_search(query, top_k):
+    results = search_similar_images(query, top_k, mode="text")
+    return process_results(results)
+
+def image_search(image, top_k):
+    results = search_similar_images(image, top_k, mode="image")
+    return process_results(results)
+
+# Set up the Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Image Search App")
+    gr.Markdown("Search for similar images using text or image input.")
+    
+    with gr.Tab("Text Search"):
+        text_input = gr.Textbox(label="Enter your search query")
+        text_button = gr.Button("Search")
+        text_output = gr.Gallery(label="Results", show_label=False, elem_id="gallery").style(columns=[2], rows=[2], object_fit="contain", height="auto")
+        text_scores = gr.JSON(label="Scores")
+
+    with gr.Tab("Image Search"):
+        image_input = gr.Image(type="pil", label="Upload an image")
+        image_button = gr.Button("Search")
+        image_output = gr.Gallery(label="Results", show_label=False, elem_id="gallery").style(columns=[2], rows=[2], object_fit="contain", height="auto")
+        image_scores = gr.JSON(label="Scores")
+
+    top_k_slider = gr.Slider(minimum=1, maximum=20, value=5, step=1, label="Number of results")
+
+    text_button.click(text_search, inputs=[text_input, top_k_slider], outputs=[text_output, text_scores])
+    image_button.click(image_search, inputs=[image_input, top_k_slider], outputs=[image_output, image_scores])
+
+# Run the setup (this should be done only once when deploying the app)
+setup_qdrant()
+
+# Launch the app
+demo.launch()