Spaces:
Running
on
Zero
Running
on
Zero
Update src/app.py
Browse files- src/app.py +4 -1
src/app.py
CHANGED
@@ -1,3 +1,4 @@
|
|
|
|
1 |
from collections.abc import Iterator
|
2 |
from datetime import datetime
|
3 |
from pathlib import Path
|
@@ -57,7 +58,9 @@ text_tokenizer.use_default_system_prompt = False
|
|
57 |
# Vision model loading
|
58 |
vision_model_path = "ibm-granite/granite-vision-3.1-2b-preview"
|
59 |
vision_processor = LlavaNextProcessor.from_pretrained(vision_model_path, use_fast=True)
|
60 |
-
vision_model = LlavaNextForConditionalGeneration.from_pretrained(
|
|
|
|
|
61 |
|
62 |
|
63 |
@spaces.GPU
|
|
|
1 |
+
# Code developed by: Ruslan Magana Vsevolodovna
|
2 |
from collections.abc import Iterator
|
3 |
from datetime import datetime
|
4 |
from pathlib import Path
|
|
|
58 |
# Vision model loading
|
59 |
vision_model_path = "ibm-granite/granite-vision-3.1-2b-preview"
|
60 |
vision_processor = LlavaNextProcessor.from_pretrained(vision_model_path, use_fast=True)
|
61 |
+
vision_model = LlavaNextForConditionalGeneration.from_pretrained(
|
62 |
+
vision_model_path, torch_dtype=torch.float16, device_map="auto", low_cpu_mem_usage=True
|
63 |
+
) # Added low_cpu_mem_usage and explicit torch_dtype
|
64 |
|
65 |
|
66 |
@spaces.GPU
|