TaahKay commited on
Commit
ce35779
·
verified ·
1 Parent(s): 90fe9f1

Update models/model_loader.py

Browse files
Files changed (1) hide show
  1. models/model_loader.py +7 -0
models/model_loader.py CHANGED
@@ -3,6 +3,8 @@ import streamlit as st
3
 
4
  CACHE_DIR = "/tmp/pretrained_models"
5
 
 
 
6
 
7
  from speechbrain.pretrained.interfaces import foreign_class
8
  from faster_whisper import WhisperModel
@@ -34,3 +36,8 @@ def load_accent_model():
34
  @st.cache_resource(show_spinner="Loading Whisper...")
35
  def load_whisper():
36
  return WhisperModel("tiny", device="cpu", compute_type="int8_float32")
 
 
 
 
 
 
3
 
4
  CACHE_DIR = "/tmp/pretrained_models"
5
 
6
+ import torch
7
+ import gc
8
 
9
  from speechbrain.pretrained.interfaces import foreign_class
10
  from faster_whisper import WhisperModel
 
36
  @st.cache_resource(show_spinner="Loading Whisper...")
37
  def load_whisper():
38
  return WhisperModel("tiny", device="cpu", compute_type="int8_float32")
39
+
40
+ def unload_model(model):
41
+ del model
42
+ torch.cuda.empty_cache()
43
+ gc.collect()