#!/usr/bin/env python3 from __future__ import annotations import argparse import dataclasses import fractions import io import json import os import re import shutil import subprocess import sys import tarfile import tempfile import types import warnings from pathlib import Path from typing import Any, Iterable from urllib.error import HTTPError from urllib.parse import quote from urllib.request import Request, urlopen # Heavy install/convert deps are imported lazily so ``list --json`` / ``info`` # work without Torch in the environment. torch: Any = None safe_open: Any = None load_file: Any = None save_file: Any = None yaml: Any = None def _ensure_install_deps() -> None: """Import Torch / safetensors / yaml on first install/convert use.""" global torch, safe_open, load_file, save_file, yaml if torch is not None: return import torch as _torch from safetensors import safe_open as _safe_open from safetensors.torch import load_file as _load_file from safetensors.torch import save_file as _save_file import yaml as _yaml torch = _torch safe_open = _safe_open load_file = _load_file save_file = _save_file yaml = _yaml REPO_ROOT = Path(__file__).resolve().parents[1] MODEL_MANAGER_ASSETS = REPO_ROOT / "assets" / "model_manager" DEMUCS_ROOT_URL = "https://dl.fbaipublicfiles.com/demucs/" WHISPER_MEDIUM_PT_URL = ( "https://openaipublic.azureedge.net/main/whisper/models/" "345ae4da62f9b3d59415adc60127b97c714f32e89e936602e85993674d08dcb1/medium.pt" ) DEMUCS_REMOTE_MODELS = { "955717e8": DEMUCS_ROOT_URL + "hybrid_transformer/955717e8-8726e21a.th", "f7e0c4bc": DEMUCS_ROOT_URL + "hybrid_transformer/f7e0c4bc-ba3fe64a.th", "d12395a8": DEMUCS_ROOT_URL + "hybrid_transformer/d12395a8-e57c48e6.th", "92cfc3b6": DEMUCS_ROOT_URL + "hybrid_transformer/92cfc3b6-ef3bcb9c.th", "04573f0d": DEMUCS_ROOT_URL + "hybrid_transformer/04573f0d-f3cf25b2.th", "75fc33f5": DEMUCS_ROOT_URL + "hybrid_transformer/75fc33f5-1941ce65.th", "5c90dfd2": DEMUCS_ROOT_URL + "5c90dfd2-34c22ccb.th", } DEMUCS_BAGS: dict[str, dict[str, Any]] = { "htdemucs": {"models": ["955717e8"]}, "htdemucs_ft": { "models": ["f7e0c4bc", "d12395a8", "92cfc3b6", "04573f0d"], "weights": [ [1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0], ], }, "htdemucs_6s": {"models": ["5c90dfd2"]}, "hdemucs_mmi": {"models": ["75fc33f5"], "segment": 44}, } def huggingface_token() -> str | None: token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") if token: return token.strip() token_path = Path.home() / ".cache" / "huggingface" / "token" if token_path.is_file(): cached = token_path.read_text(encoding="utf-8").strip() if cached: return cached return None def hf_endpoint() -> str: """Base URL for Hugging Face Hub requests. Honors the standard HF_ENDPOINT environment variable (e.g. HF_ENDPOINT=https://hf-mirror.com for the hf-mirror mirror). Falls back to https://huggingface.co. Empty values and trailing slashes are tolerated. """ endpoint = os.environ.get("HF_ENDPOINT", "").strip().rstrip("/") return endpoint or "https://huggingface.co" def http_headers() -> dict[str, str]: headers = {"User-Agent": "audio.cpp model_manager.py"} token = huggingface_token() if token: headers["Authorization"] = f"Bearer {token}" return headers @dataclasses.dataclass(frozen=True) class SnapshotSource: repo_id: str revision: str = "main" include_prefixes: tuple[str, ...] = () include_suffixes: tuple[str, ...] = () exclude_prefixes: tuple[str, ...] = () strip_prefix: str = "" @dataclasses.dataclass(frozen=True) class SnapshotPlacement: source: SnapshotSource target_subdir: str = "" required_files: tuple[str, ...] = () @dataclasses.dataclass(frozen=True) class CompositeSnapshotSource: placements: tuple[SnapshotPlacement, ...] @dataclasses.dataclass(frozen=True) class ConverterSource: kind: str description: str url: str = "" output_weights_file: str = "" config_kind: str = "" @dataclasses.dataclass(frozen=True) class UnsupportedSource: reason: str @dataclasses.dataclass(frozen=True) class ModelPackage: id: str display_name: str target_directory: str required_files: tuple[str, ...] source: SnapshotSource | CompositeSnapshotSource | ConverterSource | UnsupportedSource description: str = "" # Optional identity fields for ``list --json`` consumers. When unset, # family/tasks/standalone/gated are derived from local package metadata only. family: str | None = None standalone: bool | None = None parent_package_id: str | None = None tasks: tuple[str, ...] = () modes: tuple[str, ...] = () gated: bool | None = None UTILITY_CONVERTER_KINDS = {"pytorch_to_safetensors"} POSTPROCESS_SNAPSHOT_PACKAGE_IDS = { "voxcpm2", } # HF repos that require accepting a license / access grant before download. _GATED_REPO_MARKERS = ( "kyutai/pocket-tts", "stabilityai/", ) def package_install_kind(package: ModelPackage) -> str: source = package.source if isinstance(source, SnapshotSource): return "composite" if package.id in POSTPROCESS_SNAPSHOT_PACKAGE_IDS else "snapshot" if isinstance(source, CompositeSnapshotSource): return "composite" if isinstance(source, ConverterSource): return "utility" if source.kind in UTILITY_CONVERTER_KINDS else "composite" return "unsupported" def package_needs_install_deps(package: ModelPackage) -> bool: if package.id in POSTPROCESS_SNAPSHOT_PACKAGE_IDS: return True source = package.source if isinstance(source, ConverterSource): return True return False def package_usage_examples(package: ModelPackage) -> list[str]: if package.id == "voxcpm2_audiovae": return [ "python tools/model_manager.py install voxcpm2_audiovae --source-file models/VoxCPM2/audiovae.pth --models-root models --overwrite", "python tools/model_manager.py install voxcpm2_audiovae --source-file /path/to/audiovae.pth --output-file /path/to/audiovae.safetensors --overwrite", ] return [] CATALOG: tuple[ModelPackage, ...] = ( ModelPackage( id="ace_step", display_name="ACE-Step 1.5", target_directory="Ace-Step1.5", source=CompositeSnapshotSource( placements=( SnapshotPlacement( source=SnapshotSource(repo_id="ACE-Step/Ace-Step1.5"), required_files=( "config.json", "Qwen3-Embedding-0.6B/model.safetensors", "acestep-5Hz-lm-1.7B/model.safetensors", "acestep-v15-turbo/model.safetensors", "acestep-v15-turbo/silence_latent.pt", "vae/diffusion_pytorch_model.safetensors", ), ), SnapshotPlacement( source=SnapshotSource(repo_id="ACE-Step/acestep-v15-base"), target_subdir="acestep-v15-base", required_files=("config.json", "model.safetensors", "silence_latent.pt"), ), ), ), required_files=( "Qwen3-Embedding-0.6B/model.safetensors", "acestep-5Hz-lm-1.7B/model.safetensors", "acestep-v15-base/model.safetensors", "acestep-v15-base/silence_latent.safetensors", "acestep-v15-turbo/model.safetensors", "acestep-v15-turbo/silence_latent.safetensors", "vae/diffusion_pytorch_model.safetensors", ), ), ModelPackage( id="kokoro_82m_bf16", display_name="Kokoro 82M GGML", target_directory="kokoro-82m-v1_0-ggml", source=SnapshotSource(repo_id="mlx-community/kokoro_mlx"), required_files=( "config.json", "kokoro-v1_0.safetensors", "voices.json", "voices/af_heart.f32", ), ), ModelPackage( id="moss_tts_nano_100m", display_name="MOSS-TTS-Nano 100M", target_directory="MOSS-TTS-Nano-100M", source=CompositeSnapshotSource( placements=( SnapshotPlacement( source=SnapshotSource(repo_id="OpenMOSS-Team/MOSS-TTS-Nano-100M"), required_files=("config.json", "pytorch_model.bin", "tokenizer.model", "tokenizer_config.json"), ), SnapshotPlacement( source=SnapshotSource(repo_id="OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano"), target_subdir="audio_tokenizer", required_files=("config.json", "model-00001-of-00001.safetensors", "model.safetensors.index.json"), ), ), ), required_files=( "config.json", "model.safetensors", "tokenizer.model", "tokenizer_config.json", "audio_tokenizer/config.json", "audio_tokenizer/model-00001-of-00001.safetensors", "audio_tokenizer/model.safetensors.index.json", ), description="Recommended MOSS-TTS-Nano package; assembles the TTS model and MOSS-Audio-Tokenizer-Nano dependency.", ), ModelPackage( id="moss_tts_nano_100m_model", display_name="MOSS-TTS-Nano 100M model subcomponent", target_directory="MOSS-TTS-Nano-100M", source=SnapshotSource(repo_id="OpenMOSS-Team/MOSS-TTS-Nano-100M"), required_files=("config.json", "pytorch_model.bin", "tokenizer.model", "tokenizer_config.json"), description="Subcomponent only. Use moss_tts_nano_100m for the full framework-ready MOSS runtime layout.", ), ModelPackage( id="moss_audio_tokenizer_nano", display_name="MOSS Audio Tokenizer Nano subcomponent", target_directory="MOSS-Audio-Tokenizer-Nano", source=SnapshotSource(repo_id="OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano"), required_files=("config.json", "model-00001-of-00001.safetensors", "model.safetensors.index.json"), description="Subcomponent only. Use moss_tts_nano_100m for the full framework-ready MOSS Nano runtime layout.", ), ModelPackage( id="moss_audio_tokenizer_v2", display_name="MOSS Audio Tokenizer v2 subcomponent", target_directory="MOSS-Audio-Tokenizer-v2", source=SnapshotSource(repo_id="OpenMOSS-Team/MOSS-Audio-Tokenizer-v2"), required_files=( "config.json", "model.safetensors.index.json", "model-00001-of-00003.safetensors", "model-00002-of-00003.safetensors", "model-00003-of-00003.safetensors", ), description="Subcomponent only. Use moss_tts_local_v1_5 for the full framework-ready MOSS-TTS-Local runtime layout.", ), ModelPackage( id="moss_tts_local_v1_5", display_name="MOSS-TTS-Local Transformer v1.5", target_directory="MOSS-TTS-Local-Transformer-v1.5", source=CompositeSnapshotSource( placements=( SnapshotPlacement( source=SnapshotSource(repo_id="OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5"), required_files=( "config.json", "model.safetensors", "tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt", "special_tokens_map.json", "added_tokens.json", "chat_template.jinja", ), ), SnapshotPlacement( source=SnapshotSource(repo_id="OpenMOSS-Team/MOSS-Audio-Tokenizer-v2"), target_subdir="audio_tokenizer", required_files=( "config.json", "model.safetensors.index.json", "model-00001-of-00003.safetensors", "model-00002-of-00003.safetensors", "model-00003-of-00003.safetensors", ), ), ), ), required_files=( "config.json", "model.safetensors", "tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt", "special_tokens_map.json", "added_tokens.json", "chat_template.jinja", "audio_tokenizer/config.json", "audio_tokenizer/model.safetensors.index.json", "audio_tokenizer/model-00001-of-00003.safetensors", "audio_tokenizer/model-00002-of-00003.safetensors", "audio_tokenizer/model-00003-of-00003.safetensors", ), description="MOSS-TTS-Local Transformer v1.5 with the MOSS-Audio-Tokenizer-v2 codec dependency.", ), ModelPackage( id="omnivoice", display_name="OmniVoice", target_directory="OmniVoice", source=SnapshotSource(repo_id="k2-fsa/OmniVoice"), required_files=( "config.json", "model.safetensors", "tokenizer.json", "audio_tokenizer/config.json", "audio_tokenizer/model.safetensors", ), ), ModelPackage( id="fun_asr_nano_2512_q8_0", display_name="Fun-ASR-Nano-2512 Q8_0 GGUF", target_directory="Fun-ASR-Nano-2512-GGUF", source=SnapshotSource( repo_id="FunAudioLLM/Fun-ASR-Nano-2512-GGUF", revision="ce72677f84900f0dc57f498ace253bfb3c9155b6", include_prefixes=("fun-asr-nano-2512-q8_0.gguf",), ), required_files=("fun-asr-nano-2512-q8_0.gguf",), family="fun_asr_nano", tasks=("asr",), modes=("offline",), description=( "Standalone Fun-ASR-Nano-2512 Q8_0 GGUF for offline multilingual ASR; " "governed by the FunASR Model Open Source License Agreement v1.1." ), ), ModelPackage( id="fun_asr_nano_2512_f16", display_name="Fun-ASR-Nano-2512 F16 GGUF", target_directory="Fun-ASR-Nano-2512-GGUF", source=SnapshotSource( repo_id="FunAudioLLM/Fun-ASR-Nano-2512-GGUF", revision="ce72677f84900f0dc57f498ace253bfb3c9155b6", include_prefixes=("fun-asr-nano-2512-f16.gguf",), ), required_files=("fun-asr-nano-2512-f16.gguf",), family="fun_asr_nano", tasks=("asr",), modes=("offline",), description=( "Standalone Fun-ASR-Nano-2512 F16 GGUF for offline multilingual ASR; " "governed by the FunASR Model Open Source License Agreement v1.1." ), ), ModelPackage( id="qwen3_asr_0_6b", display_name="Qwen3 ASR 0.6B", target_directory="Qwen3-ASR-0.6B", source=SnapshotSource(repo_id="Qwen/Qwen3-ASR-0.6B"), required_files=( "config.json", "generation_config.json", "model.safetensors", "preprocessor_config.json", "tokenizer_config.json", "vocab.json", "merges.txt", ), ), ModelPackage( id="qwen3_asr_1_7b_hf", display_name="Qwen3 ASR 1.7B HF", target_directory="Qwen3-ASR-1.7B-hf", source=SnapshotSource( repo_id="Qwen/Qwen3-ASR-1.7B-hf", include_prefixes=( "config.json", "generation_config.json", "model.safetensors", "processor_config.json", "tokenizer_config.json", "tokenizer.json", ), ), required_files=( "config.json", "generation_config.json", "model.safetensors", "processor_config.json", "tokenizer_config.json", "tokenizer.json", ), description="Native Hugging Face Transformers checkpoint; no conversion is required.", ), ModelPackage( id="voxtral_realtime", display_name="Voxtral Mini 4B Realtime GGUF", target_directory="Voxtral-Mini-4B-Realtime-2602-GGUF", source=SnapshotSource( repo_id="audio-cpp/audio.cpp-gguf", include_prefixes=("Voxtral-Mini-4B-Realtime-2602-GGUF/voxtral-mini-4b-realtime-2602-q8_0.gguf",), strip_prefix="Voxtral-Mini-4B-Realtime-2602-GGUF/", ), required_files=("voxtral-mini-4b-realtime-2602-q8_0.gguf",), family="voxtral_realtime", tasks=("asr",), modes=("offline", "streaming"), description="Standalone audio.cpp Q8_0 GGUF package for Voxtral realtime ASR.", ), ModelPackage( id="fish_audio_s2_pro", display_name="Fish Audio S2 Pro GGUF", target_directory="Fish-Audio-S2-Pro-GGUF", source=SnapshotSource( repo_id="audio-cpp/audio.cpp-gguf", include_prefixes=("Fish-Audio-S2-Pro-GGUF/fish-audio-s2-pro-q8_0.gguf",), strip_prefix="Fish-Audio-S2-Pro-GGUF/", ), required_files=("fish-audio-s2-pro-q8_0.gguf",), family="fish_audio", tasks=("tts",), description="Standalone audio.cpp Q8_0 GGUF package for Fish Audio S2 Pro.", ), ModelPackage( id="higgs_audio_stt", display_name="Higgs Audio STT", target_directory="higgs-audio-v3-stt", source=CompositeSnapshotSource( placements=( SnapshotPlacement( source=SnapshotSource(repo_id="bosonai/higgs-audio-v3-stt"), required_files=( "config.json", "generation_config.json", "model.safetensors.index.json", "model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors", "tokenizer_config.json", "vocab.json", "merges.txt", ), ), SnapshotPlacement( source=SnapshotSource( repo_id="openai/whisper-large-v3", include_prefixes=("preprocessor_config.json",), ), target_subdir="../whisper-large-v3", required_files=("preprocessor_config.json",), ), ), ), required_files=( "config.json", "generation_config.json", "model.safetensors.index.json", "model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors", "tokenizer_config.json", "vocab.json", "merges.txt", "../whisper-large-v3/preprocessor_config.json", ), description="Installs Higgs Audio STT plus the sibling Whisper Large v3 preprocessor config required by the framework runtime.", ), ModelPackage( id="hviske_asr", display_name="Hviske ASR", target_directory="hviske-v5.3", source=SnapshotSource(repo_id="syvai/hviske-v5.3"), required_files=("config.json", "generation_config.json", "model.safetensors", "tokenizer.model"), ), ModelPackage( id="nemotron_asr", display_name="Nemotron ASR", target_directory="nemotron-3.5-asr-streaming-0.6b", source=SnapshotSource(repo_id="nvidia/nemotron-3.5-asr-streaming-0.6b"), required_files=("config.json", "model.safetensors", "processor_config.json", "tokenizer.json"), ), ModelPackage( id="parakeet_tdt_q8_0", display_name="Parakeet-TDT 0.6B v3 Q8_0 GGUF", target_directory="Parakeet-TDT-0.6B-v3-GGUF", source=SnapshotSource( repo_id="audio-cpp/audio.cpp-gguf", include_prefixes=("Parakeet-TDT-0.6B-v3-GGUF/parakeet-tdt-0.6b-v3-q8_0.gguf",), strip_prefix="Parakeet-TDT-0.6B-v3-GGUF/", ), required_files=("parakeet-tdt-0.6b-v3-q8_0.gguf",), family="parakeet_tdt", tasks=("asr",), modes=("offline", "streaming"), standalone=True, description="Standalone Parakeet-TDT 0.6B v3 Q8_0 GGUF for offline and buffered streaming ASR.", ), ModelPackage( id="parakeet_tdt_f16", display_name="Parakeet-TDT 0.6B v3 F16 GGUF", target_directory="Parakeet-TDT-0.6B-v3-GGUF", source=SnapshotSource( repo_id="audio-cpp/audio.cpp-gguf", include_prefixes=("Parakeet-TDT-0.6B-v3-GGUF/parakeet-tdt-0.6b-v3-f16.gguf",), strip_prefix="Parakeet-TDT-0.6B-v3-GGUF/", ), required_files=("parakeet-tdt-0.6b-v3-f16.gguf",), family="parakeet_tdt", tasks=("asr",), modes=("offline", "streaming"), standalone=True, description="Standalone Parakeet-TDT 0.6B v3 F16 GGUF for offline and buffered streaming ASR.", ), ModelPackage( id="qwen3_forced_aligner_0_6b", display_name="Qwen3 Forced Aligner 0.6B", target_directory="Qwen3-ForcedAligner-0.6B", source=SnapshotSource(repo_id="Qwen/Qwen3-ForcedAligner-0.6B"), required_files=( "config.json", "generation_config.json", "model.safetensors", "preprocessor_config.json", "tokenizer_config.json", "vocab.json", "merges.txt", ), ), ModelPackage( id="qwen3_tts_0_6b_base", display_name="Qwen3 TTS 12Hz 0.6B Base", target_directory="Qwen3-TTS-12Hz-0.6B-Base", source=SnapshotSource(repo_id="Qwen/Qwen3-TTS-12Hz-0.6B-Base"), required_files=( "config.json", "generation_config.json", "model.safetensors", "speech_tokenizer/config.json", "speech_tokenizer/model.safetensors", "tokenizer_config.json", "vocab.json", "merges.txt", ), ), ModelPackage( id="vieneu_v3_turbo_v3_turbo", display_name="VieNeu-TTS v3 Turbo Base", target_directory="VieNeu-TTS-v3-Turbo", source=SnapshotSource(repo_id="phuocnguyen90/VieNeu-TTS-v3-Turbo-GGUF"), required_files=( "config.json", "model.gguf", "speech_tokenizer/config.json", "tokenizer_config.json", "tokenizer.json", "special_tokens_map.json", ), description="Installs VieNeu-TTS v3 Turbo GGUF model and configuration sidecars for C++ inference.", family="vieneu_v3_turbo", tasks=("tts",), ), ModelPackage( id="qwen3_tts_1_7b_base", display_name="Qwen3 TTS 12Hz 1.7B Base", target_directory="Qwen3-TTS-12Hz-1.7B-Base", source=SnapshotSource(repo_id="Qwen/Qwen3-TTS-12Hz-1.7B-Base"), required_files=( "config.json", "generation_config.json", "model.safetensors", "speech_tokenizer/config.json", "speech_tokenizer/model.safetensors", "tokenizer_config.json", "vocab.json", "merges.txt", ), ), ModelPackage( id="qwen3_tts_1_7b_custom_voice", display_name="Qwen3 TTS 12Hz 1.7B Custom Voice", target_directory="Qwen3-TTS-12Hz-1.7B-CustomVoice", source=SnapshotSource(repo_id="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"), required_files=( "config.json", "generation_config.json", "model.safetensors", "speech_tokenizer/config.json", "speech_tokenizer/model.safetensors", "tokenizer_config.json", "vocab.json", "merges.txt", ), ), ModelPackage( id="qwen3_tts_1_7b_voice_design", display_name="Qwen3 TTS 12Hz 1.7B Voice Design", target_directory="Qwen3-TTS-12Hz-1.7B-VoiceDesign", source=SnapshotSource(repo_id="Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"), required_files=( "config.json", "generation_config.json", "model.safetensors", "speech_tokenizer/config.json", "speech_tokenizer/model.safetensors", "tokenizer_config.json", "vocab.json", "merges.txt", ), ), ModelPackage( id="qwen3_tts_tokenizer_12hz", display_name="Qwen3 TTS Tokenizer 12Hz subcomponent", target_directory="Qwen3-TTS-Tokenizer-12Hz", source=SnapshotSource(repo_id="Qwen/Qwen3-TTS-Tokenizer-12Hz"), required_files=("config.json", "model.safetensors"), description="Subcomponent only. Normal Qwen3 TTS packages already include the speech tokenizer needed by the framework.", ), ModelPackage( id="chatterbox", display_name="Chatterbox", target_directory="chatterbox", source=SnapshotSource(repo_id="ResembleAI/chatterbox"), required_files=( "ve.safetensors", "t3_cfg.safetensors", "t3_mtl23ls_v2.safetensors", "t3_mtl23ls_v3.safetensors", "s3gen.safetensors", "tokenizer.json", "grapheme_mtl_merged_expanded_v1.json", "Cangjie5_TC.json", "conds.pt", ), ), ModelPackage( id="sortformer_diar_4spk_v1", display_name="Sortformer diarization 4 speaker v1", target_directory="diar_sortformer_4spk-v1", source=SnapshotSource(repo_id="nvidia/diar_sortformer_4spk-v1"), required_files=("config.json", "model.safetensors", "processor_config.json"), ), ModelPackage( id="pocket_tts", display_name="PocketTTS", target_directory="pocket-tts", source=SnapshotSource( repo_id="kyutai/pocket-tts", include_prefixes=("languages/english/",), ), required_files=( "languages/english/model.safetensors", "languages/english/tokenizer.model", "languages/english/embeddings/alba.safetensors", ), ), ModelPackage( id="miocodec_25hz_44k_v2", display_name="MioCodec 25Hz 44.1kHz v2", target_directory="MioCodec-25Hz-44.1kHz-v2", source=CompositeSnapshotSource( placements=( SnapshotPlacement( source=SnapshotSource( repo_id="Aratako/MioCodec-25Hz-44.1kHz-v2", include_prefixes=("config.yaml", "model.safetensors"), ), required_files=("config.yaml", "model.safetensors"), ), SnapshotPlacement( source=SnapshotSource( repo_id="mlx-community/wavlm-base-plus-mlx", include_prefixes=("config.json", "weights.safetensors"), ), target_subdir="wavlm-base-plus-mlx", required_files=("config.json", "weights.safetensors"), ), ), ), required_files=( "config.yaml", "model.safetensors", "wavlm-base-plus-mlx/config.json", "wavlm-base-plus-mlx/weights.safetensors", ), ), ModelPackage( id="miotts_1_7b", display_name="MioTTS 1.7B", target_directory="MioTTS-1.7B", source=CompositeSnapshotSource( placements=( SnapshotPlacement( source=SnapshotSource( repo_id="Aratako/MioTTS-1.7B", include_prefixes=( "config.json", "generation_config.json", "tokenizer_config.json", "tokenizer.json", "vocab.json", "merges.txt", "model.safetensors", ), ), required_files=( "config.json", "generation_config.json", "tokenizer_config.json", "tokenizer.json", "vocab.json", "merges.txt", "model.safetensors", ), ), SnapshotPlacement( source=SnapshotSource( repo_id="Aratako/MioCodec-25Hz-44.1kHz-v2", include_prefixes=("config.yaml", "model.safetensors"), ), target_subdir="../MioCodec-25Hz-44.1kHz-v2", required_files=("config.yaml", "model.safetensors"), ), SnapshotPlacement( source=SnapshotSource( repo_id="mlx-community/wavlm-base-plus-mlx", include_prefixes=("config.json", "weights.safetensors"), ), target_subdir="../MioCodec-25Hz-44.1kHz-v2/wavlm-base-plus-mlx", required_files=("config.json", "weights.safetensors"), ), ), ), required_files=( "config.json", "generation_config.json", "tokenizer_config.json", "tokenizer.json", "vocab.json", "merges.txt", "model.safetensors", ), description="Installs MioTTS plus the sibling MioCodec dependency required by the framework runtime.", ), ModelPackage( id="vibevoice_asr", display_name="VibeVoice ASR", target_directory="VibeVoice-ASR", source=CompositeSnapshotSource( placements=( SnapshotPlacement( source=SnapshotSource(repo_id="microsoft/VibeVoice-ASR"), required_files=( "config.json", "model.safetensors.index.json", "model-00001-of-00008.safetensors", "model-00002-of-00008.safetensors", "model-00003-of-00008.safetensors", "model-00004-of-00008.safetensors", "model-00005-of-00008.safetensors", "model-00006-of-00008.safetensors", "model-00007-of-00008.safetensors", "model-00008-of-00008.safetensors", ), ), ), ), required_files=( "config.json", "model.safetensors.index.json", "model-00001-of-00008.safetensors", "model-00002-of-00008.safetensors", "model-00003-of-00008.safetensors", "model-00004-of-00008.safetensors", "model-00005-of-00008.safetensors", "model-00006-of-00008.safetensors", "model-00007-of-00008.safetensors", "model-00008-of-00008.safetensors", "tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt", ), description="Installs VibeVoice ASR plus the Qwen2.5 tokenizer files required by the framework runtime.", ), ModelPackage( id="vibevoice_1_5b", display_name="VibeVoice 1.5B", target_directory="VibeVoice-1.5B", source=CompositeSnapshotSource( placements=( SnapshotPlacement( source=SnapshotSource(repo_id="microsoft/VibeVoice-1.5B"), required_files=( "config.json", "model.safetensors.index.json", "model-00001-of-00003.safetensors", "model-00002-of-00003.safetensors", "model-00003-of-00003.safetensors", "preprocessor_config.json", ), ), ), ), required_files=( "config.json", "model.safetensors.index.json", "model-00001-of-00003.safetensors", "model-00002-of-00003.safetensors", "model-00003-of-00003.safetensors", "preprocessor_config.json", "tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt", ), ), ModelPackage( id="vibevoice_7b", display_name="VibeVoice 7B", target_directory="VibeVoice-7B", source=CompositeSnapshotSource( placements=( SnapshotPlacement( source=SnapshotSource(repo_id="vibevoice/VibeVoice-7B"), required_files=( "config.json", "model.safetensors.index.json", "model-00001-of-00010.safetensors", "model-00002-of-00010.safetensors", "model-00003-of-00010.safetensors", "model-00004-of-00010.safetensors", "model-00005-of-00010.safetensors", "model-00006-of-00010.safetensors", "model-00007-of-00010.safetensors", "model-00008-of-00010.safetensors", "model-00009-of-00010.safetensors", "model-00010-of-00010.safetensors", "preprocessor_config.json", ), ), ), ), required_files=( "config.json", "model.safetensors.index.json", "model-00001-of-00010.safetensors", "model-00002-of-00010.safetensors", "model-00003-of-00010.safetensors", "model-00004-of-00010.safetensors", "model-00005-of-00010.safetensors", "model-00006-of-00010.safetensors", "model-00007-of-00010.safetensors", "model-00008-of-00010.safetensors", "model-00009-of-00010.safetensors", "model-00010-of-00010.safetensors", "preprocessor_config.json", "tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt", ), ), ModelPackage( id="higgs_audio_v3_tts_4b", display_name="Higgs Audio v3 TTS 4B GGUF", target_directory="Higgs-Audio-v3-TTS-4B-GGUF", source=SnapshotSource( repo_id="audio-cpp/audio.cpp-gguf", include_prefixes=("Higgs-Audio-v3-TTS-4B-GGUF/higgs-audio-v3-tts-4b-q8_0.gguf",), strip_prefix="Higgs-Audio-v3-TTS-4B-GGUF/", ), required_files=("higgs-audio-v3-tts-4b-q8_0.gguf",), family="higgs_audio_tts", tasks=("tts",), description="Standalone audio.cpp Q8_0 GGUF package for Higgs Audio v3 TTS 4B.", ), ModelPackage( id="heartmula", display_name="HeartMuLa", target_directory="HeartMuLa", source=CompositeSnapshotSource( placements=( SnapshotPlacement( source=SnapshotSource(repo_id="HeartMuLa/HeartMuLaGen"), required_files=("gen_config.json", "tokenizer.json"), ), SnapshotPlacement( source=SnapshotSource(repo_id="HeartMuLa/HeartMuLa-oss-3B"), target_subdir="HeartMuLa-oss-3B", required_files=( "config.json", "model.safetensors.index.json", "model-00001-of-00004.safetensors", "model-00002-of-00004.safetensors", "model-00003-of-00004.safetensors", "model-00004-of-00004.safetensors", ), ), SnapshotPlacement( source=SnapshotSource(repo_id="HeartMuLa/HeartCodec-oss-20260123"), target_subdir="HeartCodec-oss", required_files=( "config.json", "model.safetensors.index.json", "model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors", ), ), ), ), required_files=( "tokenizer.json", "gen_config.json", "HeartMuLa-oss-3B/config.json", "HeartMuLa-oss-3B/model.safetensors.index.json", "HeartMuLa-oss-3B/model-00001-of-00004.safetensors", "HeartMuLa-oss-3B/model-00002-of-00004.safetensors", "HeartMuLa-oss-3B/model-00003-of-00004.safetensors", "HeartMuLa-oss-3B/model-00004-of-00004.safetensors", "HeartCodec-oss/config.json", "HeartCodec-oss/model.safetensors.index.json", "HeartCodec-oss/model-00001-of-00002.safetensors", "HeartCodec-oss/model-00002-of-00002.safetensors", ), ), ModelPackage( id="irodori_tts_500m_v3", display_name="Irodori-TTS 500M v3", target_directory="Irodori-TTS-500M-v3", source=CompositeSnapshotSource( placements=( SnapshotPlacement( source=SnapshotSource(repo_id="Aratako/Irodori-TTS-500M-v3"), required_files=("model.safetensors",), ), SnapshotPlacement( source=SnapshotSource(repo_id="llm-jp/llm-jp-3-150m"), target_subdir="../llm-jp-3-150m", required_files=("tokenizer.json", "model.safetensors", "config.json"), ), SnapshotPlacement( source=SnapshotSource(repo_id="Aratako/Semantic-DACVAE-Japanese-32dim"), target_subdir="../Semantic-DACVAE-Japanese-32dim", required_files=("weights.pth",), ), ), ), required_files=( "model.safetensors", "model_config.json", "../llm-jp-3-150m/tokenizer.json", "../Semantic-DACVAE-Japanese-32dim/weights.safetensors", ), description="Installs Irodori-TTS plus the sibling llm-jp tokenizer and DACVAE codec dependencies required by the framework runtime.", ), ModelPackage( id="irodori_tts_600m_v3_voice_design", display_name="Irodori-TTS 600M v3 VoiceDesign", target_directory="Irodori-TTS-600M-v3-VoiceDesign", source=CompositeSnapshotSource( placements=( SnapshotPlacement( source=SnapshotSource(repo_id="Aratako/Irodori-TTS-600M-v3-VoiceDesign"), required_files=("model.safetensors",), ), SnapshotPlacement( source=SnapshotSource(repo_id="llm-jp/llm-jp-3-150m"), target_subdir="../llm-jp-3-150m", required_files=("tokenizer.json", "model.safetensors", "config.json"), ), SnapshotPlacement( source=SnapshotSource(repo_id="Aratako/Semantic-DACVAE-Japanese-32dim"), target_subdir="../Semantic-DACVAE-Japanese-32dim", required_files=("weights.pth",), ), ), ), required_files=( "model.safetensors", "model_config.json", "../llm-jp-3-150m/tokenizer.json", "../Semantic-DACVAE-Japanese-32dim/weights.safetensors", ), description="Installs Irodori-TTS VoiceDesign plus the sibling llm-jp tokenizer and DACVAE codec dependencies required by the framework runtime.", ), ModelPackage( id="glm_tts", display_name="GLM-TTS", target_directory="GLM-TTS", source=CompositeSnapshotSource( placements=( SnapshotPlacement( source=SnapshotSource(repo_id="zai-org/GLM-TTS"), required_files=( "flow/config.yaml", "flow/flow.pt", "hift/hift.pt", "llm/config.json", "llm/generation_config.json", "llm/model-00001-of-00002.safetensors", "llm/model-00002-of-00002.safetensors", "llm/model.safetensors.index.json", "speech_tokenizer/config.json", "speech_tokenizer/model.safetensors", "speech_tokenizer/preprocessor_config.json", "vq32k-phoneme-tokenizer/tokenizer.model", "vq32k-phoneme-tokenizer/tokenizer_config.json", ), ), SnapshotPlacement( source=SnapshotSource( repo_id="mlx-community/index-tts2-mlx", include_prefixes=("campplus.safetensors",), ), target_subdir="frontend", required_files=("campplus.safetensors",), ), ), ), required_files=( "audio_cpp_config.json", "flow/config.yaml", "flow/model.safetensors", "frontend/campplus.safetensors", "hift/model.safetensors", "llm/config.json", "llm/generation_config.json", "llm/model-00001-of-00002.safetensors", "llm/model-00002-of-00002.safetensors", "llm/model.safetensors.index.json", "speech_tokenizer/config.json", "speech_tokenizer/model.safetensors", "speech_tokenizer/preprocessor_config.json", "vq32k-phoneme-tokenizer/tokenizer_config.json", "vq32k-phoneme-tokenizer/tokenizer_merges.txt", "vq32k-phoneme-tokenizer/tokenizer_vocab.json", ), family="glm_tts", tasks=("tts", "clon"), modes=("offline",), standalone=True, description=( "Installs GLM-TTS and prepares its Flow, HiFT, ChatGLM tokenizer, " "and CAMPPlus assets for zero-shot Chinese and English synthesis." ), ), ModelPackage( id="inflect_micro_v2_orig", display_name="Inflect Micro v2 Original-Dtype GGUF", target_directory="Inflect-Micro-v2-GGUF", source=SnapshotSource( repo_id="audio-cpp/audio.cpp-gguf", include_prefixes=("Inflect-Micro-v2-GGUF/inflect-micro-v2-orig.gguf",), strip_prefix="Inflect-Micro-v2-GGUF/", ), required_files=("inflect-micro-v2-orig.gguf",), family="inflect_v2", tasks=("tts",), modes=("offline",), standalone=True, description=( "Standalone Inflect Micro v2 original-dtype GGUF. Inflect v2 is FP32-only " "today, so quantized GGUF packages are not advertised." ), ), ModelPackage( id="outetts_1_0_1b", display_name="OuteTTS 1.0 1B", target_directory="Llama-OuteTTS-1.0-1B", source=CompositeSnapshotSource( placements=( SnapshotPlacement( source=SnapshotSource(repo_id="OuteAI/Llama-OuteTTS-1.0-1B"), required_files=( "config.json", "generation_config.json", "model.safetensors", "special_tokens_map.json", "tokenizer.json", "tokenizer_config.json", ), ), SnapshotPlacement( source=SnapshotSource(repo_id="ibm-research/DAC.speech.v1.0"), target_subdir="../DAC.speech.v1.0", required_files=("config.json", "weights_24khz_1.5kbps_v1.0.pth"), ), SnapshotPlacement( source=SnapshotSource(repo_id="Qwen/Qwen3-ForcedAligner-0.6B"), target_subdir="../Qwen3-ForcedAligner-0.6B", required_files=( "config.json", "generation_config.json", "model.safetensors", "preprocessor_config.json", "tokenizer_config.json", "vocab.json", "merges.txt", ), ), ), ), required_files=( "config.json", "generation_config.json", "model.safetensors", "special_tokens_map.json", "tokenizer.json", "tokenizer_config.json", "../DAC.speech.v1.0/config.json", "../DAC.speech.v1.0/model.safetensors", "../Qwen3-ForcedAligner-0.6B/config.json", "../Qwen3-ForcedAligner-0.6B/generation_config.json", "../Qwen3-ForcedAligner-0.6B/model.safetensors", "../Qwen3-ForcedAligner-0.6B/preprocessor_config.json", "../Qwen3-ForcedAligner-0.6B/tokenizer_config.json", "../Qwen3-ForcedAligner-0.6B/vocab.json", "../Qwen3-ForcedAligner-0.6B/merges.txt", ), description="Installs OuteTTS, its IBM DAC 1.5 kbps codec, and Qwen3 Forced Aligner for reliable voice cloning.", ), ModelPackage( id="stable_audio_3_small_music", display_name="Stable Audio 3 Small Music", target_directory="stable-audio-3-small-music", source=SnapshotSource(repo_id="stabilityai/stable-audio-3-small-music"), required_files=( "model_config.json", "model.safetensors", "t5gemma-b-b-ul2/config.json", "t5gemma-b-b-ul2/model.safetensors", "t5gemma-b-b-ul2/tokenizer.json", "t5gemma-b-b-ul2/tokenizer.model", ), ), ModelPackage( id="stable_audio_3_small_sfx", display_name="Stable Audio 3 Small SFX", target_directory="stable-audio-3-small-sfx", source=SnapshotSource(repo_id="stabilityai/stable-audio-3-small-sfx"), required_files=( "model_config.json", "model.safetensors", "t5gemma-b-b-ul2/config.json", "t5gemma-b-b-ul2/model.safetensors", "t5gemma-b-b-ul2/tokenizer.json", "t5gemma-b-b-ul2/tokenizer.model", ), ), ModelPackage( id="stable_audio_3_medium", display_name="Stable Audio 3 Medium", target_directory="stable-audio-3-medium", source=SnapshotSource(repo_id="stabilityai/stable-audio-3-medium"), required_files=( "model_config.json", "model.safetensors", "t5gemma-b-b-ul2/config.json", "t5gemma-b-b-ul2/model.safetensors", "t5gemma-b-b-ul2/tokenizer.json", "t5gemma-b-b-ul2/tokenizer.model", ), ), ModelPackage( id="supertonic_3", display_name="Supertonic 3", target_directory="supertonic-3", source=SnapshotSource(repo_id="mlx-community/supertonic-3-mlx"), required_files=( "config/tts.json", "config/unicode_indexer.json", "ggml/supertonic.safetensors", "voice_styles/M1.json", ), ), ModelPackage( id="index_tts2", display_name="IndexTTS2", target_directory="IndexTTS-2", source=SnapshotSource(repo_id="mlx-community/index-tts2-mlx"), required_files=( "config.yaml", "bpe.model", "gpt.safetensors", "s2mel.safetensors", "feat1.safetensors", "feat2.safetensors", "wav2vec2bert_stats.safetensors", "semantic_codec_model.safetensors", "campplus.safetensors", "w2v-bert-2.0/config.json", "w2v-bert-2.0/preprocessor_config.json", "w2v-bert-2.0/model.safetensors", "bigvgan/config.json", "bigvgan/model.safetensors", "qwen0.6bemo4-merge/config.json", "qwen0.6bemo4-merge/generation_config.json", "qwen0.6bemo4-merge/tokenizer.json", "qwen0.6bemo4-merge/tokenizer_config.json", "qwen0.6bemo4-merge/vocab.json", "qwen0.6bemo4-merge/merges.txt", "qwen0.6bemo4-merge/model.safetensors", ), description="Framework-ready IndexTTS2 safetensors layout with shared components at the model root.", ), ModelPackage( id="mel_band_roformer", display_name="Mel RoFormer MLX", target_directory="mel-roformer-mlx", source=SnapshotSource(repo_id="mlx-community/mel-roformer-mlx"), required_files=("config.json", "model.safetensors"), ), ModelPackage( id="vevo2", display_name="Vevo2", target_directory="Vevo2", source=CompositeSnapshotSource( placements=( SnapshotPlacement( source=SnapshotSource( repo_id="RMSnow/Vevo2", include_prefixes=( "acoustic_modeling/fm_emilia101k_singnet7k_repa/", "acoustic_modeling/fm_emilia101k_singnet7k_repa_text/", "contentstyle_modeling/posttrained/", "contentstyle_modeling/pretrained/", "tokenizer/contentstyle_fvq16384_12.5hz/", "tokenizer/prosody_fvq512_6.25hz/", "vocoder/", ), exclude_prefixes=( "contentstyle_modeling/posttrained/optimizer.pt", "contentstyle_modeling/posttrained/rng_state_", "contentstyle_modeling/posttrained/scheduler.pt", "contentstyle_modeling/posttrained/trainer_state.json", "contentstyle_modeling/posttrained/training_args.bin", "contentstyle_modeling/pretrained/optimizer.bin", "contentstyle_modeling/pretrained/random_states_", "contentstyle_modeling/pretrained/scheduler.bin", ), ), required_files=( "acoustic_modeling/fm_emilia101k_singnet7k_repa/config.json", "acoustic_modeling/fm_emilia101k_singnet7k_repa/model.safetensors", "acoustic_modeling/fm_emilia101k_singnet7k_repa/whisper_stats.pt", "acoustic_modeling/fm_emilia101k_singnet7k_repa_text/config.json", "acoustic_modeling/fm_emilia101k_singnet7k_repa_text/model.safetensors", "acoustic_modeling/fm_emilia101k_singnet7k_repa_text/whisper_stats.pt", "contentstyle_modeling/posttrained/amphion_config.json", "contentstyle_modeling/posttrained/config.json", "contentstyle_modeling/posttrained/generation_config.json", "contentstyle_modeling/posttrained/merges.txt", "contentstyle_modeling/posttrained/model.safetensors", "contentstyle_modeling/posttrained/tokenizer.json", "contentstyle_modeling/posttrained/tokenizer_config.json", "contentstyle_modeling/posttrained/vocab.json", "contentstyle_modeling/pretrained/config.json", "contentstyle_modeling/pretrained/generation_config.json", "contentstyle_modeling/pretrained/merges.txt", "contentstyle_modeling/pretrained/model.safetensors", "contentstyle_modeling/pretrained/tokenizer.json", "contentstyle_modeling/pretrained/tokenizer_config.json", "contentstyle_modeling/pretrained/vocab.json", "tokenizer/contentstyle_fvq16384_12.5hz/model.safetensors", "tokenizer/prosody_fvq512_6.25hz/model.safetensors", "vocoder/config.json", "vocoder/model.safetensors", "vocoder/model_1.safetensors", "vocoder/model_2.safetensors", ), ), ), ), required_files=( "acoustic_modeling/fm_emilia101k_singnet7k_repa/config.json", "acoustic_modeling/fm_emilia101k_singnet7k_repa/model.safetensors", "acoustic_modeling/fm_emilia101k_singnet7k_repa/whisper_stats.safetensors", "acoustic_modeling/fm_emilia101k_singnet7k_repa_text/config.json", "acoustic_modeling/fm_emilia101k_singnet7k_repa_text/model.safetensors", "acoustic_modeling/fm_emilia101k_singnet7k_repa_text/whisper_stats.safetensors", "contentstyle_modeling/posttrained/amphion_config.json", "contentstyle_modeling/posttrained/config.json", "contentstyle_modeling/posttrained/generation_config.json", "contentstyle_modeling/posttrained/merges.txt", "contentstyle_modeling/posttrained/model.safetensors", "contentstyle_modeling/posttrained/tokenizer.json", "contentstyle_modeling/posttrained/tokenizer_config.json", "contentstyle_modeling/posttrained/vocab.json", "contentstyle_modeling/pretrained/config.json", "contentstyle_modeling/pretrained/generation_config.json", "contentstyle_modeling/pretrained/merges.txt", "contentstyle_modeling/pretrained/model.safetensors", "contentstyle_modeling/pretrained/tokenizer.json", "contentstyle_modeling/pretrained/tokenizer_config.json", "contentstyle_modeling/pretrained/vocab.json", "tokenizer/contentstyle_fvq16384_12.5hz/model.safetensors", "tokenizer/prosody_fvq512_6.25hz/model.safetensors", "vocoder/config.json", "vocoder/model.safetensors", "vocoder/model_1.safetensors", "vocoder/model_2.safetensors", ), description="Installs Vevo2 plus the sibling whisper-medium dependency required by the framework runtime.", ), ModelPackage( id="vevo2_gguf", display_name="Vevo2 Q8_0 GGUF", target_directory="Vevo2-GGUF", source=SnapshotSource( repo_id="audio-cpp/audio.cpp-gguf", include_prefixes=("Vevo2-GGUF/vevo2-q8_0.gguf",), strip_prefix="Vevo2-GGUF/", ), required_files=("vevo2-q8_0.gguf",), family="vevo2", tasks=("tts", "music", "vc", "edit", "svc", "s2s"), description=( "Standalone audio.cpp Q8_0 GGUF package for Vevo2. It embeds every config the runtime " "reads, whisper-medium included, so unlike the safetensors package it needs no sibling " "whisper-medium download." ), ), ModelPackage( id="seed_vc", display_name="SeedVC-MLX", target_directory="SeedVC-MLX", source=SnapshotSource(repo_id="mlx-community/SeedVC-MLX"), required_files=( "seed_vc_manifest.json", "v2/vc_wrapper.json", "v2/ar.safetensors", "v2/cfm.safetensors", "v1/svc.json", "v1/svc.safetensors", "v1/whisper_bigvgan.json", "v1/whisper_bigvgan.safetensors", "v1/xlsr_hift.json", "v1/xlsr_hift.safetensors", "astral/bsq32.json", "astral/bsq32.safetensors", "astral/bsq2048.json", "astral/bsq2048.safetensors", "campplus/model.safetensors", "rmvpe/model.safetensors", "hift/config.json", "hift/model.safetensors", "bigvgan/v2_22khz_80band_256x/config.json", "bigvgan/v2_22khz_80band_256x/model.safetensors", "bigvgan/v2_44khz_128band_512x/config.json", "bigvgan/v2_44khz_128band_512x/model.safetensors", "whisper-small/config.json", "whisper-small/model.safetensors", "hubert-large-ll60k/config.json", "hubert-large-ll60k/model.safetensors", "wav2vec2-xls-r-300m/config.json", "wav2vec2-xls-r-300m/model.safetensors", ), description="Framework-ready SeedVC-MLX bundle. The Hugging Face snapshot already contains the converted safetensors layout.", ), ModelPackage( id="citrinet_asr", display_name="Citrinet ASR converted layout", target_directory="citrinet", source=ConverterSource( kind="nemo_archive", description="Download and convert the official NeMo archive into framework-friendly safetensors and sidecars.", url="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_en_citrinet_256/versions/1.0.0rc1/files/stt_en_citrinet_256.nemo", output_weights_file="citrinet_256.safetensors", config_kind="citrinet", ), required_files=( "citrinet_256.safetensors", "citrinet_256_config.json", "citrinet_256_tokenizer.model", "citrinet_256_vocab.txt", ), ), ModelPackage( id="kroko_asr_community_q8_0", display_name="Kroko Community ASR Q8_0 GGUF", target_directory="Kroko-ASR-GGUF", source=SnapshotSource( repo_id="audio-cpp/audio.cpp-gguf", include_prefixes=("Kroko-ASR-GGUF/kroko-en-community-64-l-q8_0.gguf",), strip_prefix="Kroko-ASR-GGUF/", ), required_files=("kroko-en-community-64-l-q8_0.gguf",), family="kroko_asr", standalone=True, tasks=("asr",), modes=("offline", "streaming"), description="Standalone Kroko Community English Q8_0 GGUF for offline and streaming ASR.", ), ModelPackage( id="voxcpm2", display_name="VoxCPM2", target_directory="VoxCPM2", source=SnapshotSource(repo_id="OpenBMB/VoxCPM2"), required_files=( "config.json", "model.safetensors", "tokenizer.json", "tokenizer_config.json", "audiovae.pth", "audiovae.safetensors", ), ), ModelPackage( id="voxcpm2_audiovae", display_name="VoxCPM2 AudioVAE local conversion utility", target_directory="VoxCPM2", source=ConverterSource( kind="pytorch_to_safetensors", description="Convert a local PyTorch checkpoint into safetensors.", ), required_files=("audiovae.safetensors",), description="Utility only. Use voxcpm2 for the full framework-ready VoxCPM2 package.", ), ModelPackage( id="htdemucs", display_name="HTDemucs", target_directory="htdemucs", source=ConverterSource( kind="demucs_reference", description="Download official Demucs .th checkpoints, or use --source-dir, and assemble framework safetensors plus a manifest.", ), required_files=( "manifest.json", "955717e8/config.json", "955717e8/model.safetensors", ), ), ) PACKAGE_BY_ID = {package.id: package for package in CATALOG} def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Download or convert official model packages into a local models layout." ) subparsers = parser.add_subparsers(dest="command", required=True) list_parser = subparsers.add_parser("list", help="List available model packages.") list_parser.add_argument("--json", action="store_true", help="Print machine-readable JSON.") info_parser = subparsers.add_parser("info", help="Show one package.") info_parser.add_argument("package_id") info_parser.add_argument("--json", action="store_true", help="Print machine-readable JSON.") install_parser = subparsers.add_parser("install", help="Install one package into a models root.") install_parser.add_argument("package_id") install_parser.add_argument("--models-root", default="models", help="Target models root.") install_parser.add_argument("--overwrite", action="store_true", help="Replace an existing installed directory.") install_parser.add_argument("--source-file", help="Local source checkpoint for utility packages.") install_parser.add_argument("--output-file", help="Optional output file for single-file utility packages.") install_parser.add_argument("--source-dir", help="Optional local source directory for packages that can use checkpoint folders.") install_parser.add_argument("--variant", help="Variant or model name for packages that support variants.") return parser.parse_args() def resolve_path(path: str) -> Path: candidate = Path(path).expanduser() if candidate.is_absolute(): return candidate return REPO_ROOT / candidate _FAMILY_ID_STRIP_PATTERNS = ( re.compile(r"_\d+_\d+b(?:_base|_custom_voice|_voice_design)?$"), re.compile(r"_\d+b(?:_base|_custom_voice|_voice_design)?$"), re.compile(r"_\d+m(?:_bf16)?$"), re.compile(r"_bf16$"), re.compile(r"_\d+spk_v\d+$"), re.compile(r"_\d+hz_\d+k_v\d+$"), re.compile(r"_3_small_(?:music|sfx)$"), re.compile(r"_3_medium$"), re.compile(r"_v\d+(?:_\d+)?$"), re.compile(r"_\d+$"), re.compile(r"_model$"), re.compile(r"_hf$"), re.compile(r"_12hz$"), re.compile(r"_(?:base|custom_voice|voice_design)$"), ) def _default_family_from_package_id(package_id: str) -> str: """Strip size/version suffixes from a package id. Prefer ``ModelPackage.family`` for exceptions.""" key = str(package_id or "").strip().lower() if not key: return "" changed = True while changed and key: changed = False for pattern in _FAMILY_ID_STRIP_PATTERNS: updated = pattern.sub("", key) if updated and updated != key: key = updated changed = True break return key def _default_tasks_from_family(family: str) -> list[str]: key = str(family or "").strip().lower() if not key: return [] if "forced_aligner" in key or key.endswith("_aligner") or key.endswith("_align"): return ["align"] if key.endswith("_asr") or key.endswith("_stt") or key in {"whisper", "voxtral_realtime"}: return ["asr"] if "vad" in key: return ["vad"] if "diar" in key: return ["diar"] if any(token in key for token in ("demucs", "roformer", "separator")): return ["sep"] if key in {"stable_audio", "ace_step", "heartmula"} or key.endswith("_gen"): return ["gen"] if key.endswith("_vc") or key in {"seed_vc", "vevo2"}: return ["vc"] if key.endswith("_codec") or key == "miocodec": return ["codec"] if key.endswith("_asr") or key.endswith("_stt"): return ["asr"] if "tts" in key or key in { "chatterbox", "voxcpm2", "omnivoice", "supertonic", "miotts", "pocket_tts", "irodori_tts", "index_tts2", "vibevoice", "moss_tts_nano", "moss_tts_local", }: return ["tts"] return [] def _collect_package_repo_ids(package: ModelPackage) -> list[str]: source = package.source repos: list[str] = [] if isinstance(source, SnapshotSource): repos.append(source.repo_id) elif isinstance(source, CompositeSnapshotSource): for placement in source.placements: repos.append(placement.source.repo_id) return [repo for repo in repos if isinstance(repo, str) and repo] def _package_is_gated(package: ModelPackage) -> bool: if isinstance(package.gated, bool): return package.gated repos = [repo.lower() for repo in _collect_package_repo_ids(package)] return any(any(marker in repo for marker in _GATED_REPO_MARKERS) for repo in repos) def _package_standalone_fields(package: ModelPackage) -> tuple[bool, str | None]: if isinstance(package.standalone, bool): return package.standalone, package.parent_package_id desc = str(package.description or "").strip().lower() kind = package_install_kind(package) if desc.startswith("subcomponent only.") or desc.startswith("utility only.") or kind == "utility": parent = package.parent_package_id if not parent: for token in desc.replace(".", " ").replace(",", " ").split(): if token in PACKAGE_BY_ID and token != package.id: parent = token break return False, parent if "tokenizer" in package.id.lower() or "audiovae" in package.id.lower(): return False, package.parent_package_id return True, None def package_payload(package: ModelPackage) -> dict[str, object]: source = package.source if isinstance(source, SnapshotSource): source_payload: dict[str, object] = { "kind": "huggingface_snapshot", "repo_id": source.repo_id, "revision": source.revision, "include_prefixes": list(source.include_prefixes), "include_suffixes": list(source.include_suffixes), "exclude_prefixes": list(source.exclude_prefixes), "strip_prefix": source.strip_prefix, } installable = True elif isinstance(source, CompositeSnapshotSource): source_payload = { "kind": "composite_snapshot", "placements": [ { "repo_id": placement.source.repo_id, "revision": placement.source.revision, "target_subdir": placement.target_subdir, "required_files": list(placement.required_files), "include_prefixes": list(placement.source.include_prefixes), "include_suffixes": list(placement.source.include_suffixes), "exclude_prefixes": list(placement.source.exclude_prefixes), "strip_prefix": placement.source.strip_prefix, } for placement in source.placements ], "repo_ids": _collect_package_repo_ids(package), } installable = True elif isinstance(source, ConverterSource): utility = source.kind in UTILITY_CONVERTER_KINDS source_payload = { "kind": "utility" if utility else "composite", "operation_kind": source.kind, "description": source.description, } if source.url: source_payload["url"] = source.url if source.output_weights_file: source_payload["output_weights_file"] = source.output_weights_file if source.config_kind: source_payload["config_kind"] = source.config_kind installable = True else: source_payload = { "kind": "unsupported", "reason": source.reason, } installable = False family = package.family or _default_family_from_package_id(package.id) tasks = list(package.tasks) if package.tasks else _default_tasks_from_family(family) standalone, parent_package_id = _package_standalone_fields(package) return { "id": package.id, "display_name": package.display_name, "target_directory": package.target_directory, "description": package.description, "installable": installable, "install_kind": package_install_kind(package), "usage_examples": package_usage_examples(package), "required_files": list(package.required_files), "source": source_payload, "family": family, "tasks": tasks, "modes": list(package.modes) if package.modes else (["offline"] if tasks else []), "standalone": standalone, "parent_package_id": parent_package_id, "gated": _package_is_gated(package), } def print_package_table(packages: Iterable[ModelPackage]) -> None: rows = [] for package in packages: status = package_install_kind(package) rows.append((package.id, status, package.target_directory, package.display_name)) width_id = max(len("id"), *(len(row[0]) for row in rows)) width_status = max(len("type"), *(len(row[1]) for row in rows)) width_target = max(len("target"), *(len(row[2]) for row in rows)) print(f"{'id':<{width_id}} {'type':<{width_status}} {'target':<{width_target}} name") for row in rows: print(f"{row[0]:<{width_id}} {row[1]:<{width_status}} {row[2]:<{width_target}} {row[3]}") def hf_tree_url(source: SnapshotSource) -> str: revision = quote(source.revision, safe="") repo_id = source.repo_id return f"{hf_endpoint()}/api/models/{repo_id}/tree/{revision}?recursive=true" def hf_resolve_url(source: SnapshotSource, relative_path: str) -> str: revision = quote(source.revision, safe="") path = quote(relative_path, safe="/") return f"{hf_endpoint()}/{source.repo_id}/resolve/{revision}/{path}" def http_json(url: str, timeout: float = 30.0) -> object: # Bounded by default: the WebUI probes package sizes from a request thread before # starting a download, and an unbounded urlopen there would hang the page. request = Request(url, headers=http_headers()) with urlopen(request, timeout=timeout) as response: return json.load(response) def local_snapshot_path(source: SnapshotSource, remote_path: str) -> str: if not source.strip_prefix: return remote_path if not remote_path.startswith(source.strip_prefix): raise RuntimeError(f"snapshot path does not start with strip_prefix: {remote_path}") local_path = remote_path[len(source.strip_prefix):] if not local_path: raise RuntimeError(f"snapshot strip_prefix removed full path: {remote_path}") return local_path def list_hf_files(source: SnapshotSource) -> list[tuple[str, str, int | None]]: payload = http_json(hf_tree_url(source)) if not isinstance(payload, list): raise RuntimeError(f"unexpected HuggingFace tree payload for {source.repo_id}") files: list[tuple[str, str, int | None]] = [] for entry in payload: if not isinstance(entry, dict): continue path = entry.get("path") entry_type = entry.get("type") if not isinstance(path, str) or entry_type != "file": continue if source.include_prefixes and not any(path.startswith(prefix) for prefix in source.include_prefixes): continue if source.include_suffixes and not any(path.endswith(suffix) for suffix in source.include_suffixes): continue if any(path.startswith(prefix) for prefix in source.exclude_prefixes): continue size = entry.get("size") files.append((path, local_snapshot_path(source, path), size if isinstance(size, int) else None)) if not files: raise RuntimeError(f"no installable files found for {source.repo_id}") return files def download_file(url: str, target: Path, expected_size: int | None) -> int: request = Request(url, headers=http_headers()) with urlopen(request) as response, target.open("wb") as handle: written = 0 while True: chunk = response.read(1 << 20) if not chunk: break handle.write(chunk) written += len(chunk) if expected_size is not None and written != expected_size: raise RuntimeError(f"downloaded size mismatch for {target}: {written} != {expected_size}") return written def validate_required_files(package: ModelPackage, root: Path) -> None: missing = [relative for relative in package.required_files if not (root / relative).exists()] if missing: raise RuntimeError(f"installed package is missing required files: {missing}") def validate_composite_required_files(package: ModelPackage, staged_root: Path, final_root: Path) -> None: missing = [ relative for relative in package.required_files if not normalized_join(staged_root, relative).exists() and not normalized_join(final_root, relative).exists() ] if missing: raise RuntimeError(f"installed package is missing required files: {missing}") def validate_required_files_list(required_files: Iterable[str], root: Path, label: str) -> None: missing = [relative for relative in required_files if not (root / relative).exists()] if missing: raise RuntimeError(f"{label} is missing required files: {missing}") def install_snapshot_into_dir( source: SnapshotSource, destination_root: Path, required_files: Iterable[str], *, validate: bool = True, ) -> None: files = list_hf_files(source) for remote, relative, expected_size in files: destination = destination_root / relative destination.parent.mkdir(parents=True, exist_ok=True) print(f"download {remote}") download_file(hf_resolve_url(source, remote), destination, expected_size) if validate: validate_required_files_list(required_files, destination_root, source.repo_id) def write_vevo2_whisper_stats(target_path: Path, source_file: Path) -> None: payload = torch.load(source_file, map_location="cpu", weights_only=False) state = checkpoint_state_dict(payload) tensors = tensor_state_dict(state) if set(tensors.keys()) != {"mean", "std"}: raise RuntimeError(f"Vevo2 whisper stats must contain mean/std only: {sorted(tensors.keys())}") save_file(tensors, str(target_path)) def whisper_dims(payload: dict[str, Any]) -> dict[str, int]: raw_dims = payload.get("dims") if not isinstance(raw_dims, dict): raise RuntimeError("Whisper checkpoint is missing dims") dims: dict[str, int] = {} for key, value in raw_dims.items(): if not isinstance(key, str) or not isinstance(value, int): raise RuntimeError(f"invalid Whisper dim entry: {key!r}={value!r}") dims[key] = value return dims def whisper_state(payload: dict[str, Any]) -> dict[str, torch.Tensor]: raw_state = payload.get("model_state_dict") if not isinstance(raw_state, dict): raise RuntimeError("Whisper checkpoint is missing model_state_dict") state: dict[str, torch.Tensor] = {} for key, value in raw_state.items(): if not isinstance(key, str): raise RuntimeError(f"invalid non-string Whisper tensor key: {key!r}") if not isinstance(value, torch.Tensor): raise RuntimeError(f"invalid non-tensor Whisper state entry: {key}") state[key] = value.detach().cpu().contiguous() return state def whisper_medium_config(dims: dict[str, int], tensor_count: int) -> dict[str, Any]: return { "model_type": "whisper", "variant": "medium", "source_format": "openai_whisper_pt", "weight_file": "model.safetensors", "tensor_count": tensor_count, "dims": dims, "audio_encoder": { "n_mels": dims["n_mels"], "n_audio_ctx": dims["n_audio_ctx"], "n_audio_state": dims["n_audio_state"], "n_audio_head": dims["n_audio_head"], "n_audio_layer": dims["n_audio_layer"], }, "text_decoder": { "n_vocab": dims["n_vocab"], "n_text_ctx": dims["n_text_ctx"], "n_text_state": dims["n_text_state"], "n_text_head": dims["n_text_head"], "n_text_layer": dims["n_text_layer"], }, } def install_whisper_medium_dependency(root: Path) -> None: root.mkdir(parents=True, exist_ok=True) ckpt_path = root / "medium.pt" model_path = root / "model.safetensors" config_path = root / "config.json" download_file(WHISPER_MEDIUM_PT_URL, ckpt_path, None) try: payload = torch.load(ckpt_path, map_location="cpu", weights_only=False) if not isinstance(payload, dict): raise RuntimeError(f"unexpected Whisper checkpoint payload type: {type(payload)}") dims = whisper_dims(payload) state = whisper_state(payload) save_file(state, str(model_path)) reloaded = load_file(str(model_path)) if set(reloaded.keys()) != set(state.keys()): raise RuntimeError("saved Whisper safetensors key set does not match source checkpoint") for key, value in state.items(): if not torch.equal(value, reloaded[key]): raise RuntimeError(f"saved Whisper safetensors changed tensor: {key}") config_path.write_text( json.dumps(whisper_medium_config(dims, len(state)), indent=2, ensure_ascii=False) + "\n", encoding="utf-8", ) finally: ckpt_path.unlink(missing_ok=True) def prepare_vevo2_snapshot_layout(root: Path) -> None: stats_dirs = ( root / "acoustic_modeling" / "fm_emilia101k_singnet7k_repa", root / "acoustic_modeling" / "fm_emilia101k_singnet7k_repa_text", ) if all((stats_dir / "whisper_stats.safetensors").is_file() for stats_dir in stats_dirs): return for stats_dir in stats_dirs: stats_pt = stats_dir / "whisper_stats.pt" if not stats_pt.is_file(): raise RuntimeError(f"Vevo2 snapshot is missing {stats_pt}") write_vevo2_whisper_stats(stats_dir / "whisper_stats.safetensors", stats_pt) def convert_voxcpm2_audiovae(root: Path) -> None: input_path = root / "audiovae.pth" output_path = root / "audiovae.safetensors" payload = torch.load(input_path, map_location="cpu", weights_only=True) state = checkpoint_state_dict(payload) tensors = tensor_state_dict(state) write_checked_safetensors(tensors, output_path, input_path, overwrite=True) def convert_moss_tts_weights(root: Path) -> None: input_path = root / "pytorch_model.bin" output_path = root / "model.safetensors" payload = torch.load(input_path, map_location="cpu", weights_only=True) state = checkpoint_state_dict(payload) tensors = tensor_state_dict(state) write_checked_safetensors(tensors, output_path, input_path, overwrite=True) def convert_ace_step_silence_latent(root: Path) -> None: input_path = root / "silence_latent.pt" output_path = root / "silence_latent.safetensors" tensor = torch.load(input_path, map_location="cpu", weights_only=True) if not isinstance(tensor, torch.Tensor): raise RuntimeError(f"ACE-Step silence latent must be a tensor: {input_path}") write_checked_safetensors({"silence_latent": tensor.detach().cpu().contiguous()}, output_path, input_path, overwrite=True) def convert_irodori_dacvae_weights(root: Path) -> None: input_path = root / "weights.pth" output_path = root / "weights.safetensors" payload = torch.load(input_path, map_location="cpu", weights_only=True) state = checkpoint_state_dict(payload) tensors = tensor_state_dict(state) write_checked_safetensors(tensors, output_path, input_path, overwrite=True) def convert_outetts_dac_weights(root: Path) -> None: input_path = root / "weights_24khz_1.5kbps_v1.0.pth" output_path = root / "model.safetensors" payload = torch.load(input_path, map_location="cpu", weights_only=True) state = checkpoint_state_dict(payload) tensors = tensor_state_dict(state) expected = { "quantizer.quantizers.0.codebook.weight": (1024, 8), "quantizer.quantizers.1.codebook.weight": (1024, 8), "decoder.model.0.weight_v": (1536, 1024, 7), "decoder.model.6.weight_v": (1, 96, 7), } for name, shape in expected.items(): if name not in tensors or tuple(tensors[name].shape) != shape: raise RuntimeError(f"unexpected OuteTTS DAC tensor {name}: {getattr(tensors.get(name), 'shape', None)}") write_checked_safetensors(tensors, output_path, input_path, overwrite=True) def write_irodori_model_config(root: Path) -> None: input_path = root / "model.safetensors" output_path = root / "model_config.json" with safe_open(input_path, framework="pt", device="cpu") as handle: metadata = handle.metadata() config_json = metadata.get("config_json") if config_json is None: raise RuntimeError(f"Irodori-TTS model.safetensors is missing config_json metadata: {input_path}") output_path.write_text(config_json, encoding="utf-8") def copy_bundled_model_manager_assets(asset_subdir: str, destination_root: Path, required_files: Iterable[str]) -> None: source_root = MODEL_MANAGER_ASSETS / asset_subdir validate_required_files_list(required_files, source_root, f"bundled model-manager assets {asset_subdir}") for relative in required_files: source_path = source_root / relative destination_path = destination_root / relative destination_path.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(source_path, destination_path) def install_snapshot(package: ModelPackage, source: SnapshotSource, models_root: Path, overwrite: bool) -> Path: target_dir = models_root / package.target_directory staging_root = models_root / ".engine_model_staging" staging_root.mkdir(parents=True, exist_ok=True) staging_dir = Path(tempfile.mkdtemp(prefix=f"{package.target_directory}.", dir=staging_root)) try: generated_files = {"audiovae.safetensors"} pre_validate_files = tuple( relative for relative in package.required_files if relative not in generated_files ) install_snapshot_into_dir( source, staging_dir, pre_validate_files, validate=package.id not in POSTPROCESS_SNAPSHOT_PACKAGE_IDS, ) if package.id == "voxcpm2": convert_voxcpm2_audiovae(staging_dir) validate_required_files(package, staging_dir) elif package.id == "moss_tts_nano_100m_model": convert_moss_tts_weights(staging_dir) if target_dir.exists(): if not overwrite: raise RuntimeError(f"model directory already exists: {target_dir}") shutil.rmtree(target_dir) target_dir.parent.mkdir(parents=True, exist_ok=True) staging_dir.rename(target_dir) return target_dir except Exception: shutil.rmtree(staging_dir, ignore_errors=True) raise finally: try: if staging_root.exists() and not any(staging_root.iterdir()): staging_root.rmdir() except OSError: pass def normalized_join(base: Path, relative: str) -> Path: if not relative: return base return Path(os.path.normpath(str(base / relative))) def install_composite_snapshot( package: ModelPackage, source: CompositeSnapshotSource, models_root: Path, overwrite: bool, ) -> Path: package_root = models_root / package.target_directory staging_root = models_root / ".engine_model_staging" staging_root.mkdir(parents=True, exist_ok=True) staging_bundle = Path(tempfile.mkdtemp(prefix=f"{package.target_directory}.", dir=staging_root)) staged_roots: dict[Path, Path] = {} try: staged_package_root = staging_bundle / package.target_directory for placement in source.placements: destination_root = normalized_join(staged_package_root, placement.target_subdir) final_root = normalized_join(package_root, placement.target_subdir) if final_root.exists() and not overwrite: validate_required_files_list(placement.required_files, final_root, str(final_root)) continue destination_root.mkdir(parents=True, exist_ok=True) install_snapshot_into_dir(placement.source, destination_root, placement.required_files) staged_roots[final_root] = destination_root if package.id == "vevo2": whisper_root = staging_bundle / "whisper-medium" install_whisper_medium_dependency(whisper_root) staged_roots[package_root.parent / "whisper-medium"] = whisper_root prepare_vevo2_snapshot_layout(staged_package_root) elif package.id == "ace_step": convert_ace_step_silence_latent(staged_package_root / "acestep-v15-turbo") convert_ace_step_silence_latent(staged_package_root / "acestep-v15-base") elif package.id == "moss_tts_nano_100m": convert_moss_tts_weights(staged_package_root) elif package.id in {"irodori_tts_500m_v3", "irodori_tts_600m_v3_voice_design"}: write_irodori_model_config(staged_package_root) dacvae_root = staged_package_root.parent / "Semantic-DACVAE-Japanese-32dim" if dacvae_root.exists(): convert_irodori_dacvae_weights(dacvae_root) elif package.id == "glm_tts": subprocess.run( [ sys.executable, str(REPO_ROOT / "tools" / "community_models" / "convert_glm_tts.py"), "--model-dir", str(staged_package_root), "--overwrite", ], check=True, ) elif package.id == "outetts_1_0_1b": dac_root = staged_package_root.parent / "DAC.speech.v1.0" if dac_root.exists(): convert_outetts_dac_weights(dac_root) elif package.id == "vibevoice_asr": copy_bundled_model_manager_assets( "vibevoice_1_5b", staged_package_root, ("tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt"), ) elif package.id in {"vibevoice_1_5b", "vibevoice_7b"}: # VibeVoice 1.5B and 7B share the same Qwen2.5 tokenizer, and neither # upstream repo ships the tokenizer files, so both reuse one bundle. copy_bundled_model_manager_assets( "vibevoice_1_5b", staged_package_root, ("tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt"), ) validate_composite_required_files(package, staged_package_root, package_root) top_level_roots: list[Path] = [] for final_root in sorted(staged_roots.keys(), key=lambda path: len(path.parts)): if any(final_root.is_relative_to(existing) for existing in top_level_roots): continue top_level_roots.append(final_root) for final_root in sorted(top_level_roots, key=lambda path: len(path.parts), reverse=True): if final_root.exists(): if not overwrite: raise RuntimeError(f"model directory already exists: {final_root}") shutil.rmtree(final_root) for final_root in sorted(top_level_roots, key=lambda path: len(path.parts)): destination_root = staged_roots[final_root] final_root.parent.mkdir(parents=True, exist_ok=True) destination_root.rename(final_root) return package_root except Exception: shutil.rmtree(staging_bundle, ignore_errors=True) raise finally: try: if staging_root.exists() and not any(staging_root.iterdir()): staging_root.rmdir() except OSError: pass def checkpoint_state_dict(payload: Any) -> dict[str, Any]: if not isinstance(payload, dict): raise RuntimeError(f"checkpoint payload must be a dict, got {type(payload).__name__}") state = payload.get("state_dict", payload) if not isinstance(state, dict): raise RuntimeError(f"checkpoint state_dict must be a dict, got {type(state).__name__}") return state def install_demucs_import_stubs() -> None: if "demucs.htdemucs" not in sys.modules: demucs_pkg = sys.modules.get("demucs") if demucs_pkg is None: demucs_pkg = types.ModuleType("demucs") demucs_pkg.__path__ = [] # type: ignore[attr-defined] sys.modules["demucs"] = demucs_pkg htdemucs = types.ModuleType("demucs.htdemucs") class HTDemucs: pass htdemucs.HTDemucs = HTDemucs sys.modules["demucs.htdemucs"] = htdemucs setattr(demucs_pkg, "htdemucs", htdemucs) if "openunmix.filtering" not in sys.modules: filtering = types.ModuleType("openunmix.filtering") def _wiener_stub(*_args: Any, **_kwargs: Any) -> Any: raise RuntimeError("openunmix.filtering.wiener stub should not be called during Demucs conversion") filtering.wiener = _wiener_stub openunmix = types.ModuleType("openunmix") openunmix.filtering = filtering sys.modules["openunmix"] = openunmix sys.modules["openunmix.filtering"] = filtering if "julius" not in sys.modules: julius = types.ModuleType("julius") def _resample_frac_stub(*_args: Any, **_kwargs: Any) -> Any: raise RuntimeError("julius.resample_frac stub should not be called during Demucs conversion") julius.resample_frac = _resample_frac_stub sys.modules["julius"] = julius if "dora.log" not in sys.modules: log = types.ModuleType("dora.log") def _fatal_stub(message: str) -> None: raise RuntimeError(message) log.fatal = _fatal_stub dora = types.ModuleType("dora") dora.log = log sys.modules["dora"] = dora sys.modules["dora.log"] = log def load_demucs_package(path: Path) -> dict[str, Any]: install_demucs_import_stubs() with warnings.catch_warnings(): warnings.simplefilter("ignore") package = torch.load(path, map_location="cpu", weights_only=False) if not isinstance(package, dict): raise RuntimeError(f"unexpected Demucs package type in {path}: {type(package)}") required = {"klass", "args", "kwargs", "state"} missing = required.difference(package) if missing: raise RuntimeError(f"Demucs package {path} is missing keys: {sorted(missing)}") return package def find_demucs_checkpoint(repo_dir: Path, signature: str) -> Path: matches = sorted(repo_dir.glob(f"{signature}*.th")) if not matches: raise RuntimeError(f"could not find checkpoint for signature '{signature}' in {repo_dir}") if len(matches) > 1: exact = [path for path in matches if path.stem == signature or path.stem.startswith(signature + "-")] if len(exact) == 1: return exact[0] raise RuntimeError(f"ambiguous checkpoints for signature '{signature}' in {repo_dir}: {matches}") return matches[0] def normalize_manifest_value(value: Any) -> Any: if isinstance(value, (str, int, float, bool)) or value is None: return value if isinstance(value, fractions.Fraction): return float(value) if isinstance(value, Path): return str(value) if isinstance(value, (list, tuple)): return [normalize_manifest_value(item) for item in value] if isinstance(value, dict): return {str(key): normalize_manifest_value(item) for key, item in value.items()} return repr(value) def export_demucs_model_config(package: dict[str, Any], checkpoint_path: Path, signature: str) -> dict[str, Any]: klass = package["klass"] kwargs = {str(key): normalize_manifest_value(value) for key, value in package["kwargs"].items()} return { "model_type": "demucs", "class_name": getattr(klass, "__name__", str(klass)), "signature": signature, "checkpoint_file": checkpoint_path.name, "args": [normalize_manifest_value(value) for value in package["args"]], "kwargs": kwargs, "sources": normalize_manifest_value(kwargs.get("sources")), "samplerate": normalize_manifest_value(kwargs.get("samplerate")), "audio_channels": normalize_manifest_value(kwargs.get("audio_channels")), "segment": normalize_manifest_value(kwargs.get("segment")), } def export_demucs_state_tensors(package: dict[str, Any]) -> dict[str, torch.Tensor]: state = package["state"] if not isinstance(state, dict): raise RuntimeError(f"unexpected Demucs state payload type: {type(state)}") converted: dict[str, torch.Tensor] = {} for key, value in state.items(): if key == "__quantized": raise RuntimeError("Demucs quantized checkpoints are not supported by this converter yet") if not isinstance(value, torch.Tensor): raise RuntimeError(f"unexpected non-tensor entry in Demucs state: {key} -> {type(value)}") converted[str(key)] = value.detach().cpu().contiguous() return converted def load_demucs_bag_manifest(name: str) -> dict[str, Any] | None: payload = DEMUCS_BAGS.get(name) if payload is None: return None return dict(payload) def demucs_variant_signatures(name: str) -> list[str]: bag = load_demucs_bag_manifest(name) if bag is None: return [name] signatures = bag.get("models") if not isinstance(signatures, list) or not signatures: raise RuntimeError(f"Demucs bag manifest has invalid model list: {name}") return [str(signature) for signature in signatures] def ensure_demucs_checkpoints(repo_dir: Path, name: str) -> None: repo_dir.mkdir(parents=True, exist_ok=True) for signature in demucs_variant_signatures(name): if any(repo_dir.glob(f"{signature}*.th")): continue url = DEMUCS_REMOTE_MODELS.get(signature) if url is None: raise RuntimeError(f"no remote checkpoint URL found for Demucs signature: {signature}") filename = Path(url).name print(f"download {filename}") download_file(url, repo_dir / filename, None) def convert_demucs_single(signature: str, repo_dir: Path, output_root: Path) -> dict[str, Any]: checkpoint_path = find_demucs_checkpoint(repo_dir, signature) package = load_demucs_package(checkpoint_path) state = export_demucs_state_tensors(package) model_dir = output_root / signature model_dir.mkdir(parents=True, exist_ok=True) save_file(state, str(model_dir / "model.safetensors")) write_json(model_dir / "config.json", export_demucs_model_config(package, checkpoint_path, signature)) return { "signature": signature, "checkpoint_file": checkpoint_path.name, "output_dir": str(model_dir.relative_to(output_root)), "tensor_count": len(state), } def install_demucs_reference( package: ModelPackage, models_root: Path, overwrite: bool, source_dir_arg: str | None, variant_arg: str | None, ) -> Path: name = variant_arg or "htdemucs" target_dir = models_root / package.target_directory staging_root = models_root / ".engine_model_staging" staging_root.mkdir(parents=True, exist_ok=True) staging_dir = Path(tempfile.mkdtemp(prefix="htdemucs.", dir=staging_root)) try: if source_dir_arg: repo_dir = resolve_path(source_dir_arg) if not repo_dir.is_dir(): raise RuntimeError(f"checkpoint repo directory does not exist: {repo_dir}") else: repo_dir = staging_dir / "downloaded_checkpoints" ensure_demucs_checkpoints(repo_dir, name) bag = load_demucs_bag_manifest(name) if bag is None: converted = [convert_demucs_single(name, repo_dir, staging_dir)] manifest = { "model_type": "demucs_single", "name": name, "model": converted[0], } else: signatures = bag["models"] if not isinstance(signatures, list) or not signatures: raise RuntimeError(f"Demucs bag manifest has invalid model list: {name}") converted = [convert_demucs_single(str(signature), repo_dir, staging_dir) for signature in signatures] weights = normalize_manifest_value(bag.get("weights")) segment = normalize_manifest_value(bag.get("segment")) if len(converted) == 1 and weights is None and segment is None: manifest = { "model_type": "demucs_single_alias", "name": name, "alias_of": converted[0]["signature"], "model": converted[0], } else: manifest = { "model_type": "demucs_bag", "name": name, "models": converted, "weights": weights, "segment": segment, } write_json(staging_dir / "manifest.json", manifest) validate_required_files(package, staging_dir) if target_dir.exists(): if not overwrite: raise RuntimeError(f"model directory already exists: {target_dir}") shutil.rmtree(target_dir) target_dir.parent.mkdir(parents=True, exist_ok=True) staging_dir.rename(target_dir) return target_dir except Exception: shutil.rmtree(staging_dir, ignore_errors=True) raise finally: try: if staging_root.exists() and not any(staging_root.iterdir()): staging_root.rmdir() except OSError: pass def archive_member(archive: tarfile.TarFile, name: str) -> bytes: for member in archive.getmembers(): if not member.isfile(): continue normalized = member.name.removeprefix("./") if normalized != name: continue extracted = archive.extractfile(member) if extracted is None: break return extracted.read() raise RuntimeError(f"missing archive member: {name}") def tensor_state_dict(state: dict[str, Any]) -> dict[str, torch.Tensor]: tensors: dict[str, torch.Tensor] = {} seen_storages: set[int] = set() for key, value in state.items(): if not isinstance(key, str): raise RuntimeError(f"checkpoint contains a non-string tensor key: {key!r}") if not torch.is_tensor(value): raise RuntimeError(f"checkpoint state contains a non-tensor value at key: {key}") tensor = value.detach().cpu().contiguous() storage_ptr = tensor.untyped_storage().data_ptr() if storage_ptr in seen_storages: tensor = tensor.clone() storage_ptr = tensor.untyped_storage().data_ptr() seen_storages.add(storage_ptr) tensors[key] = tensor if not tensors: raise RuntimeError("refusing to write an empty safetensors file") return tensors def write_checked_safetensors(tensors: dict[str, torch.Tensor], output_path: Path, source_file: Path, overwrite: bool) -> None: if output_path.exists() and not overwrite: raise RuntimeError(f"output file already exists: {output_path}") output_path.parent.mkdir(parents=True, exist_ok=True) metadata = { "source_format": "pytorch", "source_file": str(source_file), "tensor_count": str(len(tensors)), } save_file(tensors, str(output_path), metadata=metadata) reloaded = load_file(str(output_path), device="cpu") if set(reloaded.keys()) != set(tensors.keys()): raise RuntimeError(f"saved safetensors key set does not match source: {output_path}") for key, value in tensors.items(): actual = reloaded[key] if actual.shape != value.shape: raise RuntimeError(f"saved tensor shape changed for {key}: {tuple(value.shape)} -> {tuple(actual.shape)}") if actual.dtype != value.dtype: raise RuntimeError(f"saved tensor dtype changed for {key}: {value.dtype} -> {actual.dtype}") if not torch.equal(actual, value): raise RuntimeError(f"saved tensor value changed for {key}") def yaml_to_json(value: Any) -> Any: if isinstance(value, dict): return {str(key): yaml_to_json(item) for key, item in value.items()} if isinstance(value, list): return [yaml_to_json(item) for item in value] return value def expect_map(value: Any, label: str) -> dict[str, Any]: if not isinstance(value, dict): raise RuntimeError(f"{label} must be a map") return {str(key): item for key, item in value.items()} def map_at(root: dict[str, Any], key: str) -> dict[str, Any]: return expect_map(root.get(key), key) def list_at(root: dict[str, Any], key: str) -> list[Any]: value = root.get(key) if not isinstance(value, list): raise RuntimeError(f"{key} must be a list") return value def int_at(root: dict[str, Any], key: str) -> int: value = root.get(key) if isinstance(value, int): return value if isinstance(value, float) and value.is_integer(): return int(value) raise RuntimeError(f"{key} must be an integer") def number_at(root: dict[str, Any], key: str, default: float | None = None) -> float: value = root.get(key) if value is None and default is not None: return default if isinstance(value, (int, float)): return float(value) raise RuntimeError(f"{key} must be a number") def string_at(root: dict[str, Any], key: str) -> str: value = root.get(key) if isinstance(value, str): return value raise RuntimeError(f"{key} must be a string") def string_list(value: Any, label: str) -> list[str]: if not isinstance(value, list): raise RuntimeError(f"{label} must be a list") return [str(item) for item in value] def jasper_block( block: dict[str, Any], *, flatten_triples: bool, include_dropout: bool, include_residual_mode: bool, se_metadata: str, ) -> dict[str, Any]: def maybe_flatten(key: str) -> Any: raw = block.get(key) if not flatten_triples: return raw if isinstance(raw, list) and len(raw) == 1: return raw[0] return raw out: dict[str, Any] = { "dilation": maybe_flatten("dilation") if block.get("dilation") is not None else 1, "filters": block.get("filters"), "kernel": maybe_flatten("kernel"), "repeat": block.get("repeat"), "residual": block.get("residual", False), "separable": block.get("separable", False), "stride": maybe_flatten("stride") if block.get("stride") is not None else 1, } if include_dropout and "dropout" in block: out["dropout"] = block["dropout"] if include_residual_mode: out["residual_mode"] = block.get("residual_mode") if "se" in block: out["se"] = block["se"] if se_metadata == "reduction_ratio" and "se" in block: out["se_reduction_ratio"] = block.get("se_reduction_ratio", 8) if se_metadata == "context_size" and "se_context_size" in block: out["se_context_size"] = block["se_context_size"] return out def jasper_blocks( blocks: list[Any], *, flatten_triples: bool, include_dropout: bool, include_residual_mode: bool, se_metadata: str, ) -> list[dict[str, Any]]: return [ jasper_block( expect_map(block, "jasper block"), flatten_triples=flatten_triples, include_dropout=include_dropout, include_residual_mode=include_residual_mode, se_metadata=se_metadata, ) for block in blocks ] def write_json(path: Path, payload: Any) -> None: path.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") def write_citrinet_sidecars(root: dict[str, Any], archive: tarfile.TarFile, output_dir: Path) -> None: tokenizer_name = "citrinet_256_tokenizer.model" vocab_name = "citrinet_256_vocab.txt" (output_dir / tokenizer_name).write_bytes(archive_member(archive, "tokenizer.model")) (output_dir / vocab_name).write_bytes(archive_member(archive, "vocab.txt")) preprocessor = map_at(root, "preprocessor") encoder = map_at(root, "encoder") decoder = map_at(root, "decoder") vocab_size = int_at(decoder, "num_classes") config = { "blank_id": vocab_size, "jasper": jasper_blocks( list_at(encoder, "jasper"), flatten_triples=True, include_dropout=True, include_residual_mode=True, se_metadata="reduction_ratio", ), "n_fft": int_at(preprocessor, "n_fft"), "n_mels": int_at(preprocessor, "features"), "normalize": string_at(preprocessor, "normalize"), "num_classes": vocab_size + 1, "pad_to": int_at(preprocessor, "pad_to"), "sample_rate": int_at(preprocessor, "sample_rate"), "target": str(root.get("target", "")), "tokenizer_model_file": tokenizer_name, "vocab_count": vocab_size, "vocab_file": vocab_name, "vocab_size": vocab_size, "window": string_at(preprocessor, "window"), "window_size": number_at(preprocessor, "window_size"), "window_stride": number_at(preprocessor, "window_stride"), } write_json(output_dir / "citrinet_256_config.json", config) def write_nemo_sidecars(config_kind: str, yaml_text: str, archive: tarfile.TarFile, output_dir: Path) -> None: root = yaml_to_json(yaml.safe_load(yaml_text)) if not isinstance(root, dict): raise RuntimeError("NeMo model_config.yaml root must be a map") output_dir.mkdir(parents=True, exist_ok=True) if config_kind == "citrinet": write_citrinet_sidecars(root, archive, output_dir) return raise RuntimeError(f"unsupported NeMo config kind: {config_kind}") def install_nemo_archive(package: ModelPackage, source: ConverterSource, models_root: Path, overwrite: bool) -> Path: target_dir = models_root / package.target_directory staging_root = models_root / ".engine_model_staging" staging_root.mkdir(parents=True, exist_ok=True) staging_dir = Path(tempfile.mkdtemp(prefix=f"{package.target_directory}.", dir=staging_root)) try: archive_path = staging_dir / Path(source.url).name print(f"download {archive_path.name}") download_file(source.url, archive_path, None) raw_bytes = archive_path.read_bytes() with tarfile.open(fileobj=io.BytesIO(raw_bytes), mode="r:*") as archive: weight_bytes = archive_member(archive, "model_weights.ckpt") config_bytes = archive_member(archive, "model_config.yaml") payload = torch.load(io.BytesIO(weight_bytes), map_location="cpu", weights_only=False) state = checkpoint_state_dict(payload) tensors = tensor_state_dict(state) write_checked_safetensors( tensors, staging_dir / source.output_weights_file, archive_path, overwrite=True, ) write_nemo_sidecars(source.config_kind, config_bytes.decode("utf-8"), archive, staging_dir) archive_path.unlink() validate_required_files(package, staging_dir) if target_dir.exists(): if not overwrite: raise RuntimeError(f"model directory already exists: {target_dir}") shutil.rmtree(target_dir) target_dir.parent.mkdir(parents=True, exist_ok=True) staging_dir.rename(target_dir) return target_dir except Exception: shutil.rmtree(staging_dir, ignore_errors=True) raise finally: try: if staging_root.exists() and not any(staging_root.iterdir()): staging_root.rmdir() except OSError: pass def install_converter( package: ModelPackage, source: ConverterSource, models_root: Path, overwrite: bool, source_file_arg: str | None, output_file_arg: str | None, source_dir_arg: str | None, variant_arg: str | None, ) -> Path: target_dir = models_root / package.target_directory if source.kind == "nemo_archive": return install_nemo_archive(package, source, models_root, overwrite) if source.kind == "demucs_reference": return install_demucs_reference(package, models_root, overwrite, source_dir_arg, variant_arg) if source.kind != "pytorch_to_safetensors": raise RuntimeError(f"unsupported converter kind: {source.kind}") if not source_file_arg: raise RuntimeError(f"{package.id} requires --source-file") source_file = resolve_path(source_file_arg) if not source_file.is_file(): raise RuntimeError(f"source checkpoint does not exist: {source_file}") if output_file_arg: output_path = resolve_path(output_file_arg) else: output_path = target_dir / package.required_files[0] payload = torch.load(source_file, map_location="cpu", weights_only=True) state = checkpoint_state_dict(payload) tensors = tensor_state_dict(state) write_checked_safetensors(tensors, output_path, source_file, overwrite) if output_path.parent == target_dir: validate_required_files(package, target_dir) return target_dir def command_list(args: argparse.Namespace) -> int: if args.json: print(json.dumps([package_payload(package) for package in CATALOG], indent=2)) else: print_package_table(CATALOG) return 0 def command_info(args: argparse.Namespace) -> int: package = PACKAGE_BY_ID.get(args.package_id) if package is None: raise RuntimeError(f"unknown package id: {args.package_id}") payload = package_payload(package) if args.json: print(json.dumps(payload, indent=2)) else: print(json.dumps(payload, indent=2)) return 0 def command_install(args: argparse.Namespace) -> int: package = PACKAGE_BY_ID.get(args.package_id) if package is None: raise RuntimeError(f"unknown package id: {args.package_id}") if package_needs_install_deps(package): _ensure_install_deps() models_root = resolve_path(args.models_root) models_root.mkdir(parents=True, exist_ok=True) source = package.source if isinstance(source, UnsupportedSource): raise RuntimeError(f"{package.id} is not installable: {source.reason}") if isinstance(source, SnapshotSource): install_path = install_snapshot(package, source, models_root, args.overwrite) elif isinstance(source, CompositeSnapshotSource): install_path = install_composite_snapshot(package, source, models_root, args.overwrite) else: install_path = install_converter( package, source, models_root, args.overwrite, args.source_file, args.output_file, args.source_dir, args.variant, ) print(f"installed {package.id} -> {install_path}") return 0 def main() -> int: args = parse_args() try: if args.command == "list": return command_list(args) if args.command == "info": return command_info(args) if args.command == "install": return command_install(args) raise RuntimeError(f"unsupported command: {args.command}") except HTTPError as ex: print(f"http error: {ex.code} {ex.reason}", file=sys.stderr) return 1 except Exception as ex: print(str(ex), file=sys.stderr) return 1 if __name__ == "__main__": raise SystemExit(main())