Adding Zero decorators

app.py

@@ -1,13 +1,21 @@
+import subprocess
 import uuid
 
 import gradio as gr
+import spaces
 import torch
 from qwen_vl_utils import process_vision_info
 from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
 from voyager_index import Voyager
 
+subprocess.run(
+    "pip install flash-attn --no-build-isolation",
+    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    shell=True,
+)
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
-device = "cpu"
+
 
 # Initialize the model and processor
 model = (
@@ -32,6 +40,7 @@ def add_to_index(files, index):
     return f"Added {len(files)} files to the index."
 
 
+@spaces.GPU
 def query_index(query, index):
     res = index(query, k=1)
     retrieved_image = res["documents"][0][0]["image"]

model.py

@@ -1,13 +1,13 @@
+import spaces
 import torch
 from PIL import Image
 from qwen_vl_utils import process_vision_info
 from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-# device = "cpu"
 
 min_pixels = 1 * 28 * 28
-max_pixels = 256 * 28 * 28  # 2560 * 28 * 28
+max_pixels = 2560 * 28 * 28
 
 
 processor = AutoProcessor.from_pretrained(
@@ -35,6 +35,7 @@ def get_embedding(last_hidden_state: torch.Tensor, dimension: int):
     return reps.to(torch.float32).cpu().numpy()
 
 
+@spaces.GPU
 def encode_queries(queries: list):
     if isinstance(queries, str):
         queries = [queries]
@@ -77,6 +78,7 @@ def encode_queries(queries: list):
     return query_embeddings
 
 
+@spaces.GPU
 def encode_images(images: list):
     if isinstance(images, Image.Image):
         images = [images]