fix bug

app.py

@@ -15,6 +15,7 @@ import torch
 import fitz
 import threading
 import gradio as gr
+import spaces
 
 
 def get_image_md5(img: Image.Image):
@@ -78,6 +79,7 @@ class PDFVisualRetrieval:
             self.images[knowledge_base_name][image_md5] = image
         return
 
+    @spaces.GPU
     def add_pdf_gradio(self, pdf_file_binary, progress=gr.Progress()):
         knowledge_base_name = calculate_md5_from_binary(pdf_file_binary)
         if knowledge_base_name not in self.reps:
@@ -99,6 +101,7 @@ class PDFVisualRetrieval:
                 self.images[knowledge_base_name][image_md5] = image
         return knowledge_base_name
 
+    @spaces.GPU
     def retrieve_gradio(self, knowledge_base: str, query: str, topk: int):
         doc_reps = list(self.reps[knowledge_base].values())
         query_with_instruction = "Represent this query for retrieving relavant document: " + query

requirements.txt

@@ -6,4 +6,5 @@ torch==2.2.0
 torchvision
 sentencepiece==0.1.99
 numpy==1.26.0
-transformers==4.40.2
+transformers==4.40.2
+timm