Update vision_model.py
Browse files- vision_model.py +4 -0
vision_model.py
CHANGED
|
@@ -2,6 +2,10 @@ from transformers import AutoProcessor, AutoModelForVision2Seq
|
|
| 2 |
from PIL import Image
|
| 3 |
import torch
|
| 4 |
|
|
|
|
|
|
|
|
|
|
|
|
|
| 5 |
model_id = "llava-hf/llava-1.5-7b-hf"
|
| 6 |
processor = AutoProcessor.from_pretrained(model_id)
|
| 7 |
model = AutoModelForVision2Seq.from_pretrained(model_id, torch_dtype=torch.float16).cuda()
|
|
|
|
| 2 |
from PIL import Image
|
| 3 |
import torch
|
| 4 |
|
| 5 |
+
import os
|
| 6 |
+
os.environ["HF_HOME"] = "/app/.cache"
|
| 7 |
+
|
| 8 |
+
|
| 9 |
model_id = "llava-hf/llava-1.5-7b-hf"
|
| 10 |
processor = AutoProcessor.from_pretrained(model_id)
|
| 11 |
model = AutoModelForVision2Seq.from_pretrained(model_id, torch_dtype=torch.float16).cuda()
|