manu commited on
Commit
f1d7f41
·
verified ·
1 Parent(s): 32776d6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -6
app.py CHANGED
@@ -14,18 +14,18 @@ from tqdm import tqdm
14
  from colpali_engine.models import ColQwen2, ColQwen2Processor
15
 
16
 
17
- @spaces.GPU
18
  def install_fa2():
19
  print("Install FA2")
20
  os.system("pip install flash-attn --no-build-isolation")
21
- # install_fa2()
 
22
 
23
 
24
  model = ColQwen2.from_pretrained(
25
  "vidore/colqwen2-v1.0",
26
  torch_dtype=torch.bfloat16,
27
  device_map="cuda:0", # or "mps" if on Apple Silicon
28
- # attn_implementation="flash_attention_2", # should work on A100
29
  ).eval()
30
  processor = ColQwen2Processor.from_pretrained("vidore/colqwen2-v1.0")
31
 
@@ -84,7 +84,6 @@ def query_gpt4o_mini(query, images, api_key):
84
  return "Enter your OpenAI API key to get a custom response"
85
 
86
 
87
- @spaces.GPU
88
  def search(query: str, ds, images, k, api_key):
89
  k = min(k, len(ds))
90
  device = "cuda:0" if torch.cuda.is_available() else "cpu"
@@ -124,8 +123,8 @@ def convert_files(files):
124
  for f in files:
125
  images.extend(convert_from_path(f, thread_count=4))
126
 
127
- if len(images) >= 150:
128
- raise gr.Error("The number of images in the dataset should be less than 150.")
129
  return images
130
 
131
 
 
14
  from colpali_engine.models import ColQwen2, ColQwen2Processor
15
 
16
 
 
17
  def install_fa2():
18
  print("Install FA2")
19
  os.system("pip install flash-attn --no-build-isolation")
20
+
21
+ install_fa2()
22
 
23
 
24
  model = ColQwen2.from_pretrained(
25
  "vidore/colqwen2-v1.0",
26
  torch_dtype=torch.bfloat16,
27
  device_map="cuda:0", # or "mps" if on Apple Silicon
28
+ attn_implementation="flash_attention_2", # should work on A100
29
  ).eval()
30
  processor = ColQwen2Processor.from_pretrained("vidore/colqwen2-v1.0")
31
 
 
84
  return "Enter your OpenAI API key to get a custom response"
85
 
86
 
 
87
  def search(query: str, ds, images, k, api_key):
88
  k = min(k, len(ds))
89
  device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
123
  for f in files:
124
  images.extend(convert_from_path(f, thread_count=4))
125
 
126
+ if len(images) >= 500:
127
+ raise gr.Error("The number of images in the dataset should be less than 500.")
128
  return images
129
 
130