Update hubert/inference.py

hubert/inference.py

@@ -4,6 +4,8 @@ import numpy as np
 import argparse
 import torch
 import librosa
+import requests
+from tqdm import tqdm
 
 from hubert import hubert_model
 
@@ -22,6 +24,37 @@ def load_model(path, device):
     return model
 
 
+def check_and_download_model():
+    temp_dir = "/tmp"
+    model_path = os.path.join(temp_dir, "hubert-soft-0d54a1f4.pt")
+    
+    if os.path.exists(model_path):
+        return f"モデルは既に存在します: {model_path}"
+    
+    url = "https://github.com/bshall/hubert/releases/download/v0.1/hubert-soft-0d54a1f4.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_vec(model, wavPath, vecPath, device):
     audio = load_audio(wavPath)
     audln = audio.shape[0]
@@ -62,6 +95,6 @@ if __name__ == "__main__":
     vecPath = args.vec
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    hubert = load_model(os.path.join(
-        "hubert_pretrain", "hubert-soft-0d54a1f4.pt"), device)
+    _ = check_and_download_model()
+    hubert = load_model("/tmp/hubert-soft-0d54a1f4.pt", device)
     pred_vec(hubert, wavPath, vecPath, device)