diff --git a/vox_box/downloader/downloaders.py b/vox_box/downloader/downloaders.py index 92cea67..b5781ef 100644 --- a/vox_box/downloader/downloaders.py +++ b/vox_box/downloader/downloaders.py @@ -20,7 +20,7 @@ logger = logging.getLogger(__name__) -def download_model( +def download_file( huggingface_repo_id: Optional[str] = None, huggingface_filename: Optional[str] = None, model_scope_model_id: Optional[str] = None, @@ -28,20 +28,40 @@ def download_model( cache_dir: Optional[str] = None, huggingface_token: Optional[str] = None, ) -> str: + result_path = None + key = None + if huggingface_repo_id is not None: - return HfDownloader.download( + key = ( + f"huggingface:{huggingface_repo_id}" + if huggingface_filename is None + else f"huggingface:{huggingface_repo_id}:{huggingface_filename}" + ) + logger.debug(f"Downloading {key}") + + result_path = HfDownloader.download( repo_id=huggingface_repo_id, filename=huggingface_filename, token=huggingface_token, cache_dir=os.path.join(cache_dir, "huggingface"), ) elif model_scope_model_id is not None: - return ModelScopeDownloader.download( + key = ( + f"modelscope:{model_scope_model_id}" + if model_scope_file_path is None + else f"modelscope:{model_scope_model_id}:{model_scope_file_path}" + ) + logger.debug(f"Downloading {key}") + + result_path = ModelScopeDownloader.download( model_id=model_scope_model_id, file_path=model_scope_file_path, cache_dir=os.path.join(cache_dir, "model_scope"), ) + logger.debug(f"Downloaded {key}") + return result_path + def get_file_size( huggingface_repo_id: Optional[str] = None, @@ -150,8 +170,6 @@ def download_file( if len(matching_files) == 0: raise ValueError(f"No file found in {repo_id} that match {filename}") - logger.info(f"Downloading model {repo_id}/{filename}") - subfolder, first_filename = ( str(Path(matching_files[0]).parent), Path(matching_files[0]).name, @@ -193,7 +211,6 @@ def _inner_hf_hub_download(repo_file: str): else: model_path = os.path.join(local_dir, first_filename) - logger.info(f"Downloaded model {repo_id}/{filename}") return model_path def __call__(self): @@ -242,7 +259,6 @@ def download( name = name.replace(".", "___") lock_filename = os.path.join(cache_dir, group_or_owner, f"{name}.lock") - logger.info("Retriving file lock") with FileLock(lock_filename): if file_path is not None: matching_files = match_model_scope_file_paths(model_id, file_path) diff --git a/vox_box/estimator/bark.py b/vox_box/estimator/bark.py index db227e4..30052c8 100644 --- a/vox_box/estimator/bark.py +++ b/vox_box/estimator/bark.py @@ -3,7 +3,7 @@ import os from typing import Dict from vox_box.config.config import BackendEnum, Config, TaskTypeEnum -from vox_box.downloader.downloaders import download_model +from vox_box.downloader.downloaders import download_file from vox_box.estimator.base import Estimator from vox_box.utils.model import create_model_dict @@ -18,10 +18,8 @@ def __init__( self._cfg = cfg self._required_files = [ "config.json", - "speaker_embeddings_path.json", ] self._config_json = None - self._speaker_json = None def model_info(self) -> Dict: model = ( @@ -61,17 +59,13 @@ def _check_local_model(self, base_dir: str) -> bool: if architectures is not None and "BarkModel" in architectures: supported = True - speaker_path = os.path.join(base_dir, "speaker_embeddings_path.json") - with open(speaker_path, "r", encoding="utf-8") as f: - self._speaker_json = json.load(f) - return supported def _check_remote_model(self) -> bool: downloaded_files = [] for f in self._required_files: try: - downloaded_file_path = download_model( + downloaded_file_path = download_file( huggingface_repo_id=self._cfg.huggingface_repo_id, huggingface_filename=f, model_scope_model_id=self._cfg.model_scope_model_id, @@ -80,7 +74,7 @@ def _check_remote_model(self) -> bool: ) downloaded_files.append(downloaded_file_path) except Exception as e: - logger.error(f"Failed to download {f}, {e}") + logger.debug(f"File {f} does not exist, {e}") continue if len(downloaded_files) != 0: diff --git a/vox_box/estimator/cosyvoice.py b/vox_box/estimator/cosyvoice.py index 1d6c773..9678ee3 100644 --- a/vox_box/estimator/cosyvoice.py +++ b/vox_box/estimator/cosyvoice.py @@ -2,7 +2,7 @@ import os from typing import Dict -from vox_box.downloader.downloaders import download_model +from vox_box.downloader.downloaders import download_file from vox_box.estimator.base import Estimator from vox_box.config.config import BackendEnum, Config, TaskTypeEnum @@ -57,7 +57,7 @@ def _check_remote_model(self) -> bool: downloaded_files = [] for f in self._required_files: try: - download_file_path = download_model( + download_file_path = download_file( huggingface_repo_id=self._cfg.huggingface_repo_id, huggingface_filename=f, model_scope_model_id=self._cfg.model_scope_model_id, @@ -65,7 +65,7 @@ def _check_remote_model(self) -> bool: cache_dir=self._cfg.cache_dir, ) except Exception as e: - logger.error(f"Failed to download {f}, {e}") + logger.debug(f"File {f} does not exist, {e}") continue downloaded_files.append(download_file_path) diff --git a/vox_box/estimator/estimate.py b/vox_box/estimator/estimate.py index f7428a2..e1b7f4a 100644 --- a/vox_box/estimator/estimate.py +++ b/vox_box/estimator/estimate.py @@ -1,3 +1,4 @@ +import logging from typing import Dict, List from vox_box.config.config import Config from vox_box.estimator.bark import Bark @@ -6,14 +7,15 @@ from vox_box.estimator.faster_whisper import FasterWhisper from vox_box.estimator.funasr import FunASR from vox_box.utils.model import create_model_dict -from concurrent.futures import ThreadPoolExecutor, as_completed + +logger = logging.getLogger(__name__) def estimate_model(cfg: Config) -> Dict: estimators: List[Estimator] = [ + CosyVoice(cfg), FasterWhisper(cfg), FunASR(cfg), - CosyVoice(cfg), Bark(cfg), ] @@ -23,17 +25,9 @@ def estimate_model(cfg: Config) -> Dict: supported=False, ) - def get_model_info(estimator: Estimator) -> Dict: - return estimator.model_info() - - with ThreadPoolExecutor() as executor: - futures = {executor.submit(get_model_info, e): e for e in estimators} - for future in as_completed(futures): - result = future.result() - if result["supported"]: - for f in futures: - if not f.done(): - f.cancel() - return result + for estimator in estimators: + model_info = estimator.model_info() + if model_info["supported"]: + return model_info return model_info diff --git a/vox_box/estimator/faster_whisper.py b/vox_box/estimator/faster_whisper.py index 2235431..e1969df 100644 --- a/vox_box/estimator/faster_whisper.py +++ b/vox_box/estimator/faster_whisper.py @@ -3,7 +3,7 @@ import os from typing import Dict, List from vox_box.config.config import BackendEnum, Config, TaskTypeEnum -from vox_box.downloader.downloaders import download_model +from vox_box.downloader.downloaders import download_file from vox_box.downloader.hub import match_files from vox_box.estimator.base import Estimator from vox_box.utils.model import create_model_dict @@ -107,14 +107,14 @@ def _check_remote_model(self) -> bool: # noqa: C901 if "model.bin" not in matching_files: return False except Exception as e: - logger.error(f"Failed to download model file for estimating, {e}") + logger.debug(f"File model.bin does not exist, {e}") return False downloaded_files = [] download_files = ["tokenizer.json", "preprocessor_config.json"] for f in download_files: try: - downloaded_file_path = download_model( + downloaded_file_path = download_file( huggingface_repo_id=self._cfg.huggingface_repo_id, huggingface_filename=f, model_scope_model_id=self._cfg.model_scope_model_id, @@ -122,7 +122,7 @@ def _check_remote_model(self) -> bool: # noqa: C901 cache_dir=self._cfg.cache_dir, ) except Exception as e: - logger.error(f"Failed to download {f} for model estimate, {e}") + logger.debug(f"File {f} does not exist, {e}") continue downloaded_files.append(downloaded_file_path) diff --git a/vox_box/estimator/funasr.py b/vox_box/estimator/funasr.py index 392d393..04dfe87 100644 --- a/vox_box/estimator/funasr.py +++ b/vox_box/estimator/funasr.py @@ -5,7 +5,7 @@ import yaml from vox_box.config.config import BackendEnum, Config, TaskTypeEnum -from vox_box.downloader.downloaders import download_model +from vox_box.downloader.downloaders import download_file from vox_box.estimator.base import Estimator from vox_box.utils.model import create_model_dict @@ -91,7 +91,7 @@ def _check_remote_model(self) -> Tuple[bool, str]: downloaded_files = [] for f in self._optional_files: try: - download_file_path = download_model( + download_file_path = download_file( huggingface_repo_id=self._cfg.huggingface_repo_id, huggingface_filename=f, model_scope_model_id=self._cfg.model_scope_model_id, @@ -99,7 +99,7 @@ def _check_remote_model(self) -> Tuple[bool, str]: cache_dir=self._cfg.cache_dir, ) except Exception as e: - logger.error(f"Failed to download {f} for model estimate, {e}") + logger.debug(f"File {f} does not exist, {e}") continue downloaded_files.append(download_file_path) diff --git a/vox_box/server/model.py b/vox_box/server/model.py index 8a0f048..102d76d 100644 --- a/vox_box/server/model.py +++ b/vox_box/server/model.py @@ -1,3 +1,4 @@ +import logging from typing import Union from vox_box.backends.stt.base import STTBackend from vox_box.backends.stt.faster_whisper import FasterWhisper @@ -11,12 +12,17 @@ _instance = None +logger = logging.getLogger(__name__) + class ModelInstance: def __init__(self, cfg: Config): self._cfg = cfg self._backend_framework = None + + logger.info("Estimating model") self._estimate = estimate_model(cfg) + logger.info("Finished estimating model") if ( self._estimate is None or not self._estimate.get("supported", False) @@ -30,7 +36,8 @@ def __init__(self, cfg: Config): or self._cfg.model_scope_model_id is not None ): try: - mode_path = downloaders.download_model( + logger.info("Downloading model") + mode_path = downloaders.download_file( huggingface_repo_id=self._cfg.huggingface_repo_id, model_scope_model_id=self._cfg.model_scope_model_id, cache_dir=self._cfg.cache_dir, @@ -54,6 +61,7 @@ def run(self): if _instance is None: try: + logger.info("Loading model") _instance = self._backend_framework.load() except Exception as e: raise Exception(f"Faild to load model, {e}")