Spaces:
Sleeping
Sleeping
File size: 1,724 Bytes
60efa5a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
from transformers import AutoModelForImageClassification, AutoImageProcessor
from transformers import AutoModelForAudioClassification, AutoFeatureExtractor
import torch
class ModelLoader:
_instance = None
_video_model = None
_video_processor = None
_audio_model = None
_audio_processor = None
def __new__(cls):
if cls._instance is None:
cls._instance = super(ModelLoader, cls).__new__(cls)
return cls._instance
def load_video_model(self):
if self._video_model is None:
self._video_model = AutoModelForImageClassification.from_pretrained("./models/video_model")
self._video_processor = AutoImageProcessor.from_pretrained("./models/video_model")
self._video_model.eval()
if torch.cuda.is_available():
self._video_model = self._video_model.cuda()
print("Video model loaded!")
return self._video_model, self._video_processor
def load_audio_model(self):
if self._audio_model is None:
self._audio_model = AutoModelForAudioClassification.from_pretrained("./models/audio_model")
self._audio_processor = AutoFeatureExtractor.from_pretrained("./models/audio_model")
self._audio_model.eval()
if torch.cuda.is_available():
self._audio_model = self._audio_model.cuda()
print("Audio model loaded!")
return self._audio_model, self._audio_processor
def get_device(self):
return "cuda" if torch.cuda.is_available() else "cpu"
model_loader = ModelLoader() |