chats-bug commited on
Commit
307952a
·
1 Parent(s): 1d290e9

Updated Blip device map

Browse files
Files changed (1) hide show
  1. app.py +5 -1
app.py CHANGED
@@ -5,10 +5,14 @@ import open_clip
5
 
6
  from huggingface_hub import hf_hub_download
7
 
 
 
 
 
8
 
9
  # Load the Blip2 model
10
  preprocessor_blip2_8_bit = AutoProcessor.from_pretrained("Salesforce/blip2-opt-6.7b")
11
- model_blip2_8_bit = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-6.7b", device_map="auto", load_in_8bit=True)
12
 
13
  # Load the Blip base model
14
  preprocessor_blip_base = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
 
5
 
6
  from huggingface_hub import hf_hub_download
7
 
8
+ # Use when running on a CPU
9
+ device_map = {
10
+ "embedding": "cpu",
11
+ }
12
 
13
  # Load the Blip2 model
14
  preprocessor_blip2_8_bit = AutoProcessor.from_pretrained("Salesforce/blip2-opt-6.7b")
15
+ model_blip2_8_bit = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-6.7b", device_map=device_map, load_in_8bit_fp32_cpu_offload=True)
16
 
17
  # Load the Blip base model
18
  preprocessor_blip_base = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")