yashsharmaa commited on
Commit
f631cb4
·
verified ·
1 Parent(s): 956b1e3

Update models/caption.py

Browse files
Files changed (1) hide show
  1. models/caption.py +2 -0
models/caption.py CHANGED
@@ -1,8 +1,10 @@
1
  from transformers import BlipProcessor, BlipForConditionalGeneration
2
  import torch
3
 
 
4
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
5
 
 
6
  processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
7
  model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
8
 
 
1
  from transformers import BlipProcessor, BlipForConditionalGeneration
2
  import torch
3
 
4
+ # Use GPU if available, otherwise fallback to CPU
5
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
6
 
7
+ # Load processor and model
8
  processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
9
  model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
10