Advait3009 commited on
Commit
d06533b
·
verified ·
1 Parent(s): d931baa

Create model_loader.py

Browse files
Files changed (1) hide show
  1. utils/model_loader.py +18 -0
utils/model_loader.py ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline
2
+ import torch
3
+
4
+ def load_llava_model():
5
+ return pipeline(
6
+ "image-to-text",
7
+ model="llava-hf/llava-1.5-7b-hf",
8
+ torch_dtype=torch.float16,
9
+ device_map="auto",
10
+ max_new_tokens=200
11
+ )
12
+
13
+ def load_clip_model():
14
+ return pipeline(
15
+ "feature-extraction",
16
+ model="openai/clip-vit-base-patch32",
17
+ device_map="auto"
18
+ )