ruslanmv commited on
Commit
6256537
·
verified ·
1 Parent(s): 5f00eab

Update src/app.py

Browse files
Files changed (1) hide show
  1. 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(vision_model_path, torch_dtype="auto", device_map="auto")
 
 
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