ThreadAbort/IndexTTS-Rust

1import requests2import zipfile3import os4import argparse5 6def get_confirm_token(response):7    """8    从响应中检查是否存在下载确认令牌（cookie）9 10    Args:11        response (requests.Response): 响应对象12 13    Returns:14        str: 确认令牌的值（如果存在），否则为None15    """16    for key, value in response.cookies.items():17        if key.startswith('download_warning'): # 确认令牌的cookie通常以这个开头18            return value19    return None20 21def save_response_content(response, destination, chunk_size=32768):22    """23    以流式方式将响应内容写入文件，支持大文件下载。24 25    Args:26        response (requests.Response): 流式响应对象27        destination (str): 本地保存路径28        chunk_size (int, optional): 每次迭代写入的块大小. Defaults to 32768.29    """30    with open(destination, "wb") as f:31        for chunk in response.iter_content(chunk_size):32            if chunk: # 过滤掉保持连接的空白块33                f.write(chunk)34 35def download_model_from_modelscope(destination,hf_cache_dir):36    print(f"[ModelScope] Downloading models to {destination},model cache dir={hf_cache_dir}")37    from modelscope import snapshot_download38    snapshot_download("IndexTeam/IndexTTS-2", local_dir=destination)39    snapshot_download("amphion/MaskGCT", local_dir=os.path.join(hf_cache_dir,"models--amphion--MaskGCT"))40    snapshot_download("facebook/w2v-bert-2.0",local_dir=os.path.join(hf_cache_dir,"models--facebook--w2v-bert-2.0"))41    snapshot_download("nv-community/bigvgan_v2_22khz_80band_256x",local_dir=os.path.join(hf_cache_dir,"models--nvidia--bigvgan_v2_22khz_80band_256x"))42    snapshot_download("iic/speech_campplus_sv_zh-cn_16k-common",local_dir=os.path.join(hf_cache_dir,"models--funasr--campplus"))43 44def download_model_from_huggingface(destination,hf_cache_dir):45    print(f"[HuggingFace] Downloading models to {destination},model cache dir={hf_cache_dir}")46    from huggingface_hub import snapshot_download47    snapshot_download("IndexTeam/IndexTTS-2", local_dir=destination)48    # snapshot_download("amphion/MaskGCT", local_dir=os.path.join(hf_cache_dir,"models--amphion--MaskGCT"))49    # snapshot_download("facebook/w2v-bert-2.0",local_dir=os.path.join(hf_cache_dir,"models--facebook--w2v-bert-2.0"))50    # snapshot_download("nvidia/bigvgan_v2_22khz_80band_256x",local_dir=os.path.join(hf_cache_dir, "models--nvidia--bigvgan_v2_22khz_80band_256x"))51    # snapshot_download("funasr/campplus",local_dir=os.path.join(hf_cache_dir,"models--funasr--campplus"))52 53# 使用示例54if __name__ == "__main__":55    parser = argparse.ArgumentParser(description="Download models and example files")56    parser.add_argument('-s','--model_source', choices=['modelscope', 'huggingface'], default=None, help='Model source')57    args = parser.parse_args()58 59    if args.model_source:60        if args.model_source == 'modelscope':61            download_model_from_modelscope("checkpoints",os.path.join("checkpoints","hf_cache"))62        elif args.model_source == 'huggingface':63            download_model_from_huggingface("checkpoints",os.path.join("checkpoints","hf_cache"))64