Spaces:
Sleeping
Sleeping
Update app.py
Browse files
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 |
-
|
|
|
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 |
-
|
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) >=
|
128 |
-
raise gr.Error("The number of images in the dataset should be less than
|
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 |
|