| |
| """Validate catalog integrity and exercise every model and every voice on CPU.""" |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| from pathlib import Path |
| from typing import Any, Iterator |
|
|
| import numpy as np |
| import onnx |
| import onnxruntime as ort |
| import soundfile as sf |
| from jsonschema import Draft202012Validator |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| VOICE_SHAPE = (510, 256) |
|
|
|
|
| def sha256(path: Path) -> str: |
| digest = hashlib.sha256() |
| with path.open("rb") as source: |
| for block in iter(lambda: source.read(1024 * 1024), b""): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def artifacts(value: Any) -> Iterator[dict[str, Any]]: |
| if isinstance(value, list): |
| for item in value: |
| yield from artifacts(item) |
| elif isinstance(value, dict): |
| if {"path", "sizeBytes", "sha256"} <= value.keys(): |
| yield value |
| for item in value.values(): |
| yield from artifacts(item) |
|
|
|
|
| def unique_by_id(values: list[dict[str, Any]], description: str) -> dict[str, dict[str, Any]]: |
| result = {value["id"]: value for value in values} |
| if len(result) != len(values): |
| raise RuntimeError(f"Duplicate {description} id") |
| return result |
|
|
|
|
| catalog = json.loads((ROOT / "catalog.json").read_text(encoding="utf-8")) |
| source_lock = json.loads((ROOT / "SOURCES.lock.json").read_text(encoding="utf-8")) |
| schema = json.loads( |
| (ROOT / "schema/kokoro-catalog.schema.json").read_text(encoding="utf-8") |
| ) |
| Draft202012Validator.check_schema(schema) |
| Draft202012Validator(schema).validate(catalog) |
|
|
| artifact_by_path: dict[str, dict[str, Any]] = {} |
| for artifact in artifacts(catalog): |
| previous = artifact_by_path.setdefault(artifact["path"], artifact) |
| if previous != artifact: |
| raise RuntimeError(f"Conflicting metadata for artifact {artifact['path']}") |
| for artifact in artifact_by_path.values(): |
| path = ROOT / artifact["path"] |
| if not path.is_file(): |
| raise FileNotFoundError(path) |
| if path.stat().st_size != artifact["sizeBytes"]: |
| raise RuntimeError(f"Size mismatch: {path}") |
| if sha256(path) != artifact["sha256"]: |
| raise RuntimeError(f"SHA-256 mismatch: {path}") |
|
|
| tokenizers = unique_by_id(catalog["tokenizers"], "tokenizer") |
| models = unique_by_id(catalog["models"], "model") |
| languages = unique_by_id(catalog["languages"], "language") |
| locked_model_sources = { |
| entry["repository"]: entry |
| for key in ("kokoroBase", "kokoroChinese", "languageModels") |
| for entry in [source_lock[key]] |
| } |
| configs: dict[str, dict[str, Any]] = {} |
| for tokenizer_id, tokenizer in tokenizers.items(): |
| configs[tokenizer_id] = json.loads( |
| (ROOT / tokenizer["artifact"]["path"]).read_text(encoding="utf-8") |
| ) |
| for model in models.values(): |
| if model["tokenizerId"] not in tokenizers: |
| raise RuntimeError( |
| f"{model['id']}: unknown tokenizer {model['tokenizerId']}" |
| ) |
| graph = onnx.load(ROOT / model["artifact"]["path"], load_external_data=False) |
| metadata = {item.key: item.value for item in graph.metadata_props} |
| if metadata.get("license") != "Apache-2.0": |
| raise RuntimeError(f"{model['id']}: ONNX license metadata is not Apache-2.0") |
| if metadata.get("source_repository") != model["source"]["repository"]: |
| raise RuntimeError(f"{model['id']}: source repository metadata mismatch") |
| if metadata.get("source_revision") != model["source"]["revision"]: |
| raise RuntimeError(f"{model['id']}: source revision metadata mismatch") |
| if metadata.get("source_checkpoint") != model["source"]["checkpoint"]: |
| raise RuntimeError(f"{model['id']}: source checkpoint metadata mismatch") |
| locked_source = locked_model_sources[model["source"]["repository"]] |
| expected_checkpoint_sha = locked_source["files"][model["source"]["checkpoint"]] |
| if metadata.get("source_checkpoint_sha256") != expected_checkpoint_sha: |
| raise RuntimeError(f"{model['id']}: source checkpoint checksum metadata mismatch") |
| validation = json.loads( |
| (ROOT / model["validation"]["path"]).read_text(encoding="utf-8") |
| ) |
| if validation["modelSha256"] != model["artifact"]["sha256"]: |
| raise RuntimeError(f"{model['id']}: validation model checksum mismatch") |
|
|
| voices: dict[tuple[str, str], dict[str, Any]] = {} |
| voice_pack_cache: dict[str, np.ndarray] = {} |
| for language in languages.values(): |
| language_voices = unique_by_id(language["voices"], f"voice in {language['id']}") |
| if language["defaultVoiceId"] not in language_voices: |
| raise RuntimeError(f"{language['id']}: unknown default voice") |
| if language["textFrontend"]["status"] == "bundled": |
| roles = [item["role"] for item in language["textFrontend"]["artifacts"]] |
| if roles.count("neural-model") != 1 or len(roles) != len(set(roles)): |
| raise RuntimeError(f"{language['id']}: invalid frontend artifact roles") |
| for voice_id, voice in language_voices.items(): |
| if voice["modelId"] not in models: |
| raise RuntimeError(f"{language['id']}/{voice_id}: unknown model") |
| path = ROOT / voice["artifact"]["path"] |
| if path.stat().st_size != VOICE_SHAPE[0] * VOICE_SHAPE[1] * 4: |
| raise RuntimeError(f"{language['id']}/{voice_id}: invalid voice byte size") |
| voice_pack = np.fromfile(path, dtype="<f4").reshape(VOICE_SHAPE) |
| if not np.isfinite(voice_pack).all(): |
| raise RuntimeError(f"{language['id']}/{voice_id}: non-finite voice values") |
| voice_pack_cache[voice["artifact"]["path"]] = voice_pack |
| voices[(language["id"], voice_id)] = voice |
|
|
| source_inventory = json.loads( |
| (ROOT / "sources/upstream-assets.json").read_text(encoding="utf-8") |
| ) |
| locked_manifest = source_lock["upstreamAssetManifest"] |
| if ( |
| locked_manifest["path"] != "sources/upstream-assets.json" |
| or locked_manifest["sha256"] != sha256(ROOT / locked_manifest["path"]) |
| ): |
| raise RuntimeError("SOURCES.lock.json does not match the upstream asset manifest") |
| if source_inventory["voiceCount"] != 157 or source_inventory["catalogVoiceCount"] != 159: |
| raise RuntimeError("Pinned upstream inventory must contain exactly 157 voices") |
|
|
| phonemis_onnx_lock = source_lock["phonemisOnnx"] |
| phonemis_onnx_manifest_path = ROOT / phonemis_onnx_lock["manifestPath"] |
| if phonemis_onnx_lock["manifestSha256"] != sha256(phonemis_onnx_manifest_path): |
| raise RuntimeError("SOURCES.lock.json does not match the Phonemis ONNX manifest") |
| phonemis_onnx_manifest = json.loads( |
| phonemis_onnx_manifest_path.read_text(encoding="utf-8") |
| ) |
| if len(phonemis_onnx_manifest["artifacts"]) != phonemis_onnx_lock["artifactCount"]: |
| raise RuntimeError("Phonemis ONNX manifest artifact count mismatch") |
| phonemis_sources = { |
| item["artifactPath"]: item for item in source_inventory["phonemis"] |
| } |
| for converted in phonemis_onnx_manifest["artifacts"]: |
| source = phonemis_sources.get(converted["sourcePath"]) |
| if source is None or source["sourceSha256"] != converted["sourceSha256"]: |
| raise RuntimeError( |
| f"Phonemis ONNX source mismatch for {converted['sourcePath']}" |
| ) |
| artifact = artifact_by_path.get(converted["path"]) |
| if artifact is None or any( |
| artifact[field] != converted[field] |
| for field in ("path", "sizeBytes", "sha256") |
| ): |
| raise RuntimeError( |
| f"Phonemis ONNX catalog mismatch for {converted['path']}" |
| ) |
| for source_voice in source_inventory["voices"]: |
| key = (source_voice["languageId"], source_voice["id"]) |
| if key not in voices: |
| raise RuntimeError(f"Source inventory voice missing from catalog: {key}") |
| artifact = voices[key]["artifact"] |
| if ( |
| artifact["path"] != source_voice["artifactPath"] |
| or artifact["sizeBytes"] != source_voice["artifactSizeBytes"] |
| or artifact["sha256"] != source_voice["artifactSha256"] |
| or voices[key]["modelId"] != source_voice["modelId"] |
| ): |
| raise RuntimeError(f"Source inventory mismatch for voice {key}") |
| for language_asset in source_inventory["languageModels"]: |
| if language_asset["role"] != "voice": |
| continue |
| artifact = artifact_by_path[language_asset["artifactPath"]] |
| if ( |
| artifact["sizeBytes"] != language_asset["artifactSizeBytes"] |
| or artifact["sha256"] != language_asset["artifactSha256"] |
| ): |
| raise RuntimeError( |
| f"Language model inventory mismatch for {language_asset['artifactPath']}" |
| ) |
| if len(voices) != 159: |
| raise RuntimeError(f"Expected 159 catalog voices, found {len(voices)}") |
|
|
| frontend_validation = json.loads( |
| (ROOT / "validation/frontend-validation.json").read_text(encoding="utf-8") |
| ) |
| frontend_cases = [ |
| *frontend_validation["phonemis"]["cases"], |
| *frontend_validation["misaki"]["cases"], |
| ] |
| bundled_languages = { |
| language["textFrontend"]["language"]: language |
| for language in languages.values() |
| if language["textFrontend"]["status"] == "bundled" |
| } |
| frontend_onnx_results = [] |
| frontend_session_options = ort.SessionOptions() |
| frontend_session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| for language_id, language in bundled_languages.items(): |
| neural_model = next( |
| item["artifact"] |
| for item in language["textFrontend"]["artifacts"] |
| if item["role"] == "neural-model" |
| ) |
| if neural_model.get("mediaType") != "application/onnx": |
| raise RuntimeError(f"{language_id}: neural frontend is not a prebuilt ONNX model") |
| frontend_session = ort.InferenceSession( |
| str(ROOT / neural_model["path"]), |
| sess_options=frontend_session_options, |
| providers=["CPUExecutionProvider"], |
| ) |
| logits = frontend_session.run( |
| ["logits"], |
| {"input_ids": np.asarray([[2]], dtype=np.int64)}, |
| )[0] |
| if logits.ndim != 3 or logits.shape[0] != 1 or logits.shape[2] != 116: |
| raise RuntimeError(f"{language_id}: unexpected frontend logits shape {logits.shape}") |
| if not np.isfinite(logits).all(): |
| raise RuntimeError(f"{language_id}: non-finite frontend logits") |
| frontend_onnx_results.append( |
| { |
| "language": language_id, |
| "path": neural_model["path"], |
| "sha256": neural_model["sha256"], |
| "logitsShape": list(logits.shape), |
| } |
| ) |
| phonemis_languages = { |
| case["language"] for case in frontend_validation["phonemis"]["cases"] |
| } |
| if phonemis_languages != set(bundled_languages): |
| raise RuntimeError("Frontend validation does not cover every bundled Phonemis language") |
| for case in frontend_cases: |
| if "tokenizerId" in case: |
| tokenizer_ids = {case["tokenizerId"]} |
| else: |
| language = bundled_languages[case["language"]] |
| tokenizer_ids = { |
| models[voice["modelId"]]["tokenizerId"] for voice in language["voices"] |
| } |
| for tokenizer_id in tokenizer_ids: |
| vocab = configs[tokenizer_id][catalog["runtime"]["tokenEncoding"]["vocabularyField"]] |
| missing = sorted(set(case["phonemes"]) - set(vocab)) |
| if missing or not case["tokenizerCompatible"]: |
| raise RuntimeError( |
| f"Frontend case {case['language']} is incompatible with {tokenizer_id}: {missing}" |
| ) |
|
|
| session_options = ort.SessionOptions() |
| session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| sessions: dict[str, ort.InferenceSession] = {} |
|
|
|
|
| def session_for(model_id: str) -> ort.InferenceSession: |
| if model_id not in sessions: |
| sessions[model_id] = ort.InferenceSession( |
| str(ROOT / models[model_id]["artifact"]["path"]), |
| sess_options=session_options, |
| providers=["CPUExecutionProvider"], |
| ) |
| return sessions[model_id] |
|
|
|
|
| def infer( |
| language_id: str, |
| voice_id: str, |
| phonemes: str, |
| ) -> tuple[np.ndarray, np.ndarray, int]: |
| voice = voices[(language_id, voice_id)] |
| model = models[voice["modelId"]] |
| config = configs[model["tokenizerId"]] |
| vocab = config[catalog["runtime"]["tokenEncoding"]["vocabularyField"]] |
| missing = sorted({phoneme for phoneme in phonemes if phoneme not in vocab}) |
| if missing: |
| raise RuntimeError( |
| f"{language_id}/{voice_id}: phonemes missing from {model['tokenizerId']}: {missing}" |
| ) |
| if not 1 <= len(phonemes) <= catalog["runtime"]["maxPhonemeCodePoints"]: |
| raise RuntimeError(f"{language_id}/{voice_id}: invalid phoneme length") |
| token_ids = [vocab[phoneme] for phoneme in phonemes] |
| input_ids = np.asarray([[0, *token_ids, 0]], dtype=np.int64) |
| voice_pack = voice_pack_cache[voice["artifact"]["path"]] |
| waveform, duration = session_for(model["id"]).run( |
| None, |
| { |
| "input_ids": input_ids, |
| "style": voice_pack[len(phonemes) - 1].reshape(1, 256), |
| "speed": np.asarray([1.0], dtype=np.float32), |
| }, |
| ) |
| if not np.isfinite(waveform).all(): |
| raise RuntimeError(f"{language_id}/{voice_id}: non-finite waveform") |
| if duration.shape != (input_ids.shape[1],): |
| raise RuntimeError( |
| f"{language_id}/{voice_id}: unexpected duration shape {duration.shape}" |
| ) |
| expected_samples = int(duration.sum()) * catalog["runtime"]["durationToSamplesFactor"] |
| if waveform.shape != (expected_samples,): |
| raise RuntimeError( |
| f"{language_id}/{voice_id}: expected {expected_samples} samples, got {waveform.shape}" |
| ) |
| return waveform, duration, int(input_ids.shape[1]) |
|
|
|
|
| cases = json.loads((ROOT / "validation/cases.json").read_text(encoding="utf-8")) |
| case_results = [] |
| for case in cases: |
| if (case["language"], case["voice"]) not in voices: |
| raise RuntimeError(f"{case['id']}: unknown language or voice") |
| waveform, _duration, input_tokens = infer( |
| case["language"], case["voice"], case["phonemes"] |
| ) |
| if "previewPath" in case: |
| preview_path = ROOT / case["previewPath"] |
| preview_path.parent.mkdir(parents=True, exist_ok=True) |
| sf.write(preview_path, waveform, catalog["audio"]["sampleRateHz"], subtype="PCM_16") |
| case_results.append( |
| { |
| "id": case["id"], |
| "language": case["language"], |
| "voice": case["voice"], |
| "modelId": voices[(case["language"], case["voice"])]["modelId"], |
| "phonemeCodePoints": len(case["phonemes"]), |
| "inputTokens": input_tokens, |
| "outputSamples": int(waveform.size), |
| "durationSeconds": waveform.size / catalog["audio"]["sampleRateHz"], |
| "peakAbsoluteAmplitude": float(np.abs(waveform).max(initial=0.0)), |
| } |
| ) |
|
|
| voice_results = [] |
| for language in languages.values(): |
| for voice in language["voices"]: |
| waveform, _duration, input_tokens = infer(language["id"], voice["id"], "a") |
| voice_results.append( |
| { |
| "language": language["id"], |
| "voice": voice["id"], |
| "modelId": voice["modelId"], |
| "inputTokens": input_tokens, |
| "outputSamples": int(waveform.size), |
| "peakAbsoluteAmplitude": float(np.abs(waveform).max(initial=0.0)), |
| } |
| ) |
|
|
| boundary_results = [] |
| for model_id, model in models.items(): |
| compatible = next( |
| (language_id, voice_id) |
| for (language_id, voice_id), voice in voices.items() |
| if voice["modelId"] == model_id |
| ) |
| config = configs[model["tokenizerId"]] |
| vocab = config[catalog["runtime"]["tokenEncoding"]["vocabularyField"]] |
| if "a" not in vocab: |
| raise RuntimeError(f"{model_id}: tokenizer has no boundary-test token 'a'") |
| for phoneme_count in (1, 128, catalog["runtime"]["maxPhonemeCodePoints"]): |
| waveform, _duration, input_tokens = infer( |
| compatible[0], compatible[1], "a" * phoneme_count |
| ) |
| boundary_results.append( |
| { |
| "modelId": model_id, |
| "phonemeCodePoints": phoneme_count, |
| "inputTokens": input_tokens, |
| "outputSamples": int(waveform.size), |
| } |
| ) |
|
|
| report = { |
| "catalogVersion": catalog["catalogVersion"], |
| "catalogSchemaValidated": True, |
| "onnxRuntimeVersion": ort.__version__, |
| "artifactsVerified": len(artifact_by_path), |
| "modelsValidated": len(models), |
| "languageProfilesValidated": len(languages), |
| "voicesValidated": len(voice_results), |
| "frontendCasesValidated": len(frontend_cases), |
| "frontendOnnxModelsValidated": len(frontend_onnx_results), |
| "frontendOnnxSmokeTests": frontend_onnx_results, |
| "cases": case_results, |
| "voiceSmokeTests": voice_results, |
| "boundaryCases": boundary_results, |
| } |
| (ROOT / "validation/package-validation.json").write_text( |
| json.dumps(report, ensure_ascii=False, indent=2) + "\n", |
| encoding="utf-8", |
| ) |
| print( |
| json.dumps( |
| { |
| "catalogVersion": report["catalogVersion"], |
| "artifactsVerified": report["artifactsVerified"], |
| "modelsValidated": report["modelsValidated"], |
| "languageProfilesValidated": report["languageProfilesValidated"], |
| "voicesValidated": report["voicesValidated"], |
| "frontendCasesValidated": report["frontendCasesValidated"], |
| "frontendOnnxModelsValidated": report["frontendOnnxModelsValidated"], |
| "naturalCasesValidated": len(case_results), |
| "boundaryCasesValidated": len(boundary_results), |
| }, |
| indent=2, |
| ) |
| ) |
|
|
