Update my_model/captioner/image_captioning.py
Browse files
my_model/captioner/image_captioning.py
CHANGED
@@ -6,7 +6,8 @@ from transformers import InstructBlipProcessor, InstructBlipForConditionalGenera
|
|
6 |
import bitsandbytes
|
7 |
import accelerate
|
8 |
from my_model.captioner import captioning_config as config
|
9 |
-
|
|
|
10 |
|
11 |
class ImageCaptioningModel:
|
12 |
def __init__(self):
|
@@ -76,6 +77,8 @@ def get_caption(img):
|
|
76 |
captioner.load_model()
|
77 |
caption = captioner.generate_caption(img)
|
78 |
|
|
|
|
|
79 |
return caption
|
80 |
|
81 |
|
|
|
6 |
import bitsandbytes
|
7 |
import accelerate
|
8 |
from my_model.captioner import captioning_config as config
|
9 |
+
from my_model.utilities import free_gpu_resources
|
10 |
+
|
11 |
|
12 |
class ImageCaptioningModel:
|
13 |
def __init__(self):
|
|
|
77 |
captioner.load_model()
|
78 |
caption = captioner.generate_caption(img)
|
79 |
|
80 |
+
free_gpu_resources()
|
81 |
+
|
82 |
return caption
|
83 |
|
84 |
|