draft

app.py

@@ -1,122 +1,159 @@
+import base64
+import io
+import os
+import random
+
 import gradio as gr
+import numpy as np
 import torch
+from colpali_engine.models import ColPali, ColPaliProcessor
 from datasets import load_dataset
+from dotenv import load_dotenv
+from PIL import Image, ImageDraw
 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
-
-# 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"
-
-# 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"):
+from tqdm import tqdm
+from gradio.themes.base import Base
+from gradio.themes.utils import colors, fonts, sizes
+from typing import Iterable
+
+# Load environment variables
+load_dotenv()
+
+# Set up device
+if torch.cuda.is_available():
+    device = "cuda:0"
+elif torch.backends.mps.is_available():
+    device = "mps"
+else:
+    device = "cpu"
+print(f"Using device: {device}")
+
+# Set up Qdrant client
+QDRANT_URL = os.getenv("QDRANT_URL")
+QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
+qdrant_client = QdrantClient(
+    url=QDRANT_URL,
+    api_key=QDRANT_API_KEY,
+    prefer_grpc=True,
+)
+
+
+# Load dataset and set up model
+dataset = load_dataset("davanstrien/ufo-ColPali", split="train")
+collection_name = "ufo"
+
+model_name = "davanstrien/finetune_colpali_v1_2-ufo-4bit"
+colpali_model = ColPali.from_pretrained(
+    model_name,
+    torch_dtype=torch.bfloat16,
+    device_map=device,
+)
+colpali_processor = ColPaliProcessor.from_pretrained(
+    "vidore/colpaligemma-3b-pt-448-base"
+)
+
+
+def search_images_by_text(query_text, top_k=5):
     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)
+        batch_query = colpali_processor.process_queries([query_text]).to(
+            colpali_model.device
+        )
         query_embedding = colpali_model(**batch_query)
 
     multivector_query = query_embedding[0].cpu().float().numpy().tolist()
-
-    search_result = qdrant_client.search(
+    results = qdrant_client.query_points(
         collection_name=collection_name,
-        query_vector=multivector_query,
-        limit=top_k
+        query=multivector_query,
+        limit=top_k,
+        timeout=60,
     )
+    print(results)
+    return results
+
+
+def search_by_text_and_return_images(query_text, top_k=5):
+    results = search_images_by_text(query_text, top_k)
+    print(results)
+    row_ids = [r.id for r in results.points]
+    subset = dataset.select(row_ids)
+    return list(subset["image"])
+
+
+class Geocities90s(Base):
+    def __init__(
+        self,
+        *,
+        primary_hue: colors.Color | str = colors.yellow,
+        secondary_hue: colors.Color | str = colors.purple,
+        neutral_hue: colors.Color | str = colors.gray,
+        font: fonts.Font | str = fonts.GoogleFont("Comic Neue"),
+        font_mono: fonts.Font | str = fonts.GoogleFont("VT323"),
+    ):
+        super().__init__(
+            primary_hue=primary_hue,
+            secondary_hue=secondary_hue,
+            neutral_hue=neutral_hue,
+            font=(font, "Comic Sans MS", "ui-sans-serif", "sans-serif"),
+            font_mono=(font_mono, "Courier New", "monospace"),
+        )
+        self.set(
+            body_background_fill="url('https://web.archive.org/web/20091020152706/http://hk.geocities.com/neonlightfantasy/image/stars.gif')",
+            button_primary_background_fill="linear-gradient(90deg, *primary_500, *secondary_500)",
+            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_500, *primary_500)",
+            button_primary_text_color="*neutral_50",
+        )
 
-    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()
+
+geocities90s = Geocities90s()
+
+css = """
+body {
+    margin: 0;
+    padding: 0;
+    color: #00ff00;
+    font-family: 'Comic Sans MS', cursive;
+}
+.gradio-container {
+    background-image: url('https://i.ytimg.com/vi/5WapcCXEcXA/maxresdefault.jpg');
+    background-repeat: repeat;
+    background-size: 300px 300px;
+}
+h1 {
+    text-align: center;
+    color: #ff00ff;
+    text-shadow: 2px 2px #000000;
+    font-size: 36px;
+}
+.yellow-text {
+    color: #ffff00;
+    text-shadow: 2px 2px #000000;
+    font-weight: bold;
+}
+"""
+
+demo = gr.Interface(
+    fn=search_by_text_and_return_images,
+    inputs=[
+        gr.Textbox(
+            label="Enter your cosmic query",
+            placeholder="e.g., alien abduction, crop circles",
+        ),
+        gr.Slider(
+            minimum=1,
+            maximum=10,
+            step=1,
+            label="Number of classified documents",
+            value=5,
+        ),
+    ],
+    outputs=gr.Gallery(label="Declassified UFO Sightings", elem_id="gallery"),
+    title="🛸 Top Secret UFO Document Search 🛸",
+    description="<marquee direction='left' scrollamount='5' class='yellow-text'>Uncover the truth that's out there! The government doesn't want you to know!</marquee>",
+    css=css,
+    allow_flagging="never",
+    theme=geocities90s,
+)
+
+if __name__ == "__main__":
+    demo.launch()