DESUCLUB commited on
Commit
ffa13c1
1 Parent(s): e16d4f8

Upload BLIPIntepret.py

Browse files
Files changed (1) hide show
  1. BLIPIntepret.py +6 -3
BLIPIntepret.py CHANGED
@@ -7,10 +7,13 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
7
  print(device)
8
 
9
  def init_BLIP(device):
10
-
 
 
 
11
  processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
12
  model = Blip2ForConditionalGeneration.from_pretrained(
13
- "Salesforce/blip2-opt-2.7b", load_in_8bit=True,torch_dtype=torch.float16, device_map = 'auto'
14
  )
15
  model.eval()
16
  if torch.__version__ >= "2":
@@ -42,4 +45,4 @@ Testing
42
 
43
  model,processor = init_BLIP(device)
44
  image = Image.open('/home/spooky/Downloads/IMG20221214012021.jpg')
45
- infer_BLIP2(model,processor,image,device)'''
 
7
  print(device)
8
 
9
  def init_BLIP(device):
10
+ if device == 'cuda':
11
+ bit_load = True
12
+ else:
13
+ bit_load = False
14
  processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
15
  model = Blip2ForConditionalGeneration.from_pretrained(
16
+ "Salesforce/blip2-opt-2.7b", load_in_8bit= bit_load,torch_dtype=torch.float16, device_map = 'auto'
17
  )
18
  model.eval()
19
  if torch.__version__ >= "2":
 
45
 
46
  model,processor = init_BLIP(device)
47
  image = Image.open('/home/spooky/Downloads/IMG20221214012021.jpg')
48
+ infer_BLIP2(model,processor,image,device)'''