debug

app.py

@@ -79,7 +79,6 @@ 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:
@@ -90,18 +89,20 @@ class PDFVisualRetrieval:
             self.images[knowledge_base_name] = {}
         dpi = 100
         doc = fitz.open("pdf", pdf_file_binary)
-        for page in progress.tqdm(doc):
-            with self.lock: # because we hope one 16G gpu only process one image at the same time
-                pix = page.get_pixmap(dpi=dpi)
-                image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-                image_md5 = get_image_md5(image)
-                with torch.no_grad():
-                    reps = self.model(text=[''], image=[image], tokenizer=self.tokenizer).reps
-                self.reps[knowledge_base_name][image_md5] = reps.squeeze(0)
-                self.images[knowledge_base_name][image_md5] = image
+        
+        with spaces.GPU():
+            for page in progress.tqdm(doc):
+                with self.lock: # because we hope one 16G gpu only process one image at the same time
+                    pix = page.get_pixmap(dpi=dpi)
+                    image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+                    image_md5 = get_image_md5(image)
+                    with torch.no_grad():
+                        reps = self.model(text=[''], image=[image], tokenizer=self.tokenizer).reps
+                    self.reps[knowledge_base_name][image_md5] = reps.squeeze(0)
+                    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