Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add cache_dir and local_dir parameters #864

Open
wants to merge 2 commits into
base: dev
Choose a base branch
from
Open

add cache_dir and local_dir parameters #864

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
913bdcc
Update core.py
BBC-Esq Dec 31, 2024
3a0d866
chore(format): run black on main
github-actions[bot] Dec 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
121 changes: 88 additions & 33 deletions ChatTTS/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,63 +61,114 @@ def has_loaded(self, use_decoder=False):

return not not_finish

# Modified
def download_models(
self,
source: Literal["huggingface", "local", "custom"] = "local",
force_redownload=False,
custom_path: Optional[torch.serialization.FILE_LIKE] = None,
cache_dir: Optional[str] = None,
local_dir: Optional[str] = None,
) -> Optional[str]:
if source == "local":
download_path = os.getcwd()
if (
not check_all_assets(Path(download_path), self.sha256_map, update=True)
or force_redownload
):
download_path = (
local_dir if local_dir else (cache_dir if cache_dir else os.getcwd())
)
if local_dir:
# Skip hash checking when using local_dir
with tempfile.TemporaryDirectory() as tmp:
download_all_assets(tmpdir=tmp)
if not check_all_assets(
Path(download_path), self.sha256_map, update=False
):
self.logger.error(
"download to local path %s failed.", download_path
else:
# Do hash checking for cache_dir or default path
if (
not check_all_assets(
Path(download_path), self.sha256_map, update=True
)
return None
or force_redownload
):
with tempfile.TemporaryDirectory() as tmp:
download_all_assets(tmpdir=tmp)
if not check_all_assets(
Path(download_path), self.sha256_map, update=False
):
self.logger.error(
"download to local path %s failed.", download_path
)
return None

elif source == "huggingface":
hf_home = os.getenv("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
try:
download_path = get_latest_modified_file(
os.path.join(hf_home, "hub/models--2Noise--ChatTTS/snapshots")
)
except:
download_path = None
if download_path is None or force_redownload:
self.logger.log(
logging.INFO,
f"download from HF: https://huggingface.co/2Noise/ChatTTS",
)
try:
if local_dir:
# Skip hash checking when using local_dir
download_path = snapshot_download(
repo_id="2Noise/ChatTTS",
allow_patterns=["*.yaml", "*.json", "*.safetensors"],
local_dir=local_dir,
force_download=force_redownload,
)
except:
download_path = None
else:
self.logger.log(
logging.INFO, f"load latest snapshot from cache: {download_path}"
)
if download_path is None:
self.logger.error("download from huggingface failed.")
return None
elif cache_dir:
# Download to cache_dir and verify hashes
download_path = snapshot_download(
repo_id="2Noise/ChatTTS",
allow_patterns=["*.yaml", "*.json", "*.safetensors"],
cache_dir=cache_dir,
force_download=force_redownload,
)
if not check_all_assets(
Path(download_path), self.sha256_map, update=False
):
self.logger.error("Model verification failed")
return None
else:
# Original behavior for default HF cache
hf_home = os.getenv(
"HF_HOME", os.path.expanduser("~/.cache/huggingface")
)
try:
download_path = get_latest_modified_file(
os.path.join(
hf_home, "hub/models--2Noise--ChatTTS/snapshots"
)
)
except:
download_path = None
if download_path is None or force_redownload:
self.logger.log(
logging.INFO,
f"download from HF: https://huggingface.co/2Noise/ChatTTS",
)
download_path = snapshot_download(
repo_id="2Noise/ChatTTS",
allow_patterns=["*.yaml", "*.json", "*.safetensors"],
)
if not check_all_assets(
Path(download_path), self.sha256_map, update=False
):
self.logger.error("Model verification failed")
return None
else:
self.logger.log(
logging.INFO,
f"load latest snapshot from cache: {download_path}",
)
except Exception as e:
self.logger.error(f"Failed to download models: {str(e)}")
download_path = None

elif source == "custom":
self.logger.log(logging.INFO, f"try to load from local: {custom_path}")
if not check_all_assets(Path(custom_path), self.sha256_map, update=False):
self.logger.error("check models in custom path %s failed.", custom_path)
return None
download_path = custom_path

if download_path is None:
self.logger.error("Model download failed")
return None

return download_path

# Modified
def load(
self,
source: Literal["huggingface", "local", "custom"] = "local",
Expand All @@ -129,8 +180,12 @@ def load(
use_flash_attn=False,
use_vllm=False,
experimental: bool = False,
cache_dir: Optional[str] = None,
local_dir: Optional[str] = None,
) -> bool:
download_path = self.download_models(source, force_redownload, custom_path)
download_path = self.download_models(
source, force_redownload, custom_path, cache_dir, local_dir
)
if download_path is None:
return False
return self._load(
Expand Down
14 changes: 7 additions & 7 deletions tools/audio/av.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,11 +41,11 @@ def wav2(i: BytesIO, o: BufferedWriter, format: str):


def load_audio(
file: Union[str, BytesIO, Path],
sr: Optional[int] = None,
format: Optional[str] = None,
mono=True,
) -> Union[np.ndarray, Tuple[np.ndarray, int]]:
file: Union[str, BytesIO, Path],
sr: Optional[int] = None,
format: Optional[str] = None,
mono=True,
) -> Union[np.ndarray, Tuple[np.ndarray, int]]:
"""
https://github.com/fumiama/Retrieval-based-Voice-Conversion-WebUI/blob/412a9950a1e371a018c381d1bfb8579c4b0de329/infer/lib/audio.py#L39
"""
Expand Down Expand Up @@ -113,7 +113,7 @@ def frame_iter(container):

np.copyto(decoded_audio[..., offset:end_index], frame_data)
offset += len(frame_data[0])

container.close()

# Truncate the array to the actual size
Expand All @@ -124,4 +124,4 @@ def frame_iter(container):

if sr is not None:
return decoded_audio
return decoded_audio, rate
return decoded_audio, rate