Update neural_acoustic_distance.py

neural_acoustic_distance.py

@@ -54,6 +54,13 @@ def load_wav2vec2_featurizer(model_id: str, layer: Optional[int] = None):
             model.cuda()
     # st.success("Done!")
 
+    return model
+
+
+#@st.cache(persist=True, show_spinner=False, max_entries=3)
+def run(model_id, layer, filename_x, filename_y):
+    model = load_wav2vec2_featurizer(model_id, layer)
+    
     @torch.no_grad()
     def _featurize(path):
         input_values, rate = sf.read(path, dtype=np.float32)
@@ -83,16 +90,9 @@ def load_wav2vec2_featurizer(model_id: str, layer: Optional[int] = None):
 
         return hidden_state
 
-    return _featurize
-
-
-#@st.cache(persist=True, show_spinner=False, max_entries=3)
-def run(model_id, layer, filename_x, filename_y):
-    featurizer = load_wav2vec2_featurizer(model_id, layer)
-
     with st.spinner("Measuring distance..."):
-        feats_x = featurizer(filename_x)
-        feats_y = featurizer(filename_y)
+        feats_x = _featurize(filename_x)
+        feats_y = _featurize(filename_y)
         print('3. Features computed', datetime.now().strftime('%d-%m-%Y %H:%M:%S')) # test
         gcm = aligner(feats_x, feats_y)
         print('4. Alignments computed', datetime.now().strftime('%d-%m-%Y %H:%M:%S')) # test
@@ -102,6 +102,8 @@ def run(model_id, layer, filename_x, filename_y):
 
         c, n = compute_costs(gcm)
         print('5. Costs computed', datetime.now().strftime('%d-%m-%Y %H:%M:%S')) # test
+    
+    del model
     return d, c, n