add more

audio_to_text.py

@@ -10,6 +10,7 @@ from transformers import pipeline
 
 class AudioPipeline(object):
     def __init__(self, audio_text_path, audio_text_embeddings_path):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.model = laion_clap.CLAP_Module(enable_fusion=False)
         self.model.load_ckpt()  # download the default pretrained checkpoint.
         self.audio_text_path = audio_text_path
@@ -39,7 +40,7 @@ class AudioPipeline(object):
             texts = json.load(f)
 
         tensors = {}
-        with safe_open(self.audio_text_embeddings_path, framework="pt", device=0) as f:
+        with safe_open(self.audio_text_embeddings_path, framework="pt", device=self.device) as f:
             for k in f.keys():
                 tensors[k] = f.get_tensor(k)
         text_embed = tensors["text_embed"]