Update initialize.py

initialize.py

@@ -9,13 +9,14 @@ import logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-def create_symlink_if_not_exists(repo_id, filename, link_dir):
+def create_symlink_if_not_exists(repo_id, repo_filepath, local_subpath, link_dir):
     """
-    Hugging Face Hubからファイルをキャッシュにダウンロードし、
-    指定された場所にシンボリックリンクを作成するヘルパー関数。
+    Hugging Face Hubからファイルをダウンロードし、シンボリックリンクを作成する。
+    リポジトリ上のパスとローカルのパスを分離して扱えるように修正。
     """
-    # リンクを配置したいパスを定義 (例: pretrained/G_0.safetensors)
-    link_path = Path(link_dir) / filename
+    # リンクを配置したいローカルのフルパスを定義
+    # 例: link_dir="model_assets", local_subpath="ANNYUI/config.json" -> "model_assets/ANNYUI/config.json"
+    link_path = Path(link_dir) / local_subpath
 
     # リンクが既に存在する場合はスキップ
     if link_path.exists():
@@ -24,13 +25,14 @@ def create_symlink_if_not_exists(repo_id, filename, link_dir):
     # リンク先の親ディレクトリがなければ作成
     link_path.parent.mkdir(parents=True, exist_ok=True)
     
-    logger.info(f"Downloading {repo_id}/{filename}")
+    logger.info(f"Downloading {repo_id}/{repo_filepath}")
     
-    # Hugging Face Hubからファイルをダウンロード（キャッシュ優先）し、実際のパスを取得
+    # Hugging Face Hubからファイルをダウンロード（キャッシュ優先）
     try:
-        actual_file_path = hf_hub_download(repo_id=repo_id, filename=filename)
+        # ダウンロードにはリポジトリ上のフルパス(repo_filepath)を使用
+        actual_file_path = hf_hub_download(repo_id=repo_id, filename=repo_filepath)
     except Exception as e:
-        logger.error(f"Failed to download {filename}: {e}")
+        logger.error(f"Failed to download {repo_filepath}: {e}")
         return
 
     # シンボリックリンクを作成
@@ -47,15 +49,18 @@ def setup_bert_models():
         # BERTモデルは 'bert/モデル名/' というサブディレクトリに配置
         link_dir = Path("bert") / model_name 
         for file in model_info["files"]:
-            create_symlink_if_not_exists(repo_id, file, link_dir)
+            create_symlink_if_not_exists(repo_id, file, file, link_dir)
 
 def setup_slm_model():
     logger.info("Setting up SLM model...")
+    filename = "pytorch_model.bin"
+    link_dir = "slm/wavlm-base-plus/"
     # SLMモデルは 'slm/wavlm-base-plus/' というサブディレクトリに配置
     create_symlink_if_not_exists(
         repo_id="microsoft/wavlm-base-plus",
-        filename="pytorch_model.bin",
-        link_dir="slm/wavlm-base-plus/"
+        repo_filepath=filename,
+        local_subpath=filename,
+        link_dir=link_dir
     )
 
 def setup_pretrained_models():
@@ -63,61 +68,77 @@ def setup_pretrained_models():
     repo_id = "litagin/Style-Bert-VITS2-1.0-base"
     files = ["G_0.safetensors", "D_0.safetensors", "DUR_0.safetensors"]
     for file in files:
-        create_symlink_if_not_exists(repo_id, file, "pretrained")
+        create_symlink_if_not_exists(repo_id, file, file, "pretrained")
 
 def setup_jp_extra_pretrained_models():
     logger.info("Setting up JP-Extra Pretrained models...")
     repo_id = "litagin/Style-Bert-VITS2-2.0-base-JP-Extra"
     files = ["G_0.safetensors", "D_0.safetensors", "WD_0.safetensors"]
     for file in files:
-        create_symlink_if_not_exists(repo_id, file, "pretrained_jp_extra")
+        create_symlink_if_not_exists(repo_id, file, file, "pretrained_jp_extra")
 
 def setup_default_models():
     logger.info("Setting up default speaker models...")
-    # 先ほどの回答で示した `model_assets` 以下のモデル
+    # `models_to_link` 辞書を修正し、キーの重複をなくす
     models_to_link = {
         "litagin/style_bert_vits2_jvnv": [
             "jvnv-F1-jp/config.json", "jvnv-F1-jp/jvnv-F1-jp_e160_s14000.safetensors", "jvnv-F1-jp/style_vectors.npy",
-            # ... 他のjvnvモデルも同様に追加
         ],
-        "teradakokoro/voice_models/CO": ["ANNYUI/config.json", "ANNYUI/style_settings.json", "ANNYUI/style_vectors.npy", "ANNYUI/ANNYUI_e101_s21000.safetensors"],
-        "teradakokoro/voice_models": ["CO/ASTK/config.json", "CO/ASTK/style_settings.json", "CO/ASTK/style_vectors.npy", "CO/ASTK/ASTK_e501_s28000.safetensors"],
-        "teradakokoro/voice_models": ["CO/AZS/config.json", "CO/AZS/style_settings.json", "CO/AZS/style_vectors.npy", "CO/AZS/AZS_e442_s34000.safetensors"],
-        "teradakokoro/voice_models": ["CO/BNKRG/config.json", "CO/BNKRG/style_settings.json", "CO/BNKRG/style_vectors.npy", "CO/BNKRG/BNKRG_e222_s31000.safetensors"],
-        "teradakokoro/voice_models": ["CO/ESK/config.json", "CO/ESK/style_settings.json", "CO/ESK/style_vectors.npy", "CO/ESK/ESK_e42_s11000.safetensors"],
-        "teradakokoro/voice_models": ["CO/HNS/config.json", "CO/HNS/style_settings.json", "CO/HNS/style_vectors.npy", "CO/HNS/HNS_e1000_s11000.safetensors"],
-        "teradakokoro/voice_models": ["CO/HSI/config.json", "CO/HSI/style_settings.json", "CO/HSI/style_vectors.npy", "CO/HSI/HSI_e209_s26000.safetensors"],
-        "teradakokoro/voice_models": ["CO/KNN/config.json", "CO/KNN/style_settings.json", "CO/KNN/style_vectors.npy", "CO/KNN/KNN_e68_s15000.safetensors"],
-        "teradakokoro/voice_models": ["CO/MZ/config.json", "CO/MZ/style_settings.json", "CO/MZ/style_vectors.npy", "CO/MZ/MZ_e137_s14000.safetensors"],
-        "teradakokoro/voice_models": ["CO/NEL/config.json", "CO/NEL/style_settings.json", "CO/NEL/style_vectors.npy", "CO/NEL/NEL_e1000_s16000.safetensors"],
-        "teradakokoro/voice_models": ["CO/PSR/config.json", "CO/PSR/style_settings.json", "CO/PSR/style_vectors.npy", "CO/PSR/PSR_e142_s25000.safetensors"],
-        "teradakokoro/voice_models": ["CO/RI/config.json", "CO/RI/style_settings.json", "CO/RI/style_vectors.npy", "CO/RI/RI_e94_s20000.safetensors"],
-        "teradakokoro/voice_models": ["CO/SKRNBU/config.json", "CO/SKRNBU/style_settings.json", "CO/SKRNBU/style_vectors.npy", "CO/SKRNBU/SKRNBU_e135_s25000.safetensors"],
-        "teradakokoro/voice_models": ["CO/SNNN/config.json", "CO/SNNN/style_settings.json", "CO/SNNN/style_vectors.npy", "CO/SNNN/SNNN_e201_s11000.safetensors"],
-        "teradakokoro/voice_models": ["CO/SRMY/config.json", "CO/SRMY/style_settings.json", "CO/SRMY/style_vectors.npy", "CO/SRMY/SRMY_e1000_s20000.safetensors"],
-        "teradakokoro/voice_models": ["CO/TIS/config.json", "CO/TIS/style_settings.json", "CO/TIS/style_vectors.npy", "CO/TIS/TIS_e596_s28000.safetensors"],
-        "teradakokoro/voice_models": ["CO/UDK/config.json", "CO/UDK/style_settings.json", "CO/UDK/style_vectors.npy", "CO/UDK/UDK_e472_s25000.safetensors"],
-        "teradakokoro/voice_models": ["CO/VVAN/config.json", "CO/VVAN/style_settings.json", "CO/VVAN/style_vectors.npy", "CO/VVAN/VVAN_e728_s32000.safetensors"],
-        "teradakokoro/voice_models": ["CO/YZY/config.json", "CO/YZY/style_settings.json", "CO/YZY/style_vectors.npy", "CO/YZY/YZY_e1000_s15000.safetensors"],
-
-        "teradakokoro/voice_models": ["FN/FN1/config.json", "FN/FN1/style_settings.json", "FN/FN1/style_vectors.npy", "FN/FN1/FN1.safetensors"],
-        "teradakokoro/voice_models": ["FN/FN10/config.json", "FN/FN10/style_settings.json", "FN/FN10/style_vectors.npy", "FN/FN10/FN10.safetensors"],
-        "teradakokoro/voice_models": ["FN/FN2/config.json", "FN/FN2/style_settings.json", "FN/FN2/style_vectors.npy", "FN/FN2/FN2.safetensors"],
-        "teradakokoro/voice_models": ["FN/FN3/config.json", "FN/FN3/style_settings.json", "FN/FN3/style_vectors.npy", "FN/FN3/FN3.safetensors"],
-        "teradakokoro/voice_models": ["FN/FN4/config.json", "FN/FN4/style_settings.json", "FN/FN4/style_vectors.npy", "FN/FN4/FN4.safetensors"],
-        "teradakokoro/voice_models": ["FN/FN5/config.json", "FN/FN5/style_settings.json", "FN/FN5/style_vectors.npy", "FN/FN5/FN5.safetensors"],
-        "teradakokoro/voice_models": ["FN/FN6/config.json", "FN/FN6/style_settings.json", "FN/FN6/style_vectors.npy", "FN/FN6/FN6.safetensors"],
-        "teradakokoro/voice_models": ["FN/FN7/config.json", "FN/FN7/style_settings.json", "FN/FN7/style_vectors.npy", "FN/FN7/FN7.safetensors"],
-        "teradakokoro/voice_models": ["FN/FN8/config.json", "FN/FN8/style_settings.json", "FN/FN8/style_vectors.npy", "FN/FN8/FN8.safetensors"],
-        "teradakokoro/voice_models": ["FN/FN9/config.json", "FN/FN9/style_settings.json", "FN/FN9/style_vectors.npy", "FN/FN9/FN9.safetensors"],
-
-        "teradakokoro/voice_models/other": ["whisper/config.json", "whisper/style_vectors.npy", "whisper/whisper.safetensors"],
+        "teradakokoro/voice_models": [
+            # CO Models
+            "CO/ANNYUI/config.json", "CO/ANNYUI/style_settings.json", "CO/ANNYUI/style_vectors.npy", "CO/ANNYUI/ANNYUI_e101_s21000.safetensors",
+            "CO/ASTK/config.json", "CO/ASTK/style_settings.json", "CO/ASTK/style_vectors.npy", "CO/ASTK/ASTK_e501_s28000.safetensors",
+            "CO/AZS/config.json", "CO/AZS/style_settings.json", "CO/AZS/style_vectors.npy", "CO/AZS/AZS_e442_s34000.safetensors",
+            "CO/BNKRG/config.json", "CO/BNKRG/style_settings.json", "CO/BNKRG/style_vectors.npy", "CO/BNKRG/BNKRG_e222_s31000.safetensors",
+            "CO/ESK/config.json", "CO/ESK/style_settings.json", "CO/ESK/style_vectors.npy", "CO/ESK/ESK_e42_s11000.safetensors",
+            "CO/HNS/config.json", "CO/HNS/style_settings.json", "CO/HNS/style_vectors.npy", "CO/HNS/HNS_e1000_s11000.safetensors",
+            "CO/HSI/config.json", "CO/HSI/style_settings.json", "CO/HSI/style_vectors.npy", "CO/HSI/HSI_e209_s26000.safetensors",
+            "CO/KNN/config.json", "CO/KNN/style_settings.json", "CO/KNN/style_vectors.npy", "CO/KNN/KNN_e68_s15000.safetensors",
+            "CO/MZ/config.json", "CO/MZ/style_settings.json", "CO/MZ/style_vectors.npy", "CO/MZ/MZ_e137_s14000.safetensors",
+            "CO/NEL/config.json", "CO/NEL/style_settings.json", "CO/NEL/style_vectors.npy", "CO/NEL/NEL_e1000_s16000.safetensors",
+            "CO/PSR/config.json", "CO/PSR/style_settings.json", "CO/PSR/style_vectors.npy", "CO/PSR/PSR_e142_s25000.safetensors",
+            "CO/RI/config.json", "CO/RI/style_settings.json", "CO/RI/style_vectors.npy", "CO/RI/RI_e94_s20000.safetensors",
+            "CO/SKRNBU/config.json", "CO/SKRNBU/style_settings.json", "CO/SKRNBU/style_vectors.npy", "CO/SKRNBU/SKRNBU_e135_s25000.safetensors",
+            "CO/SNNN/config.json", "CO/SNNN/style_settings.json", "CO/SNNN/style_vectors.npy", "CO/SNNN/SNNN_e201_s11000.safetensors",
+            "CO/SRMY/config.json", "CO/SRMY/style_settings.json", "CO/SRMY/style_vectors.npy", "CO/SRMY/SRMY_e1000_s20000.safetensors",
+            "CO/TIS/config.json", "CO/TIS/style_settings.json", "CO/TIS/style_vectors.npy", "CO/TIS/TIS_e596_s28000.safetensors",
+            "CO/UDK/config.json", "CO/UDK/style_settings.json", "CO/UDK/style_vectors.npy", "CO/UDK/UDK_e472_s25000.safetensors",
+            "CO/VVAN/config.json", "CO/VVAN/style_settings.json", "CO/VVAN/style_vectors.npy", "CO/VVAN/VVAN_e728_s32000.safetensors",
+            "CO/YZY/config.json", "CO/YZY/style_settings.json", "CO/YZY/style_vectors.npy", "CO/YZY/YZY_e1000_s15000.safetensors",
+            # FN Models
+            "FN/FN1/config.json", "FN/FN1/style_settings.json", "FN/FN1/style_vectors.npy", "FN/FN1/FN1.safetensors",
+            "FN/FN10/config.json", "FN/FN10/style_settings.json", "FN/FN10/style_vectors.npy", "FN/FN10/FN10.safetensors",
+            "FN/FN2/config.json", "FN/FN2/style_settings.json", "FN/FN2/style_vectors.npy", "FN/FN2/FN2.safetensors",
+            "FN/FN3/config.json", "FN/FN3/style_settings.json", "FN/FN3/style_vectors.npy", "FN/FN3/FN3.safetensors",
+            "FN/FN4/config.json", "FN/FN4/style_settings.json", "FN/FN4/style_vectors.npy", "FN/FN4/FN4.safetensors",
+            "FN/FN5/config.json", "FN/FN5/style_settings.json", "FN/FN5/style_vectors.npy", "FN/FN5/FN5.safetensors",
+            "FN/FN6/config.json", "FN/FN6/style_settings.json", "FN/FN6/style_vectors.npy", "FN/FN6/FN6.safetensors",
+            "FN/FN7/config.json", "FN/FN7/style_settings.json", "FN/FN7/style_vectors.npy", "FN/FN7/FN7.safetensors",
+            "FN/FN8/config.json", "FN/FN8/style_settings.json", "FN/FN8/style_vectors.npy", "FN/FN8/FN8.safetensors",
+            "FN/FN9/config.json", "FN/FN9/style_settings.json", "FN/FN9/style_vectors.npy", "FN/FN9/FN9.safetensors",
+            # other Models
+            "other/whisper/config.json", "other/whisper/style_vectors.npy", "other/whisper/whisper.safetensors",
+        ]
     }
 
     for repo_id, files_in_repo in models_to_link.items():
-        for file_path in files_in_repo:
-            # ここではファイル名だけでなくディレクトリ構造も含むパスを渡す
-            create_symlink_if_not_exists(repo_id, file_path, "model_assets")
+        for repo_filepath in files_in_repo:
+            local_subpath = repo_filepath  # デフォルトではリポジトリパスをそのままローカルパスとする
+
+            # teradakokoro/voice_models の場合のみ、パスの先頭部分を加工
+            if repo_id == "teradakokoro/voice_models":
+                path_parts = repo_filepath.split('/')
+                # パスの先頭が 'CO', 'FN', 'other' のいずれかなら、それを取り除く
+                if len(path_parts) > 1 and path_parts[0] in ["CO", "FN", "other"]:
+                    local_subpath = '/'.join(path_parts[1:])
+            
+            # 修正したヘルパー関数を呼び出す
+            create_symlink_if_not_exists(
+                repo_id=repo_id,
+                repo_filepath=repo_filepath, # ダウンロード用には元のフルパス
+                local_subpath=local_subpath,   # リンク作成用には加工したパス
+                link_dir="model_assets"
+            )
 
 def main(skip_default_models=False, only_infer=False):
     """