Update whisper/inference.py

whisper/inference.py

@@ -3,6 +3,8 @@ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 import numpy as np
 import argparse
 import torch
+import requests
+from tqdm import tqdm
 
 from whisper.model import Whisper, ModelDimensions
 from whisper.audio import load_audio, pad_or_trim, log_mel_spectrogram
@@ -29,6 +31,37 @@ def load_model(path, device) -> Whisper:
     return model
 
 
+def check_and_download_model():
+    temp_dir = "/tmp"
+    model_path = os.path.join(temp_dir, "large-v2.pt")
+    
+    if os.path.exists(model_path):
+        return f"モデルは既に存在します: {model_path}"
+    
+    url = "https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt"
+    
+    try:
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+        total_size = int(response.headers.get('content-length', 0))
+        
+        with open(model_path, 'wb') as f, tqdm(
+            desc=model_path,
+            total=total_size,
+            unit='iB',
+            unit_scale=True,
+            unit_divisor=1024,
+        ) as pbar:
+            for data in response.iter_content(chunk_size=1024):
+                size = f.write(data)
+                pbar.update(size)
+                
+        return f"モデルのダウンロードが完了しました: {model_path}"
+        
+    except Exception as e:
+        return f"エラーが発生しました: {e}"
+
+
 def pred_ppg(whisper: Whisper, wavPath, ppgPath, device):
     audio = load_audio(wavPath)
     audln = audio.shape[0]
@@ -74,5 +107,7 @@ if __name__ == "__main__":
     ppgPath = args.ppg
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    whisper = load_model(os.path.join("/tmp/large-v2.pt"), device)
+
+    _ =check_and_download_model()
+    whisper = load_model("/tmp/large-v2.pt", device)
     pred_ppg(whisper, wavPath, ppgPath, device)