ok

app.py

@@ -90,16 +90,15 @@ class PDFVisualRetrieval:
         dpi = 100
         doc = fitz.open("pdf", pdf_file_binary)
         
-        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
+        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
 
@@ -137,6 +136,14 @@ if __name__ == "__main__":
     
     retriever = PDFVisualRetrieval(model=model, tokenizer=tokenizer)
     
+    @spaces.GPU
+    def add_pdf_gradio(pdf_file_binary):
+        return retriever.add_pdf_gradio(pdf_file_binary)
+    
+    @spaces.GPU
+    def retrieve_gradio(knowledge_base, query, topk):
+        return retriever.retrieve_gradio(knowledge_base, query, topk)
+
     # topk_doc_ids_np, topk_values_np, images_topk = retriever.retrieve(knowledge_base='test', query='what is the number of VQ of this kind of codec method?', topk=1)
     # # 2
     # topk_doc_ids_np, topk_values_np, images_topk = retriever.retrieve(knowledge_base='test', query='the training loss curve of this paper?', topk=1)
@@ -152,7 +159,7 @@ if __name__ == "__main__":
             file_result = gr.Text(label="Knowledge Base ID (remember this!)")
             process_button = gr.Button("Process PDF")
         
-        process_button.click(retriever.add_pdf_gradio, inputs=[file_input], outputs=file_result)
+        process_button.click(add_pdf_gradio, inputs=[file_input], outputs=file_result)
 
         with gr.Row():
             kb_id_input = gr.Text(label="Your Knowledge Base ID")
@@ -163,7 +170,7 @@ if __name__ == "__main__":
         with gr.Row():
             images_output = gr.Gallery(label="Retrieved Pages")
         
-        retrieve_button.click(retriever.retrieve_gradio, inputs=[kb_id_input, query_input, topk_input], outputs=images_output)
+        retrieve_button.click(retrieve_gradio, inputs=[kb_id_input, query_input, topk_input], outputs=images_output)
 
     app.launch()