Update initialize.py

initialize.py

@@ -1,147 +1,114 @@
-import argparse
+import os
 import json
-import shutil
 from pathlib import Path
-
-import yaml
 from huggingface_hub import hf_hub_download
+import logging
+
+# loggerをセットアップ (元のコードに合わせて)
+# from style_bert_vits2.logging import logger でも可
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def create_symlink_if_not_exists(repo_id, filename, link_dir):
+    """
+    Hugging Face Hubからファイルをキャッシュにダウンロードし、
+    指定された場所にシンボリックリンクを作成するヘルパー関数。
+    """
+    # リンクを配置したいパスを定義 (例: pretrained/G_0.safetensors)
+    link_path = Path(link_dir) / filename
+
+    # リンクが既に存在する場合はスキップ
+    if link_path.exists():
+        return
 
-from style_bert_vits2.logging import logger
+    # リンク先の親ディレクトリがなければ作成
+    link_path.parent.mkdir(parents=True, exist_ok=True)
+    
+    logger.info(f"Downloading {repo_id}/{filename}")
+    
+    # Hugging Face Hubからファイルをダウンロード（キャッシュ優先）し、実際のパスを取得
+    try:
+        actual_file_path = hf_hub_download(repo_id=repo_id, filename=filename)
+    except Exception as e:
+        logger.error(f"Failed to download {filename}: {e}")
+        return
 
+    # シンボリックリンクを作成
+    logger.info(f"Creating symlink: {link_path} -> {actual_file_path}")
+    os.symlink(actual_file_path, link_path)
 
-def download_bert_models():
+def setup_bert_models():
+    logger.info("Setting up BERT models...")
     with open("bert/bert_models.json", encoding="utf-8") as fp:
         models = json.load(fp)
-    for k, v in models.items():
-        local_path = Path("bert").joinpath(k)
-        for file in v["files"]:
-            if not Path(local_path).joinpath(file).exists():
-                logger.info(f"Downloading {k} {file}")
-                hf_hub_download(v["repo_id"], file, local_dir=local_path)
-
-
-def download_slm_model():
-    local_path = Path("slm/wavlm-base-plus/")
-    file = "pytorch_model.bin"
-    if not Path(local_path).joinpath(file).exists():
-        logger.info(f"Downloading wavlm-base-plus {file}")
-        hf_hub_download("microsoft/wavlm-base-plus", file, local_dir=local_path)
-
+    
+    for model_name, model_info in models.items():
+        repo_id = model_info["repo_id"]
+        # BERTモデルは 'bert/モデル名/' というサブディレクトリに配置
+        link_dir = Path("bert") / model_name 
+        for file in model_info["files"]:
+            create_symlink_if_not_exists(repo_id, file, link_dir)
+
+def setup_slm_model():
+    logger.info("Setting up SLM model...")
+    # 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/"
+    )
 
-def download_pretrained_models():
+def setup_pretrained_models():
+    logger.info("Setting up Pretrained models...")
+    repo_id = "litagin/Style-Bert-VITS2-1.0-base"
     files = ["G_0.safetensors", "D_0.safetensors", "DUR_0.safetensors"]
-    local_path = Path("pretrained")
     for file in files:
-        if not Path(local_path).joinpath(file).exists():
-            logger.info(f"Downloading pretrained {file}")
-            hf_hub_download(
-                "litagin/Style-Bert-VITS2-1.0-base", file, local_dir=local_path
-            )
-
+        create_symlink_if_not_exists(repo_id, file, "pretrained")
 
-def download_jp_extra_pretrained_models():
+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"]
-    local_path = Path("pretrained_jp_extra")
     for file in files:
-        if not Path(local_path).joinpath(file).exists():
-            logger.info(f"Downloading JP-Extra pretrained {file}")
-            hf_hub_download(
-                "litagin/Style-Bert-VITS2-2.0-base-JP-Extra", file, local_dir=local_path
-            )
-
-
-#def download_default_models():
-#    files = [
-#        "jvnv-F1-jp/config.json",
-#        "jvnv-F1-jp/jvnv-F1-jp_e160_s14000.safetensors",
-#        "jvnv-F1-jp/style_vectors.npy",
-#        "jvnv-F2-jp/config.json",
-#        "jvnv-F2-jp/jvnv-F2_e166_s20000.safetensors",
-#        "jvnv-F2-jp/style_vectors.npy",
-#        "jvnv-M1-jp/config.json",
-#        "jvnv-M1-jp/jvnv-M1-jp_e158_s14000.safetensors",
-#        "jvnv-M1-jp/style_vectors.npy",
-#        "jvnv-M2-jp/config.json",
-#        "jvnv-M2-jp/jvnv-M2-jp_e159_s17000.safetensors",
-#        "jvnv-M2-jp/style_vectors.npy",
-#    ]
-#    for file in files:
-#        if not Path(f"model_assets/{file}").exists():
-#            logger.info(f"Downloading {file}")
-#            hf_hub_download(
-#                "litagin/style_bert_vits2_jvnv",
-#                file,
-#                local_dir="model_assets",
-#            )
-#    additional_files = {
-#        "litagin/sbv2_koharune_ami": [
-#            "koharune-ami/config.json",
-#            "koharune-ami/style_vectors.npy",
-#            "koharune-ami/koharune-ami.safetensors",
-#        ],
-#        "litagin/sbv2_amitaro": [
-#            "amitaro/config.json",
-#            "amitaro/style_vectors.npy",
-#            "amitaro/amitaro.safetensors",
-#        ],
-#    }
-#    for repo_id, files in additional_files.items():
-#        for file in files:
-#            if not Path(f"model_assets/{file}").exists():
-#                logger.info(f"Downloading {file}")
-#                hf_hub_download(
-#                    repo_id,
-#                    file,
-#                    local_dir="model_assets",
-#                )
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--skip_default_models", action="store_true")
-    parser.add_argument("--only_infer", action="store_true")
-    parser.add_argument(
-        "--dataset_root",
-        type=str,
-        help="Dataset root path (default: Data)",
-        default=None,
-    )
-    parser.add_argument(
-        "--assets_root",
-        type=str,
-        help="Assets root path (default: model_assets)",
-        default=None,
-    )
-    args = parser.parse_args()
-
-    download_bert_models()
-
-#    if not args.skip_default_models:
-#        download_default_models()
-    if not args.only_infer:
-        download_slm_model()
-        download_pretrained_models()
-        download_jp_extra_pretrained_models()
-
-    # If configs/paths.yml not exists, create it
-    default_paths_yml = Path("configs/default_paths.yml")
-    paths_yml = Path("configs/paths.yml")
-    if not paths_yml.exists():
-        shutil.copy(default_paths_yml, paths_yml)
-
-    if args.dataset_root is None and args.assets_root is None:
-        return
-
-    # Change default paths if necessary
-    with open(paths_yml, encoding="utf-8") as f:
-        yml_data = yaml.safe_load(f)
-    if args.assets_root is not None:
-        yml_data["assets_root"] = args.assets_root
-    if args.dataset_root is not None:
-        yml_data["dataset_root"] = args.dataset_root
-    with open(paths_yml, "w", encoding="utf-8") as f:
-        yaml.dump(yml_data, f, allow_unicode=True)
-
+        create_symlink_if_not_exists(repo_id, file, "pretrained_jp_extra")
+
+def setup_default_models():
+    logger.info("Setting up default speaker models...")
+    # 先ほどの回答で示した `model_assets` 以下のモデル
+    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モデルも同様に追加
+        ],
+        "litagin/sbv2_koharune_ami": ["koharune-ami/config.json", "koharune-ami/style_vectors.npy", "koharune-ami/koharune-ami.safetensors"],
+        "litagin/sbv2_amitaro": ["amitaro/config.json", "amitaro/style_vectors.npy", "amitaro/amitaro.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")
+
+def main(skip_default_models=False, only_infer=False):
+    """
+    アプリケーションが必要とする全てのモデルのシンボリックリンクをセットアップする。
+    元のスクリプトの引数も模倣。
+    """
+    logger.info("Starting model setup...")
+    
+    # 常に必要なモデル
+    setup_bert_models()
+    if not skip_default_models:
+        setup_default_models()
+
+    # 推論時のみ不要なモデル
+    if not only_infer:
+        setup_slm_model()
+        setup_pretrained_models()
+        setup_jp_extra_pretrained_models()
+
+    logger.info("All models are set up and ready.")
 
 if __name__ == "__main__":
-    main()
+    # Gradioアプリから呼び出す場合は引数なしで実行
+    main()