From 98bd2e1d38ed6fe47e2256f30ee135a647673642 Mon Sep 17 00:00:00 2001 From: owenawsong Date: Sun, 26 Jul 2026 14:34:45 -0700 Subject: [PATCH] Add public Inflect adaptation toolkit --- finetune/.gitignore | 11 + finetune/CONTRACT.md | 90 + finetune/README.md | 227 +++ finetune/docs/CUSTOM_G2P.md | 68 + finetune/docs/DATA_QUALITY.md | 88 + finetune/docs/LANGUAGES.md | 81 + finetune/docs/MANIFESTS.md | 90 + finetune/docs/RESPONSIBLE_USE.md | 76 + finetune/docs/SCOPE.md | 74 + finetune/docs/TRAINING.md | 173 ++ finetune/docs/TROUBLESHOOTING.md | 117 ++ finetune/examples/custom_frontend_hook.py | 34 + finetune/examples/dataset_manifest.csv | 4 + finetune/examples/evaluation_manifest.jsonl | 3 + finetune/examples/export_language_package.py | 47 + finetune/examples/export_options.json | 13 + finetune/examples/language_adaptation.json | 16 + .../prephonemized_evaluation_manifest.jsonl | 1 + finetune/examples/self_test_export_eval.py | 374 ++++ .../examples/transcript_evaluator_plugin.py | 20 + finetune/inflect_finetune/__init__.py | 3 + finetune/inflect_finetune/__main__.py | 5 + finetune/inflect_finetune/audio.py | 198 ++ finetune/inflect_finetune/audit.py | 332 ++++ finetune/inflect_finetune/checkpoint.py | 522 +++++ finetune/inflect_finetune/cli.py | 373 ++++ finetune/inflect_finetune/evaluation.py | 521 +++++ finetune/inflect_finetune/exporting.py | 1686 +++++++++++++++++ finetune/inflect_finetune/frontend.py | 499 +++++ finetune/inflect_finetune/manifest.py | 250 +++ finetune/inflect_finetune/modeling.py | 280 +++ finetune/inflect_finetune/monotonic_align.py | 78 + finetune/inflect_finetune/prepare.py | 436 +++++ finetune/inflect_finetune/presets/__init__.py | 57 + finetune/inflect_finetune/reporting.py | 147 ++ finetune/inflect_finetune/symbols.py | 174 ++ finetune/inflect_finetune/training.py | 869 +++++++++ finetune/inflect_finetune/training_data.py | 184 ++ finetune/presets/balanced.json | 11 + finetune/presets/micro-12gb.json | 11 + finetune/presets/nano-8gb.json | 11 + finetune/pyproject.toml | 45 + finetune/tests/test_checkpoint_identity.py | 269 +++ finetune/tests/test_cli.py | 150 ++ finetune/tests/test_data_frontend_prepare.py | 313 +++ finetune/tests/test_model_resolution.py | 30 + finetune/tests/test_monotonic_align.py | 42 + finetune/tests/test_packaged_presets.py | 68 + finetune/tests/test_public_safety.py | 81 + finetune/tests/test_training_stages.py | 29 + 50 files changed, 9281 insertions(+) create mode 100644 finetune/.gitignore create mode 100644 finetune/CONTRACT.md create mode 100644 finetune/README.md create mode 100644 finetune/docs/CUSTOM_G2P.md create mode 100644 finetune/docs/DATA_QUALITY.md create mode 100644 finetune/docs/LANGUAGES.md create mode 100644 finetune/docs/MANIFESTS.md create mode 100644 finetune/docs/RESPONSIBLE_USE.md create mode 100644 finetune/docs/SCOPE.md create mode 100644 finetune/docs/TRAINING.md create mode 100644 finetune/docs/TROUBLESHOOTING.md create mode 100644 finetune/examples/custom_frontend_hook.py create mode 100644 finetune/examples/dataset_manifest.csv create mode 100644 finetune/examples/evaluation_manifest.jsonl create mode 100644 finetune/examples/export_language_package.py create mode 100644 finetune/examples/export_options.json create mode 100644 finetune/examples/language_adaptation.json create mode 100644 finetune/examples/prephonemized_evaluation_manifest.jsonl create mode 100644 finetune/examples/self_test_export_eval.py create mode 100644 finetune/examples/transcript_evaluator_plugin.py create mode 100644 finetune/inflect_finetune/__init__.py create mode 100644 finetune/inflect_finetune/__main__.py create mode 100644 finetune/inflect_finetune/audio.py create mode 100644 finetune/inflect_finetune/audit.py create mode 100644 finetune/inflect_finetune/checkpoint.py create mode 100644 finetune/inflect_finetune/cli.py create mode 100644 finetune/inflect_finetune/evaluation.py create mode 100644 finetune/inflect_finetune/exporting.py create mode 100644 finetune/inflect_finetune/frontend.py create mode 100644 finetune/inflect_finetune/manifest.py create mode 100644 finetune/inflect_finetune/modeling.py create mode 100644 finetune/inflect_finetune/monotonic_align.py create mode 100644 finetune/inflect_finetune/prepare.py create mode 100644 finetune/inflect_finetune/presets/__init__.py create mode 100644 finetune/inflect_finetune/reporting.py create mode 100644 finetune/inflect_finetune/symbols.py create mode 100644 finetune/inflect_finetune/training.py create mode 100644 finetune/inflect_finetune/training_data.py create mode 100644 finetune/presets/balanced.json create mode 100644 finetune/presets/micro-12gb.json create mode 100644 finetune/presets/nano-8gb.json create mode 100644 finetune/pyproject.toml create mode 100644 finetune/tests/test_checkpoint_identity.py create mode 100644 finetune/tests/test_cli.py create mode 100644 finetune/tests/test_data_frontend_prepare.py create mode 100644 finetune/tests/test_model_resolution.py create mode 100644 finetune/tests/test_monotonic_align.py create mode 100644 finetune/tests/test_packaged_presets.py create mode 100644 finetune/tests/test_public_safety.py create mode 100644 finetune/tests/test_training_stages.py diff --git a/finetune/.gitignore b/finetune/.gitignore new file mode 100644 index 0000000..6633189 --- /dev/null +++ b/finetune/.gitignore @@ -0,0 +1,11 @@ +__pycache__/ +*.py[cod] +.pytest_cache/ +.ruff_cache/ +build/ +dist/ +*.egg-info/ +.venv/ +prepared/ +runs/ +exports/ diff --git a/finetune/CONTRACT.md b/finetune/CONTRACT.md new file mode 100644 index 0000000..430b8c8 --- /dev/null +++ b/finetune/CONTRACT.md @@ -0,0 +1,90 @@ +# Inflect adaptation toolkit contract + +This directory is a public, generic warm-start adaptation toolkit. It is not +the private recipe used to produce the Inflect v2 release checkpoints. + +## Supported result + +The toolkit produces a separate, fixed-voice checkpoint for one configured +language. The dataset speaker becomes the checkpoint voice. It does not add +runtime-selectable speakers or languages to the original checkpoint. + +## Public workflow + +```text +python -m inflect_finetune prepare ... +python -m inflect_finetune audit ... +python -m inflect_finetune train ... +python -m inflect_finetune export ... +``` + +Every command must support `--help`, fail with an actionable message, and write +machine-readable reports alongside human-readable summaries. + +## Prepared dataset layout + +```text +prepared/ + dataset.json + symbols.json + train.jsonl + validation.jsonl + audio/ +``` + +Each JSONL row contains: + +```json +{ + "audio": "audio/example.wav", + "text": "Original transcript.", + "normalized_text": "Normalized transcript.", + "phonemes": "phoneme string", + "duration_seconds": 2.34 +} +``` + +Prepared rows may also preserve `id`, `speaker`, `group_id`, and +`group_field`. Source manifests may provide `session` as a split boundary and +`phonemes` for prephonemized preparation. + +`dataset.json` records the language, sample rate, frontend, source-manifest +hash, frontend source/metadata hashes where applicable, speaker, split seed, +row counts, and aggregate diagnostics. `symbols.json` records the ordered +symbol inventory and its relationship to the base inventory. + +## Checkpoint migration + +Checkpoint migration copies all shape-compatible generator weights. Text +embedding rows are copied by symbol string, not by numeric index. Newly added +symbols receive deterministic initialization. Training-only discriminators, +optimizers, and schedulers are initialized from public generic defaults. + +## Safety boundaries + +Public code must not contain or infer: + +- private corpus paths, transcripts, or generated audio; +- teacher-model names or corpus-generation prompts; +- internal filtering thresholds unrelated to generic audio validity; +- the original base-model curriculum, search history, or checkpoint ranking; +- credentials, rental identifiers, or private storage locations. + +## Validation gates + +At minimum, automated tests must cover: + +- manifest parsing and path traversal rejection; +- deterministic splitting; +- group, duplicate-audio, and duplicate-transcript leakage prevention; +- audio-format validation; +- language/frontend failures; +- phoneme coverage and unknown-symbol reporting; +- embedding migration by symbol identity; +- checkpoint save/resume behavior; +- inference-only export loadability; +- language-aware frontend packaging and refusal of silent English fallback; +- ONNX Runtime parity when ONNX export is requested. + +A release may call language adaptation experimental until at least one +non-English end-to-end run passes preparation, training, export, and inference. diff --git a/finetune/README.md b/finetune/README.md new file mode 100644 index 0000000..8bdfdb6 --- /dev/null +++ b/finetune/README.md @@ -0,0 +1,227 @@ +# Inflect adaptation toolkit + +Public tools for warm-starting Inflect v2 with user-owned speech data. One run +creates one fixed voice for one configured language. + +This is a generic compatible trainer, not the private process used to create +the official Inflect v2 checkpoints. It does not include the private corpus, +corpus-generation process, original curriculum, hyperparameter search, or +internal checkpoint-selection process. + +## What is implemented + +- JSONL and CSV manifests with strict path and audio validation +- deterministic 24 kHz mono preparation +- leakage-safe train/validation splitting by transcript and recording group +- eSpeak, prephonemized, and explicit custom Python frontends +- symbol-aware embedding migration from Micro or Nano +- staged generator/discriminator training with AMP and accumulation +- atomic checkpoints and strict same-run resume validation +- held-out waveform diagnostics and optional transcript evaluators +- inference-only PyTorch and ONNX packages +- language-aware deployment frontends with no silent English fallback + +The software path has been tested end to end, including a real Nano CUDA +training step, resume, strict PyTorch load, and ONNX Runtime parity. New +language and new voice quality remain experimental: data quality, phoneme +coverage, and fluent-speaker review determine whether an adaptation is useful. + +Read [CONTRACT.md](CONTRACT.md) and [scope](docs/SCOPE.md) before starting. + +## Install + +From this directory: + +```bash +python -m venv .venv +python -m pip install --upgrade pip +python -m pip install -e . +``` + +Activate on PowerShell with: + +```powershell +.\.venv\Scripts\Activate.ps1 +``` + +For ONNX export and parity checks: + +```bash +python -m pip install -e ".[onnx]" +``` + +Every command is available as either `inflect-adapt` or +`python -m inflect_finetune`. + +## 1. Create a manifest + +JSONL and CSV are supported. A minimal JSONL row is: + +```json +{"audio":"session01/000001.wav","text":"Buenos dias.","speaker":"voice-a","session":"session01"} +``` + +`audio` and `text` are required. `speaker`, `session`, `group_id`, `id`, and +`phonemes` are documented in [MANIFESTS.md](docs/MANIFESTS.md). + +Use one consenting speaker per dataset. Put clips cut from the same source +recording in the same `group_id` or `session` so they cannot leak across +splits. + +## 2. Prepare and audit + +For an eSpeak-supported language: + +```bash +inflect-adapt prepare \ + --manifest data/metadata.jsonl \ + --audio-root data/audio \ + --language es \ + --frontend espeak \ + --output prepared/es + +inflect-adapt audit --dataset prepared/es +``` + +Preparation writes converted audio, ordered symbols, frontend identity, hashes, +and deterministic nonempty train/validation splits. Do not train until audit +passes and a fluent speaker has inspected representative normalized text and +phonemes. + +If eSpeak is unsuitable, use prephonemized rows or the documented +[custom frontend hook](docs/CUSTOM_G2P.md). + +## 3. Train + +Start with Micro unless minimum footprint is the primary goal: + +```bash +inflect-adapt train \ + --base owensong/Inflect-Micro-v2 \ + --dataset prepared/es \ + --preset micro-12gb \ + --output runs/es-micro +``` + +Available presets are `balanced`, `micro-12gb`, and `nano-8gb`. They are +starting points, not memory or quality guarantees. CLI flags override preset +values exactly: + +```bash +inflect-adapt train \ + --base nano \ + --dataset prepared/es \ + --preset nano-8gb \ + --batch-size 1 \ + --gradient-accumulation-steps 12 \ + --output runs/es-nano +``` + +The public release checkpoint contains inference weights only. Training-only +posterior and discriminator components are initialized by this toolkit, and +new symbol embeddings are initialized deterministically. + +## 4. Resume safely + +```bash +inflect-adapt train \ + --base owensong/Inflect-Micro-v2 \ + --dataset prepared/es \ + --preset micro-12gb \ + --output runs/es-micro \ + --resume runs/es-micro/checkpoints/latest.pth +``` + +Resume is accepted only for the same run identity. Changes to the base model, +prepared data, symbols, frontend, public optimizer schema, or relevant +configuration are rejected. + +The final resumable checkpoint is +`runs/es-micro/checkpoints/adaptation-final.pth`. Step checkpoints and held-out +audio are written at the configured intervals. The toolkit does not label a +checkpoint "best"; select one using a declared held-out process. + +## 5. Export + +Export the selected training checkpoint together with the exact prepared +frontend metadata: + +```bash +inflect-adapt export \ + --checkpoint runs/es-micro/checkpoints/adaptation-final.pth \ + --prepared-dataset prepared/es \ + --format pytorch \ + --output exports/es-micro +``` + +For ONNX: + +```bash +inflect-adapt export \ + --checkpoint runs/es-micro/checkpoints/adaptation-final.pth \ + --prepared-dataset prepared/es \ + --format onnx \ + --output exports/es-micro-onnx +``` + +For a custom frontend, also pass the same source file: + +```bash +inflect-adapt export \ + --checkpoint runs/custom/checkpoints/adaptation-final.pth \ + --prepared-dataset prepared/custom \ + --frontend-hook my_frontend.py \ + --format onnx \ + --output exports/custom +``` + +Export strips posterior, discriminator, optimizer, scheduler, scaler, RNG, and +other training-only state. It writes `frontend.json`, `symbols.json`, +checksums, a report, a deployable runtime, and optional ONNX graphs. Adapted +checkpoints without frontend metadata are rejected instead of silently using +the release English frontend. + +## 6. Evaluate the exported package + +Use original text for eSpeak/custom packages or include `phonemes` in each +evaluation row for a prephonemized package: + +```bash +inflect-adapt evaluate \ + --model-dir exports/es-micro \ + --manifest prepared/es/validation.jsonl \ + --output evaluations/es-micro \ + --max-samples 100 +``` + +The report covers waveform duration, silence, clipping, peak, RMS, DC offset, +and non-finite values. An optional transcript-evaluator hook can add ASR or +other metrics. These diagnostics do not replace blind listening or +fluent-speaker review. + +## Documentation + +- [Supported scope](docs/SCOPE.md) +- [Manifests and split safety](docs/MANIFESTS.md) +- [Data quality](docs/DATA_QUALITY.md) +- [Languages and symbols](docs/LANGUAGES.md) +- [Training, checkpoints, evaluation, and export](docs/TRAINING.md) +- [Custom G2P/frontend hooks](docs/CUSTOM_G2P.md) +- [Troubleshooting](docs/TROUBLESHOOTING.md) +- [Consent and responsible use](docs/RESPONSIBLE_USE.md) + +## Release gate + +Before publishing an adapted checkpoint: + +1. Re-run preparation and audit from the source manifest. +2. Confirm groups and duplicate transcripts do not cross the split. +3. Compare several checkpoints on untouched held-out text. +4. Load the exported package in a clean environment. +5. Run its checksum, inference, and ONNX parity checks where applicable. +6. Have fluent speakers review pronunciation and naturalness. +7. Publish data provenance, consent, language, voice, frontend, base model, + toolkit version, evaluation method, and limitations. + +Passing the software checks does not establish naturalness, voice identity, or +acceptable pronunciation. diff --git a/finetune/docs/CUSTOM_G2P.md b/finetune/docs/CUSTOM_G2P.md new file mode 100644 index 0000000..7918805 --- /dev/null +++ b/finetune/docs/CUSTOM_G2P.md @@ -0,0 +1,68 @@ +# Custom G2P and frontend hooks + +Use a custom frontend when eSpeak does not provide acceptable normalization or +phonemization. A frontend controls normalization, punctuation, word +boundaries, and the exact symbol stream consumed by the generator. + +## Hook contract + +Pass a trusted factory as either `module:callable` or `file.py:function`: + +```bash +python -m inflect_finetune prepare \ + --manifest data/metadata.jsonl \ + --audio-root data/audio \ + --language my-language \ + --frontend custom \ + --frontend-hook ./my_frontend.py:create_frontend \ + --output prepared/my-language +``` + +The factory may accept no arguments or the keyword argument `language`. It +must return an object with: + +```python +normalize(text: str) -> str +phonemize(normalized_text: str) -> str +symbols() -> list[str] +metadata() -> dict +``` + +`symbols()` must return a nonempty ordered list of unique, one-character +strings. `metadata()` must contain `name`, `version`, `language`, and +`configuration`. + +The hook is trusted Python code and executes during preparation. Do not run a +hook from an untrusted source. + +## Reproducibility checks + +Preparation calls normalization and phonemization twice and rejects +nondeterministic output. It also rejects empty text, control characters, and +symbols not declared by the hook. + +The prepared dataset stores hashes of the hook source and declared metadata. +Export requires the matching hook source for custom frontends, verifies those +hashes, and copies the hook into the deployment package. It will not silently +replace a custom or non-English frontend with the release English frontend. + +## Frontend validation + +Before training, test: + +- ordinary sentences and every target phoneme; +- punctuation and sentence boundaries; +- numbers, dates, currencies, and abbreviations; +- names, loanwords, and mixed scripts; +- unsupported and empty input; +- repeated calls for exact determinism. + +Have fluent speakers inspect both normalized text and phonemes. A technically +valid symbol stream can still encode the wrong pronunciation. + +## Deployment requirements + +An adapted package must include the exact frontend needed for inference. +Third-party dictionaries or models remain subject to their own licenses and +must be packaged or documented separately. Do not describe an export as +self-contained if its frontend requires an external artifact. diff --git a/finetune/docs/DATA_QUALITY.md b/finetune/docs/DATA_QUALITY.md new file mode 100644 index 0000000..7ff2820 --- /dev/null +++ b/finetune/docs/DATA_QUALITY.md @@ -0,0 +1,88 @@ +# Data quality + +Adaptation quality is constrained by the corpus. More hours do not compensate +for inaccurate transcripts, inconsistent speakers, clipping, or poor phoneme +coverage. + +## Required properties + +Use recordings that are: + +- legally usable for training and redistribution under the intended terms; +- spoken by one consenting speaker for a fixed-voice checkpoint; +- paired with manually verified transcripts; +- mono or safely convertible to mono; +- consistently recorded, without changing microphones or aggressive effects; +- free from clipping, dropouts, corruption, background music, and overlapping + speakers; +- trimmed without cutting initial consonants, breaths needed for natural + phrasing, or sentence endings; +- diverse in phonemes, word positions, sentence lengths, punctuation, and + prosodic patterns. + +Keep raw source audio unchanged. Preparation should write converted copies under +the prepared dataset rather than destructively replacing source files. + +## Coverage matters + +Inspect the audit report for: + +- phonemes absent or rare in training; +- symbols introduced by only one transcript; +- validation phonemes not represented in training; +- repeated sentence templates; +- narrow pitch or duration distributions; +- unusually long or short clips; +- names, numbers, abbreviations, and punctuation patterns relevant to the + intended use. + +An eSpeak frontend producing a phoneme does not mean the model has enough data +to learn it. Newly initialized symbol rows need repeated, acoustically clear +examples in varied contexts. + +## Recording consistency + +Room tone, microphone frequency response, denoising, compression, and loudness +changes can become part of the learned voice. Avoid mixing studio audio, +telephone audio, and heavily processed clips unless that variation is an +intentional target and is evaluated. + +Do not apply strong denoising or de-essing blindly. Processing can create +musical noise, phase artifacts, dull consonants, or unstable sibilance that a +small decoder reproduces prominently. + +## Automated checks are diagnostics + +The preparation and audit stages report: + +- decode and sample-rate failures; +- channel count and duration; +- peak level, clipping indicators, silence, and non-finite samples; +- duplicate audio content and transcripts crossing split boundaries; +- transcript and frontend failures; +- unknown symbols and phoneme coverage; +- split statistics and source-manifest hash. + +Configured duration and structural thresholds are recorded. A clip passing +automated checks does not prove that its transcript, speaker identity, +pronunciation, or audio quality is correct. + +## Manual review + +Listen to a random sample, every flagged clip, and all validation clips. Check +the beginning and ending of each clip, consonant clarity, background sound, +speaker consistency, and transcript agreement. + +Before a long run, train a short smoke run and listen to held-out synthesis. +Stop if the model develops severe buzz, metallic resonance, clipped endings, +identity collapse, unintelligible phonemes, or unstable duration. + +## Data volume + +This toolkit does not promise a minimum number of minutes or hours that will +work for every language and speaker. Required data depends on phoneme coverage, +recording consistency, desired quality, distance from the base language and +voice, and which parts of the generator must adapt. + +Any future presets describing data volume must be validated experimentally and +must not be interpreted as quality guarantees. diff --git a/finetune/docs/LANGUAGES.md b/finetune/docs/LANGUAGES.md new file mode 100644 index 0000000..cc34e1e --- /dev/null +++ b/finetune/docs/LANGUAGES.md @@ -0,0 +1,81 @@ +# Languages and frontend configuration + +## One configured language per checkpoint + +An adaptation run produces a checkpoint tied to the language and symbol +inventory recorded in its prepared dataset. It does not make the original +checkpoint multilingual, and it does not expose a runtime language selector. + +For example, separate English and Spanish runs produce separate English and +Spanish checkpoints. Combining their files does not produce a bilingual model. + +## eSpeak frontend + +The built-in frontend uses `phonemizer` with eSpeak NG: + +```bash +python -m inflect_finetune prepare \ + --manifest data/metadata.jsonl \ + --audio-root data/audio \ + --language es \ + --frontend espeak \ + --output prepared/es +``` + +Use an eSpeak language or voice code appropriate to the transcripts and +recordings. Language codes, stress behavior, punctuation preservation, and +normalization must be recorded in `dataset.json`. + +Before training: + +1. Confirm eSpeak supports the requested code in the installed version. +2. Inspect normalized text and phonemes for a representative sample. +3. Ask a fluent speaker to review names, loanwords, abbreviations, and numbers. +4. Run `audit` and resolve unknown symbols. +5. Verify validation phonemes are covered by training. + +## Symbol inventory migration + +The base checkpoint and prepared language may use different symbols. Migration +must: + +- copy shape-compatible generator weights; +- match text embedding rows by symbol string, never numeric position; +- deterministically initialize newly added symbol rows; +- record copied, new, and unused symbols; +- reject ambiguous or duplicate symbol definitions; +- preserve the prepared symbol order in the exported package. + +New symbol initialization makes training possible; it does not provide a +pronunciation. The corpus must teach the acoustic realization and timing. + +## Normalization is language-specific + +Numbers, currencies, dates, abbreviations, casing, punctuation, and symbols +cannot be normalized reliably with one universal rule set. Users must verify +that normalized text matches what the speaker says. + +Do not reuse English-specific normalization for a different language without +review. If the built-in frontend cannot represent the intended reading, use a +custom frontend rather than editing prepared phonemes by hand. + +## Quality expectations + +Availability of an eSpeak voice is not evidence that Inflect will train well on +that language. Languages with substantially different phonology, writing +systems, timing, or prosody may require more data, frontend work, model +capacity, and optimization changes. + +Nano has less capacity and less tolerance for poor coverage than Micro. Use +Micro for the first adaptation attempt unless footprint is the primary +constraint and the Micro workflow has already been validated. + +Every release should state: + +- language and regional variety; +- frontend name and version; +- symbol inventory; +- speaker and corpus provenance; +- training and validation quantities; +- known pronunciation limitations; +- whether fluent speakers evaluated held-out output. diff --git a/finetune/docs/MANIFESTS.md b/finetune/docs/MANIFESTS.md new file mode 100644 index 0000000..3fa71a2 --- /dev/null +++ b/finetune/docs/MANIFESTS.md @@ -0,0 +1,90 @@ +# Manifests and prepared datasets + +## Source formats + +UTF-8 JSON Lines and CSV are supported. Every row requires: + +| Field | Type | Meaning | +| --- | --- | --- | +| `audio` | string | Relative path below `--audio-root` | +| `text` | string | Exact transcript of the recording | + +Supported optional fields: + +| Field | Meaning | +| --- | --- | +| `id` | Stable user identifier preserved in prepared metadata | +| `speaker` | Consistency metadata; all nonempty values must identify one speaker | +| `group_id` | Strongest split boundary, such as source recording or chapter | +| `session` | Split boundary used when `group_id` is absent | +| `phonemes` | Required on every row with `--frontend prephonemized` | + +JSONL example: + +```json +{"audio":"session01/000001.wav","text":"Buenos dias.","speaker":"voice-a","session":"session01"} +{"audio":"session02/000002.wav","text":"El tren llega a las nueve.","speaker":"voice-a","session":"session02"} +``` + +CSV example: + +```csv +audio,text,speaker,session +session01/000001.wav,Buenos dias.,voice-a,session01 +session02/000002.wav,El tren llega a las nueve.,voice-a,session02 +``` + +## Path and content rules + +- Paths must stay below the declared audio root. +- Absolute paths and `..` traversal are rejected. +- Missing, duplicate, undecodable, empty, or non-finite audio is rejected. +- Audio is converted deterministically to 24 kHz mono WAV. +- Reused audio content is rejected, even if it appears under another filename. +- Transcripts should describe exactly what is spoken. +- Empty transcripts and unsupported frontend output are rejected. + +Do not place comments, timestamps, markup, or speaker directions in `text` +unless the chosen frontend intentionally handles them. + +## Leakage-safe splitting + +Preparation creates deterministic, nonempty train and validation splits from +the recorded split seed. + +Rows connected by the following relationships stay in the same split: + +- identical normalized transcript; +- the same `group_id`; +- the same `session` when no `group_id` is present. + +`speaker` is not a split group. A fixed-voice corpus normally has the same +speaker on every row, so treating it as a group would make a validation split +impossible. Multiple nonempty speaker values are rejected instead. + +Use `group_id` for the strongest leakage boundary. For example, clips cut from +one long recording should share a `group_id`, even if they have different row +IDs. + +## Prepared layout + +`prepare` writes: + +```text +prepared/ + dataset.json + symbols.json + train.jsonl + validation.jsonl + audio/ + preparation_report.json +``` + +Prepared rows contain the converted audio path, original and normalized text, +phoneme string, duration, split metadata, and any supported source identifiers. +`dataset.json` records the language, frontend identity, hashes, speaker, +split configuration, counts, and diagnostics. `symbols.json` contains the +ordered inventory used to migrate embedding rows by symbol identity. + +Prepared data is immutable input. Correct the source data or frontend and run +`prepare` again instead of editing prepared files by hand. diff --git a/finetune/docs/RESPONSIBLE_USE.md b/finetune/docs/RESPONSIBLE_USE.md new file mode 100644 index 0000000..662de95 --- /dev/null +++ b/finetune/docs/RESPONSIBLE_USE.md @@ -0,0 +1,76 @@ +# Consent, licensing, and responsible use + +## Speaker consent + +Use speech only when the speaker has knowingly authorized its use for model +training and the intended distribution. Consent to publish recordings does not +automatically include consent to train a synthetic voice. + +Document: + +- who owns or controls the recordings; +- what training and distribution uses were authorized; +- whether commercial use is permitted; +- whether the resulting model may be redistributed; +- how withdrawal requests are handled where applicable. + +Do not train on private calls, leaked recordings, scraped personal media, or a +person's voice without permission. + +## Voice identity + +An adapted checkpoint is a fixed synthetic voice learned from its corpus. Do +not present it as the real speaker, use it for impersonation, or use it to +mislead listeners about who said something. + +Where a generated voice resembles an identifiable person, disclose that the +audio is synthetic and follow applicable laws, platform rules, and contractual +restrictions. + +## Data licensing + +The Inflect code or base checkpoint license does not grant rights to a user's +adaptation dataset. Users are responsible for transcript, audio, dictionary, +frontend, and model-output rights. + +An exported checkpoint cannot be distributed more broadly than its data and +dependency licenses permit. Preserve required attribution and notices. + +## Language and cultural review + +For new languages, involve fluent speakers in pronunciation and acceptability +review. Automated ASR and phonemization scores do not detect all offensive, +misleading, or culturally inappropriate output. + +Document regional variety and known limitations rather than describing a model +as supporting an entire language based on one speaker or corpus. + +## Security and privacy + +Treat manifests and custom frontend hooks as untrusted inputs. Do not publish +source paths, user names, credentials, or private metadata in reports or +checkpoints. + +Review exported files for: + +- absolute paths; +- source transcripts that were not intended for release; +- speaker identifiers and personal information; +- embedded credentials or remote URLs; +- training logs containing private data. + +## Release disclosure + +An adapted model card should state: + +- that it is community-adapted and not an official Inflect checkpoint; +- base model and toolkit version; +- fixed voice and configured language; +- data source, consent, and license; +- evaluation procedure and fluent-speaker involvement; +- known failure modes; +- intended and prohibited uses; +- contact or reporting path for harmful outputs. + +Technical capability does not remove the obligation to obtain consent or use +the model responsibly. diff --git a/finetune/docs/SCOPE.md b/finetune/docs/SCOPE.md new file mode 100644 index 0000000..4dc77e2 --- /dev/null +++ b/finetune/docs/SCOPE.md @@ -0,0 +1,74 @@ +# Scope and support status + +## What this toolkit produces + +This toolkit warm-starts a new Inflect-compatible generator from an official +Inflect v2 checkpoint and a user-supplied, single-speaker corpus. One run +produces one fixed voice for one configured language. + +It does not add a speaker or language selector to the original model. Separate +voices or languages require separate prepared datasets, training runs, and +exports. + +## Capability status + +| Capability | Status | +| --- | --- | +| Prepare JSONL or CSV data | Implemented and tested | +| Resample and validate 24 kHz mono audio | Implemented and tested | +| Leakage-safe train/validation splitting | Implemented and tested | +| eSpeak language frontend | Implemented | +| Prephonemized input | Implemented | +| Explicit custom Python frontend | Implemented and source-hashed | +| Warm-start Micro or Nano | Implemented; Nano CUDA smoke-tested | +| Strict same-run resume | Implemented and tested | +| Fixed-voice adaptation | Experimental quality | +| New-language adaptation | Experimental quality | +| PyTorch deployment export | Implemented and load-tested | +| ONNX deployment export | Implemented; Nano parity-tested | +| Runtime-selectable voices or languages | Not supported | +| Zero-shot or few-shot cloning | Not supported | +| Quantized export | Not included | + +The software path has been exercised end to end with a real prepared dataset, +a CUDA training step, resume, PyTorch export, and ONNX Runtime parity. This does +not prove that an arbitrary corpus or language will produce a good voice. +Fluent-speaker review and held-out listening remain required. + +## Why adaptation is difficult + +Language adaptation changes the symbol inventory, pronunciation rules, timing, +and acoustic distribution. Voice adaptation changes pitch range, formants, +speaking rate, recording conditions, and waveform statistics. Changing both at +once asks every part of a very small generator to move. + +A model may remain intelligible while becoming thin, buzzy, metallic, +sibilant, unstable, or unlike the target speaker. Low training loss is not +evidence that adaptation succeeded. + +## Public toolkit versus private release process + +This package contains a generic compatible trainer: manifest readers, audio +checks, frontend adapters, symbol migration, generic losses and optimizer +defaults, evaluation utilities, checkpointing, and inference-only export. + +It does not publish or reconstruct: + +- private corpora, paths, transcripts, or generated audio; +- private corpus-generation or filtering methods; +- the exact curriculum used for the official checkpoints; +- private hyperparameter searches or failed experiments; +- internal checkpoint-ranking and release-selection procedures; +- credentials, rental identifiers, or private storage locations. + +The released generator is an initialization point, not a resumable copy of the +original training run. It does not include the posterior encoder, +discriminators, optimizers, schedulers, manifests, or RNG state needed to +continue that run. + +## Claims for adapted checkpoints + +Do not describe an adapted checkpoint as an official Inflect voice or +language. Do not claim language support only because a frontend can emit its +phonemes. Publish the language, data provenance, consent basis, base model, +toolkit version, frontend, known limitations, and evaluation method. diff --git a/finetune/docs/TRAINING.md b/finetune/docs/TRAINING.md new file mode 100644 index 0000000..1818a26 --- /dev/null +++ b/finetune/docs/TRAINING.md @@ -0,0 +1,173 @@ +# Training, resume, evaluation, and export + +## Warm-start behavior + +The released `model.pth` is an inference generator, not a training snapshot. +The trainer: + +1. builds the compatible training form of Micro or Nano; +2. copies all compatible released generator weights; +3. migrates text embeddings by symbol identity; +4. deterministically initializes newly added symbols; +5. initializes the training-only posterior encoder and discriminator; +6. creates fresh public optimizers, schedulers, scaler, and RNG state; +7. records hashes for the base model, prepared data, symbols, and options. + +This begins a new adaptation run. It does not continue the private release run. + +## Stages + +The generic public schedule has three stages: + +- `posterior_warmup`: establish the training-only posterior path; +- `linguistic_adaptation`: adapt timing, text, latent, and acoustic behavior; +- `decoder_polish`: optionally unfreeze the waveform decoder at a lower + learning-rate multiplier. + +Stage boundaries are explicit and resumable. `--decoder-unfreeze-step none` +keeps the decoder frozen. + +## Presets and overrides + +| Preset | Intended starting point | +| --- | --- | +| `balanced` | General CUDA starting point | +| `micro-12gb` | Conservative Micro setup near the 12 GB class | +| `nano-8gb` | Small-batch Nano setup near the 8 GB class | + +Actual memory depends on clip length, model, batch size, validation, PyTorch, +CUDA, and allocator behavior. Preset names are not hardware guarantees. + +Every explicitly supplied CLI flag overrides its preset, even when the value +matches a library default: + +```bash +inflect-adapt train \ + --base owensong/Inflect-Micro-v2 \ + --dataset prepared/es \ + --preset micro-12gb \ + --batch-size 1 \ + --gradient-accumulation-steps 8 \ + --output runs/es-micro +``` + +If memory is exhausted, reduce the batch size and increase accumulation. Do not +compare two runs as equivalent if precision, segment length, effective batch, +or optimization settings differ. + +## Outputs + +```text +runs/es-micro/ + run-identity.json + training-options.json + compatibility-report.json + metrics.jsonl + training-summary.json + checkpoints/ + adaptation-step-00001000.pth + adaptation-final.pth + latest.pth + exports/ + model-step-00001000.pth + model.pth + validation/ + step-00000500.json + step-00000500.wav +``` + +Training checkpoints contain the public training state needed for exact +same-run resume. Files in `exports/` are lightweight generator checkpoints, +not complete deployment packages. + +## Resume identity + +Resume validates before mutating model or optimizer state. It rejects +incompatible toolkit schema, run ID, base model, prepared dataset, symbols, +frontend, model shapes, optimizer schema, or relevant options. + +```bash +inflect-adapt train \ + --base owensong/Inflect-Micro-v2 \ + --dataset prepared/es \ + --preset micro-12gb \ + --output runs/es-micro \ + --resume runs/es-micro/checkpoints/latest.pth +``` + +Do not bypass a rejection. Start a new output directory when the data or +configuration changes. + +## Checkpoint selection + +Validation intervals create fixed-seed held-out synthesis and loss +diagnostics. The trainer intentionally does not create `best.pth`, because no +single training loss reliably identifies the best-sounding TTS checkpoint. + +Declare a selection rule before inspecting the final test set. Consider: + +- intelligibility on difficult held-out text; +- clipping, silence, duration, and truncated endings; +- pronunciation and phoneme coverage; +- voice consistency; +- buzz, metallic resonance, sibilance, thinness, and transients; +- blinded listening by fluent speakers. + +Predicted MOS, ASR WER, and training loss are useful diagnostics, not complete +quality measures. + +## Deployment exports + +PyTorch: + +```bash +inflect-adapt export \ + --checkpoint runs/es-micro/checkpoints/adaptation-final.pth \ + --prepared-dataset prepared/es \ + --format pytorch \ + --output exports/es-micro +``` + +ONNX: + +```bash +inflect-adapt export \ + --checkpoint runs/es-micro/checkpoints/adaptation-final.pth \ + --prepared-dataset prepared/es \ + --format onnx \ + --output exports/es-micro-onnx +``` + +The exporter carries the exact language/frontend contract into the package. +eSpeak packages use the configured language, prephonemized packages require +phoneme input, and custom packages require a source-hash-matching +`--frontend-hook`. + +The exporter removes training-only state and verifies strict inference loading. +When ONNX is requested, it exports separate duration and decode graphs and +compares them with PyTorch on fixed inputs. Quantized export is not included. + +## Evaluation + +```bash +inflect-adapt evaluate \ + --model-dir exports/es-micro \ + --manifest prepared/es/validation.jsonl \ + --output evaluations/es-micro +``` + +Evaluation writes audio, per-item diagnostics, aggregate JSON, and a readable +summary. Supply `--transcript-evaluator module:callable` to integrate a +project-specific ASR or metric implementation. + +## Clean-environment test + +Before release: + +1. build and install the wheel in a clean environment; +2. download only the public base and adapted artifacts; +3. verify `checksums.sha256`; +4. run inference without the training checkout; +5. test ordinary and difficult text; +6. test the target frontend mode; +7. repeat ONNX parity on the release artifact if ONNX is published. diff --git a/finetune/docs/TROUBLESHOOTING.md b/finetune/docs/TROUBLESHOOTING.md new file mode 100644 index 0000000..387faa6 --- /dev/null +++ b/finetune/docs/TROUBLESHOOTING.md @@ -0,0 +1,117 @@ +# Troubleshooting + +Start with the first failing stage. Do not work around preparation or audit +errors by editing prepared files. + +## Command or module not found + +Install from the `finetune` directory in an activated environment: + +```bash +python -m pip install -e . +python -m inflect_finetune --help +``` + +If the editable install is not wanted, build and install the wheel instead. + +## Manifest rejected + +Check that: + +- the file is UTF-8 JSONL or CSV; +- every non-empty JSONL line is one JSON object; +- `audio` and `text` are strings; +- audio paths are relative to `--audio-root`; +- paths do not contain traversal outside the root; +- referenced files exist; +- JSON quoting and escaping are valid. + +## Audio cannot be decoded + +Preserve the original file, decode it with a trusted audio tool, and convert a +copy to a standard uncompressed WAV. Confirm sample rate, channels, duration, +and finite samples. Do not rename a compressed file to `.wav`. + +## eSpeak or phonemizer failure + +Confirm the installed eSpeak NG library is available, the language code exists, +and `phonemizer` can process a short sentence independently. Record package and +eSpeak versions when reporting the problem. + +If only specific text fails, reduce it to the smallest reproducible input and +inspect unsupported symbols, language mixing, and normalization. + +## Unknown symbols + +Do not delete unknown symbols from prepared text. Determine whether: + +- the frontend emitted an undeclared symbol; +- Unicode normalization produced a different code point; +- the base inventory lacks a valid target-language phoneme; +- punctuation or a word-boundary marker was omitted from the inventory. + +Correct the frontend or symbol declaration, rerun `prepare`, and audit again. +New valid symbols should be initialized and learned; they should not be mapped +silently to unrelated base symbols. + +## Out of GPU memory + +Use a smaller validated memory preset, reduce the configured batch workload, +shorten training segments only if the trainer supports that change, or increase +gradient accumulation. Restart after clearing the failed process. + +Do not compare runs as equivalent if batch semantics, segment length, precision, +or optimizer behavior changed. + +## Resume rejected + +Compare the recorded toolkit version, base checkpoint, dataset hash, symbols +hash, frontend metadata, model configuration, and optimizer schema. Resume +rejection usually means the run inputs changed. + +Start a new run rather than forcing incompatible state to load. + +## Training loss improves but audio worsens + +Stop and inspect held-out audio. Common failures include decoder buzz, metallic +resonance, excessive sibilance, clipped endings, duration collapse, speaker +drift, and overfitting. + +Select checkpoints using matched held-out listening and diagnostics, not loss +alone. More steps can make adaptation worse. + +## Output is intelligible but pronunciation is wrong + +Verify the source transcript, normalized text, phonemes, symbol coverage, and +training examples for the affected sound. If the frontend is wrong, fix it and +re-prepare the corpus. Training longer will not reliably correct a systematically +wrong phoneme sequence. + +## Adapted voice sounds unlike the speaker + +Check speaker consistency, recording conditions, data quantity and coverage, +and whether the run changed language and voice simultaneously. Speaker +similarity is not guaranteed by the fixed-voice architecture. + +## ONNX model fails to parse or load + +Confirm the file is fully downloaded and its checksum matches. Test with the +documented ONNX Runtime and ONNX versions. Re-export from the same inference +checkpoint and validate every graph before upload. + +An ONNX file existing on disk is not evidence that export succeeded. + +## Reporting an issue + +Include: + +- exact command; +- toolkit commit or version; +- operating system, Python, PyTorch, CUDA, eSpeak, and ONNX Runtime versions; +- base model and checksum; +- redacted manifest schema and failing row; +- prepared dataset and symbols hashes; +- complete error text; +- minimal reproducible input. + +Do not attach private speech data without the speaker's permission. diff --git a/finetune/examples/custom_frontend_hook.py b/finetune/examples/custom_frontend_hook.py new file mode 100644 index 0000000..6115fac --- /dev/null +++ b/finetune/examples/custom_frontend_hook.py @@ -0,0 +1,34 @@ +"""Minimal public custom-frontend example for Inflect adaptation exports.""" + +from __future__ import annotations + +import re +import unicodedata + + +class ExampleFrontend: + def __init__(self, language: str) -> None: + self.language = language + + def normalize(self, text: str) -> str: + value = unicodedata.normalize("NFKC", text) + return re.sub(r"\s+", " ", value).strip() + + def phonemize(self, text: str) -> str: + # Replace this illustrative mapping with a real language frontend. + return " ".join(character.lower() for character in text if character.isalpha()) + + def symbols(self) -> tuple[str, ...]: + return tuple("abcdefghijklmnopqrstuvwxyz ") + + def metadata(self) -> dict[str, object]: + return { + "name": "public-example-character-frontend", + "version": "1", + "language": self.language, + "configuration": {"mapping": "lowercase Unicode alphabetic characters"}, + } + + +def create_frontend(*, language: str) -> ExampleFrontend: + return ExampleFrontend(language) diff --git a/finetune/examples/dataset_manifest.csv b/finetune/examples/dataset_manifest.csv new file mode 100644 index 0000000..89df678 --- /dev/null +++ b/finetune/examples/dataset_manifest.csv @@ -0,0 +1,4 @@ +audio,text +audio/000001.wav,"Esta es una grabacion de ejemplo para preparar el conjunto de datos." +audio/000002.wav,"Los datos reales deben incluir transcripciones revisadas y audio de un solo hablante." +audio/000003.wav,"Los numeros, nombres propios y signos de puntuacion necesitan cobertura explicita." diff --git a/finetune/examples/evaluation_manifest.jsonl b/finetune/examples/evaluation_manifest.jsonl new file mode 100644 index 0000000..316c51e --- /dev/null +++ b/finetune/examples/evaluation_manifest.jsonl @@ -0,0 +1,3 @@ +{"id":"held-out-001","text":"Esta frase se reserva para evaluar el punto de control adaptado."} +{"id":"held-out-002","text":"La prueba incluye numeros: veintisiete, ciento cuatro y dos mil veintiseis."} +{"id":"held-out-003","text":"¿Mantiene pausas naturales, ritmo estable y una pronunciacion clara?"} diff --git a/finetune/examples/export_language_package.py b/finetune/examples/export_language_package.py new file mode 100644 index 0000000..7559f4a --- /dev/null +++ b/finetune/examples/export_language_package.py @@ -0,0 +1,47 @@ +"""Export a language-aware inference package through the public Python API.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +TOOLKIT_ROOT = Path(__file__).resolve().parents[1] +if str(TOOLKIT_ROOT) not in sys.path: + sys.path.insert(0, str(TOOLKIT_ROOT)) + +from inflect_finetune.exporting import ExportOptions, export_checkpoint + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--prepared-dataset", type=Path, required=True) + parser.add_argument("--package-template", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--config", type=Path) + parser.add_argument("--symbols", type=Path) + parser.add_argument("--frontend-hook", type=Path) + parser.add_argument("--onnx", action="store_true") + parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + symbols = args.symbols or args.prepared_dataset / "symbols.json" + report = export_checkpoint( + ExportOptions( + checkpoint=args.checkpoint, + output_dir=args.output, + config=args.config, + symbols=symbols, + package_template=args.package_template, + prepared_dataset=args.prepared_dataset, + frontend_hook=args.frontend_hook, + include_onnx=args.onnx, + overwrite=args.overwrite, + ) + ) + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/finetune/examples/export_options.json b/finetune/examples/export_options.json new file mode 100644 index 0000000..07e70ea --- /dev/null +++ b/finetune/examples/export_options.json @@ -0,0 +1,13 @@ +{ + "checkpoint": "runs/my-language/checkpoints/adaptation-final.pth", + "prepared_dataset": "prepared/my-language", + "config": "path/to/Inflect-Micro-v2/config.json", + "symbols": "prepared/my-language/symbols.json", + "frontend_hook": null, + "package_template": "path/to/Inflect-Micro-v2", + "output_dir": "exports/my-language", + "include_onnx": true, + "onnx_opset": 17, + "model_name": "My-Inflect-Adaptation", + "source_revision": "local-training-run" +} diff --git a/finetune/examples/language_adaptation.json b/finetune/examples/language_adaptation.json new file mode 100644 index 0000000..6ff5c79 --- /dev/null +++ b/finetune/examples/language_adaptation.json @@ -0,0 +1,16 @@ +{ + "format": "inflect_adaptation_example_v1", + "language": "es", + "frontend": { + "type": "espeak", + "language": "es", + "preserve_punctuation": true, + "with_stress": true + }, + "audio": { + "sample_rate": 24000, + "channels": 1 + }, + "base_model": "owensong/Inflect-Micro-v2", + "notes": "Example only. Supply audio recorded with the fixed voice you want the adapted checkpoint to use." +} diff --git a/finetune/examples/prephonemized_evaluation_manifest.jsonl b/finetune/examples/prephonemized_evaluation_manifest.jsonl new file mode 100644 index 0000000..743fd6e --- /dev/null +++ b/finetune/examples/prephonemized_evaluation_manifest.jsonl @@ -0,0 +1 @@ +{"id":"prepared-001","text":"Reference transcript for optional ASR scoring.","phonemes":"təst"} diff --git a/finetune/examples/self_test_export_eval.py b/finetune/examples/self_test_export_eval.py new file mode 100644 index 0000000..06b873a --- /dev/null +++ b/finetune/examples/self_test_export_eval.py @@ -0,0 +1,374 @@ +"""Executable acceptance test for the public export/evaluation APIs. + +Run from the ``finetune`` directory: + + python examples/self_test_export_eval.py \ + --package ../release_assets/hf_clean_download/Inflect-Nano-v2 + +The script uses only a released package and temporary synthetic metadata. It +does not require or inspect an adaptation corpus. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import sys +import tempfile +from pathlib import Path +from typing import Any + +import torch + +TOOLKIT_ROOT = Path(__file__).resolve().parents[1] +if str(TOOLKIT_ROOT) not in sys.path: + sys.path.insert(0, str(TOOLKIT_ROOT)) + +from inflect_finetune.checkpoint import capture_rng_state +from inflect_finetune.evaluation import EvaluationOptions, evaluate_checkpoint +from inflect_finetune.exporting import ( + ExportOptions, + _build_model, + _extract_state, + _load_checkpoint, + export_checkpoint, +) +from inflect_finetune.frontend import FrontendOptions, custom_frontend_metadata + + +def _symbols(package: Path) -> list[str]: + source = package / "runtime" / "text" / "symbols.py" + namespace: dict[str, Any] = {} + exec(compile(source.read_text(encoding="utf-8"), str(source), "exec"), namespace) + return list(namespace["symbols"]) + + +def _training_fixture(package: Path, destination: Path) -> int: + config = json.loads((package / "config.json").read_text(encoding="utf-8")) + symbols = _symbols(package) + deployable, _, _ = _extract_state(_load_checkpoint(package / "model.pth")) + model, _ = _build_model( + package / "runtime", + config, + len(symbols), + inference_only=False, + ) + incompatible = model.load_state_dict(deployable, strict=False) + missing = list(incompatible.missing_keys) + if not missing or not all(key.startswith("enc_q.") for key in missing): + raise AssertionError(f"Training model had invalid missing keys: {missing}") + if incompatible.unexpected_keys: + raise AssertionError( + f"Training model had unexpected deployable keys: {incompatible.unexpected_keys}" + ) + torch.save( + { + "format": "inflect_vits_adaptation_training_checkpoint_v1", + "generator": model.state_dict(), + "discriminator": {"self_test": torch.zeros(1)}, + "optimizer_g": {"state": {}, "param_groups": []}, + "optimizer_d": {"state": {}, "param_groups": []}, + "scheduler_g": {"last_epoch": 0}, + "scheduler_d": {"last_epoch": 0}, + "scaler": {"scale": 1.0}, + "step": 17, + "epoch": 1, + "stage": "self-test", + "options": {"base_model": str(package)}, + "symbols": symbols, + "compatibility": {"self_test": True}, + "rng_state": capture_rng_state(), + }, + destination, + ) + return len(missing) + + +def _prepared_fixture( + destination: Path, + symbols: list[str], + *, + language: str, + mode: str, + hook_metadata: dict[str, Any] | None = None, +) -> Path: + destination.mkdir(parents=True, exist_ok=True) + (destination / "symbols.json").write_text( + json.dumps( + { + "format": "inflect_v2_symbol_inventory_v1", + "symbols": symbols, + "count": len(symbols), + }, + ensure_ascii=False, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + frontend: dict[str, Any] = { + "type": mode, + "language": language, + "preserve_punctuation": True, + "with_stress": True, + } + if hook_metadata is not None: + frontend["hook"] = hook_metadata + (destination / "dataset.json").write_text( + json.dumps( + { + "format": "inflect_prepared_dataset_v1", + "language": language, + "sample_rate": 24000, + "frontend": frontend, + "source_manifest_sha256": "0" * 64, + }, + ensure_ascii=False, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + return destination + + +def _onnx_available() -> bool: + return ( + importlib.util.find_spec("onnx") is not None + and importlib.util.find_spec("onnxruntime") is not None + ) + + +def run(package: Path, *, include_onnx: bool) -> dict[str, Any]: + package = package.resolve() + required = ("config.json", "model.pth", "runtime", "inference.py") + missing = [name for name in required if not (package / name).exists()] + if missing: + raise FileNotFoundError(f"Release package is missing: {missing}") + + with tempfile.TemporaryDirectory(prefix="inflect-export-eval-self-test-") as temporary: + root = Path(temporary) + checkpoint = root / "training.pth" + expected_enc_q = _training_fixture(package, checkpoint) + symbols = _symbols(package) + + try: + export_checkpoint( + ExportOptions( + checkpoint=checkpoint, + output_dir=root / "must-reject-missing-frontend", + package_template=package, + ) + ) + except ValueError as exc: + if "requires prepared dataset frontend metadata" not in str(exc): + raise + else: + raise AssertionError( + "Adapted checkpoint silently reused the release English frontend." + ) + + prepared = _prepared_fixture( + root / "prepared-es", + symbols, + language="es", + mode="espeak", + ) + export_dir = root / "export" + export_report = export_checkpoint( + ExportOptions( + checkpoint=checkpoint, + output_dir=export_dir, + package_template=package, + prepared_dataset=prepared, + include_onnx=include_onnx, + verify=True, + ) + ) + frontend_contract = json.loads( + (export_dir / "frontend.json").read_text(encoding="utf-8") + ) + if ( + frontend_contract["mode"] != "espeak" + or frontend_contract["language"] != "es" + or not frontend_contract["accepts_prephonemized_input"] + ): + raise AssertionError( + f"Exported language frontend contract is invalid: {frontend_contract}" + ) + payload = _load_checkpoint(export_dir / "model.pth") + deployable, _, stripped_after_reload = _extract_state(payload) + if stripped_after_reload or any(key.startswith("enc_q.") for key in deployable): + raise AssertionError("Inference export retained enc_q.* tensors.") + + required_omissions = { + "discriminator", + "optimizer_g", + "optimizer_d", + "scheduler_g", + "scheduler_d", + "scaler", + "rng_state", + } + omissions = set( + export_report["stripped_training_tensors"]["top_level_fields"] + ) + if not required_omissions <= omissions: + raise AssertionError( + f"Export report did not prove all omissions: {required_omissions - omissions}" + ) + training_check = next( + check for check in export_report["checks"] if "enc_q" in check["message"] + ) + if ( + not training_check["ok"] + or len(training_check["missing_keys"]) != expected_enc_q + or training_check["unexpected_keys"] + ): + raise AssertionError(f"Training-form compatibility failed: {training_check}") + + manifest = root / "evaluation.jsonl" + manifest.write_text( + json.dumps( + { + "id": "held-out-self-test", + "text": "The exported checkpoint produces a complete waveform.", + "phonemes": "test", + } + ) + + "\n", + encoding="utf-8", + ) + evaluation_report = evaluate_checkpoint( + EvaluationOptions( + model_dir=export_dir, + manifest=manifest, + output_dir=root / "evaluation", + max_samples=1, + ) + ) + if not evaluation_report["ok"]: + raise AssertionError(f"Evaluation failed: {evaluation_report['failures']}") + if evaluation_report["items"][0]["input_mode"] != "prephonemized": + raise AssertionError("Evaluation did not use the explicit phoneme bypass.") + + hook = Path(__file__).with_name("custom_frontend_hook.py") + hook_options = FrontendOptions( + mode="custom", + language="x-example", + hook=f"{hook}:create_frontend", + ) + hook_metadata = custom_frontend_metadata(hook_options) + assert hook_metadata is not None + custom_prepared = _prepared_fixture( + root / "prepared-custom", + symbols, + language="x-example", + mode="custom", + hook_metadata=hook_metadata, + ) + try: + export_checkpoint( + ExportOptions( + checkpoint=checkpoint, + output_dir=root / "must-reject-missing-hook", + package_template=package, + prepared_dataset=custom_prepared, + ) + ) + except ValueError as exc: + if "frontend_hook" not in str(exc): + raise + else: + raise AssertionError("Custom export succeeded without a package hook.") + + custom_export = root / "custom-export" + custom_report = export_checkpoint( + ExportOptions( + checkpoint=checkpoint, + output_dir=custom_export, + package_template=package, + prepared_dataset=custom_prepared, + frontend_hook=hook, + verify=True, + ) + ) + packaged_hook = custom_export / "frontend_hook.py" + if ( + not packaged_hook.is_file() + or packaged_hook.read_bytes() != hook.read_bytes() + ): + raise AssertionError("Custom frontend hook was not copied byte-for-byte.") + custom_manifest = root / "custom-evaluation.jsonl" + custom_manifest.write_text( + json.dumps({"id": "custom-hook-self-test", "text": "Test"}) + "\n", + encoding="utf-8", + ) + custom_evaluation = evaluate_checkpoint( + EvaluationOptions( + model_dir=custom_export, + manifest=custom_manifest, + output_dir=root / "custom-evaluation", + max_samples=1, + ) + ) + if not custom_evaluation["ok"]: + raise AssertionError( + f"Packaged custom frontend failed: {custom_evaluation['failures']}" + ) + + onnx_report = export_report["onnx"] + onnx_verification = onnx_report.get("verification", {}) + return { + "ok": True, + "deployable_parameters": export_report["deployable_parameters"], + "stripped_enc_q_tensors": export_report["stripped_training_tensors"][ + "count" + ], + "omitted_training_fields": sorted(omissions), + "strict_runtime_load": any( + check["ok"] and check.get("strict_model_load") + for check in export_report["checks"] + ), + "training_model_missing_only_enc_q": training_check["ok"], + "frontend": { + "mode": frontend_contract["mode"], + "language": frontend_contract["language"], + "prephonemized_evaluation": True, + "custom_hook_packaged": custom_report["deployment_frontend"]["mode"] + == "custom", + "missing_metadata_rejected": True, + "missing_custom_hook_rejected": True, + }, + "evaluation_duration_seconds": evaluation_report["items"][0]["signal"][ + "duration_seconds" + ], + "onnx": { + "requested": onnx_report.get("requested", False), + "status": onnx_report.get("status"), + "message": onnx_report.get("message"), + "parity": onnx_verification.get("parity"), + }, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--package", type=Path, required=True) + parser.add_argument( + "--onnx", + choices=("auto", "on", "off"), + default="auto", + help="Auto runs ONNX parity only when both optional dependencies are installed.", + ) + args = parser.parse_args() + available = _onnx_available() + if args.onnx == "on" and not available: + raise SystemExit("ONNX self-test requested but onnx/onnxruntime are unavailable.") + include_onnx = available if args.onnx == "auto" else args.onnx == "on" + print(json.dumps(run(args.package, include_onnx=include_onnx), indent=2)) + + +if __name__ == "__main__": + main() diff --git a/finetune/examples/transcript_evaluator_plugin.py b/finetune/examples/transcript_evaluator_plugin.py new file mode 100644 index 0000000..e36cfa6 --- /dev/null +++ b/finetune/examples/transcript_evaluator_plugin.py @@ -0,0 +1,20 @@ +"""Adapter example for an ASR system that the user installed separately. + +This file intentionally does not import or download an ASR model. Replace the +body with a call to an already-configured local or hosted transcription system. +""" + +from pathlib import Path +from typing import Any + + +def evaluate_transcript( + audio_path: Path, + reference_text: str, + sample_rate: int, +) -> dict[str, Any]: + raise RuntimeError( + "Configure your own transcript evaluator in " + "finetune/examples/transcript_evaluator_plugin.py. The adaptation " + "toolkit never downloads a heavyweight ASR model automatically." + ) diff --git a/finetune/inflect_finetune/__init__.py b/finetune/inflect_finetune/__init__.py new file mode 100644 index 0000000..6660402 --- /dev/null +++ b/finetune/inflect_finetune/__init__.py @@ -0,0 +1,3 @@ +"""Public adaptation utilities for Inflect v2.""" + +__version__ = "0.1.0" diff --git a/finetune/inflect_finetune/__main__.py b/finetune/inflect_finetune/__main__.py new file mode 100644 index 0000000..a049ad7 --- /dev/null +++ b/finetune/inflect_finetune/__main__.py @@ -0,0 +1,5 @@ +from .cli import main + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/finetune/inflect_finetune/audio.py b/finetune/inflect_finetune/audio.py new file mode 100644 index 0000000..834897e --- /dev/null +++ b/finetune/inflect_finetune/audio.py @@ -0,0 +1,198 @@ +"""Audio validation and deterministic conversion for Inflect adaptation datasets.""" + +from __future__ import annotations + +import math +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import soundfile as sf +from scipy.signal import resample_poly + + +class AudioValidationError(ValueError): + """Raised when an input recording cannot be used safely for preparation.""" + + +@dataclass(frozen=True) +class AudioOptions: + """Generic audio validity and conversion settings.""" + + sample_rate: int = 24_000 + min_duration_seconds: float = 0.05 + max_duration_seconds: float | None = None + max_channels: int = 8 + peak_limit: float = 1.0 + output_subtype: str = "PCM_16" + + def validate(self) -> None: + """Validate option values before reading any source files.""" + if self.sample_rate <= 0: + raise ValueError("sample_rate must be positive.") + if self.min_duration_seconds < 0: + raise ValueError("min_duration_seconds cannot be negative.") + if ( + self.max_duration_seconds is not None + and self.max_duration_seconds <= self.min_duration_seconds + ): + raise ValueError("max_duration_seconds must exceed min_duration_seconds.") + if self.max_channels < 1: + raise ValueError("max_channels must be at least one.") + if not 0 < self.peak_limit <= 1: + raise ValueError("peak_limit must be in the interval (0, 1].") + + +@dataclass(frozen=True) +class AudioDiagnostics: + """Machine-readable diagnostics for one converted recording.""" + + source_sample_rate: int + source_channels: int + source_frames: int + output_sample_rate: int + output_frames: int + duration_seconds: float + source_peak: float + output_peak: float + source_clipped_fraction: float + resampled: bool + downmixed: bool + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-serializable representation.""" + return asdict(self) + + +def inspect_wav(path: Path, options: AudioOptions | None = None) -> sf.SoundFile: + """Validate WAV container metadata and return the libsndfile descriptor. + + The returned object is metadata only; callers do not need to close it. + """ + options = options or AudioOptions() + options.validate() + path = Path(path) + if not path.is_file(): + raise AudioValidationError(f"Audio file does not exist: {path}") + if path.suffix.lower() not in {".wav", ".wave"}: + raise AudioValidationError( + f"Expected a WAV file but received '{path.name}'. Convert it to WAV first." + ) + try: + info = sf.info(path) + except (RuntimeError, TypeError) as exc: + raise AudioValidationError(f"Could not read WAV metadata for {path}: {exc}") from exc + if info.format != "WAV": + raise AudioValidationError( + f"{path} has extension .wav but container format '{info.format}' is not WAV." + ) + if info.frames <= 0: + raise AudioValidationError(f"Audio file is empty: {path}") + if info.samplerate < 1_000 or info.samplerate > 384_000: + raise AudioValidationError( + f"Unsupported sample rate {info.samplerate} Hz in {path}; " + "expected a conventional audio sample rate." + ) + if info.channels < 1 or info.channels > options.max_channels: + raise AudioValidationError( + f"Unsupported channel count {info.channels} in {path}; " + f"the configured maximum is {options.max_channels}." + ) + duration = info.frames / info.samplerate + if duration < options.min_duration_seconds: + raise AudioValidationError( + f"{path} is only {duration:.3f}s; minimum is " + f"{options.min_duration_seconds:.3f}s." + ) + if options.max_duration_seconds is not None and duration > options.max_duration_seconds: + raise AudioValidationError( + f"{path} is {duration:.3f}s; maximum is " + f"{options.max_duration_seconds:.3f}s." + ) + return info + + +def _mono(audio: np.ndarray) -> np.ndarray: + if audio.ndim == 1: + return audio + if audio.ndim != 2: + raise AudioValidationError(f"Expected one- or two-dimensional audio, got {audio.shape}.") + return np.mean(audio, axis=1, dtype=np.float64) + + +def _resample(audio: np.ndarray, source_rate: int, target_rate: int) -> np.ndarray: + if source_rate == target_rate: + return audio + divisor = math.gcd(source_rate, target_rate) + return resample_poly( + audio, + up=target_rate // divisor, + down=source_rate // divisor, + padtype="line", + ) + + +def convert_wav( + source: Path, + destination: Path, + options: AudioOptions | None = None, +) -> AudioDiagnostics: + """Validate, downmix, resample, and write one canonical 24 kHz mono WAV.""" + options = options or AudioOptions() + info = inspect_wav(source, options) + try: + audio, sample_rate = sf.read( + source, + dtype="float64", + always_2d=True, + fill_value=0.0, + ) + except (RuntimeError, TypeError) as exc: + raise AudioValidationError(f"Could not decode WAV audio from {source}: {exc}") from exc + if sample_rate != info.samplerate: + raise AudioValidationError( + f"Metadata/decode sample-rate mismatch in {source}: " + f"{info.samplerate} versus {sample_rate}." + ) + if not np.isfinite(audio).all(): + raise AudioValidationError(f"Audio contains NaN or infinite samples: {source}") + + source_peak = float(np.max(np.abs(audio), initial=0.0)) + source_clipped_fraction = float(np.mean(np.abs(audio) >= 0.999)) + mono = _mono(audio) + converted = _resample(mono, sample_rate, options.sample_rate) + if not np.isfinite(converted).all() or converted.size == 0: + raise AudioValidationError(f"Audio conversion produced invalid samples for {source}.") + converted = np.clip(converted, -options.peak_limit, options.peak_limit) + + destination = Path(destination) + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_name(f".{destination.name}.tmp") + try: + sf.write( + temporary, + converted, + options.sample_rate, + format="WAV", + subtype=options.output_subtype, + ) + temporary.replace(destination) + except (OSError, RuntimeError, TypeError) as exc: + temporary.unlink(missing_ok=True) + raise AudioValidationError(f"Could not write prepared WAV {destination}: {exc}") from exc + + output_info = sf.info(destination) + return AudioDiagnostics( + source_sample_rate=sample_rate, + source_channels=info.channels, + source_frames=info.frames, + output_sample_rate=options.sample_rate, + output_frames=output_info.frames, + duration_seconds=output_info.frames / options.sample_rate, + source_peak=source_peak, + output_peak=float(np.max(np.abs(converted), initial=0.0)), + source_clipped_fraction=source_clipped_fraction, + resampled=sample_rate != options.sample_rate, + downmixed=info.channels != 1, + ) diff --git a/finetune/inflect_finetune/audit.py b/finetune/inflect_finetune/audit.py new file mode 100644 index 0000000..7125837 --- /dev/null +++ b/finetune/inflect_finetune/audit.py @@ -0,0 +1,332 @@ +"""Integrity and coverage audits for prepared Inflect adaptation datasets.""" + +from __future__ import annotations + +import hashlib +import json +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterator + +import soundfile as sf + +from .audio import AudioOptions, AudioValidationError, inspect_wav +from .manifest import ManifestError, resolve_audio_path +from .symbols import SymbolInventoryError, audit_symbol_coverage + + +class DatasetAuditError(RuntimeError): + """Raised when strict auditing finds invalid prepared data.""" + + +@dataclass(frozen=True) +class AuditOptions: + """Options suitable for programmatic use and a future audit CLI.""" + + prepared_dir: Path + strict: bool = True + duration_tolerance_seconds: float = 0.02 + + def validate(self) -> None: + """Validate audit settings.""" + if self.duration_tolerance_seconds < 0: + raise ValueError("duration_tolerance_seconds cannot be negative.") + + +def _load_json(path: Path) -> Any: + try: + return json.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError as exc: + raise DatasetAuditError(f"Required prepared dataset file is missing: {path}") from exc + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise DatasetAuditError(f"Invalid JSON in {path}: {exc}") from exc + + +def _jsonl(path: Path) -> Iterator[tuple[int, dict[str, Any]]]: + try: + with path.open("r", encoding="utf-8") as handle: + for line_number, line in enumerate(handle, 1): + if not line.strip(): + continue + try: + value = json.loads(line) + except json.JSONDecodeError as exc: + raise DatasetAuditError(f"Invalid JSON at {path}:{line_number}: {exc}") from exc + if not isinstance(value, dict): + raise DatasetAuditError(f"Expected a JSON object at {path}:{line_number}.") + yield line_number, value + except FileNotFoundError as exc: + raise DatasetAuditError(f"Required prepared split is missing: {path}") from exc + + +def _summary(report: dict[str, Any]) -> str: + lines = [ + "Inflect dataset audit", + "=====================", + f"Status: {'PASS' if report['valid'] else 'FAIL'}", + f"Rows checked: {report['row_counts']['total']}", + f"Audio duration: {report['total_duration_seconds']:.2f} seconds", + f"Errors: {len(report['errors'])}", + f"Warnings: {len(report['warnings'])}", + ] + if report["errors"]: + lines.extend(["", "Errors:"]) + lines.extend(f"- {error}" for error in report["errors"]) + if report["warnings"]: + lines.extend(["", "Warnings:"]) + lines.extend(f"- {warning}" for warning in report["warnings"]) + return "\n".join(lines) + "\n" + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _normalized_split_key(text: str) -> str: + return re.sub(r"\s+", " ", text).strip().casefold() + + +def audit_dataset(options: AuditOptions) -> dict[str, Any]: + """Audit layout, WAV invariants, duration metadata, and symbol coverage.""" + options.validate() + root = Path(options.prepared_dir).expanduser().resolve() + if not root.is_dir(): + raise DatasetAuditError(f"Prepared dataset directory does not exist: {root}") + + dataset = _load_json(root / "dataset.json") + symbols_payload = _load_json(root / "symbols.json") + if not isinstance(dataset, dict): + raise DatasetAuditError("dataset.json must contain a JSON object.") + symbols = symbols_payload.get("symbols") if isinstance(symbols_payload, dict) else None + if not isinstance(symbols, list): + raise DatasetAuditError("symbols.json must contain a 'symbols' list.") + + sample_rate = dataset.get("sample_rate") + if sample_rate != 24_000: + raise DatasetAuditError( + f"dataset.json sample_rate must be 24000, found {sample_rate!r}." + ) + errors: list[str] = [] + warnings: list[str] = [] + seen_audio: set[str] = set() + audio_hash_locations: dict[str, tuple[str, str]] = {} + normalized_text_splits: dict[str, set[str]] = {} + normalized_text_locations: dict[str, list[str]] = {} + group_splits: dict[tuple[str, str], set[str]] = {} + group_locations: dict[tuple[str, str], list[str]] = {} + speakers: set[str] = set() + phonemes: list[str] = [] + split_counts: dict[str, int] = {} + total_duration = 0.0 + audio_options = AudioOptions(sample_rate=24_000) + + for split in ("train", "validation"): + split_count = 0 + for line_number, row in _jsonl(root / f"{split}.jsonl"): + split_count += 1 + location = f"{split}.jsonl:{line_number}" + required = { + "audio", + "text", + "normalized_text", + "phonemes", + "duration_seconds", + } + missing = sorted(required.difference(row)) + if missing: + errors.append(f"{location} is missing fields: {', '.join(missing)}") + continue + text_fields = required - {"duration_seconds"} + if not all( + isinstance(row[field], str) and row[field] for field in text_fields + ): + errors.append(f"{location} contains an empty or non-string text field.") + continue + try: + relative, audio_path = resolve_audio_path( + root, str(row["audio"]), location=location + ) + except ManifestError as exc: + errors.append(str(exc)) + continue + if relative in seen_audio: + errors.append(f"Prepared audio path appears more than once: {relative}") + continue + seen_audio.add(relative) + try: + info = inspect_wav(audio_path, audio_options) + except AudioValidationError as exc: + errors.append(str(exc)) + continue + if info.samplerate != 24_000 or info.channels != 1: + errors.append( + f"{relative} must be 24 kHz mono; found " + f"{info.samplerate} Hz and {info.channels} channel(s)." + ) + audio_sha256 = _sha256(audio_path) + recorded_sha256 = row.get("audio_sha256") + if recorded_sha256 is not None and recorded_sha256 != audio_sha256: + errors.append( + f"{location} audio_sha256 does not match the prepared WAV content." + ) + prior_audio = audio_hash_locations.get(audio_sha256) + if prior_audio is not None: + prior_split, prior_location = prior_audio + boundary = ( + " across train and validation" + if prior_split != split + else " within the prepared dataset" + ) + errors.append( + f"Duplicate audio content detected{boundary}: " + f"{prior_location} and {location}." + ) + else: + audio_hash_locations[audio_sha256] = (split, location) + + normalized_key = _normalized_split_key(row["normalized_text"]) + normalized_text_splits.setdefault(normalized_key, set()).add(split) + normalized_text_locations.setdefault(normalized_key, []).append(location) + group_id = row.get("group_id") + group_field = row.get("group_field", "group_id") + speaker = row.get("speaker") + if speaker is not None: + if not isinstance(speaker, str) or not speaker.strip(): + errors.append(f"{location} has an empty or non-string speaker.") + else: + speakers.add(speaker) + if group_id is not None: + if not isinstance(group_id, str) or not group_id.strip(): + errors.append(f"{location} has an empty or non-string group_id.") + elif group_field not in {"group_id", "session"}: + errors.append( + f"{location} has unsupported group_field {group_field!r}." + ) + else: + group_key = (str(group_field), group_id) + group_splits.setdefault(group_key, set()).add(split) + group_locations.setdefault(group_key, []).append(location) + actual_duration = info.frames / info.samplerate + try: + recorded_duration = float(row["duration_seconds"]) + except (TypeError, ValueError): + errors.append(f"{location} has a non-numeric duration_seconds value.") + continue + if abs(actual_duration - recorded_duration) > options.duration_tolerance_seconds: + errors.append( + f"{location} duration differs from WAV metadata by " + f"{abs(actual_duration - recorded_duration):.4f}s." + ) + total_duration += actual_duration + phonemes.append(row["phonemes"]) + split_counts[split] = split_count + + try: + coverage = audit_symbol_coverage(phonemes, symbols) + except SymbolInventoryError as exc: + errors.append(str(exc)) + coverage = None + if coverage and coverage.unknown_counts: + errors.append( + "Phoneme text contains symbols absent from symbols.json: " + + ", ".join(repr(symbol) for symbol in coverage.unknown_counts) + ) + + for normalized_key, splits in normalized_text_splits.items(): + if len(splits) > 1: + locations = ", ".join(normalized_text_locations[normalized_key]) + errors.append( + "The same normalized transcript crosses train and validation " + f"({normalized_key!r}): {locations}." + ) + for (group_field, group_id), splits in group_splits.items(): + if len(splits) > 1: + locations = ", ".join(group_locations[(group_field, group_id)]) + errors.append( + f"{group_field}={group_id!r} crosses train and validation: {locations}." + ) + if len(speakers) > 1: + errors.append( + "Single-speaker adaptation requires one speaker value, but prepared rows " + f"contain {len(speakers)} values: " + + ", ".join(repr(speaker) for speaker in sorted(speakers)) + ) + dataset_speaker = dataset.get("speaker") + if dataset_speaker is not None and dataset_speaker not in speakers: + errors.append( + "dataset.json speaker does not match the speaker value in prepared rows." + ) + + expected_counts = dataset.get("row_counts", {}) + if expected_counts.get("train") != split_counts["train"]: + errors.append("dataset.json train row count does not match train.jsonl.") + if expected_counts.get("validation") != split_counts["validation"]: + errors.append("dataset.json validation row count does not match validation.jsonl.") + if expected_counts.get("total") != sum(split_counts.values()): + errors.append("dataset.json total row count does not match the prepared splits.") + if not split_counts["train"]: + errors.append( + "The prepared dataset has no training rows; training-ready data requires " + "nonempty train and validation splits." + ) + if not split_counts["validation"]: + errors.append( + "The prepared dataset has no validation rows; training-ready data requires " + "at least two independent rows or groups and a nonzero validation fraction." + ) + if sum(split_counts.values()) < 2: + errors.append( + "The prepared dataset is too small for training: at least two usable rows " + "that can occupy independent train and validation splits are required." + ) + + frontend = dataset.get("frontend") + if isinstance(frontend, dict) and frontend.get("type") == "custom": + hook = frontend.get("hook") + if not isinstance(hook, dict): + errors.append("Custom frontend metadata is missing its hook identity and hashes.") + else: + for field in ("identity", "source_sha256", "metadata_sha256"): + value = hook.get(field) + if not isinstance(value, str) or not value: + errors.append(f"Custom frontend hook metadata is missing '{field}'.") + for field in ("source_sha256", "metadata_sha256"): + value = hook.get(field) + if isinstance(value, str) and not re.fullmatch(r"[0-9a-f]{64}", value): + errors.append( + f"Custom frontend hook metadata field '{field}' is not a SHA-256 hash." + ) + + report = { + "format": "inflect_dataset_audit_v1", + "valid": not errors, + "prepared_dir": str(root), + "row_counts": { + "train": split_counts["train"], + "validation": split_counts["validation"], + "total": sum(split_counts.values()), + }, + "audio_files": len(seen_audio), + "unique_audio_hashes": len(audio_hash_locations), + "explicit_groups": len(group_splits), + "normalized_text_keys": len(normalized_text_splits), + "total_duration_seconds": round(total_duration, 6), + "symbol_coverage": coverage.to_dict() if coverage else None, + "errors": errors, + "warnings": warnings, + } + (root / "audit.json").write_text( + json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" + ) + (root / "AUDIT_REPORT.txt").write_text(_summary(report), encoding="utf-8") + if errors and options.strict: + raise DatasetAuditError( + f"Prepared dataset audit failed with {len(errors)} error(s). " + f"See {root / 'audit.json'}." + ) + return report diff --git a/finetune/inflect_finetune/checkpoint.py b/finetune/inflect_finetune/checkpoint.py new file mode 100644 index 0000000..5577369 --- /dev/null +++ b/finetune/inflect_finetune/checkpoint.py @@ -0,0 +1,522 @@ +"""Strict public-checkpoint migration and toolkit checkpoint I/O.""" + +from __future__ import annotations + +import hashlib +import json +import os +import random +import shutil +from collections import defaultdict +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence + +import numpy as np +import torch +from torch import nn + +from . import __version__ + +RELEASE_FORMAT = "inflect_vits_inference_checkpoint_v1" +TRAINING_FORMAT = "inflect_adaptation_training_checkpoint_v1" +INFERENCE_FORMAT = "inflect_vits_inference_checkpoint_v1" +RUN_IDENTITY_FORMAT = "inflect_adaptation_run_identity_v1" +EMBEDDING_KEY = "enc_p.emb.weight" +FRESH_PREFIXES = ("enc_q.",) + + +@dataclass(frozen=True) +class CompatibilityReport: + source_path: str + source_format: str + source_tensor_count: int + source_parameter_count: int + copied_tensor_count: int + copied_parameter_count: int + exact_tensor_count: int + migrated_embedding_rows: int + initialized_embedding_rows: int + fresh_tensor_count: int + fresh_parameter_count: int + fresh_prefixes: tuple[str, ...] + verified_equal_after_copy: bool + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + def write(self, path: str | Path) -> Path: + destination = Path(path) + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text( + json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return destination + + +def _torch_load(path: Path) -> Any: + try: + return torch.load(path, map_location="cpu", weights_only=False) + except TypeError: + return torch.load(path, map_location="cpu") + + +def load_release_checkpoint(path: str | Path) -> tuple[dict, Mapping[str, torch.Tensor]]: + source = Path(path).resolve() + payload = _torch_load(source) + if not isinstance(payload, dict) or not isinstance(payload.get("model"), Mapping): + raise ValueError(f"{source} is not an Inflect inference checkpoint.") + checkpoint_format = str(payload.get("format", "")) + if checkpoint_format != RELEASE_FORMAT: + raise ValueError( + f"Unsupported checkpoint format {checkpoint_format!r}; expected {RELEASE_FORMAT!r}." + ) + state = payload["model"] + if not all(isinstance(key, str) and torch.is_tensor(value) for key, value in state.items()): + raise ValueError("The release model state contains non-tensor entries.") + return payload, state + + +def _symbol_seed(seed: int, symbol: str) -> int: + digest = hashlib.sha256(f"{seed}\0{symbol}".encode("utf-8")).digest() + return int.from_bytes(digest[:8], "little") & 0x7FFF_FFFF + + +def _initialize_embedding_row(row: torch.Tensor, symbol: str, seed: int) -> None: + generator = torch.Generator(device="cpu") + generator.manual_seed(_symbol_seed(seed, symbol)) + values = torch.empty(row.shape, dtype=torch.float32, device="cpu") + values.normal_(mean=0.0, std=float(row.shape[-1]) ** -0.5, generator=generator) + row.copy_(values.to(dtype=row.dtype, device=row.device)) + + +def _occurrence_indices(symbols: Sequence[str]) -> dict[tuple[str, int], int]: + """Map repeated symbols by occurrence so release duplicate rows stay distinct.""" + + counts: defaultdict[str, int] = defaultdict(int) + indices: dict[tuple[str, int], int] = {} + for index, symbol in enumerate(symbols): + occurrence = counts[symbol] + indices[(symbol, occurrence)] = index + counts[symbol] += 1 + return indices + + +def warm_start_from_release( + model: nn.Module, + checkpoint_path: str | Path, + base_symbols: Sequence[str], + target_symbols: Sequence[str], + *, + initialization_seed: int = 1234, +) -> CompatibilityReport: + """Load all released tensors exactly, allowing only a fresh enc_q. + + The embedding may grow for a new language. Existing rows are mapped by + symbol identity and remain bit-identical; only genuinely new rows are + initialized. + """ + + source_path = Path(checkpoint_path).resolve() + payload, source = load_release_checkpoint(source_path) + target = model.state_dict() + source_keys = set(source) + target_keys = set(target) + + unexpected = sorted(source_keys - target_keys) + if unexpected: + raise RuntimeError(f"Release checkpoint has unexpected generator keys: {unexpected}") + fresh = sorted(target_keys - source_keys) + invalid_fresh = [key for key in fresh if not key.startswith(FRESH_PREFIXES)] + if invalid_fresh: + raise RuntimeError( + "Training model contains non-release keys outside the allowed training-only " + f"prefixes {FRESH_PREFIXES}: {invalid_fresh}" + ) + if EMBEDDING_KEY not in source or EMBEDDING_KEY not in target: + raise RuntimeError(f"Both states must contain {EMBEDDING_KEY}.") + if len(base_symbols) != source[EMBEDDING_KEY].shape[0]: + raise RuntimeError( + "Base symbol inventory length does not match the released embedding rows." + ) + if len(target_symbols) != target[EMBEDDING_KEY].shape[0]: + raise RuntimeError( + "Target symbol inventory length does not match the training embedding rows." + ) + + mismatched = [] + for key in sorted(source_keys - {EMBEDDING_KEY}): + if source[key].shape != target[key].shape or source[key].dtype != target[key].dtype: + mismatched.append( + ( + key, + tuple(source[key].shape), + tuple(target[key].shape), + source[key].dtype, + target[key].dtype, + ) + ) + if mismatched: + raise RuntimeError(f"Released non-embedding tensors do not match exactly: {mismatched}") + + migrated = target[EMBEDDING_KEY].clone() + base_index = _occurrence_indices(base_symbols) + target_index = _occurrence_indices(target_symbols) + target_occurrences: defaultdict[str, int] = defaultdict(int) + copied_rows = 0 + initialized_rows = 0 + for row_index, symbol in enumerate(target_symbols): + occurrence = target_occurrences[symbol] + target_occurrences[symbol] += 1 + source_index = base_index.get((symbol, occurrence)) + if source_index is None: + _initialize_embedding_row( + migrated[row_index], + f"{symbol}\0occurrence={occurrence}", + initialization_seed, + ) + initialized_rows += 1 + else: + migrated[row_index].copy_(source[EMBEDDING_KEY][source_index]) + copied_rows += 1 + + loaded = dict(target) + for key, value in source.items(): + if key != EMBEDDING_KEY: + loaded[key] = value + loaded[EMBEDDING_KEY] = migrated + model.load_state_dict(loaded, strict=True) + + verified = model.state_dict() + unequal = [ + key + for key in source_keys - {EMBEDDING_KEY} + if not torch.equal(verified[key].cpu(), source[key].cpu()) + ] + for identity, source_index in base_index.items(): + if identity in target_index: + if not torch.equal( + verified[EMBEDDING_KEY][target_index[identity]].cpu(), + source[EMBEDDING_KEY][source_index].cpu(), + ): + unequal.append(f"{EMBEDDING_KEY}[{identity!r}]") + if unequal: + raise RuntimeError(f"Warm-start verification failed for released tensors: {unequal}") + + source_parameters = sum(tensor.numel() for tensor in source.values()) + fresh_parameters = sum(target[key].numel() for key in fresh) + return CompatibilityReport( + source_path=str(source_path), + source_format=str(payload["format"]), + source_tensor_count=len(source), + source_parameter_count=source_parameters, + copied_tensor_count=len(source), + copied_parameter_count=source_parameters, + exact_tensor_count=len(source) - 1, + migrated_embedding_rows=copied_rows, + initialized_embedding_rows=initialized_rows, + fresh_tensor_count=len(fresh), + fresh_parameter_count=fresh_parameters, + fresh_prefixes=FRESH_PREFIXES, + verified_equal_after_copy=True, + ) + + +def cpu_compatibility_report( + model: nn.Module, + checkpoint_path: str | Path, + base_symbols: Sequence[str], + target_symbols: Sequence[str], + *, + initialization_seed: int = 1234, +) -> CompatibilityReport: + """Run strict migration on CPU and return its machine-readable audit.""" + + model.cpu() + return warm_start_from_release( + model, + checkpoint_path, + base_symbols, + target_symbols, + initialization_seed=initialization_seed, + ) + + +def _atomic_torch_save(payload: dict, destination: Path) -> None: + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_name(destination.name + ".tmp") + torch.save(payload, temporary) + os.replace(temporary, destination) + + +def _canonical_json(payload: Mapping[str, Any]) -> str: + return json.dumps(payload, ensure_ascii=False, separators=(",", ":"), sort_keys=True) + + +def sha256_file(path: str | Path) -> str: + source = Path(path) + digest = hashlib.sha256() + with source.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def build_run_identity( + *, + run_id: str, + base_root: str | Path, + prepared_dir: str | Path, + options: Mapping[str, Any], + optimizer_schema: Mapping[str, Any], +) -> dict[str, Any]: + """Build a public, path-independent identity for one adaptation run.""" + + base = Path(base_root).resolve() + prepared = Path(prepared_dir).resolve() + required = { + "base checkpoint": base / "model.pth", + "base config": base / "config.json", + "dataset metadata": prepared / "dataset.json", + "training split": prepared / "train.jsonl", + "validation split": prepared / "validation.jsonl", + "symbol inventory": prepared / "symbols.json", + } + missing = [f"{label}: {path}" for label, path in required.items() if not path.is_file()] + if missing: + raise FileNotFoundError( + "Cannot establish a resumable run identity; required public inputs are missing: " + + "; ".join(missing) + ) + if not run_id or not isinstance(run_id, str): + raise ValueError("run_id must be a non-empty string.") + identity = { + "format": RUN_IDENTITY_FORMAT, + "toolkit_version": __version__, + "run_id": run_id, + "base": { + "identity": base.name, + "checkpoint_sha256": sha256_file(required["base checkpoint"]), + "config_sha256": sha256_file(required["base config"]), + }, + "prepared_dataset": { + "dataset_json_sha256": sha256_file(required["dataset metadata"]), + "train_jsonl_sha256": sha256_file(required["training split"]), + "validation_jsonl_sha256": sha256_file(required["validation split"]), + }, + "symbols_sha256": sha256_file(required["symbol inventory"]), + "options": dict(options), + "optimizer_schema": dict(optimizer_schema), + } + # Enforce JSON stability at construction time instead of failing at save. + return json.loads(_canonical_json(identity)) + + +def validate_run_identity( + actual: Mapping[str, Any], + expected: Mapping[str, Any], + *, + source: str = "resume checkpoint", +) -> None: + if not isinstance(actual, Mapping): + raise ValueError(f"{source} does not contain run identity metadata.") + if actual.get("format") != RUN_IDENTITY_FORMAT: + raise ValueError(f"{source} has an unsupported run identity format.") + if _canonical_json(actual) != _canonical_json(expected): + differing = sorted( + key + for key in set(actual) | set(expected) + if actual.get(key) != expected.get(key) + ) + raise ValueError( + f"{source} belongs to a different adaptation run; identity fields differ: " + f"{differing}" + ) + + +def write_run_identity(path: str | Path, identity: Mapping[str, Any]) -> Path: + destination = Path(path) + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_name(destination.name + ".tmp") + temporary.write_text( + json.dumps(identity, indent=2, ensure_ascii=False, sort_keys=True) + "\n", + encoding="utf-8", + ) + os.replace(temporary, destination) + return destination + + +def load_run_identity(path: str | Path) -> dict[str, Any]: + source = Path(path) + try: + payload = json.loads(source.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ValueError(f"Could not read run identity marker {source}: {error}") from error + if not isinstance(payload, dict) or payload.get("format") != RUN_IDENTITY_FORMAT: + raise ValueError(f"{source} is not an Inflect adaptation run identity marker.") + return payload + + +def copy_checkpoint_alias(source: str | Path, destination: str | Path) -> Path: + """Atomically refresh a stable checkpoint copy without symlink assumptions.""" + + source_path = Path(source).resolve() + destination_path = Path(destination) + if not source_path.is_file(): + raise FileNotFoundError(f"Checkpoint alias source does not exist: {source_path}") + if source_path == destination_path.resolve(): + return destination_path + destination_path.parent.mkdir(parents=True, exist_ok=True) + temporary = destination_path.with_name(destination_path.name + ".tmp") + shutil.copyfile(source_path, temporary) + os.replace(temporary, destination_path) + return destination_path + + +def capture_rng_state() -> dict[str, Any]: + state: dict[str, Any] = { + "python": random.getstate(), + "numpy": np.random.get_state(), + "torch": torch.random.get_rng_state(), + } + if torch.cuda.is_available(): + state["cuda"] = torch.cuda.get_rng_state_all() + return state + + +def restore_rng_state(state: Mapping[str, Any]) -> None: + random.setstate(state["python"]) + np.random.set_state(state["numpy"]) + torch.random.set_rng_state(state["torch"]) + if "cuda" in state and torch.cuda.is_available(): + torch.cuda.set_rng_state_all(state["cuda"]) + + +def save_training_checkpoint( + path: str | Path, + *, + generator: nn.Module, + discriminator: nn.Module, + optimizer_g: torch.optim.Optimizer, + optimizer_d: torch.optim.Optimizer, + scheduler_g: Any, + scheduler_d: Any, + scaler: Any, + step: int, + epoch: int, + stage: str, + options: Mapping[str, Any], + symbols: Sequence[str], + compatibility: CompatibilityReport, + run_identity: Mapping[str, Any], + latest_path: str | Path | None = None, +) -> Path: + destination = Path(path) + payload = { + "format": TRAINING_FORMAT, + "generator": generator.state_dict(), + "discriminator": discriminator.state_dict(), + "optimizer_g": optimizer_g.state_dict(), + "optimizer_d": optimizer_d.state_dict(), + "scheduler_g": scheduler_g.state_dict(), + "scheduler_d": scheduler_d.state_dict(), + "scaler": scaler.state_dict(), + "step": int(step), + "epoch": int(epoch), + "stage": str(stage), + "options": dict(options), + "symbols": list(symbols), + "compatibility": compatibility.to_dict(), + "run_identity": dict(run_identity), + "rng_state": capture_rng_state(), + } + _atomic_torch_save(payload, destination) + if latest_path is not None: + copy_checkpoint_alias(destination, latest_path) + return destination + + +def resume_training_checkpoint( + path: str | Path, + *, + generator: nn.Module, + discriminator: nn.Module, + optimizer_g: torch.optim.Optimizer, + optimizer_d: torch.optim.Optimizer, + scheduler_g: Any, + scheduler_d: Any, + scaler: Any, + expected_symbols: Sequence[str], + expected_run_identity: Mapping[str, Any], +) -> tuple[int, int, str]: + payload = _torch_load(Path(path)) + if not isinstance(payload, dict) or payload.get("format") != TRAINING_FORMAT: + raise ValueError( + "Resume requires an Inflect adaptation training checkpoint, not a base release " + "checkpoint." + ) + validate_run_identity( + payload.get("run_identity"), + expected_run_identity, + source="resume checkpoint", + ) + if tuple(payload.get("symbols", ())) != tuple(expected_symbols): + raise ValueError("Resume checkpoint symbol inventory differs from the prepared dataset.") + stage = payload.get("stage") + if not isinstance(stage, str) or not stage: + raise ValueError("Resume checkpoint does not record the active warm-start stage.") + required_state = { + "generator", + "discriminator", + "optimizer_g", + "optimizer_d", + "scheduler_g", + "scheduler_d", + "scaler", + "step", + "epoch", + } + missing_state = sorted(required_state - payload.keys()) + if missing_state: + raise ValueError(f"Resume checkpoint is missing mutable state: {missing_state}") + + # Identity, symbols, stage, and payload shape are validated before any + # live module, optimizer, scaler, scheduler, or RNG state is mutated. + generator.load_state_dict(payload["generator"], strict=True) + discriminator.load_state_dict(payload["discriminator"], strict=True) + optimizer_g.load_state_dict(payload["optimizer_g"]) + optimizer_d.load_state_dict(payload["optimizer_d"]) + scheduler_g.load_state_dict(payload["scheduler_g"]) + scheduler_d.load_state_dict(payload["scheduler_d"]) + scaler.load_state_dict(payload["scaler"]) + if "rng_state" in payload: + restore_rng_state(payload["rng_state"]) + return int(payload["step"]), int(payload["epoch"]), stage + + +def save_inference_checkpoint( + path: str | Path, + *, + generator: nn.Module, + iteration: int, + learning_rate: float, +) -> Path: + """Save deployable state; training-only posterior tensors are excluded.""" + + state = { + key: value.detach().cpu() + for key, value in generator.state_dict().items() + if not key.startswith(FRESH_PREFIXES) + } + deployable_parameters = sum(tensor.numel() for tensor in state.values()) + payload = { + "format": INFERENCE_FORMAT, + "model": state, + "iteration": int(iteration), + "learning_rate": float(learning_rate), + "deployable_parameters": deployable_parameters, + } + destination = Path(path) + _atomic_torch_save(payload, destination) + return destination diff --git a/finetune/inflect_finetune/cli.py b/finetune/inflect_finetune/cli.py new file mode 100644 index 0000000..78a4b4f --- /dev/null +++ b/finetune/inflect_finetune/cli.py @@ -0,0 +1,373 @@ +"""Command-line interface for the public Inflect adaptation toolkit.""" + +from __future__ import annotations + +import argparse +import json +import logging +import sys +from pathlib import Path +from typing import Any, Sequence + +from . import __version__ + + +def _path(value: str) -> Path: + return Path(value).expanduser() + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("value must be positive") + return parsed + + +def _fraction(value: str) -> float: + parsed = float(value) + if not 0.0 < parsed < 1.0: + raise argparse.ArgumentTypeError("value must be in (0, 1)") + return parsed + + +def _optional_step(value: str) -> int | None: + if value.lower() in {"none", "never", "off"}: + return None + parsed = int(value) + if parsed < 0: + raise argparse.ArgumentTypeError("step must be non-negative or 'none'") + return parsed + + +def _add_prepare(subparsers: Any) -> None: + parser = subparsers.add_parser( + "prepare", + help="Validate and prepare user-owned speech data.", + description=( + "Convert a CSV or JSONL manifest into Inflect's versioned, 24 kHz " + "prepared-dataset format." + ), + ) + parser.add_argument("--manifest", type=_path, required=True) + parser.add_argument("--audio-root", type=_path) + parser.add_argument("--language", default="en-us") + parser.add_argument( + "--frontend", + choices=("espeak", "prephonemized", "custom"), + default="espeak", + help="Use eSpeak NG, manifest phonemes, or an explicit custom frontend hook.", + ) + parser.add_argument( + "--frontend-hook", + help=( + "Trusted Python factory in module:callable or file.py:function form. " + "Required for --frontend custom; loading it executes that Python code." + ), + ) + parser.add_argument("--output", type=_path, required=True) + parser.add_argument("--validation-fraction", type=_fraction, default=0.05) + parser.add_argument("--split-seed", type=int, default=1337) + parser.add_argument("--min-duration-seconds", type=float, default=0.05) + parser.add_argument("--max-duration-seconds", type=float) + parser.add_argument("--base-symbols", type=_path) + parser.set_defaults(handler=_run_prepare) + + +def _add_audit(subparsers: Any) -> None: + parser = subparsers.add_parser( + "audit", + help="Audit a prepared dataset before training.", + ) + parser.add_argument("--dataset", type=_path, required=True) + parser.add_argument( + "--strict", + action=argparse.BooleanOptionalAction, + default=True, + help="Treat all structural warnings selected by the auditor as fatal.", + ) + parser.add_argument("--duration-tolerance-seconds", type=float, default=0.02) + parser.set_defaults(handler=_run_audit) + + +def _add_train(subparsers: Any) -> None: + parser = subparsers.add_parser( + "train", + help="Warm-start a fixed-voice, single-language checkpoint.", + description=( + "Run generic staged adaptation from a released Inflect generator. " + "The public checkpoint has no posterior encoder, optimizer, or discriminator; " + "those training-only components are initialized by this toolkit." + ), + ) + parser.add_argument( + "--base", + required=True, + help="micro, nano, a local model directory, or a Hugging Face repo ID", + ) + parser.add_argument("--dataset", type=_path, required=True) + parser.add_argument("--output", type=_path, required=True) + parser.add_argument("--preset", default="balanced") + parser.add_argument("--resume", type=_path) + parser.add_argument( + "--device", default=argparse.SUPPRESS, help="auto, cpu, cuda, or cuda:N" + ) + parser.add_argument("--seed", type=int, default=argparse.SUPPRESS) + parser.add_argument("--batch-size", type=_positive_int, default=argparse.SUPPRESS) + parser.add_argument( + "--gradient-accumulation-steps", + type=_positive_int, + default=argparse.SUPPRESS, + ) + parser.add_argument("--num-workers", type=int, default=argparse.SUPPRESS) + parser.add_argument("--max-steps", type=_positive_int, default=argparse.SUPPRESS) + parser.add_argument("--learning-rate-g", type=float, default=argparse.SUPPRESS) + parser.add_argument("--learning-rate-d", type=float, default=argparse.SUPPRESS) + parser.add_argument( + "--posterior-warmup-steps", type=int, default=argparse.SUPPRESS + ) + parser.add_argument( + "--decoder-unfreeze-step", + type=_optional_step, + default=argparse.SUPPRESS, + metavar="STEP|none", + ) + parser.add_argument( + "--amp", + action=argparse.BooleanOptionalAction, + default=argparse.SUPPRESS, + help="Use automatic mixed precision on CUDA.", + ) + parser.add_argument( + "--checkpoint-interval", type=_positive_int, default=argparse.SUPPRESS + ) + parser.add_argument( + "--validation-interval", type=_positive_int, default=argparse.SUPPRESS + ) + parser.add_argument("--log-interval", type=_positive_int, default=argparse.SUPPRESS) + parser.set_defaults(handler=_run_train) + + +def _add_evaluate(subparsers: Any) -> None: + parser = subparsers.add_parser( + "evaluate", + help="Synthesize held-out text and write machine-readable diagnostics.", + ) + parser.add_argument("--model-dir", type=_path, required=True) + parser.add_argument("--manifest", type=_path, required=True) + parser.add_argument("--output", type=_path, required=True) + parser.add_argument("--checkpoint", type=_path) + parser.add_argument("--device", default="cpu") + parser.add_argument("--max-samples", type=_positive_int) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--speed", type=float, default=1.0) + parser.add_argument("--variation", type=float, default=0.667) + parser.add_argument( + "--transcript-evaluator", + help="Optional module:attribute callable that returns a transcript or metrics.", + ) + parser.add_argument("--overwrite", action="store_true") + parser.add_argument( + "--save-audio", + action=argparse.BooleanOptionalAction, + default=True, + ) + parser.set_defaults(handler=_run_evaluate) + + +def _add_export(subparsers: Any) -> None: + parser = subparsers.add_parser( + "export", + help="Strip training state and create a verified inference package.", + ) + parser.add_argument("--checkpoint", type=_path, required=True) + parser.add_argument("--output", type=_path, required=True) + parser.add_argument("--format", choices=("pytorch", "onnx"), default="pytorch") + parser.add_argument("--config", type=_path) + parser.add_argument("--symbols", type=_path) + parser.add_argument( + "--prepared-dataset", + type=_path, + help=( + "Prepared directory or dataset.json containing the exact language " + "frontend metadata for this checkpoint." + ), + ) + parser.add_argument( + "--frontend-hook", + type=_path, + help=( + "Matching trusted .py frontend source. Required only for exports " + "prepared with --frontend custom." + ), + ) + parser.add_argument( + "--package-template", + type=_path, + help="Released Micro/Nano directory whose public runtime should be copied.", + ) + parser.add_argument("--onnx-opset", type=int, default=17) + parser.add_argument("--model-name") + parser.add_argument("--source-revision") + parser.add_argument("--overwrite", action="store_true") + parser.add_argument( + "--verify", + action=argparse.BooleanOptionalAction, + default=True, + ) + parser.set_defaults(handler=_run_export) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="inflect-adapt", + description=( + "Prepare data, audit it, warm-start Inflect v2, evaluate held-out " + "speech, and export inference-only packages." + ), + ) + parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}") + parser.add_argument( + "--verbose", + action="store_true", + help="Enable informational logging.", + ) + subparsers = parser.add_subparsers(dest="command", required=True) + _add_prepare(subparsers) + _add_audit(subparsers) + _add_train(subparsers) + _add_evaluate(subparsers) + _add_export(subparsers) + return parser + + +def _run_prepare(args: argparse.Namespace) -> dict[str, Any]: + from .prepare import PrepareOptions, prepare_dataset + + return prepare_dataset( + PrepareOptions( + manifest_path=args.manifest, + audio_root=args.audio_root, + language=args.language, + frontend=args.frontend, + frontend_hook=args.frontend_hook, + output_dir=args.output, + validation_fraction=args.validation_fraction, + split_seed=args.split_seed, + min_duration_seconds=args.min_duration_seconds, + max_duration_seconds=args.max_duration_seconds, + base_symbols_path=args.base_symbols, + ) + ) + + +def _run_audit(args: argparse.Namespace) -> dict[str, Any]: + from .audit import AuditOptions, audit_dataset + + return audit_dataset( + AuditOptions( + prepared_dir=args.dataset, + strict=args.strict, + duration_tolerance_seconds=args.duration_tolerance_seconds, + ) + ) + + +def _run_train(args: argparse.Namespace) -> dict[str, Any]: + from .training import TrainingOptions, train_adaptation + + override_names = ( + "resume", + "device", + "seed", + "batch_size", + "gradient_accumulation_steps", + "num_workers", + "max_steps", + "learning_rate_g", + "learning_rate_d", + "posterior_warmup_steps", + "decoder_unfreeze_step", + "amp", + "checkpoint_interval", + "validation_interval", + "log_interval", + ) + parsed = vars(args) + overrides = {name: parsed[name] for name in override_names if name in parsed} + options = TrainingOptions.from_preset( + args.preset, + base_model=args.base, + prepared_dir=args.dataset, + output_dir=args.output, + **overrides, + ) + return train_adaptation(options) + + +def _run_evaluate(args: argparse.Namespace) -> dict[str, Any]: + from .evaluation import EvaluationOptions, evaluate_checkpoint + + return evaluate_checkpoint( + EvaluationOptions( + model_dir=args.model_dir, + manifest=args.manifest, + output_dir=args.output, + checkpoint=args.checkpoint, + transcript_evaluator=args.transcript_evaluator, + device=args.device, + max_samples=args.max_samples, + seed=args.seed, + speed=args.speed, + variation=args.variation, + overwrite=args.overwrite, + save_audio=args.save_audio, + ) + ) + + +def _run_export(args: argparse.Namespace) -> dict[str, Any]: + from .exporting import ExportOptions, export_checkpoint + + return export_checkpoint( + ExportOptions( + checkpoint=args.checkpoint, + output_dir=args.output, + config=args.config, + symbols=args.symbols, + prepared_dataset=args.prepared_dataset, + frontend_hook=args.frontend_hook, + package_template=args.package_template, + include_onnx=args.format == "onnx", + onnx_opset=args.onnx_opset, + model_name=args.model_name, + source_revision=args.source_revision, + overwrite=args.overwrite, + verify=args.verify, + ) + ) + + +def main(argv: Sequence[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + logging.basicConfig( + level=logging.INFO if args.verbose else logging.WARNING, + format="%(levelname)s: %(message)s", + ) + try: + report = args.handler(args) + except KeyboardInterrupt: + print("Interrupted.", file=sys.stderr) + return 130 + except Exception as exc: + if args.verbose: + logging.exception("Command failed") + else: + print(f"error: {exc}", file=sys.stderr) + return 1 + print(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True, default=str)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/finetune/inflect_finetune/evaluation.py b/finetune/inflect_finetune/evaluation.py new file mode 100644 index 0000000..68300b3 --- /dev/null +++ b/finetune/inflect_finetune/evaluation.py @@ -0,0 +1,521 @@ +"""Lightweight held-out synthesis and signal diagnostics.""" + +from __future__ import annotations + +import importlib.util +import inspect +import json +import math +import re +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Protocol, Sequence + +import numpy as np +import soundfile as sf +from scipy import signal + +from .reporting import file_record, make_report, status, write_json, write_text + + +class Synthesizer(Protocol): + def __call__(self, text: str | None = None, **kwargs: Any) -> Any: ... + + +class TranscriptEvaluator(Protocol): + def __call__( + self, + audio_path: Path, + reference_text: str, + sample_rate: int, + ) -> str | Mapping[str, Any]: ... + + +@dataclass(slots=True) +class EvaluationOptions: + """Inputs for :func:`evaluate_checkpoint`. + + A custom ``synthesizer`` avoids coupling evaluation to a particular + trainer. If omitted, ``model_dir/inference.py`` is loaded. A manifest can + also reference existing audio, allowing diagnostics without model loading. + No ASR model is installed or downloaded; transcript scoring only runs when + the caller explicitly supplies ``transcript_evaluator``. + """ + + model_dir: str | Path | None + manifest: str | Path + output_dir: str | Path + checkpoint: str | Path | None = None + synthesizer: Synthesizer | None = None + transcript_evaluator: TranscriptEvaluator | str | None = None + device: str = "cpu" + max_samples: int | None = None + seed: int = 0 + speed: float = 1.0 + variation: float = 0.667 + overwrite: bool = False + save_audio: bool = True + clipping_threshold: float = 0.999 + silence_threshold_db: float = -50.0 + frame_ms: float = 25.0 + + +def _read_manifest(path: Path) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + try: + with path.open("r", encoding="utf-8") as handle: + for line_number, line in enumerate(handle, start=1): + if not line.strip(): + continue + row = json.loads(line) + if not isinstance(row, dict): + raise ValueError(f"line {line_number} is not a JSON object") + text = row.get("normalized_text") or row.get("text") + phonemes = ( + row.get("phonemes") + or row.get("phoneme_text") + or row.get("phones") + ) + has_text = isinstance(text, str) and bool(text.strip()) + has_phonemes = isinstance(phonemes, str) and bool(phonemes.strip()) + if not has_text and not has_phonemes: + raise ValueError( + f"line {line_number} has neither non-empty text nor phonemes" + ) + row["_line"] = line_number + rows.append(row) + except (OSError, json.JSONDecodeError, ValueError) as exc: + raise ValueError(f"Invalid evaluation manifest {path}: {exc}") from exc + if not rows: + raise ValueError(f"Evaluation manifest is empty: {path}") + return rows + + +def _safe_id(row: Mapping[str, Any], index: int) -> str: + raw = str(row.get("id") or row.get("key") or f"sample-{index:04d}") + cleaned = re.sub(r"[^A-Za-z0-9._-]+", "-", raw).strip(".-") + return cleaned or f"sample-{index:04d}" + + +def _load_default_synthesizer( + model_dir: Path, + device: str, + checkpoint: Path | None, +) -> Synthesizer: + inference_path = model_dir / "inference.py" + if not inference_path.is_file(): + raise FileNotFoundError( + f"No inference.py was found in {model_dir}. Pass EvaluationOptions(synthesizer=...)." + ) + module_name = f"_inflect_eval_inference_{abs(hash(inference_path))}" + module_names = ( + "inference", + "deployment_frontend", + "inflect_vits_frontend", + "models", + "commons", + "utils", + "modules", + "attentions", + "transforms", + "text", + "text.symbols", + ) + previous = {name: sys.modules.get(name) for name in module_names} + old_path = list(sys.path) + try: + sys.path.insert(0, str(model_dir)) + sys.path.insert(0, str(model_dir / "runtime")) + for name in module_names: + sys.modules.pop(name, None) + spec = importlib.util.spec_from_file_location(module_name, inference_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Could not import {inference_path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + engine = module.InflectTTS(model_dir=model_dir, device=device) + finally: + sys.path[:] = old_path + for name in module_names: + sys.modules.pop(name, None) + if previous[name] is not None: + sys.modules[name] = previous[name] + if checkpoint is not None and checkpoint.resolve() != (model_dir / "model.pth").resolve(): + from .exporting import _extract_state, _load_checkpoint + + state, _, _ = _extract_state(_load_checkpoint(checkpoint)) + incompatible = engine.model.load_state_dict(state, strict=True) + if incompatible.missing_keys or incompatible.unexpected_keys: + raise RuntimeError( + "Evaluation checkpoint did not load strictly: " + f"missing={list(incompatible.missing_keys)}, " + f"unexpected={list(incompatible.unexpected_keys)}" + ) + return engine.synthesize + + +def _load_transcript_evaluator( + value: TranscriptEvaluator | str | None, +) -> TranscriptEvaluator | None: + if value is None or callable(value): + return value + if ":" not in value: + raise ValueError("Transcript evaluator must use 'module_or_file:function' syntax.") + source, function_name = value.rsplit(":", 1) + path = Path(source) + if path.is_file(): + spec = importlib.util.spec_from_file_location( + f"_inflect_transcript_eval_{abs(hash(path.resolve()))}", path.resolve() + ) + if spec is None or spec.loader is None: + raise RuntimeError(f"Could not import transcript evaluator: {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + else: + spec = importlib.util.find_spec(source) + if spec is None or spec.loader is None: + raise ImportError(f"Transcript evaluator module is unavailable: {source}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + evaluator = getattr(module, function_name, None) + if not callable(evaluator): + raise ValueError(f"Transcript evaluator is not callable: {value}") + return evaluator + + +def _invoke_synthesizer( + synthesizer: Synthesizer, + text: str | None, + *, + phonemes: str | None, + seed: int, + speed: float, + variation: float, +) -> tuple[int, np.ndarray, dict[str, Any]]: + kwargs = {"seed": seed, "speed": speed, "variation": variation} + if phonemes is not None: + kwargs["phonemes"] = phonemes + try: + signature = inspect.signature(synthesizer) + if not any(param.kind == param.VAR_KEYWORD for param in signature.parameters.values()): + kwargs = {key: val for key, val in kwargs.items() if key in signature.parameters} + except (TypeError, ValueError): + pass + result = synthesizer(text, **kwargs) + metadata: dict[str, Any] = {} + if isinstance(result, Mapping): + sample_rate = result.get("sample_rate") + waveform = result.get("waveform", result.get("audio")) + metadata = dict(result.get("metadata", {})) + elif isinstance(result, tuple) and len(result) >= 2: + sample_rate, waveform = result[:2] + if len(result) >= 3 and isinstance(result[2], Mapping): + metadata = dict(result[2]) + else: + raise TypeError( + "Synthesizer must return (sample_rate, waveform) or a mapping with those values." + ) + audio = np.asarray(waveform, dtype=np.float32).squeeze() + if audio.ndim != 1: + raise ValueError(f"Synthesizer returned non-mono audio with shape {audio.shape}.") + return int(sample_rate), audio, metadata + + +def _load_existing_audio(path: Path) -> tuple[int, np.ndarray]: + waveform, sample_rate = sf.read(path, dtype="float32", always_2d=False) + audio = np.asarray(waveform, dtype=np.float32) + if audio.ndim == 2: + audio = np.mean(audio, axis=1, dtype=np.float32) + return int(sample_rate), audio + + +def _leading_trailing_silence(mask: np.ndarray, sample_rate: int) -> tuple[float, float]: + non_silent = np.flatnonzero(~mask) + if not non_silent.size: + duration = mask.size / sample_rate + return duration, duration + return non_silent[0] / sample_rate, (mask.size - 1 - non_silent[-1]) / sample_rate + + +def _signal_metrics( + waveform: np.ndarray, + sample_rate: int, + *, + clipping_threshold: float, + silence_threshold_db: float, + frame_ms: float, +) -> dict[str, Any]: + if waveform.size == 0: + raise ValueError("Waveform is empty.") + finite = np.isfinite(waveform) + safe = np.nan_to_num(waveform, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float64) + absolute = np.abs(safe) + peak = float(np.max(absolute)) + rms = float(np.sqrt(np.mean(np.square(safe)))) + rms_dbfs = 20.0 * math.log10(max(rms, 1e-12)) + silence_amplitude = 10.0 ** (silence_threshold_db / 20.0) + sample_silence = absolute <= silence_amplitude + leading, trailing = _leading_trailing_silence(sample_silence, sample_rate) + frame_length = max(1, round(sample_rate * frame_ms / 1000.0)) + usable = safe[: (safe.size // frame_length) * frame_length] + if usable.size: + frame_rms = np.sqrt(np.mean(np.square(usable.reshape(-1, frame_length)), axis=1)) + silent_frames = float(np.mean(frame_rms <= silence_amplitude)) + else: + silent_frames = float(np.mean(sample_silence)) + zero_crossing = ( + float(np.mean(np.signbit(safe[1:]) != np.signbit(safe[:-1]))) + if safe.size > 1 + else 0.0 + ) + nperseg = min(1024, safe.size) + frequencies, spectrum = signal.welch(safe, fs=sample_rate, nperseg=nperseg) + power_sum = float(np.sum(spectrum)) + if power_sum > 0: + centroid = float(np.sum(frequencies * spectrum) / power_sum) + high_frequency = frequencies >= min(8000, sample_rate * 0.4) + high_ratio = float(np.sum(spectrum[high_frequency]) / power_sum) + else: + centroid = 0.0 + high_ratio = 0.0 + return { + "sample_rate": sample_rate, + "samples": int(safe.size), + "duration_seconds": safe.size / sample_rate, + "all_finite": bool(np.all(finite)), + "non_finite_samples": int(np.count_nonzero(~finite)), + "peak": peak, + "rms": rms, + "rms_dbfs": rms_dbfs, + "dc_offset": float(np.mean(safe)), + "clipped_fraction": float(np.mean(absolute >= clipping_threshold)), + "silent_sample_fraction": float(np.mean(sample_silence)), + "silent_frame_fraction": silent_frames, + "leading_silence_seconds": leading, + "trailing_silence_seconds": trailing, + "zero_crossing_rate": zero_crossing, + "crest_factor_db": 20.0 * math.log10(max(peak, 1e-12) / max(rms, 1e-12)), + "spectral_centroid_hz": centroid, + "high_frequency_energy_ratio": high_ratio, + } + + +def _percentile(values: Sequence[float], percentile: float) -> float | None: + return float(np.percentile(values, percentile)) if values else None + + +def _aggregate(rows: Sequence[Mapping[str, Any]]) -> dict[str, Any]: + metric_names = ( + "duration_seconds", + "peak", + "rms_dbfs", + "dc_offset", + "clipped_fraction", + "silent_sample_fraction", + "silent_frame_fraction", + "leading_silence_seconds", + "trailing_silence_seconds", + "zero_crossing_rate", + "crest_factor_db", + "spectral_centroid_hz", + "high_frequency_energy_ratio", + "characters_per_second", + "words_per_second", + ) + result: dict[str, Any] = {} + for name in metric_names: + values = [ + float(row["signal"][name] if name in row["signal"] else row[name]) + for row in rows + if name in row.get("signal", {}) or name in row + ] + result[name] = { + "mean": float(np.mean(values)) if values else None, + "p50": _percentile(values, 50), + "p95": _percentile(values, 95), + "max": max(values) if values else None, + } + result["clips_with_clipping"] = sum( + row["signal"]["clipped_fraction"] > 0 for row in rows + ) + result["clips_all_silent"] = sum( + row["signal"]["silent_sample_fraction"] >= 0.999 for row in rows + ) + result["clips_with_non_finite_samples"] = sum( + not row["signal"]["all_finite"] for row in rows + ) + return result + + +def _run_transcript_evaluator( + evaluator: TranscriptEvaluator, + audio_path: Path, + reference: str, + sample_rate: int, +) -> dict[str, Any]: + result = evaluator(audio_path, reference, sample_rate) + if isinstance(result, str): + return {"transcript": result} + if isinstance(result, Mapping): + return dict(result) + raise TypeError("Transcript evaluator must return a string or mapping.") + + +def evaluate_checkpoint(options: EvaluationOptions) -> dict[str, Any]: + """Synthesize/evaluate held-out rows and write JSON plus a short summary.""" + + manifest = Path(options.manifest).resolve() + output = Path(options.output_dir).resolve() + if output.exists() and any(output.iterdir()) and not options.overwrite: + raise FileExistsError(f"Output directory is not empty: {output}") + output.mkdir(parents=True, exist_ok=True) + audio_dir = output / "audio" + rows = _read_manifest(manifest) + if options.max_samples is not None: + if options.max_samples <= 0: + raise ValueError("max_samples must be positive when supplied.") + rows = rows[: options.max_samples] + + model_dir = Path(options.model_dir).resolve() if options.model_dir is not None else None + checkpoint = ( + Path(options.checkpoint).resolve() if options.checkpoint is not None else None + ) + if checkpoint is not None and not checkpoint.is_file(): + raise FileNotFoundError(f"Evaluation checkpoint does not exist: {checkpoint}") + synthesizer = options.synthesizer + if synthesizer is None and not all(row.get("audio") for row in rows): + if model_dir is None: + raise ValueError("model_dir or a synthesizer is required for rows without audio.") + synthesizer = _load_default_synthesizer(model_dir, options.device, checkpoint) + transcript_evaluator = _load_transcript_evaluator(options.transcript_evaluator) + if options.save_audio or transcript_evaluator is not None: + audio_dir.mkdir(parents=True, exist_ok=True) + + evaluated: list[dict[str, Any]] = [] + failures: list[dict[str, Any]] = [] + for index, row in enumerate(rows): + sample_id = _safe_id(row, index) + raw_text = row.get("normalized_text") or row.get("text") + text = str(raw_text).strip() if isinstance(raw_text, str) else "" + raw_phonemes = ( + row.get("phonemes") + or row.get("phoneme_text") + or row.get("phones") + ) + phonemes = ( + str(raw_phonemes).strip() + if isinstance(raw_phonemes, str) and raw_phonemes.strip() + else None + ) + try: + synthesis_metadata: dict[str, Any] = {} + if synthesizer is not None: + sample_rate, waveform, synthesis_metadata = _invoke_synthesizer( + synthesizer, + text or None, + phonemes=phonemes, + seed=options.seed + index, + speed=options.speed, + variation=options.variation, + ) + else: + source_audio = (manifest.parent / str(row["audio"])).resolve() + try: + source_audio.relative_to(manifest.parent) + except ValueError as exc: + raise ValueError( + f"Audio path escapes manifest directory: {row['audio']}" + ) from exc + sample_rate, waveform = _load_existing_audio(source_audio) + metrics = _signal_metrics( + waveform, + sample_rate, + clipping_threshold=options.clipping_threshold, + silence_threshold_db=options.silence_threshold_db, + frame_ms=options.frame_ms, + ) + duration = metrics["duration_seconds"] + scoring_text = text or phonemes or "" + words = len(scoring_text.split()) + result: dict[str, Any] = { + "id": sample_id, + "text": text, + "input_mode": "prephonemized" if phonemes is not None else "text", + "phoneme_characters": len(phonemes) if phonemes is not None else None, + "seed": options.seed + index if synthesizer is not None else None, + "signal": metrics, + "characters_per_second": ( + len(scoring_text) / duration if duration else None + ), + "words_per_second": words / duration if duration else None, + "synthesis_metadata": synthesis_metadata, + } + destination = audio_dir / f"{sample_id}.wav" + if options.save_audio or transcript_evaluator is not None: + sf.write(destination, waveform, sample_rate, subtype="PCM_16") + result["audio"] = file_record(destination, relative_to=output) + if transcript_evaluator is not None: + if not text: + raise ValueError( + "Transcript evaluation requires reference text even when " + "synthesis uses prephonemized input." + ) + result["transcript_evaluation"] = _run_transcript_evaluator( + transcript_evaluator, + destination, + str(row.get("text") or text), + sample_rate, + ) + evaluated.append(result) + except Exception as exc: + failures.append( + { + "id": sample_id, + "line": row.get("_line"), + "error_type": type(exc).__name__, + "error": str(exc), + } + ) + + checks = [ + status(bool(evaluated), "At least one held-out item evaluated successfully"), + status(not failures, "All requested held-out items evaluated", failures=len(failures)), + ] + report = make_report( + "evaluation_report", + ok=bool(evaluated) and not failures, + source={ + "manifest": manifest.name, + "model_dir": model_dir.name if model_dir else None, + "checkpoint": checkpoint.name if checkpoint else None, + "device": options.device, + }, + settings={ + "seed": options.seed, + "speed": options.speed, + "variation": options.variation, + "clipping_threshold": options.clipping_threshold, + "silence_threshold_db": options.silence_threshold_db, + "frame_ms": options.frame_ms, + "transcript_evaluator_enabled": transcript_evaluator is not None, + }, + counts={"requested": len(rows), "evaluated": len(evaluated), "failed": len(failures)}, + checks=checks, + aggregate=_aggregate(evaluated), + items=evaluated, + failures=failures, + ) + write_json(output / "evaluation_report.json", report) + summary = [ + "Inflect adaptation evaluation", + f"Evaluated: {len(evaluated)}/{len(rows)}", + f"Failures: {len(failures)}", + f"Clips with clipping: {report['aggregate'].get('clips_with_clipping', 0)}", + f"All-silent clips: {report['aggregate'].get('clips_all_silent', 0)}", + "Transcript evaluator: " + + ("enabled (caller supplied)" if transcript_evaluator else "disabled"), + f"Result: {'PASS' if report['ok'] else 'FAIL'}", + ] + write_text(output / "evaluation_summary.txt", "\n".join(summary)) + return report diff --git a/finetune/inflect_finetune/exporting.py b/finetune/inflect_finetune/exporting.py new file mode 100644 index 0000000..83e4054 --- /dev/null +++ b/finetune/inflect_finetune/exporting.py @@ -0,0 +1,1686 @@ +"""Inference-only PyTorch and optional ONNX exports for adapted checkpoints.""" + +from __future__ import annotations + +import copy +import hashlib +import importlib +import json +import shutil +import sys +import textwrap +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterator, Mapping, Sequence + +import numpy as np +import torch +from torch import nn + +from .reporting import file_record, make_report, sha256_file, status, write_checksums, write_json + + +_TRAINING_ONLY_PREFIXES = ( + "enc_q.", + "discriminator.", + "discriminators.", + "disc.", + "mpd.", + "net_d.", + "optimizer.", + "optim.", + "scheduler.", + "scaler.", + "amp_scaler.", +) + +_TRAINING_ONLY_TOP_LEVEL_FIELDS = frozenset( + { + "compatibility", + "discriminator", + "discriminators", + "disc", + "epoch", + "mpd", + "net_d", + "optim", + "optimizer", + "optimizer_d", + "optimizer_g", + "options", + "rng_state", + "scaler", + "scheduler", + "scheduler_d", + "scheduler_g", + "stage", + "step", + } +) + + +def _is_training_only_name(name: str) -> bool: + lowered = name.lower() + return lowered in _TRAINING_ONLY_TOP_LEVEL_FIELDS or any( + lowered == prefix[:-1] or lowered.startswith(prefix) + for prefix in _TRAINING_ONLY_PREFIXES + ) + + +@dataclass(slots=True) +class ExportOptions: + """Inputs for :func:`export_checkpoint`. + + ``package_template`` may point at a released Inflect package. Its public + runtime and inference files are copied into the result and used for strict + model-load verification. The source checkpoint may be either an inference + checkpoint or a training checkpoint containing ``model``/``state_dict``. + """ + + checkpoint: str | Path + output_dir: str | Path + config: str | Path | None = None + symbols: str | Path | Sequence[str] | Mapping[str, Any] | None = None + package_template: str | Path | None = None + prepared_dataset: str | Path | Mapping[str, Any] | None = None + frontend_hook: str | Path | None = None + include_onnx: bool = False + onnx_opset: int = 17 + model_name: str | None = None + source_revision: str | None = None + overwrite: bool = False + verify: bool = True + + +def _load_json(path: Path) -> Any: + try: + return json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"Could not read valid JSON from {path}: {exc}") from exc + + +def _canonical_json_sha256(value: Any) -> str: + encoded = json.dumps( + value, + ensure_ascii=False, + allow_nan=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def _resolve_config( + options: ExportOptions, + checkpoint: Path, + package_template: Path | None, +) -> tuple[dict[str, Any], Path | None]: + candidates = [] + if options.config is not None: + candidates.append(Path(options.config)) + candidates.extend((checkpoint.parent / "config.json", checkpoint.parent.parent / "config.json")) + if package_template is not None: + candidates.append(package_template / "config.json") + for candidate in candidates: + if candidate.is_file(): + config = _load_json(candidate) + if not isinstance(config, dict): + raise ValueError(f"Config must contain a JSON object: {candidate}") + return config, candidate.resolve() + raise FileNotFoundError( + "No config.json was found. Pass ExportOptions(config=...) or place it beside " + "the checkpoint." + ) + + +def _symbols_from_payload(payload: Any) -> list[str]: + if isinstance(payload, (list, tuple)): + symbols = list(payload) + elif isinstance(payload, Mapping): + for key in ("symbols", "ordered_symbols", "inventory"): + if key in payload: + return _symbols_from_payload(payload[key]) + raise ValueError("symbols.json must contain 'symbols', 'ordered_symbols', or 'inventory'.") + else: + raise ValueError("Symbols must be an ordered JSON list or an object containing one.") + if not symbols or any(not isinstance(symbol, str) or not symbol for symbol in symbols): + raise ValueError("Every symbol must be a non-empty string.") + # The published v2 inventory contains a legacy duplicate apostrophe row. + # Preserve ordered inventories exactly because embedding migration is by + # symbol identity and numeric compatibility depends on every row. + return symbols + + +def _load_symbols( + options: ExportOptions, + checkpoint: Path, + checkpoint_payload: Mapping[str, Any], + package_template: Path | None, +) -> tuple[list[str], Path | None]: + source = options.symbols + if source is None: + for key in ("symbols", "symbol_inventory"): + if key in checkpoint_payload: + return _symbols_from_payload(checkpoint_payload[key]), None + candidates = [ + checkpoint.parent / "symbols.json", + checkpoint.parent.parent / "symbols.json", + ] + if package_template is not None: + candidates.append(package_template / "symbols.json") + for candidate in candidates: + if candidate.is_file(): + return _symbols_from_payload(_load_json(candidate)), candidate.resolve() + if package_template is not None: + symbols_py = package_template / "runtime" / "text" / "symbols.py" + if symbols_py.is_file(): + namespace: dict[str, Any] = {} + source = symbols_py.read_text(encoding="utf-8") + exec(compile(source, str(symbols_py), "exec"), namespace) + return _symbols_from_payload(namespace["symbols"]), symbols_py.resolve() + raise FileNotFoundError( + "No ordered symbol inventory was found. Pass ExportOptions(symbols=...) or place " + "symbols.json beside the checkpoint." + ) + if isinstance(source, (str, Path)): + path = Path(source) + return _symbols_from_payload(_load_json(path)), path.resolve() + return _symbols_from_payload(source), None + + +def _extract_state( + payload: Any, +) -> tuple[dict[str, torch.Tensor], dict[str, Any], list[str]]: + if not isinstance(payload, Mapping): + raise ValueError("Checkpoint must contain a mapping.") + state: Any = payload + for key in ("model", "state_dict", "generator", "net_g"): + if key in payload and isinstance(payload[key], Mapping): + state = payload[key] + break + if not isinstance(state, Mapping): + raise ValueError("Checkpoint does not contain a model state dictionary.") + tensors: dict[str, torch.Tensor] = {} + stripped: list[str] = [] + for raw_key, value in state.items(): + if not isinstance(value, torch.Tensor): + continue + key = str(raw_key) + for prefix in ("module.", "model.", "generator."): + if key.startswith(prefix): + key = key[len(prefix) :] + break + if _is_training_only_name(key): + stripped.append(key) + continue + tensors[key] = value.detach().cpu().contiguous() + if not tensors: + raise ValueError("Checkpoint state dictionary contains no tensors.") + metadata = {str(key): value for key, value in payload.items() if value is not state} + return tensors, metadata, sorted(stripped) + + +def _load_checkpoint(path: Path) -> Mapping[str, Any]: + try: + payload = torch.load(path, map_location="cpu", weights_only=True) + except Exception as safe_error: + try: + payload = torch.load(path, map_location="cpu", weights_only=False) + except Exception: + raise safe_error + if not isinstance(payload, Mapping): + raise ValueError(f"Checkpoint must contain a mapping: {path}") + return payload + + +def _resolve_package_template( + options: ExportOptions, + checkpoint_payload: Mapping[str, Any], +) -> Path | None: + if options.package_template is not None: + candidate: str | Path | None = options.package_template + else: + training_options = checkpoint_payload.get("options") + candidate = ( + training_options.get("base_model") + if isinstance(training_options, Mapping) + else None + ) + if candidate is None: + return None + try: + from .modeling import resolve_base_model + + return resolve_base_model(candidate) + except FileNotFoundError: + path = Path(candidate).expanduser() + if path.is_dir(): + return path.resolve() + if options.package_template is not None: + raise + return None + + +def _is_unmodified_template_checkpoint(checkpoint: Path, template: Path | None) -> bool: + template_checkpoint = template / "model.pth" if template is not None else None + return bool( + template_checkpoint + and template_checkpoint.is_file() + and sha256_file(checkpoint) == sha256_file(template_checkpoint) + ) + + +def _prepared_dataset_candidates( + options: ExportOptions, + checkpoint: Path, +) -> list[Path]: + candidates: list[Path] = [] + if isinstance(options.symbols, (str, Path)): + symbols_path = Path(options.symbols).expanduser() + candidates.append(symbols_path.parent / "dataset.json") + candidates.extend( + ( + checkpoint.parent / "dataset.json", + checkpoint.parent.parent / "dataset.json", + checkpoint.parent.parent / "prepared" / "dataset.json", + ) + ) + return candidates + + +def _load_prepared_dataset( + options: ExportOptions, + checkpoint: Path, + package_template: Path | None, +) -> tuple[dict[str, Any], Path | None, bool]: + source = options.prepared_dataset + if isinstance(source, Mapping): + return dict(source), None, False + candidates: list[Path] = [] + if source is not None: + explicit = Path(source).expanduser() + candidates.append(explicit / "dataset.json" if explicit.is_dir() else explicit) + else: + candidates.extend(_prepared_dataset_candidates(options, checkpoint)) + for candidate in candidates: + if candidate.is_file(): + payload = _load_json(candidate) + if not isinstance(payload, dict): + raise ValueError(f"Prepared dataset metadata must be an object: {candidate}") + return payload, candidate.resolve(), False + if source is not None: + raise FileNotFoundError( + "Prepared dataset metadata does not exist. Pass a prepared directory or " + f"dataset.json, not {source!r}." + ) + if _is_unmodified_template_checkpoint(checkpoint, package_template): + return { + "format": "inflect_release_frontend_compatibility_v1", + "language": "en-us", + "sample_rate": 24000, + "frontend": { + "type": "espeak", + "language": "en-us", + "preserve_punctuation": True, + "with_stress": True, + }, + }, None, True + raise ValueError( + "A non-template checkpoint requires prepared dataset frontend metadata. Pass " + "ExportOptions(prepared_dataset=...) or, through the existing CLI, pass " + "--symbols PREPARED_DIR/symbols.json so sibling dataset.json can be resolved. " + "Export refuses to reuse the release English frontend for an adapted checkpoint." + ) + + +def _validate_prepared_symbols( + dataset_path: Path | None, + symbols: Sequence[str], +) -> dict[str, Any]: + if dataset_path is None: + return status( + True, + "Prepared symbol inventory check skipped for inline or release metadata", + skipped=True, + ) + symbols_path = dataset_path.parent / "symbols.json" + if not symbols_path.is_file(): + raise FileNotFoundError( + f"Prepared dataset is missing its ordered symbol inventory: {symbols_path}" + ) + prepared_symbols = _symbols_from_payload(_load_json(symbols_path)) + return status( + list(prepared_symbols) == list(symbols), + "Export symbols exactly match the prepared dataset inventory", + prepared_symbol_count=len(prepared_symbols), + export_symbol_count=len(symbols), + ) + + +def _custom_hook_contract( + options: ExportOptions, + frontend: Mapping[str, Any], + language: str, + symbols: Sequence[str], +) -> tuple[dict[str, Any], Path]: + hook_metadata = frontend.get("hook") + if not isinstance(hook_metadata, Mapping): + raise ValueError( + "Custom prepared frontend metadata has no reproducibility record. " + "Re-run dataset preparation with an explicit custom hook." + ) + if options.frontend_hook is None: + raise ValueError( + "Custom frontend export requires ExportOptions(frontend_hook=PATH_TO_HOOK.py). " + "The hook is copied into the package; export will never substitute English." + ) + hook_file = Path(options.frontend_hook).expanduser().resolve() + if not hook_file.is_file() or hook_file.suffix.lower() != ".py": + raise ValueError(f"Custom frontend hook must be an existing .py file: {hook_file}") + expected_source_hash = hook_metadata.get("source_sha256") + actual_source_hash = sha256_file(hook_file) + if expected_source_hash != actual_source_hash: + raise ValueError( + "Custom frontend hook source does not match prepared dataset metadata: " + f"expected {expected_source_hash}, got {actual_source_hash}." + ) + identity = str(hook_metadata.get("identity", "")) + _, separator, factory_name = identity.rpartition(":") + if not separator or not factory_name.isidentifier(): + raise ValueError( + "Prepared custom frontend identity does not name a valid top-level factory." + ) + from .frontend import ( + FrontendOptions, + custom_frontend_metadata, + custom_frontend_symbols, + ) + + frontend_options = FrontendOptions( + mode="custom", + language=language, + hook=f"{hook_file}:{factory_name}", + ) + current = custom_frontend_metadata(frontend_options) + assert current is not None + for field in ( + "source_sha256", + "metadata_sha256", + "factory_invocation", + "declared_metadata", + "declared_symbol_count", + ): + if current.get(field) != hook_metadata.get(field): + raise ValueError( + f"Custom frontend {field} differs from prepared dataset metadata." + ) + declared_symbols = custom_frontend_symbols(frontend_options) or () + missing_symbols = sorted(set(declared_symbols).difference(symbols)) + if missing_symbols: + raise ValueError( + "Custom frontend declares symbols absent from the prepared/export " + f"inventory: {missing_symbols[:16]!r}." + ) + return { + "path": "frontend_hook.py", + "factory": factory_name, + "source_sha256": actual_source_hash, + "metadata_sha256": current["metadata_sha256"], + "declared_metadata": current["declared_metadata"], + }, hook_file + + +def _deployment_frontend_contract( + options: ExportOptions, + checkpoint: Path, + package_template: Path | None, + config: Mapping[str, Any], + symbols: Sequence[str], +) -> tuple[dict[str, Any], Path | None, list[dict[str, Any]]]: + dataset, dataset_path, release_compatibility = _load_prepared_dataset( + options, + checkpoint, + package_template, + ) + if dataset.get("format") not in { + "inflect_prepared_dataset_v1", + "inflect_release_frontend_compatibility_v1", + }: + raise ValueError( + "Unsupported prepared dataset format. Expected inflect_prepared_dataset_v1." + ) + language = dataset.get("language") + frontend = dataset.get("frontend") + if not isinstance(language, str) or not language.strip(): + raise ValueError("Prepared dataset must declare a non-empty language.") + if not isinstance(frontend, Mapping): + raise ValueError("Prepared dataset must contain a frontend metadata object.") + mode = frontend.get("type", frontend.get("mode")) + if mode not in {"espeak", "prephonemized", "custom"}: + raise ValueError( + "Prepared frontend type must be espeak, prephonemized, or custom." + ) + frontend_language = frontend.get("language", language) + if frontend_language != language: + raise ValueError( + "Prepared dataset language and frontend language differ: " + f"{language!r} != {frontend_language!r}." + ) + sample_rate = dataset.get("sample_rate") + configured_rate = config.get("data", {}).get("sampling_rate") + if int(sample_rate or 0) != int(configured_rate or 0): + raise ValueError( + "Prepared dataset sample rate does not match model config: " + f"{sample_rate!r} != {configured_rate!r}." + ) + symbol_check = _validate_prepared_symbols(dataset_path, symbols) + if not symbol_check["ok"]: + raise ValueError(symbol_check["message"]) + hook_contract: dict[str, Any] | None = None + hook_file: Path | None = None + if mode == "custom": + hook_contract, hook_file = _custom_hook_contract( + options, + frontend, + language, + symbols, + ) + elif options.frontend_hook is not None: + raise ValueError("frontend_hook may only be supplied for a custom frontend.") + dataset_hash = ( + sha256_file(dataset_path) + if dataset_path is not None + else _canonical_json_sha256(dataset) + ) + contract = { + "format": "inflect_deployment_frontend_v1", + "mode": mode, + "language": language, + "preserve_punctuation": bool(frontend.get("preserve_punctuation", True)), + "with_stress": bool(frontend.get("with_stress", True)), + "accepts_prephonemized_input": True, + "prepared_frontend": { + "type": mode, + "language": language, + "preserve_punctuation": bool( + frontend.get("preserve_punctuation", True) + ), + "with_stress": bool(frontend.get("with_stress", True)), + "custom_hook": ( + { + "source_kind": frontend["hook"].get("source_kind"), + "source_sha256": frontend["hook"].get("source_sha256"), + "metadata_sha256": frontend["hook"].get("metadata_sha256"), + "factory_invocation": frontend["hook"].get( + "factory_invocation" + ), + "declared_metadata": frontend["hook"].get( + "declared_metadata" + ), + "declared_symbol_count": frontend["hook"].get( + "declared_symbol_count" + ), + } + if mode == "custom" + else None + ), + }, + "prepared_dataset": { + "format": dataset["format"], + "dataset_json_sha256": dataset_hash, + "source_manifest_sha256": dataset.get("source_manifest_sha256"), + "release_compatibility": release_compatibility, + }, + "custom_hook": hook_contract, + } + checks = [ + symbol_check, + status( + True, + "Deployment frontend is explicit and language-aware", + mode=mode, + language=language, + accepts_prephonemized_input=True, + release_compatibility=release_compatibility, + ), + ] + return contract, hook_file, checks + + +def _validate_config_and_symbols( + config: Mapping[str, Any], + symbols: Sequence[str], + state: Mapping[str, torch.Tensor], +) -> list[dict[str, Any]]: + checks: list[dict[str, Any]] = [] + data = config.get("data") + model = config.get("model") + if not isinstance(data, Mapping) or not isinstance(model, Mapping): + raise ValueError("Config must contain 'data' and 'model' objects.") + sample_rate = data.get("sampling_rate") + duplicates = sorted({symbol for symbol in symbols if symbols.count(symbol) > 1}) + checks.append( + status( + True, + "ordered symbol inventory preserved exactly", + duplicate_symbols=duplicates, + note=( + "Duplicate rows are retained for published-checkpoint compatibility." + if duplicates + else None + ), + ) + ) + checks.append( + status( + isinstance(sample_rate, int) and sample_rate > 0, + "sampling_rate is a positive integer", + value=sample_rate, + ) + ) + embedding = state.get("enc_p.emb.weight") + if embedding is None: + checks.append(status(False, "enc_p.emb.weight is missing from the checkpoint")) + else: + checks.append( + status( + embedding.ndim == 2 and embedding.shape[0] == len(symbols), + "embedding rows match the ordered symbol inventory", + embedding_shape=list(embedding.shape), + symbol_count=len(symbols), + ) + ) + expected_hidden = model.get("hidden_channels") + checks.append( + status( + embedding.ndim == 2 and embedding.shape[1] == expected_hidden, + "embedding width matches model.hidden_channels", + embedding_width=embedding.shape[1] if embedding.ndim == 2 else None, + configured_hidden_channels=expected_hidden, + ) + ) + conv_pre = state.get("dec.conv_pre.weight") + expected_inter = model.get("inter_channels") + if conv_pre is not None: + checks.append( + status( + conv_pre.ndim == 3 and conv_pre.shape[1] == expected_inter, + "decoder input width matches model.inter_channels", + decoder_shape=list(conv_pre.shape), + configured_inter_channels=expected_inter, + ) + ) + rates = model.get("upsample_rates") + hop_length = data.get("hop_length") + if isinstance(rates, list) and all(isinstance(item, int) for item in rates): + product = 1 + for rate in rates: + product *= rate + checks.append( + status( + product == hop_length, + "decoder upsample product matches data.hop_length", + upsample_product=product, + hop_length=hop_length, + ) + ) + return checks + + +def _copy_public_runtime(template: Path, destination: Path) -> list[Path]: + copied: list[Path] = [] + runtime = template / "runtime" + if not runtime.is_dir(): + raise FileNotFoundError(f"Package template has no runtime directory: {runtime}") + shutil.copytree( + runtime, + destination / "runtime", + ignore=shutil.ignore_patterns("__pycache__", "*.pyc", "*.pyo"), + ) + copied.extend(path for path in (destination / "runtime").rglob("*") if path.is_file()) + for name in ( + "inference.py", + "inflect_vits_frontend.py", + "inflect_nano_v2_frontend.py", + "requirements.txt", + "requirements-tested.txt", + "LICENSE", + "THIRD_PARTY_NOTICES.md", + ): + source = template / name + if source.is_file(): + target = destination / name + shutil.copy2(source, target) + copied.append(target) + return copied + + +_DEPLOYMENT_FRONTEND_SOURCE = r''' +from __future__ import annotations + +import hashlib +import importlib.util +import inspect +import json +import os +import re +import sys +import unicodedata +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +PACKAGE_ROOT = Path(__file__).resolve().parent +RUNTIME_ROOT = PACKAGE_ROOT / "runtime" +sys.path.insert(0, str(RUNTIME_ROOT)) + +from text.symbols import symbols + + +class DeploymentFrontendError(RuntimeError): + pass + + +@dataclass(frozen=True) +class FrontendOutput: + raw_text: str + normalized_text: str + phoneme_text: str + + +def _load_contract() -> dict[str, Any]: + path = PACKAGE_ROOT / "frontend.json" + try: + contract = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise DeploymentFrontendError(f"Could not load frontend contract: {exc}") from exc + if contract.get("format") != "inflect_deployment_frontend_v1": + raise DeploymentFrontendError("Unsupported or missing deployment frontend contract.") + if contract.get("mode") not in {"espeak", "prephonemized", "custom"}: + raise DeploymentFrontendError("Deployment frontend mode is invalid.") + if not isinstance(contract.get("language"), str) or not contract["language"].strip(): + raise DeploymentFrontendError("Deployment frontend language is missing.") + return contract + + +CONTRACT = _load_contract() +SYMBOLS = frozenset(symbols) +_ESPEAK_BACKEND: Any = None +_CUSTOM_FRONTEND: Any = None + + +def _normalize_generic(text: str) -> str: + if not isinstance(text, str): + raise DeploymentFrontendError("Text input must be a Unicode string.") + if "\x00" in text: + raise DeploymentFrontendError("Text input contains a null byte.") + value = unicodedata.normalize("NFKC", text) + value = "".join(" " if char in "\r\n\t" else char for char in value) + value = "".join( + char for char in value if not unicodedata.category(char).startswith("C") + ) + value = re.sub(r"\s+", " ", value).strip() + if not value: + raise DeploymentFrontendError("Text input is empty after normalization.") + return value + + +def _clean_phonemes(value: str) -> str: + if not isinstance(value, str): + raise DeploymentFrontendError("Phoneme input must be a Unicode string.") + value = unicodedata.normalize("NFC", value) + if "\x00" in value or any(char in "\r\n" for char in value): + raise DeploymentFrontendError("Phoneme input contains unsupported controls.") + value = re.sub(r"\s+", " ", value).strip() + if not value: + raise DeploymentFrontendError("Phoneme input is empty.") + unknown = sorted(set(value).difference(SYMBOLS)) + if unknown: + rendered = ", ".join(repr(char) for char in unknown[:16]) + raise DeploymentFrontendError( + "Phoneme input contains symbols absent from this checkpoint: " + rendered + ) + return value + + +def _configure_espeak() -> None: + candidates = ( + Path("/usr/lib/x86_64-linux-gnu/libespeak-ng.so.1"), + Path("/usr/lib/aarch64-linux-gnu/libespeak-ng.so.1"), + Path("/usr/lib64/libespeak-ng.so.1"), + ) + system = next((path for path in candidates if path.is_file()), None) + if system is not None: + os.environ.setdefault("PHONEMIZER_ESPEAK_LIBRARY", str(system)) + return + try: + import espeakng_loader + + os.environ.setdefault( + "PHONEMIZER_ESPEAK_LIBRARY", espeakng_loader.get_library_path() + ) + os.environ.setdefault("ESPEAK_DATA_PATH", espeakng_loader.get_data_path()) + espeakng_loader.make_library_available() + espeakng_loader.load_library() + except (ImportError, OSError, RuntimeError) as exc: + raise DeploymentFrontendError( + "Could not initialize eSpeak NG. Install espeakng-loader or a system " + "eSpeak NG library." + ) from exc + + +def _espeak() -> Any: + global _ESPEAK_BACKEND + if _ESPEAK_BACKEND is not None: + return _ESPEAK_BACKEND + _configure_espeak() + try: + from phonemizer.backend import EspeakBackend + + _ESPEAK_BACKEND = EspeakBackend( + language=CONTRACT["language"], + preserve_punctuation=bool(CONTRACT["preserve_punctuation"]), + with_stress=bool(CONTRACT["with_stress"]), + language_switch="remove-flags", + ) + except (ImportError, RuntimeError, ValueError) as exc: + raise DeploymentFrontendError( + f"Could not create eSpeak frontend for {CONTRACT['language']!r}: {exc}" + ) from exc + return _ESPEAK_BACKEND + + +def _canonical_hash(value: Any) -> str: + payload = json.dumps( + value, + ensure_ascii=False, + allow_nan=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _custom() -> Any: + global _CUSTOM_FRONTEND + if _CUSTOM_FRONTEND is not None: + return _CUSTOM_FRONTEND + record = CONTRACT.get("custom_hook") + if not isinstance(record, dict): + raise DeploymentFrontendError("Custom frontend package metadata is missing.") + hook_path = (PACKAGE_ROOT / str(record.get("path", ""))).resolve() + try: + hook_path.relative_to(PACKAGE_ROOT) + except ValueError as exc: + raise DeploymentFrontendError("Custom hook path escapes the package.") from exc + if hook_path.suffix.lower() != ".py" or not hook_path.is_file(): + raise DeploymentFrontendError(f"Packaged custom hook is missing: {hook_path.name}") + digest = hashlib.sha256(hook_path.read_bytes()).hexdigest() + if digest != record.get("source_sha256"): + raise DeploymentFrontendError("Packaged custom hook hash verification failed.") + spec = importlib.util.spec_from_file_location("_inflect_package_frontend", hook_path) + if spec is None or spec.loader is None: + raise DeploymentFrontendError("Could not load packaged custom frontend.") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + factory = getattr(module, str(record.get("factory", "")), None) + if not callable(factory): + raise DeploymentFrontendError("Packaged custom frontend factory is unavailable.") + try: + signature = inspect.signature(factory) + try: + signature.bind(language=CONTRACT["language"]) + implementation = factory(language=CONTRACT["language"]) + except TypeError: + signature.bind() + implementation = factory() + except Exception as exc: + raise DeploymentFrontendError(f"Custom frontend factory failed: {exc}") from exc + for method in ("normalize", "phonemize", "metadata"): + if not callable(getattr(implementation, method, None)): + raise DeploymentFrontendError( + f"Custom frontend does not provide callable {method}()." + ) + metadata = implementation.metadata() + if _canonical_hash(metadata) != record.get("metadata_sha256"): + raise DeploymentFrontendError("Custom frontend metadata hash verification failed.") + _CUSTOM_FRONTEND = implementation + return implementation + + +def process_input( + text: str | None = None, + *, + phonemes: str | None = None, +) -> FrontendOutput: + raw_text = text or "" + if phonemes is not None: + normalized = _normalize_generic(text) if text and text.strip() else "" + return FrontendOutput(raw_text, normalized, _clean_phonemes(phonemes)) + mode = CONTRACT["mode"] + if mode == "prephonemized": + raise DeploymentFrontendError( + "This checkpoint uses a prephonemized frontend. Supply phonemes=... " + "or use inference.py --phonemes." + ) + if text is None: + raise DeploymentFrontendError("Text input is required.") + if mode == "custom": + implementation = _custom() + try: + normalized = _normalize_generic(implementation.normalize(text)) + phoneme_text = implementation.phonemize(normalized) + except Exception as exc: + raise DeploymentFrontendError(f"Custom frontend failed: {exc}") from exc + else: + normalized = _normalize_generic(text) + try: + from phonemizer.separator import Separator + + phoneme_text = _espeak().phonemize( + [normalized], + separator=Separator(phone="", word=" ", syllable=""), + strip=True, + njobs=1, + )[0] + except (RuntimeError, ValueError, OSError) as exc: + raise DeploymentFrontendError( + f"eSpeak failed for language {CONTRACT['language']!r}: {exc}" + ) from exc + return FrontendOutput(raw_text, normalized, _clean_phonemes(phoneme_text)) +''' + + +_VITS_FRONTEND_SOURCE = r''' +from __future__ import annotations + +from deployment_frontend import FrontendOutput, process_input + + +VitsFrontendOutput = FrontendOutput + + +def run_vits_frontend( + text: str | None = None, + *, + phonemes: str | None = None, +) -> VitsFrontendOutput: + return process_input(text, phonemes=phonemes) + + +def run_vits_frontend_batch( + texts: list[str], + *, + jobs: int = 1, +) -> list[VitsFrontendOutput]: + del jobs + return [process_input(text) for text in texts] +''' + + +_INFERENCE_SOURCE = r''' +from __future__ import annotations + +import argparse +import logging +import re +import sys +from pathlib import Path + +import numpy as np +import soundfile as sf +import torch + + +PACKAGE_ROOT = Path(__file__).resolve().parent +RUNTIME_ROOT = PACKAGE_ROOT / "runtime" +sys.path.insert(0, str(RUNTIME_ROOT)) +sys.path.insert(0, str(PACKAGE_ROOT)) + +import commons +import utils +from inflect_vits_frontend import run_vits_frontend +from models import SynthesizerTrn +from text import cleaned_text_to_sequence +from text.symbols import symbols + + +def split_text(text: str, limit: int = 280) -> list[str]: + normalized = " ".join(text.split()) + sentences = [ + part.strip() + for part in re.split(r"(?<=[.!?;:。！？；：])\s*", normalized) + if part.strip() + ] + chunks: list[str] = [] + for sentence in sentences or [normalized]: + while len(sentence) > limit: + search = sentence[: limit + 1] + punctuation = max(search.rfind(mark) for mark in (",", ";", ":", "，", "；", "：")) + split_at = ( + punctuation + 1 + if punctuation >= limit // 2 + else sentence.rfind(" ", 0, limit + 1) + ) + if split_at < limit // 2: + split_at = limit + chunks.append(sentence[:split_at].strip()) + sentence = sentence[split_at:].strip() + if sentence: + chunks.append(sentence) + return chunks + + +def boundary_pause_seconds(chunk: str) -> float: + ending = chunk.rstrip()[-1:] if chunk.strip() else "" + return { + "?": 0.28, "？": 0.28, "!": 0.24, "！": 0.24, + ".": 0.22, "。": 0.22, ";": 0.16, "；": 0.16, + ":": 0.13, "：": 0.13, ",": 0.09, "，": 0.09, + }.get(ending, 0.08) + + +def edge_fade( + waveform: np.ndarray, + sample_rate: int, + milliseconds: float = 5.0, +) -> np.ndarray: + frames = min(round(sample_rate * milliseconds / 1000.0), waveform.size // 2) + if frames <= 0: + return waveform + output = waveform.copy() + ramp = np.linspace(0.0, 1.0, frames, endpoint=True, dtype=np.float32) + output[:frames] *= ramp + output[-frames:] *= ramp[::-1] + return output + + +class InflectTTS: + def __init__( + self, + model_dir: str | Path = PACKAGE_ROOT, + device: str = "cpu", + ) -> None: + self.root = Path(model_dir).resolve() + self.device = torch.device(device) + self.hps = utils.get_hparams_from_file(str(self.root / "config.json")) + self.model = SynthesizerTrn( + len(symbols), + self.hps.data.filter_length // 2 + 1, + self.hps.train.segment_size // self.hps.data.hop_length, + **self.hps.model, + ).to(self.device).eval() + root_logger = logging.getLogger() + previous_level = root_logger.level + try: + root_logger.setLevel(logging.WARNING) + utils.load_checkpoint(str(self.root / "model.pth"), self.model, None) + finally: + root_logger.setLevel(previous_level) + self.sample_rate = int(self.hps.data.sampling_rate) + + def _tokens( + self, + text: str | None = None, + *, + phonemes: str | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + output = run_vits_frontend(text, phonemes=phonemes) + sequence = cleaned_text_to_sequence(output.phoneme_text) + if self.hps.data.add_blank: + sequence = commons.intersperse(sequence, 0) + if not sequence: + raise ValueError("The deployment frontend produced no speakable tokens.") + tokens = torch.LongTensor(sequence).to(self.device).unsqueeze(0) + lengths = torch.LongTensor([tokens.size(1)]).to(self.device) + return tokens, lengths + + @torch.inference_mode() + def synthesize( + self, + text: str | None = None, + *, + phonemes: str | None = None, + speed: float = 1.0, + variation: float = 0.667, + seed: int = 0, + ) -> tuple[int, np.ndarray]: + if phonemes is None: + normalized = " ".join((text or "").split()) + if not normalized: + raise ValueError("Text must not be empty.") + chunks: list[tuple[str | None, str | None]] = [ + (chunk, None) for chunk in split_text(normalized) + ] + else: + chunks = [(text, phonemes)] + if not 0.5 <= speed <= 2.0: + raise ValueError("speed must be between 0.5 and 2.0") + if not 0.0 <= variation <= 1.0: + raise ValueError("variation must be between 0.0 and 1.0") + pieces: list[np.ndarray] = [] + for index, (chunk_text, chunk_phonemes) in enumerate(chunks): + if index: + previous = chunks[index - 1][0] or "" + pieces.append( + np.zeros( + round(self.sample_rate * boundary_pause_seconds(previous)), + dtype=np.float32, + ) + ) + tokens, lengths = self._tokens(chunk_text, phonemes=chunk_phonemes) + torch.manual_seed(seed + index) + if self.device.type == "cuda": + torch.cuda.manual_seed_all(seed + index) + waveform = self.model.infer( + tokens, + lengths, + noise_scale=variation, + noise_scale_w=0.8, + length_scale=1.0 / speed, + max_len=4000, + )[0][0, 0].float().cpu().numpy() + pieces.append(edge_fade(waveform, self.sample_rate)) + return self.sample_rate, np.clip(np.concatenate(pieces), -1.0, 1.0) + + def save( + self, + text: str | None, + output: str | Path, + **kwargs: object, + ) -> Path: + destination = Path(output) + destination.parent.mkdir(parents=True, exist_ok=True) + sample_rate, waveform = self.synthesize(text, **kwargs) + sf.write(destination, waveform, sample_rate) + return destination + + +def main() -> None: + parser = argparse.ArgumentParser(description="Run standalone Inflect synthesis.") + parser.add_argument("--model-dir", type=Path, default=PACKAGE_ROOT) + inputs = parser.add_mutually_exclusive_group(required=True) + inputs.add_argument("--text") + inputs.add_argument("--phonemes") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--device", default="cpu") + parser.add_argument("--speed", type=float, default=1.0) + parser.add_argument("--variation", type=float, default=0.667) + parser.add_argument("--seed", type=int, default=0) + args = parser.parse_args() + engine = InflectTTS(args.model_dir, args.device) + engine.save( + args.text, + args.output, + phonemes=args.phonemes, + speed=args.speed, + variation=args.variation, + seed=args.seed, + ) + print(f"wrote {args.output} at {engine.sample_rate} Hz") + + +if __name__ == "__main__": + main() +''' + + +def _write_deployment_runtime( + output: Path, + hook_file: Path | None, +) -> list[Path]: + paths = [ + output / "deployment_frontend.py", + output / "inflect_vits_frontend.py", + output / "inference.py", + ] + paths[0].write_text( + textwrap.dedent(_DEPLOYMENT_FRONTEND_SOURCE).lstrip(), + encoding="utf-8", + ) + paths[1].write_text( + textwrap.dedent(_VITS_FRONTEND_SOURCE).lstrip(), + encoding="utf-8", + ) + paths[2].write_text( + textwrap.dedent(_INFERENCE_SOURCE).lstrip(), + encoding="utf-8", + ) + if hook_file is not None: + destination = output / "frontend_hook.py" + shutil.copy2(hook_file, destination) + paths.append(destination) + return paths + + +def _verify_deployment_runtime( + output: Path, + expected_contract: Mapping[str, Any], + symbols: Sequence[str], +) -> dict[str, Any]: + """Import the packaged frontend and prove its explicit-phoneme bypass works.""" + + source_paths = ( + output / "deployment_frontend.py", + output / "inflect_vits_frontend.py", + output / "inference.py", + ) + for source in source_paths: + compile(source.read_text(encoding="utf-8"), str(source), "exec") + + module_names = ( + "deployment_frontend", + "inflect_vits_frontend", + "text", + "text.symbols", + ) + previous = {name: sys.modules.get(name) for name in module_names} + old_path = list(sys.path) + try: + sys.path.insert(0, str(output)) + sys.path.insert(0, str(output / "runtime")) + for name in module_names: + sys.modules.pop(name, None) + generated_module_name = f"_inflect_deployment_frontend_{abs(hash(output))}" + spec = importlib.util.spec_from_file_location( + generated_module_name, + output / "deployment_frontend.py", + ) + if spec is None or spec.loader is None: + raise RuntimeError("Could not import generated deployment_frontend.py.") + module = importlib.util.module_from_spec(spec) + sys.modules[generated_module_name] = module + try: + spec.loader.exec_module(module) + finally: + sys.modules.pop(generated_module_name, None) + if module.CONTRACT != dict(expected_contract): + raise RuntimeError("Generated runtime loaded a different frontend contract.") + probe = next( + ( + symbol + for symbol in symbols + if len(symbol) == 1 and not symbol.isspace() + ), + None, + ) + if probe is None: + raise RuntimeError("No single-character symbol is available for bypass test.") + result = module.process_input(phonemes=probe) + if result.phoneme_text != probe: + raise RuntimeError("Prephonemized runtime bypass changed the supplied input.") + finally: + sys.path[:] = old_path + for name in module_names: + sys.modules.pop(name, None) + if previous[name] is not None: + sys.modules[name] = previous[name] + + return status( + True, + "Generated deployment runtime imports its language contract and accepts " + "prephonemized input without invoking a text frontend", + mode=expected_contract["mode"], + language=expected_contract["language"], + compiled_files=[path.name for path in source_paths], + ) + + +@contextmanager +def _runtime_imports(runtime_root: Path) -> Iterator[tuple[Any, Any]]: + old_path = list(sys.path) + names = ("models", "commons", "utils", "modules", "attentions", "transforms") + previous = {name: sys.modules.get(name) for name in names} + try: + sys.path.insert(0, str(runtime_root)) + for name in names: + sys.modules.pop(name, None) + models = importlib.import_module("models") + commons = importlib.import_module("commons") + yield models, commons + finally: + sys.path[:] = old_path + for name in names: + sys.modules.pop(name, None) + if previous[name] is not None: + sys.modules[name] = previous[name] + + +def _build_model( + runtime_root: Path, + config: Mapping[str, Any], + symbol_count: int, + *, + inference_only: bool, +) -> tuple[nn.Module, Any]: + train = config["train"] + data = config["data"] + model_config = dict(config["model"]) + model_config["inference_only"] = inference_only + with _runtime_imports(runtime_root) as (models, commons): + model = models.SynthesizerTrn( + symbol_count, + int(data["filter_length"]) // 2 + 1, + int(train["segment_size"]) // int(data["hop_length"]), + **model_config, + ).cpu().eval() + return model, commons + + +class _DurationGraph(nn.Module): + def __init__(self, model: nn.Module, commons_module: Any) -> None: + super().__init__() + self.enc_p = model.enc_p + self.dp = model.dp + self.commons = commons_module + + def forward( + self, + tokens: torch.Tensor, + lengths: torch.Tensor, + length_scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + hidden, means, logs, text_mask = self.enc_p(tokens, lengths) + log_duration = self.dp(hidden, text_mask, g=None) + durations = torch.ceil(torch.exp(log_duration) * text_mask * length_scale) + output_lengths = torch.clamp_min(torch.sum(durations, [1, 2]), 1).long() + output_mask = torch.unsqueeze( + self.commons.sequence_mask(output_lengths, None), 1 + ).to(text_mask.dtype) + attention_mask = torch.unsqueeze(text_mask, 2) * torch.unsqueeze(output_mask, -1) + attention = self.commons.generate_path(durations, attention_mask) + expanded_means = torch.matmul(attention.squeeze(1), means.transpose(1, 2)).transpose(1, 2) + expanded_logs = torch.matmul(attention.squeeze(1), logs.transpose(1, 2)).transpose(1, 2) + return expanded_means, expanded_logs, output_mask + + +class _DecodeGraph(nn.Module): + def __init__(self, model: nn.Module) -> None: + super().__init__() + self.flow = model.flow + self.decoder = model.dec + + def forward( + self, + means: torch.Tensor, + logs: torch.Tensor, + output_mask: torch.Tensor, + noise: torch.Tensor, + noise_scale: torch.Tensor, + ) -> torch.Tensor: + latent = means + noise * torch.exp(logs) * noise_scale + decoded = self.flow(latent, output_mask, g=None, reverse=True) + return self.decoder(decoded * output_mask, g=None) + + +def _export_onnx( + model: nn.Module, + commons_module: Any, + destination: Path, + *, + opset: int, +) -> tuple[list[Path], dict[str, Any]]: + try: + import onnx + except ImportError as exc: + raise RuntimeError( + "ONNX export was requested, but 'onnx' is not installed. " + "Install the toolkit with the 'onnx' extra." + ) from exc + + destination.mkdir(parents=True, exist_ok=True) + duration_path = destination / "duration.onnx" + decode_path = destination / "decode.onnx" + duration = _DurationGraph(model, commons_module).eval() + decode = _DecodeGraph(model).eval() + tokens = torch.tensor([[0, 1, 0, 2, 0, 3, 0]], dtype=torch.long) + lengths = torch.tensor([tokens.shape[1]], dtype=torch.long) + length_scale = torch.tensor(1.0, dtype=torch.float32) + torch.onnx.export( + duration, + (tokens, lengths, length_scale), + duration_path, + input_names=["tokens", "lengths", "length_scale"], + output_names=["m_p_exp", "logs_p_exp", "y_mask"], + dynamic_axes={ + "tokens": {1: "text_len"}, + "m_p_exp": {2: "audio_frames"}, + "logs_p_exp": {2: "audio_frames"}, + "y_mask": {2: "audio_frames"}, + }, + opset_version=opset, + do_constant_folding=True, + dynamo=False, + ) + # These tensors are traced again as decode-graph inputs. inference_mode() + # marks them as inference tensors, which older supported PyTorch versions + # cannot save for backward while constructing the second ONNX graph. + with torch.no_grad(): + means, logs, output_mask = duration(tokens, lengths, length_scale) + noise = torch.zeros_like(means) + noise_scale = torch.tensor(0.667, dtype=torch.float32) + with torch.no_grad(): + reference_waveform = decode( + means, + logs, + output_mask, + noise, + noise_scale, + ).detach().cpu().numpy() + torch.onnx.export( + decode, + (means, logs, output_mask, noise, noise_scale), + decode_path, + input_names=["m_p_exp", "logs_p_exp", "y_mask", "zp_noise", "noise_scale"], + output_names=["waveform"], + dynamic_axes={ + "m_p_exp": {2: "audio_frames"}, + "logs_p_exp": {2: "audio_frames"}, + "y_mask": {2: "audio_frames"}, + "zp_noise": {2: "audio_frames"}, + "waveform": {2: "waveform_samples"}, + }, + opset_version=opset, + do_constant_folding=True, + dynamo=False, + ) + onnx.checker.check_model(onnx.load(duration_path)) + onnx.checker.check_model(onnx.load(decode_path)) + provider_check: dict[str, Any] + try: + import onnxruntime as ort + + sessions = [ + ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) + for path in (duration_path, decode_path) + ] + duration_outputs = sessions[0].run( + None, + { + "tokens": tokens.numpy(), + "lengths": lengths.numpy(), + "length_scale": length_scale.numpy(), + }, + ) + onnx_noise = np.zeros_like(duration_outputs[0], dtype=np.float32) + onnx_waveform = sessions[1].run( + None, + { + "m_p_exp": duration_outputs[0], + "logs_p_exp": duration_outputs[1], + "y_mask": duration_outputs[2], + "zp_noise": onnx_noise, + "noise_scale": noise_scale.numpy(), + }, + )[0] + reference_duration = [ + value.detach().cpu().numpy() for value in (means, logs, output_mask) + ] + duration_errors = [ + float(np.max(np.abs(actual - expected))) + for actual, expected in zip(duration_outputs[:2], reference_duration[:2]) + ] + mask_exact = bool(np.array_equal(duration_outputs[2], reference_duration[2])) + waveform_delta = np.abs(onnx_waveform - reference_waveform) + waveform_max_error = float(np.max(waveform_delta)) + waveform_mean_error = float(np.mean(waveform_delta)) + correlation = float( + np.corrcoef(onnx_waveform.reshape(-1), reference_waveform.reshape(-1))[0, 1] + ) + parity_ok = ( + all(error <= 1.0e-4 for error in duration_errors) + and mask_exact + and waveform_max_error <= 5.0e-4 + and correlation >= 0.9999 + ) + provider_check = status( + parity_ok, + "ONNX Runtime executed both graphs with PyTorch parity", + inputs=[[item.name for item in session.get_inputs()] for session in sessions], + outputs=[[item.name for item in session.get_outputs()] for session in sessions], + parity={ + "duration_max_abs_error": duration_errors, + "duration_mask_exact": mask_exact, + "waveform_max_abs_error": waveform_max_error, + "waveform_mean_abs_error": waveform_mean_error, + "waveform_correlation": correlation, + "tolerances": { + "duration_max_abs_error": 1.0e-4, + "waveform_max_abs_error": 5.0e-4, + "waveform_correlation_minimum": 0.9999, + }, + }, + ) + except ImportError: + provider_check = status( + True, + "ONNX checker passed; ONNX Runtime load check skipped because it is not installed", + skipped=True, + ) + return [duration_path, decode_path], provider_check + + +def export_checkpoint(options: ExportOptions) -> dict[str, Any]: + """Export and verify an adapted checkpoint. + + Returns the same report written to ``export_report.json``. Existing output + directories are rejected unless ``overwrite=True``. + """ + + checkpoint = Path(options.checkpoint).resolve() + output = Path(options.output_dir).resolve() + if not checkpoint.is_file(): + raise FileNotFoundError(f"Checkpoint does not exist: {checkpoint}") + if output.exists() and any(output.iterdir()): + if not options.overwrite: + raise FileExistsError(f"Output directory is not empty: {output}") + shutil.rmtree(output) + output.mkdir(parents=True, exist_ok=True) + + payload = _load_checkpoint(checkpoint) + state, source_metadata, stripped_keys = _extract_state(payload) + package_template = _resolve_package_template(options, payload) + config, config_source = _resolve_config(options, checkpoint, package_template) + symbols, symbols_source = _load_symbols( + options, + checkpoint, + payload, + package_template, + ) + checks = _validate_config_and_symbols(config, symbols, state) + frontend_contract, frontend_hook, frontend_checks = _deployment_frontend_contract( + options, + checkpoint, + package_template, + config, + symbols, + ) + checks.extend(frontend_checks) + remaining_training_keys = [key for key in state if _is_training_only_name(key)] + omitted_training_fields = sorted( + str(key) for key in payload if _is_training_only_name(str(key)) + ) + checks.append( + status( + not remaining_training_keys, + "Inference state contains no posterior encoder, discriminator, optimizer, " + "scheduler, or scaler tensors", + stripped_tensor_count=len(stripped_keys), + stripped_tensor_keys=stripped_keys, + omitted_top_level_fields=omitted_training_fields, + remaining_training_only_keys=remaining_training_keys, + ) + ) + failed = [check for check in checks if not check["ok"]] + if failed: + raise ValueError("Checkpoint/config/symbol validation failed: " + "; ".join( + check["message"] for check in failed + )) + + export_config = copy.deepcopy(config) + export_config.setdefault("model", {})["inference_only"] = True + export_config["deployment_frontend"] = copy.deepcopy(frontend_contract) + deployable_parameters = sum(tensor.numel() for tensor in state.values()) + inference_payload = { + "format": "inflect_vits_inference_checkpoint_v1", + "model": state, + "iteration": int( + source_metadata.get("iteration", source_metadata.get("step", 0)) or 0 + ), + "learning_rate": 0.0, + "deployable_parameters": deployable_parameters, + "adaptation": { + "model_name": options.model_name, + "source_revision": options.source_revision, + "symbol_count": len(symbols), + "frontend": copy.deepcopy(frontend_contract), + }, + } + model_path = output / "model.pth" + torch.save(inference_payload, model_path) + config_path = write_json(output / "config.json", export_config) + symbols_path = write_json( + output / "symbols.json", + { + "format": "inflect_v2_symbol_inventory_v1", + "symbols": symbols, + "count": len(symbols), + }, + ) + frontend_path = write_json(output / "frontend.json", frontend_contract) + produced = [model_path, config_path, symbols_path, frontend_path] + + strict_load_check = status( + True, + "Tensor checkpoint reloaded and symbol/config consistency passed", + strict_model_load=False, + ) + model: nn.Module | None = None + commons_module: Any = None + if package_template is not None: + template = package_template + produced.extend(_copy_public_runtime(template, output)) + produced.extend(_write_deployment_runtime(output, frontend_hook)) + runtime_symbols = output / "runtime" / "text" / "symbols.py" + if runtime_symbols.is_file(): + runtime_symbols.write_text( + "# Generated by inflect-finetune. Keep this order unchanged.\n" + f"symbols = {symbols!r}\n" + 'SPACE_ID = symbols.index(" ") if " " in symbols else -1\n', + encoding="utf-8", + ) + checks.append( + _verify_deployment_runtime(output, frontend_contract, symbols) + ) + model, commons_module = _build_model( + output / "runtime", + export_config, + len(symbols), + inference_only=True, + ) + incompatible = model.load_state_dict(state, strict=True) + strict_load_check = status( + not incompatible.missing_keys and not incompatible.unexpected_keys, + "Runtime model instantiated and checkpoint loaded strictly", + strict_model_load=True, + missing_keys=list(incompatible.missing_keys), + unexpected_keys=list(incompatible.unexpected_keys), + ) + checks.append(strict_load_check) + training_model, _ = _build_model( + output / "runtime", + export_config, + len(symbols), + inference_only=False, + ) + training_incompatible = training_model.load_state_dict(state, strict=False) + training_missing = list(training_incompatible.missing_keys) + training_unexpected = list(training_incompatible.unexpected_keys) + checks.append( + status( + bool(training_missing) + and all(key.startswith("enc_q.") for key in training_missing) + and not training_unexpected, + "Training-form generator differs only by intentionally omitted enc_q.* tensors", + missing_keys=training_missing, + unexpected_keys=training_unexpected, + ) + ) + del training_model + elif options.verify: + raise ValueError( + "Verified export requires a released package runtime. Pass " + "ExportOptions(package_template=...) or export a toolkit training checkpoint " + "whose saved options still identify an available base model." + ) + + onnx_report: dict[str, Any] = {"requested": options.include_onnx} + if options.include_onnx: + if model is None or commons_module is None: + raise ValueError("ONNX export requires ExportOptions(package_template=...).") + onnx_files, onnx_check = _export_onnx( + model, + commons_module, + output / "onnx", + opset=options.onnx_opset, + ) + produced.extend(onnx_files) + checks.append(onnx_check) + onnx_report.update( + { + "opset": options.onnx_opset, + "files": [file_record(path, relative_to=output) for path in onnx_files], + "verification": onnx_check, + } + ) + + reloaded = torch.load(model_path, map_location="cpu", weights_only=True) + reload_state, _, reloaded_stripped = _extract_state(reloaded) + checks.append( + status( + list(reload_state) == list(state) + and all(torch.equal(reload_state[key], state[key]) for key in state), + "Saved inference checkpoint reloads with exact tensor parity", + tensor_count=len(state), + training_only_tensors_after_reload=reloaded_stripped, + ) + ) + if any(not check["ok"] for check in checks): + raise RuntimeError("One or more export verification checks failed.") + + checksum_path = write_checksums( + output / "checksums.sha256", + list(dict.fromkeys(produced)), + relative_to=output, + ) + produced = list(dict.fromkeys(produced)) + produced.append(checksum_path) + report = make_report( + "export_report", + ok=True, + model_name=options.model_name, + source={ + "checkpoint": checkpoint.name, + "checkpoint_sha256": sha256_file(checkpoint), + "config": config_source.name if config_source else None, + "symbols": symbols_source.name if symbols_source else "checkpoint_or_inline", + "revision": options.source_revision, + }, + output_dir=".", + deployable_parameters=deployable_parameters, + stripped_training_tensors={ + "count": len(stripped_keys), + "keys": stripped_keys, + "top_level_fields": omitted_training_fields, + }, + symbol_count=len(symbols), + sample_rate=int(export_config["data"]["sampling_rate"]), + deployment_frontend=frontend_contract, + checks=checks, + onnx=onnx_report, + files=[file_record(path, relative_to=output) for path in produced], + ) + write_json(output / "export_report.json", report) + return report diff --git a/finetune/inflect_finetune/frontend.py b/finetune/inflect_finetune/frontend.py new file mode 100644 index 0000000..04cdefe --- /dev/null +++ b/finetune/inflect_finetune/frontend.py @@ -0,0 +1,499 @@ +"""Configurable eSpeak, prephonemized, and explicit custom frontends.""" + +from __future__ import annotations + +import hashlib +import importlib +import importlib.util +import inspect +import json +import os +import re +import sys +import unicodedata +from dataclasses import asdict, dataclass +from pathlib import Path +from types import ModuleType +from typing import Any, Literal, Mapping, Sequence + + +class FrontendError(RuntimeError): + """Raised when text normalization or phonemization cannot be completed.""" + + +FrontendMode = Literal["espeak", "prephonemized", "custom"] + + +@dataclass(frozen=True) +class FrontendOptions: + """Frontend configuration shared by preparation and future CLI commands.""" + + mode: FrontendMode = "espeak" + language: str = "en-us" + preserve_punctuation: bool = True + with_stress: bool = True + hook: str | None = None + + def validate(self) -> None: + """Validate options without loading eSpeak.""" + if self.mode not in {"espeak", "prephonemized", "custom"}: + raise ValueError("mode must be 'espeak', 'prephonemized', or 'custom'.") + if not self.language.strip(): + raise ValueError("language cannot be empty.") + if self.mode == "custom" and not (self.hook and self.hook.strip()): + raise ValueError( + "Custom frontend mode requires an explicit hook in " + "'module:callable' or 'file.py:function' form." + ) + if self.mode != "custom" and self.hook: + raise ValueError("frontend hook may only be supplied when mode is 'custom'.") + + +@dataclass(frozen=True) +class FrontendResult: + """Normalized transcript and the character-level model input string.""" + + raw_text: str + normalized_text: str + phonemes: str + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-serializable representation.""" + return asdict(self) + + +@dataclass(frozen=True) +class CustomFrontend: + """Validated custom frontend instance plus its reproducibility record.""" + + identity: str + source_kind: str + source_sha256: str + metadata_sha256: str + invocation: str + metadata: dict[str, Any] + symbols: tuple[str, ...] + implementation: Any + + def public_metadata(self) -> dict[str, Any]: + """Return metadata safe to persist in a prepared dataset.""" + return { + "identity": self.identity, + "source_kind": self.source_kind, + "source_sha256": self.source_sha256, + "metadata_sha256": self.metadata_sha256, + "factory_invocation": self.invocation, + "declared_metadata": self.metadata, + "declared_symbol_count": len(self.symbols), + } + + +_CONFIGURED = False +_BACKENDS: dict[tuple[str, bool, bool], Any] = {} +_CUSTOM_FRONTENDS: dict[tuple[str, str], CustomFrontend] = {} +_IMPORT_PATH = re.compile(r"^[A-Za-z_]\w*(?:\.[A-Za-z_]\w*)*$") +_CALLABLE_NAME = re.compile(r"^[A-Za-z_]\w*$") + + +def normalize_text(text: str) -> str: + """Apply language-neutral Unicode and whitespace normalization. + + Language-specific number and abbreviation expansion should be supplied by + a custom frontend rather than silently applying English rules. + """ + if not isinstance(text, str): + raise FrontendError(f"Transcript must be a string, got {type(text).__name__}.") + if "\x00" in text: + raise FrontendError("Transcript contains a null byte.") + normalized = unicodedata.normalize("NFKC", text) + normalized = "".join( + " " if character in {"\r", "\n", "\t"} else character for character in normalized + ) + normalized = "".join( + character + for character in normalized + if not unicodedata.category(character).startswith("C") + ) + normalized = re.sub(r"\s+", " ", normalized).strip() + if not normalized: + raise FrontendError("Transcript is empty after Unicode normalization.") + return normalized + + +def _sha256_bytes(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def _json_hash(value: Mapping[str, Any]) -> str: + try: + encoded = json.dumps( + value, + ensure_ascii=False, + allow_nan=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + except (TypeError, ValueError) as exc: + raise FrontendError( + f"Custom frontend metadata must be finite and JSON-serializable: {exc}" + ) from exc + return _sha256_bytes(encoded) + + +def _source_hash(factory: Any, module: ModuleType, explicit_file: Path | None) -> str: + source_path = explicit_file + if source_path is None: + candidate = inspect.getsourcefile(factory) or getattr(module, "__file__", None) + source_path = Path(candidate).resolve() if candidate else None + if source_path and source_path.is_file(): + return _sha256_bytes(source_path.read_bytes()) + try: + return _sha256_bytes(inspect.getsource(factory).encode("utf-8")) + except (OSError, TypeError) as exc: + raise FrontendError( + "The custom frontend callable has no readable source to fingerprint. " + "Use a Python module or .py file with available source." + ) from exc + + +def _load_hook_callable(specification: str) -> tuple[Any, str, str, str]: + source, separator, callable_name = specification.rpartition(":") + source = source.strip() + callable_name = callable_name.strip() + if not separator: + raise FrontendError( + "Custom frontend hook must include ':' in " + "'module:callable' or 'file.py:function' form." + ) + if not source or not _CALLABLE_NAME.fullmatch(callable_name): + raise FrontendError( + "Custom frontend hook must name one top-level callable, for example " + "'my_frontend:create_frontend' or './frontend.py:create_frontend'." + ) + + explicit_file: Path | None = None + if source.lower().endswith(".py"): + explicit_file = Path(source).expanduser().resolve() + if not explicit_file.is_file(): + raise FrontendError(f"Custom frontend file does not exist: {explicit_file}") + module_name = ( + "_inflect_custom_frontend_" + + hashlib.sha256(str(explicit_file).encode("utf-8")).hexdigest()[:16] + ) + module_spec = importlib.util.spec_from_file_location(module_name, explicit_file) + if module_spec is None or module_spec.loader is None: + raise FrontendError(f"Could not load custom frontend file: {explicit_file}") + module = importlib.util.module_from_spec(module_spec) + sys.modules[module_name] = module + try: + module_spec.loader.exec_module(module) + except Exception as exc: + sys.modules.pop(module_name, None) + raise FrontendError( + f"Custom frontend file raised while loading: {explicit_file}: {exc}" + ) from exc + identity = f"file:{explicit_file.name}:{callable_name}" + source_kind = "file" + else: + if not _IMPORT_PATH.fullmatch(source): + raise FrontendError( + "Custom frontend module must be a dotted Python import path, or the " + "source must be a .py file." + ) + try: + module = importlib.import_module(source) + except Exception as exc: + raise FrontendError( + f"Could not import custom frontend module '{source}': {exc}" + ) from exc + identity = f"{source}:{callable_name}" + source_kind = "module" + + factory = getattr(module, callable_name, None) + if not callable(factory): + raise FrontendError(f"Custom frontend hook '{identity}' is not callable.") + return factory, identity, source_kind, _source_hash(factory, module, explicit_file) + + +def _create_custom_frontend(options: FrontendOptions) -> CustomFrontend: + options.validate() + assert options.hook is not None + cache_key = (options.hook, options.language) + if cache_key in _CUSTOM_FRONTENDS: + return _CUSTOM_FRONTENDS[cache_key] + + factory, identity, source_kind, source_sha256 = _load_hook_callable(options.hook) + try: + signature = inspect.signature(factory) + except (TypeError, ValueError) as exc: + raise FrontendError( + f"Could not inspect custom frontend factory '{identity}': {exc}" + ) from exc + try: + signature.bind(language=options.language) + invocation = "factory(language=)" + except TypeError: + try: + signature.bind() + except TypeError as exc: + raise FrontendError( + f"Custom frontend factory '{identity}' must accept either no arguments " + "or the keyword argument 'language'." + ) from exc + invocation = "factory()" + try: + implementation = ( + factory(language=options.language) + if invocation.startswith("factory(language=") + else factory() + ) + except Exception as exc: + raise FrontendError( + f"Custom frontend factory '{identity}' failed: {exc}" + ) from exc + + for method_name in ("normalize", "phonemize", "symbols", "metadata"): + if not callable(getattr(implementation, method_name, None)): + raise FrontendError( + f"Custom frontend '{identity}' must provide callable " + f"{method_name}(...)." + ) + try: + metadata = implementation.metadata() + declared_symbols = implementation.symbols() + except Exception as exc: + raise FrontendError( + f"Custom frontend '{identity}' failed while declaring metadata or symbols: {exc}" + ) from exc + if not isinstance(metadata, Mapping): + raise FrontendError( + f"Custom frontend '{identity}' metadata() must return a mapping." + ) + metadata_dict = dict(metadata) + required_metadata = ("name", "version", "language", "configuration") + missing = [field for field in required_metadata if field not in metadata_dict] + if missing: + raise FrontendError( + f"Custom frontend '{identity}' metadata() is missing: {', '.join(missing)}." + ) + if not isinstance(declared_symbols, Sequence) or isinstance( + declared_symbols, (str, bytes) + ): + raise FrontendError( + f"Custom frontend '{identity}' symbols() must return an ordered sequence." + ) + symbols = tuple(declared_symbols) + if not symbols or any( + not isinstance(symbol, str) or len(symbol) != 1 for symbol in symbols + ): + raise FrontendError( + f"Custom frontend '{identity}' symbols() must contain one-character strings." + ) + if len(symbols) != len(set(symbols)): + raise FrontendError(f"Custom frontend '{identity}' symbols() contains duplicates.") + normalized_symbols = [unicodedata.normalize("NFC", symbol) for symbol in symbols] + if len(normalized_symbols) != len(set(normalized_symbols)): + raise FrontendError( + f"Custom frontend '{identity}' symbols() has Unicode normalization collisions." + ) + + result = CustomFrontend( + identity=identity, + source_kind=source_kind, + source_sha256=source_sha256, + metadata_sha256=_json_hash(metadata_dict), + invocation=invocation, + metadata=metadata_dict, + symbols=symbols, + implementation=implementation, + ) + _CUSTOM_FRONTENDS[cache_key] = result + return result + + +def _hook_text(value: Any, *, field: str, identity: str) -> str: + if not isinstance(value, str): + raise FrontendError( + f"Custom frontend '{identity}' {field} must return a Unicode string." + ) + normalized = unicodedata.normalize("NFC", value) + if "\x00" in normalized or any( + unicodedata.category(character).startswith("C") and character != "\t" + for character in normalized + ): + raise FrontendError( + f"Custom frontend '{identity}' {field} returned unsupported control characters." + ) + normalized = re.sub(r"\s+", " ", normalized).strip() + if not normalized: + raise FrontendError( + f"Custom frontend '{identity}' {field} returned an empty string." + ) + return normalized + + +def custom_frontend_metadata(options: FrontendOptions) -> dict[str, Any] | None: + """Return a reproducibility record for an explicitly configured custom hook.""" + if options.mode != "custom": + return None + return _create_custom_frontend(options).public_metadata() + + +def custom_frontend_symbols(options: FrontendOptions) -> tuple[str, ...] | None: + """Return the ordered symbol inventory declared by a custom frontend.""" + if options.mode != "custom": + return None + return _create_custom_frontend(options).symbols + + +def _configure_espeak() -> None: + global _CONFIGURED + if _CONFIGURED: + return + candidates = ( + Path("/usr/lib/x86_64-linux-gnu/libespeak-ng.so.1"), + Path("/usr/lib/aarch64-linux-gnu/libespeak-ng.so.1"), + Path("/usr/lib64/libespeak-ng.so.1"), + ) + system_library = next((candidate for candidate in candidates if candidate.is_file()), None) + try: + if system_library: + os.environ.setdefault("PHONEMIZER_ESPEAK_LIBRARY", str(system_library)) + else: + import espeakng_loader + + os.environ.setdefault( + "PHONEMIZER_ESPEAK_LIBRARY", espeakng_loader.get_library_path() + ) + os.environ.setdefault("ESPEAK_DATA_PATH", espeakng_loader.get_data_path()) + espeakng_loader.make_library_available() + espeakng_loader.load_library() + except (ImportError, OSError, RuntimeError) as exc: + raise FrontendError( + "Could not initialize eSpeak NG. Install the system eSpeak NG library or " + "the declared espeakng-loader dependency." + ) from exc + _CONFIGURED = True + + +def _backend(options: FrontendOptions) -> Any: + options.validate() + _configure_espeak() + key = (options.language, options.preserve_punctuation, options.with_stress) + if key in _BACKENDS: + return _BACKENDS[key] + try: + from phonemizer.backend import EspeakBackend + + backend = EspeakBackend( + language=options.language, + preserve_punctuation=options.preserve_punctuation, + with_stress=options.with_stress, + language_switch="remove-flags", + ) + except (ImportError, RuntimeError, ValueError) as exc: + raise FrontendError( + f"Could not create the eSpeak frontend for language '{options.language}'. " + "Confirm the language with `espeak-ng --voices` and install phonemizer." + ) from exc + _BACKENDS[key] = backend + return backend + + +def phonemize_text(text: str, options: FrontendOptions | None = None) -> str: + """Phonemize normalized text into the character sequence consumed by Inflect.""" + options = options or FrontendOptions() + if options.mode != "espeak": + raise FrontendError("phonemize_text requires frontend mode 'espeak'.") + normalized = normalize_text(text) + try: + from phonemizer.separator import Separator + + result = _backend(options).phonemize( + [normalized], + separator=Separator(phone="", word=" ", syllable=""), + strip=True, + njobs=1, + )[0] + except (RuntimeError, ValueError, OSError) as exc: + raise FrontendError( + f"eSpeak failed to phonemize text using language '{options.language}': {exc}" + ) from exc + result = unicodedata.normalize("NFC", result) + result = re.sub(r"\s+", " ", result).strip() + if not result: + raise FrontendError( + f"eSpeak produced an empty phoneme sequence for language '{options.language}'." + ) + return result + + +def process_text( + text: str, + *, + options: FrontendOptions | None = None, + prephonemized: str | None = None, +) -> FrontendResult: + """Normalize a transcript and obtain phonemes from eSpeak or supplied input.""" + options = options or FrontendOptions() + options.validate() + normalized = normalize_text(text) + if options.mode == "prephonemized": + if prephonemized is None: + raise FrontendError( + "Prephonemized mode requires a non-empty phoneme string in every manifest row." + ) + phonemes = unicodedata.normalize("NFC", prephonemized) + phonemes = re.sub(r"\s+", " ", phonemes).strip() + if not phonemes: + raise FrontendError("Prephonemized input is empty after normalization.") + if "\x00" in phonemes or any(character in "\r\n" for character in phonemes): + raise FrontendError("Prephonemized input contains an unsupported control character.") + elif options.mode == "custom": + custom = _create_custom_frontend(options) + try: + first_normalized = custom.implementation.normalize(text) + second_normalized = custom.implementation.normalize(text) + except Exception as exc: + raise FrontendError( + f"Custom frontend '{custom.identity}' normalize() failed: {exc}" + ) from exc + if first_normalized != second_normalized: + raise FrontendError( + f"Custom frontend '{custom.identity}' normalize() is not deterministic." + ) + normalized = _hook_text( + first_normalized, field="normalize()", identity=custom.identity + ) + try: + first_phonemes = custom.implementation.phonemize(normalized) + second_phonemes = custom.implementation.phonemize(normalized) + except Exception as exc: + raise FrontendError( + f"Custom frontend '{custom.identity}' phonemize() failed: {exc}" + ) from exc + if first_phonemes != second_phonemes: + raise FrontendError( + f"Custom frontend '{custom.identity}' phonemize() is not deterministic." + ) + phonemes = _hook_text( + first_phonemes, field="phonemize()", identity=custom.identity + ) + undeclared = sorted(set(phonemes).difference(custom.symbols)) + if undeclared: + raise FrontendError( + f"Custom frontend '{custom.identity}' emitted undeclared symbols: " + + ", ".join(repr(symbol) for symbol in undeclared) + ) + else: + phonemes = phonemize_text(normalized, options) + return FrontendResult(raw_text=text, normalized_text=normalized, phonemes=phonemes) + + +def validate_frontend(options: FrontendOptions) -> None: + """Fail early when a configured frontend cannot be initialized.""" + options.validate() + if options.mode == "espeak": + _backend(options) + elif options.mode == "custom": + _create_custom_frontend(options) diff --git a/finetune/inflect_finetune/manifest.py b/finetune/inflect_finetune/manifest.py new file mode 100644 index 0000000..b4ef9e9 --- /dev/null +++ b/finetune/inflect_finetune/manifest.py @@ -0,0 +1,250 @@ +"""Strict CSV and JSONL manifest parsing with confined audio paths.""" + +from __future__ import annotations + +import csv +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterator, Mapping + + +class ManifestError(ValueError): + """Raised when a source manifest is malformed or unsafe.""" + + +@dataclass(frozen=True) +class ManifestRow: + """Canonical source row consumed by dataset preparation.""" + + index: int + source_location: str + audio_relative: str + audio_path: Path + text: str + phonemes: str | None = None + row_id: str | None = None + speaker: str | None = None + group_id: str | None = None + group_field: str | None = None + + +_AUDIO_FIELDS = ("audio", "audio_path", "wav", "wav_path", "path") +_TEXT_FIELDS = ("text", "transcript", "sentence") +_PHONEME_FIELDS = ("phonemes", "phoneme_text", "phones") +_ID_FIELDS = ("id", "utt_id", "utterance_id", "name") +_GROUP_FIELDS = ("group_id", "session") +_SPEAKER_FIELDS = ("speaker",) +_MAX_MANIFEST_BYTES = 128 * 1024 * 1024 +_MAX_LINE_BYTES = 4 * 1024 * 1024 +_MAX_FIELD_CHARS = 1_000_000 + + +def _first(row: Mapping[str, Any], fields: tuple[str, ...]) -> Any: + for field in fields: + value = row.get(field) + if value is not None and str(value).strip(): + return value + return None + + +def _first_named( + row: Mapping[str, Any], fields: tuple[str, ...] +) -> tuple[str | None, Any]: + for field in fields: + value = row.get(field) + if value is not None and str(value).strip(): + return field, value + return None, None + + +def resolve_audio_path(audio_root: Path, value: str, *, location: str) -> tuple[str, Path]: + """Resolve a relative manifest path while rejecting traversal and symlink escapes.""" + if "\x00" in value: + raise ManifestError(f"Audio path contains a null byte at {location}.") + raw = Path(value.strip()) + if not value.strip(): + raise ManifestError(f"Audio path is empty at {location}.") + if raw.is_absolute() or raw.drive or raw.root: + raise ManifestError( + f"Audio path must be relative to the configured audio root at {location}: {value!r}" + ) + if any(part == ".." for part in raw.parts): + raise ManifestError(f"Audio path traversal is not allowed at {location}: {value!r}") + + root = Path(audio_root).expanduser().resolve() + candidate = (root / raw).resolve() + try: + relative = candidate.relative_to(root) + except ValueError as exc: + raise ManifestError( + f"Audio path escapes the configured audio root at {location}: {value!r}" + ) from exc + if not candidate.is_file(): + raise ManifestError(f"Audio file referenced at {location} does not exist: {candidate}") + return relative.as_posix(), candidate + + +def _validate_scalar(value: Any, field: str, location: str) -> str: + if isinstance(value, (dict, list, tuple)): + raise ManifestError(f"Field '{field}' must be a string at {location}.") + text = str(value) + if "\x00" in text: + raise ManifestError(f"Field '{field}' contains a null byte at {location}.") + if len(text) > _MAX_FIELD_CHARS: + raise ManifestError( + f"Field '{field}' exceeds {_MAX_FIELD_CHARS:,} characters at {location}." + ) + return text.strip() + + +def _canonicalize( + raw: Mapping[str, Any], + *, + index: int, + location: str, + audio_root: Path, + require_phonemes: bool, +) -> ManifestRow: + audio_value = _first(raw, _AUDIO_FIELDS) + text_value = _first(raw, _TEXT_FIELDS) + phone_value = _first(raw, _PHONEME_FIELDS) + row_id_value = _first(raw, _ID_FIELDS) + speaker_value = _first(raw, _SPEAKER_FIELDS) + group_field, group_value = _first_named(raw, _GROUP_FIELDS) + + if audio_value is None: + raise ManifestError( + f"No audio field found at {location}; expected one of {', '.join(_AUDIO_FIELDS)}." + ) + if text_value is None: + raise ManifestError( + f"No transcript field found at {location}; expected one of {', '.join(_TEXT_FIELDS)}." + ) + if require_phonemes and phone_value is None: + raise ManifestError( + f"No phoneme field found at {location}; prephonemized mode expects one of " + f"{', '.join(_PHONEME_FIELDS)}." + ) + + audio_text = _validate_scalar(audio_value, "audio", location) + text = _validate_scalar(text_value, "text", location) + phonemes = ( + _validate_scalar(phone_value, "phonemes", location) + if phone_value is not None + else None + ) + row_id = ( + _validate_scalar(row_id_value, "id", location) if row_id_value is not None else None + ) + speaker = ( + _validate_scalar(speaker_value, "speaker", location) + if speaker_value is not None + else None + ) + group_id = ( + _validate_scalar(group_value, group_field or "group_id", location) + if group_value is not None + else None + ) + if not text: + raise ManifestError(f"Transcript is empty at {location}.") + if require_phonemes and not phonemes: + raise ManifestError(f"Phoneme text is empty at {location}.") + relative, resolved = resolve_audio_path(audio_root, audio_text, location=location) + return ManifestRow( + index=index, + source_location=location, + audio_relative=relative, + audio_path=resolved, + text=text, + phonemes=phonemes, + row_id=row_id, + speaker=speaker, + group_id=group_id, + group_field=group_field, + ) + + +def _jsonl_rows(path: Path) -> Iterator[tuple[int, Mapping[str, Any]]]: + with path.open("rb") as handle: + for line_number, raw_line in enumerate(handle, 1): + if len(raw_line) > _MAX_LINE_BYTES: + raise ManifestError( + f"JSONL line {line_number} exceeds {_MAX_LINE_BYTES:,} bytes in {path}." + ) + if not raw_line.strip(): + continue + try: + line = raw_line.decode("utf-8-sig" if line_number == 1 else "utf-8") + value = json.loads(line) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ManifestError(f"Invalid JSON at {path}:{line_number}: {exc}") from exc + if not isinstance(value, dict): + raise ManifestError(f"Expected a JSON object at {path}:{line_number}.") + yield line_number, value + + +def _csv_rows(path: Path) -> Iterator[tuple[int, Mapping[str, Any]]]: + try: + with path.open("r", encoding="utf-8-sig", newline="") as handle: + reader = csv.DictReader(handle) + if not reader.fieldnames: + raise ManifestError(f"CSV manifest has no header: {path}") + normalized = [str(field).strip() for field in reader.fieldnames] + if not all(normalized) or len(set(normalized)) != len(normalized): + raise ManifestError(f"CSV header contains empty or duplicate fields: {path}") + reader.fieldnames = normalized + for line_number, row in enumerate(reader, 2): + if None in row: + raise ManifestError( + f"CSV row has more values than header columns at {path}:{line_number}." + ) + yield line_number, row + except UnicodeDecodeError as exc: + raise ManifestError(f"CSV manifest must be UTF-8 encoded: {path}") from exc + except csv.Error as exc: + raise ManifestError(f"Invalid CSV manifest {path}: {exc}") from exc + + +def parse_manifest( + path: Path, + *, + audio_root: Path | None = None, + require_phonemes: bool = False, +) -> list[ManifestRow]: + """Parse a UTF-8 CSV or JSONL manifest into validated canonical rows.""" + path = Path(path).expanduser().resolve() + if not path.is_file(): + raise ManifestError(f"Manifest does not exist: {path}") + if path.stat().st_size > _MAX_MANIFEST_BYTES: + raise ManifestError( + f"Manifest exceeds the {_MAX_MANIFEST_BYTES // (1024 * 1024)} MiB safety limit: {path}" + ) + root = Path(audio_root).expanduser().resolve() if audio_root else path.parent + if not root.is_dir(): + raise ManifestError(f"Configured audio root is not a directory: {root}") + + suffix = path.suffix.lower() + if suffix == ".jsonl": + source_rows = _jsonl_rows(path) + elif suffix == ".csv": + source_rows = _csv_rows(path) + else: + raise ManifestError( + f"Unsupported manifest extension '{path.suffix}'. Use .csv or .jsonl." + ) + + rows = [ + _canonicalize( + raw, + index=index, + location=f"{path}:{line_number}", + audio_root=root, + require_phonemes=require_phonemes, + ) + for index, (line_number, raw) in enumerate(source_rows) + ] + if not rows: + raise ManifestError(f"Manifest contains no data rows: {path}") + return rows diff --git a/finetune/inflect_finetune/modeling.py b/finetune/inflect_finetune/modeling.py new file mode 100644 index 0000000..af27898 --- /dev/null +++ b/finetune/inflect_finetune/modeling.py @@ -0,0 +1,280 @@ +"""Construction helpers for public Inflect v2 warm-start training.""" + +from __future__ import annotations + +import importlib +import json +import sys +from collections import Counter +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType +from typing import Iterator, Sequence + +import torch +from torch import nn + +from . import monotonic_align as training_monotonic_align +from .symbols import BASE_SYMBOLS as RELEASE_BASE_SYMBOLS + + +BASE_SYMBOL_COUNT = 178 +_RELEASE_DUPLICATE_POSITIONS = { + symbol: tuple( + index for index, value in enumerate(RELEASE_BASE_SYMBOLS) if value == symbol + ) + for symbol, count in Counter(RELEASE_BASE_SYMBOLS).items() + if count > 1 +} +_RUNTIME_MODULE_NAMES = ( + "attentions", + "commons", + "inflect_alias_free", + "models", + "modules", + "monotonic_align", + "transforms", + "utils", + "text", + "text.cleaners", + "text.symbols", +) + + +@dataclass(frozen=True) +class RuntimeComponents: + """Objects imported from one released model's self-contained runtime.""" + + synthesizer_class: type[nn.Module] + discriminator_class: type[nn.Module] + commons: ModuleType + base_symbols: tuple[str, ...] + runtime_root: Path + + +@dataclass(frozen=True) +class ModelBundle: + generator: nn.Module + discriminator: nn.Module + config: dict + symbols: tuple[str, ...] + base_symbols: tuple[str, ...] + components: RuntimeComponents + + +def repository_root() -> Path: + """Return the repository root without relying on the process CWD.""" + + return Path(__file__).resolve().parents[2] + + +def resolve_base_model(model: str | Path) -> Path: + """Resolve a local release, Hugging Face repository ID, or model shorthand.""" + + candidate = Path(model).expanduser() + if candidate.is_dir(): + resolved = candidate.resolve() + else: + model_text = str(model).strip() + name = model_text.lower().replace("\\", "/").rstrip("/").split("/")[-1] + aliases = { + "micro": "Inflect-Micro-v2", + "inflect-micro-v2": "Inflect-Micro-v2", + "nano": "Inflect-Nano-v2", + "inflect-nano-v2": "Inflect-Nano-v2", + } + release_name = aliases.get(name) + checkout = ( + repository_root() / "release_assets" / "hf_clean_download" / release_name + if release_name is not None + else None + ) + if checkout is not None and checkout.is_dir(): + resolved = checkout.resolve() + else: + repo_id = f"owensong/{release_name}" if release_name is not None else model_text + revision = None + if "@" in repo_id: + repo_id, revision = repo_id.rsplit("@", 1) + if "/" not in repo_id or not all(part for part in repo_id.split("/", 1)): + raise FileNotFoundError( + f"Base model {model!r} is neither a local directory, a Micro/Nano " + "shorthand, nor a Hugging Face repository ID such as " + "'owensong/Inflect-Micro-v2'." + ) + try: + from huggingface_hub import snapshot_download + except ImportError as exc: + raise RuntimeError( + "Downloading a base model requires huggingface-hub. Install the " + "toolkit dependencies or pass a local model directory." + ) from exc + resolved = Path( + snapshot_download( + repo_id=repo_id, + revision=revision, + allow_patterns=[ + "config.json", + "model.pth", + "runtime/**", + "inference.py", + "inflect*_frontend.py", + "requirements*.txt", + "LICENSE", + "THIRD_PARTY_NOTICES.md", + ], + ) + ).resolve() + required = ("config.json", "model.pth", "runtime") + missing = [name for name in required if not (resolved / name).exists()] + if missing: + raise FileNotFoundError( + f"{resolved} is not a complete Inflect release directory; missing {missing}." + ) + return resolved + + +def load_symbols(path: str | Path) -> tuple[str, ...]: + """Load an ordered symbol inventory written by the preparation workflow.""" + + source = Path(path) + payload = json.loads(source.read_text(encoding="utf-8")) + if isinstance(payload, list): + values = payload + elif isinstance(payload, dict): + values = payload.get("symbols", payload.get("ordered_symbols")) + else: + values = None + if not isinstance(values, list) or not all(isinstance(item, str) for item in values): + raise ValueError(f"{source} must contain a JSON symbol list or a 'symbols' list.") + if not values or values[0] != "_": + raise ValueError("The ordered symbol inventory must begin with '_' at index 0.") + if len(values) < BASE_SYMBOL_COUNT: + raise ValueError( + f"The symbol inventory must preserve the {BASE_SYMBOL_COUNT}-symbol release prefix." + ) + if tuple(values[:BASE_SYMBOL_COUNT]) != tuple(RELEASE_BASE_SYMBOLS): + raise ValueError( + "The first 178 symbols must exactly preserve the published Inflect v2 release " + "inventory and indices." + ) + duplicate_positions = { + symbol: tuple(index for index, value in enumerate(values) if value == symbol) + for symbol, count in Counter(values).items() + if count > 1 + } + if duplicate_positions != _RELEASE_DUPLICATE_POSITIONS: + raise ValueError( + "The symbol inventory contains duplicate/custom-added symbols outside the " + "release-compatible duplicate apostrophe at base indices 174 and 176: " + f"{duplicate_positions}" + ) + return tuple(values) + + +@contextmanager +def _isolated_runtime_import(runtime_root: Path) -> Iterator[None]: + """Temporarily isolate VITS' historical absolute module imports.""" + + previous = {name: sys.modules.get(name) for name in _RUNTIME_MODULE_NAMES} + for name in _RUNTIME_MODULE_NAMES: + sys.modules.pop(name, None) + sys.path.insert(0, str(runtime_root)) + try: + yield + finally: + if sys.path and sys.path[0] == str(runtime_root): + sys.path.pop(0) + for name in _RUNTIME_MODULE_NAMES: + sys.modules.pop(name, None) + for name, module in previous.items(): + if module is not None: + sys.modules[name] = module + + +def load_runtime_components(base_model: str | Path) -> RuntimeComponents: + """Import training-capable classes from the exact released runtime.""" + + model_root = resolve_base_model(base_model) + runtime_root = model_root / "runtime" + with _isolated_runtime_import(runtime_root): + models = importlib.import_module("models") + commons = importlib.import_module("commons") + symbol_module = importlib.import_module("text.symbols") + # The release bundle intentionally contains an inference-only stub. + # Replace only that module global with the public training operation. + models.monotonic_align = training_monotonic_align + synthesizer = models.SynthesizerTrn + discriminator = models.MultiPeriodDiscriminator + symbols = tuple(symbol_module.symbols) + if len(symbols) != BASE_SYMBOL_COUNT: + raise RuntimeError( + f"Expected {BASE_SYMBOL_COUNT} release symbols, found {len(symbols)}." + ) + return RuntimeComponents( + synthesizer_class=synthesizer, + discriminator_class=discriminator, + commons=commons, + base_symbols=symbols, + runtime_root=runtime_root, + ) + + +def load_release_config(base_model: str | Path) -> dict: + model_root = resolve_base_model(base_model) + payload = json.loads((model_root / "config.json").read_text(encoding="utf-8")) + for section in ("train", "data", "model"): + if not isinstance(payload.get(section), dict): + raise ValueError(f"Release config is missing the {section!r} object.") + return payload + + +def build_training_models( + base_model: str | Path, + symbols: Sequence[str], + *, + seed: int = 1234, +) -> ModelBundle: + """Construct the released generator in training mode plus a fresh MPD.""" + + model_root = resolve_base_model(base_model) + config = load_release_config(model_root) + components = load_runtime_components(model_root) + model_kwargs = dict(config["model"]) + model_kwargs["inference_only"] = False + # These release fields are retained even where a runtime version consumes + # them through **kwargs; they are part of the architecture contract. + model_kwargs["n_layers_q"] = int(model_kwargs.get("n_layers_q", 3)) + model_kwargs["n_speakers"] = int(config["data"].get("n_speakers", 0)) + + with torch.random.fork_rng(devices=[]): + torch.manual_seed(seed) + generator = components.synthesizer_class( + len(symbols), + int(config["data"]["filter_length"]) // 2 + 1, + int(config["train"]["segment_size"]) // int(config["data"]["hop_length"]), + **model_kwargs, + ) + discriminator = components.discriminator_class( + bool(model_kwargs.get("use_spectral_norm", False)) + ) + return ModelBundle( + generator=generator, + discriminator=discriminator, + config=config, + symbols=tuple(symbols), + base_symbols=components.base_symbols, + components=components, + ) + + +def trainable_parameter_count(module: nn.Module) -> int: + return sum(parameter.numel() for parameter in module.parameters() if parameter.requires_grad) + + +def optimizer_parameters(module: nn.Module) -> list[nn.Parameter]: + parameters = [parameter for parameter in module.parameters() if parameter.requires_grad] + if not parameters: + raise ValueError("No trainable parameters remain after applying the freeze policy.") + return parameters diff --git a/finetune/inflect_finetune/monotonic_align.py b/finetune/inflect_finetune/monotonic_align.py new file mode 100644 index 0000000..07494ce --- /dev/null +++ b/finetune/inflect_finetune/monotonic_align.py @@ -0,0 +1,78 @@ +"""Dependency-free monotonic alignment for public Inflect adaptation.""" + +from __future__ import annotations + +import numpy as np +import torch + + +def maximum_path(neg_cent: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + """Return the maximum-score monotonic path through ``[batch, audio, text]``. + + Each valid audio frame is assigned to one text token. The token index starts + at zero, ends at the final valid token, and either stays fixed or advances + by one at each frame. This is a training-only operation; deployable + inference checkpoints do not contain or call it. + """ + + if neg_cent.ndim != 3 or mask.ndim != 3 or neg_cent.shape != mask.shape: + raise ValueError( + "neg_cent and mask must have the same [batch, audio, text] shape." + ) + scores = neg_cent.detach().to(device="cpu", dtype=torch.float32).numpy() + valid = mask.detach().to(device="cpu").numpy() > 0 + paths = np.zeros(scores.shape, dtype=np.float32) + + for batch_index in range(scores.shape[0]): + valid_audio = np.any(valid[batch_index], axis=1) + valid_text = np.any(valid[batch_index], axis=0) + audio_length = int(valid_audio.sum()) + text_length = int(valid_text.sum()) + if audio_length == 0 or text_length == 0: + raise ValueError("Monotonic alignment received an empty valid sequence.") + if audio_length < text_length: + raise ValueError( + "Monotonic alignment requires at least one audio frame per text token; " + f"received {audio_length} frames and {text_length} tokens." + ) + expected_mask = np.zeros_like(valid[batch_index]) + expected_mask[:audio_length, :text_length] = True + if not np.array_equal(valid[batch_index], expected_mask): + raise ValueError( + "Monotonic alignment mask must be one top-left rectangular valid region." + ) + + values = scores[batch_index, :audio_length, :text_length] + accumulated = np.full(values.shape, -np.inf, dtype=np.float32) + advanced = np.zeros(values.shape, dtype=np.bool_) + accumulated[0, 0] = values[0, 0] + + for audio_index in range(1, audio_length): + minimum_text = max(0, text_length + audio_index - audio_length) + maximum_text = min(text_length - 1, audio_index) + for text_index in range(minimum_text, maximum_text + 1): + stay = accumulated[audio_index - 1, text_index] + move = ( + accumulated[audio_index - 1, text_index - 1] + if text_index > 0 + else -np.inf + ) + use_move = text_index == audio_index or move > stay + predecessor = move if use_move else stay + accumulated[audio_index, text_index] = ( + values[audio_index, text_index] + predecessor + ) + advanced[audio_index, text_index] = use_move + + text_index = text_length - 1 + for audio_index in range(audio_length - 1, -1, -1): + paths[batch_index, audio_index, text_index] = 1.0 + if audio_index > 0 and advanced[audio_index, text_index]: + text_index -= 1 + if text_index != 0: + raise RuntimeError("Monotonic alignment backtracking did not reach the first token.") + + return torch.from_numpy(paths).to(device=neg_cent.device, dtype=neg_cent.dtype) + + +__all__ = ["maximum_path"] diff --git a/finetune/inflect_finetune/prepare.py b/finetune/inflect_finetune/prepare.py new file mode 100644 index 0000000..8a35299 --- /dev/null +++ b/finetune/inflect_finetune/prepare.py @@ -0,0 +1,436 @@ +"""Atomic preparation of public Inflect adaptation datasets.""" + +from __future__ import annotations + +import hashlib +import json +import re +import shutil +import uuid +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Sequence + +from .audio import AudioOptions, convert_wav +from .frontend import ( + FrontendOptions, + custom_frontend_metadata, + custom_frontend_symbols, + process_text, + validate_frontend, +) +from .manifest import ManifestRow, parse_manifest +from .symbols import ( + BASE_SYMBOLS, + audit_symbol_coverage, + build_symbol_inventory, + load_base_symbols, + write_symbol_inventory, +) + + +class PreparationError(RuntimeError): + """Raised when a source dataset cannot be prepared completely.""" + + +@dataclass(frozen=True) +class PrepareOptions: + """Options suitable for programmatic use and a future CLI.""" + + manifest_path: Path + output_dir: Path + audio_root: Path | None = None + language: str = "en-us" + frontend: str = "espeak" + frontend_hook: str | None = None + validation_fraction: float = 0.05 + split_seed: int = 1337 + sample_rate: int = 24_000 + min_duration_seconds: float = 0.05 + max_duration_seconds: float | None = None + base_symbols_path: Path | None = None + + def validate(self) -> None: + """Validate preparation settings before any output is written.""" + if self.frontend not in {"espeak", "prephonemized", "custom"}: + raise ValueError("frontend must be 'espeak', 'prephonemized', or 'custom'.") + if not 0 < self.validation_fraction < 1: + raise ValueError("validation_fraction must be in the interval (0, 1).") + if self.frontend == "custom" and not self.frontend_hook: + raise ValueError("frontend_hook is required when frontend='custom'.") + if self.frontend != "custom" and self.frontend_hook: + raise ValueError("frontend_hook may only be used when frontend='custom'.") + if self.sample_rate != 24_000: + raise ValueError("Inflect prepared datasets must use a 24,000 Hz sample rate.") + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with Path(path).open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _safe_name(row: ManifestRow) -> str: + source = row.row_id or Path(row.audio_relative).stem or f"row-{row.index}" + slug = re.sub(r"[^A-Za-z0-9._-]+", "-", source).strip(".-_")[:48] or "utterance" + suffix = hashlib.sha256( + f"{row.index}\0{row.audio_relative}\0{row.text}".encode("utf-8") + ).hexdigest()[:10] + return f"{row.index:07d}-{slug}-{suffix}.wav" + + +class _UnionFind: + """Small deterministic union-find used to construct leakage-safe split units.""" + + def __init__(self, count: int) -> None: + self.parent = list(range(count)) + + def find(self, index: int) -> int: + """Return a component root with path compression.""" + while self.parent[index] != index: + self.parent[index] = self.parent[self.parent[index]] + index = self.parent[index] + return index + + def union(self, left: int, right: int) -> None: + """Join two components using the smaller root as a stable representative.""" + left_root = self.find(left) + right_root = self.find(right) + if left_root == right_root: + return + smaller, larger = sorted((left_root, right_root)) + self.parent[larger] = smaller + + +def _normalized_split_key(text: str) -> str: + """Return the canonical key used to prevent transcript leakage.""" + return re.sub(r"\s+", " ", text).strip().casefold() + + +def _split_components( + prepared_rows: Sequence[dict[str, Any]], + source_rows: Sequence[ManifestRow], + *, + fraction: float, + seed: int, +) -> tuple[set[int], dict[str, Any]]: + """Split deterministic atomic components linked by group or normalized text.""" + count = len(prepared_rows) + if count < 2: + raise PreparationError( + "Training-ready preparation requires at least two usable rows so both " + "train and validation are nonempty." + ) + + union_find = _UnionFind(count) + first_group: dict[tuple[str, str], int] = {} + first_text: dict[str, int] = {} + for index, (prepared, source) in enumerate(zip(prepared_rows, source_rows)): + if source.group_id: + group_key = (source.group_field or "group_id", source.group_id) + if group_key in first_group: + union_find.union(index, first_group[group_key]) + else: + first_group[group_key] = index + text_key = _normalized_split_key(prepared["normalized_text"]) + if text_key in first_text: + union_find.union(index, first_text[text_key]) + else: + first_text[text_key] = index + + components: dict[int, list[int]] = {} + for index in range(count): + components.setdefault(union_find.find(index), []).append(index) + if len(components) < 2: + raise PreparationError( + "The dataset is too small to create leakage-safe train and validation " + "splits: all rows are connected by the same group or normalized transcript. " + "Add at least one independent group with different text." + ) + + def component_rank(indices: list[int]) -> bytes: + identities = [ + "\0".join( + ( + source_rows[index].group_field or "", + source_rows[index].group_id or "", + prepared_rows[index]["normalized_text"], + prepared_rows[index]["audio_sha256"], + ) + ) + for index in indices + ] + payload = f"{seed}\0" + "\0".join(sorted(identities)) + return hashlib.sha256(payload.encode("utf-8")).digest() + + ranked = sorted(components.values(), key=component_rank) + target = max(1, min(count - 1, int(round(count * fraction)))) + prefix_sizes: list[int] = [] + running = 0 + for component in ranked[:-1]: + running += len(component) + prefix_sizes.append(running) + chosen_prefix = min( + range(1, len(ranked)), + key=lambda length: ( + abs(prefix_sizes[length - 1] - target), + prefix_sizes[length - 1] > target, + length, + ), + ) + validation_indices = { + index for component in ranked[:chosen_prefix] for index in component + } + if not validation_indices or len(validation_indices) == count: + raise PreparationError( + "Could not produce nonempty train and validation splits. Add more independent " + "groups or transcripts." + ) + return validation_indices, { + "strategy": "deterministic_group_aware_v1", + "atomic_component_count": len(components), + "explicit_group_count": len(first_group), + "normalized_text_key_count": len(first_text), + "target_validation_rows": target, + "actual_validation_rows": len(validation_indices), + "group_fields": [ + field + for field in ("group_id", "session") + if any(row.group_field == field for row in source_rows) + ], + "normalized_text_duplicates_co_located": True, + } + + +def _write_jsonl(path: Path, rows: Sequence[dict[str, Any]]) -> None: + with path.open("w", encoding="utf-8", newline="\n") as handle: + for row in rows: + handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n") + + +def _human_summary(dataset: dict[str, Any]) -> str: + diagnostics = dataset["diagnostics"] + return "\n".join( + [ + "Inflect prepared dataset", + "========================", + f"Language: {dataset['language']}", + f"Frontend: {dataset['frontend']}", + f"Sample rate: {dataset['sample_rate']} Hz mono", + f"Rows: {dataset['row_counts']['total']} " + f"(train {dataset['row_counts']['train']}, " + f"validation {dataset['row_counts']['validation']})", + f"Audio duration: {diagnostics['total_duration_seconds']:.2f} seconds", + f"Resampled files: {diagnostics['resampled_files']}", + f"Downmixed files: {diagnostics['downmixed_files']}", + f"Added symbols: {diagnostics['added_symbol_count']}", + f"Base-symbol coverage before extension: " + f"{diagnostics['base_symbol_coverage_fraction']:.6f}", + "", + ] + ) + + +def prepare_dataset(options: PrepareOptions) -> dict[str, Any]: + """Prepare a manifest into the versioned public dataset layout. + + Output is assembled in a sibling staging directory and moved into place + only after every row succeeds. + """ + options.validate() + manifest_path = Path(options.manifest_path).expanduser().resolve() + output_dir = Path(options.output_dir).expanduser().resolve() + if output_dir.exists() and not output_dir.is_dir(): + raise PreparationError(f"Output path exists and is not a directory: {output_dir}") + if output_dir.exists() and any(output_dir.iterdir()): + raise PreparationError( + f"Output directory is not empty: {output_dir}. Choose a new directory." + ) + + frontend_options = FrontendOptions( + mode=options.frontend, + language=options.language, + hook=options.frontend_hook, + ) + validate_frontend(frontend_options) + rows = parse_manifest( + manifest_path, + audio_root=options.audio_root, + require_phonemes=options.frontend == "prephonemized", + ) + speakers = sorted({row.speaker for row in rows if row.speaker}) + if len(speakers) > 1: + raise PreparationError( + "Single-speaker adaptation requires one consistent nonempty speaker value, " + f"but the manifest contains {len(speakers)} values: " + + ", ".join(repr(speaker) for speaker in speakers) + ) + dataset_speaker = speakers[0] if speakers else None + base_symbols = ( + load_base_symbols(options.base_symbols_path) + if options.base_symbols_path + else BASE_SYMBOLS + ) + audio_options = AudioOptions( + sample_rate=options.sample_rate, + min_duration_seconds=options.min_duration_seconds, + max_duration_seconds=options.max_duration_seconds, + ) + + stage = output_dir.with_name(f".{output_dir.name}.preparing-{uuid.uuid4().hex}") + if stage.exists(): + raise PreparationError(f"Unexpected preparation staging collision: {stage}") + (stage / "audio").mkdir(parents=True) + + prepared_rows: list[dict[str, Any]] = [] + audio_diagnostics: list[dict[str, Any]] = [] + seen_audio_hashes: dict[str, ManifestRow] = {} + try: + for row in rows: + filename = _safe_name(row) + destination = stage / "audio" / filename + try: + audio = convert_wav(row.audio_path, destination, audio_options) + frontend = process_text( + row.text, + options=frontend_options, + prephonemized=row.phonemes, + ) + except Exception as exc: + raise PreparationError(f"Failed to prepare {row.source_location}: {exc}") from exc + audio_sha256 = _sha256(destination) + duplicate = seen_audio_hashes.get(audio_sha256) + if duplicate is not None: + raise PreparationError( + "Duplicate audio content detected after canonical 24 kHz conversion " + f"between {duplicate.source_location} and {row.source_location}. " + "Remove or replace one duplicate before training." + ) + seen_audio_hashes[audio_sha256] = row + prepared = { + "audio": f"audio/{filename}", + "audio_sha256": audio_sha256, + "text": frontend.raw_text, + "normalized_text": frontend.normalized_text, + "phonemes": frontend.phonemes, + "duration_seconds": round(audio.duration_seconds, 6), + } + if row.group_id: + prepared["group_id"] = row.group_id + prepared["group_field"] = row.group_field or "group_id" + if row.speaker: + prepared["speaker"] = row.speaker + prepared_rows.append(prepared) + audio_diagnostics.append(audio.to_dict()) + + phoneme_texts = [row["phonemes"] for row in prepared_rows] + base_coverage = audit_symbol_coverage(phoneme_texts, base_symbols) + declared_symbols = custom_frontend_symbols(frontend_options) or () + inventory = build_symbol_inventory( + phoneme_texts, + base_symbols=base_symbols, + extension_symbols=declared_symbols, + ) + validation_indices, split_metadata = _split_components( + prepared_rows, + rows, + fraction=options.validation_fraction, + seed=options.split_seed, + ) + train_rows = [ + row for index, row in enumerate(prepared_rows) if index not in validation_indices + ] + validation_rows = [ + row for index, row in enumerate(prepared_rows) if index in validation_indices + ] + _write_jsonl(stage / "train.jsonl", train_rows) + _write_jsonl(stage / "validation.jsonl", validation_rows) + write_symbol_inventory(stage / "symbols.json", inventory) + + frontend_metadata = { + "type": options.frontend, + "language": options.language, + "preserve_punctuation": frontend_options.preserve_punctuation, + "with_stress": frontend_options.with_stress, + } + custom_metadata = custom_frontend_metadata(frontend_options) + if custom_metadata is not None: + frontend_metadata["hook"] = custom_metadata + + dataset = { + "format": "inflect_prepared_dataset_v1", + "language": options.language, + "sample_rate": options.sample_rate, + "channels": 1, + "speaker": dataset_speaker, + "frontend": frontend_metadata, + "source_manifest_sha256": _sha256(manifest_path), + "split_seed": options.split_seed, + "validation_fraction": options.validation_fraction, + "split": split_metadata, + "row_counts": { + "total": len(prepared_rows), + "train": len(train_rows), + "validation": len(validation_rows), + }, + "diagnostics": { + "total_duration_seconds": round( + sum(item["duration_seconds"] for item in audio_diagnostics), 6 + ), + "min_duration_seconds": min( + item["duration_seconds"] for item in audio_diagnostics + ), + "max_duration_seconds": max( + item["duration_seconds"] for item in audio_diagnostics + ), + "resampled_files": sum(bool(item["resampled"]) for item in audio_diagnostics), + "downmixed_files": sum(bool(item["downmixed"]) for item in audio_diagnostics), + "source_clipped_files": sum( + item["source_clipped_fraction"] > 0 for item in audio_diagnostics + ), + "base_symbol_coverage_fraction": base_coverage.coverage_fraction, + "base_unknown_symbols": base_coverage.unknown_counts, + "added_symbol_count": len(inventory.added_symbols), + }, + } + (stage / "dataset.json").write_text( + json.dumps(dataset, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + report = { + "format": "inflect_preparation_report_v1", + "dataset": dataset, + "audio": audio_diagnostics, + "symbol_coverage_before_extension": base_coverage.to_dict(), + } + (stage / "preparation_report.json").write_text( + json.dumps(report, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + (stage / "phoneme_coverage.json").write_text( + json.dumps( + { + "format": "inflect_phoneme_coverage_v1", + "base_inventory": base_coverage.to_dict(), + "added_symbols": list(inventory.added_symbols), + "adapted_inventory_size": len(inventory.symbols), + }, + ensure_ascii=False, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + (stage / "PREPARATION_REPORT.txt").write_text( + _human_summary(dataset), encoding="utf-8" + ) + + if output_dir.exists(): + output_dir.rmdir() + stage.replace(output_dir) + return dataset + except Exception: + shutil.rmtree(stage, ignore_errors=True) + raise diff --git a/finetune/inflect_finetune/presets/__init__.py b/finetune/inflect_finetune/presets/__init__.py new file mode 100644 index 0000000..ea586bd --- /dev/null +++ b/finetune/inflect_finetune/presets/__init__.py @@ -0,0 +1,57 @@ +"""Wheel-safe built-in adaptation presets.""" + +from __future__ import annotations + +from copy import deepcopy +from typing import Any + + +_PRESETS: dict[str, dict[str, Any]] = { + "balanced": { + "batch_size": 4, + "gradient_accumulation_steps": 2, + "learning_rate_g": 0.0001, + "learning_rate_d": 0.0001, + "max_steps": 20_000, + "num_workers": 4, + "amp": True, + "checkpoint_interval": 1_000, + "validation_interval": 500, + }, + "micro-12gb": { + "batch_size": 2, + "gradient_accumulation_steps": 4, + "learning_rate_g": 0.00008, + "learning_rate_d": 0.00008, + "max_steps": 20_000, + "num_workers": 3, + "amp": True, + "checkpoint_interval": 1_000, + "validation_interval": 500, + }, + "nano-8gb": { + "batch_size": 1, + "gradient_accumulation_steps": 8, + "learning_rate_g": 0.00008, + "learning_rate_d": 0.00008, + "max_steps": 20_000, + "num_workers": 2, + "amp": True, + "checkpoint_interval": 1_000, + "validation_interval": 500, + }, +} + + +def available_presets() -> tuple[str, ...]: + return tuple(sorted(_PRESETS)) + + +def load_packaged_preset(name: str) -> dict[str, Any]: + try: + return deepcopy(_PRESETS[name]) + except KeyError as error: + raise KeyError(name) from error + + +__all__ = ["available_presets", "load_packaged_preset"] diff --git a/finetune/inflect_finetune/reporting.py b/finetune/inflect_finetune/reporting.py new file mode 100644 index 0000000..fafec55 --- /dev/null +++ b/finetune/inflect_finetune/reporting.py @@ -0,0 +1,147 @@ +"""Small, dependency-free helpers for reproducible public reports.""" + +from __future__ import annotations + +import dataclasses +import hashlib +import json +import math +import os +import tempfile +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Iterable, Mapping + + +REPORT_SCHEMA_VERSION = 1 + + +def utc_now() -> str: + """Return an RFC 3339 timestamp without host-specific locale data.""" + + return datetime.now(timezone.utc).isoformat(timespec="seconds").replace("+00:00", "Z") + + +def sha256_file(path: str | Path, *, block_size: int = 1024 * 1024) -> str: + """Hash a file without reading it all into memory.""" + + digest = hashlib.sha256() + with Path(path).open("rb") as handle: + for block in iter(lambda: handle.read(block_size), b""): + digest.update(block) + return digest.hexdigest() + + +def file_record(path: str | Path, *, relative_to: str | Path | None = None) -> dict[str, Any]: + """Describe a file using a stable relative path, size, and SHA-256.""" + + source = Path(path).resolve() + if relative_to is None: + display_path = source.name + else: + display_path = source.relative_to(Path(relative_to).resolve()).as_posix() + return { + "path": display_path, + "bytes": source.stat().st_size, + "sha256": sha256_file(source), + } + + +def _json_value(value: Any) -> Any: + if dataclasses.is_dataclass(value): + return _json_value(dataclasses.asdict(value)) + if isinstance(value, Path): + return str(value) + if isinstance(value, Mapping): + return {str(key): _json_value(item) for key, item in value.items()} + if isinstance(value, (list, tuple, set)): + return [_json_value(item) for item in value] + if hasattr(value, "item") and callable(value.item): + try: + return _json_value(value.item()) + except (TypeError, ValueError): + pass + if isinstance(value, float) and not math.isfinite(value): + return None + if isinstance(value, (str, int, float, bool)) or value is None: + return value + return repr(value) + + +def write_json(path: str | Path, payload: Any) -> Path: + """Atomically write deterministic, UTF-8 JSON.""" + + destination = Path(path) + destination.parent.mkdir(parents=True, exist_ok=True) + encoded = (json.dumps(_json_value(payload), indent=2, sort_keys=True) + "\n").encode("utf-8") + fd, temporary = tempfile.mkstemp(prefix=f".{destination.name}.", dir=destination.parent) + try: + with os.fdopen(fd, "wb") as handle: + handle.write(encoded) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, destination) + except BaseException: + try: + os.unlink(temporary) + except FileNotFoundError: + pass + raise + return destination + + +def write_text(path: str | Path, text: str) -> Path: + """Atomically write a human-readable UTF-8 report.""" + + destination = Path(path) + destination.parent.mkdir(parents=True, exist_ok=True) + fd, temporary = tempfile.mkstemp(prefix=f".{destination.name}.", dir=destination.parent) + try: + with os.fdopen(fd, "w", encoding="utf-8", newline="\n") as handle: + handle.write(text.rstrip() + "\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, destination) + except BaseException: + try: + os.unlink(temporary) + except FileNotFoundError: + pass + raise + return destination + + +def write_checksums( + path: str | Path, + files: Iterable[str | Path], + *, + relative_to: str | Path, +) -> Path: + """Write a standard SHA-256 checksum list.""" + + root = Path(relative_to).resolve() + rows = [] + resolved = (Path(file).resolve() for file in files) + for item in sorted(resolved, key=lambda value: value.as_posix()): + rows.append(f"{sha256_file(item)} {item.relative_to(root).as_posix()}") + return write_text(path, "\n".join(rows)) + + +def make_report(kind: str, **payload: Any) -> dict[str, Any]: + """Create the shared envelope used by export and evaluation reports.""" + + return { + "format": f"inflect_adaptation_{kind}_v{REPORT_SCHEMA_VERSION}", + "created_at": utc_now(), + **payload, + } + + +def status(ok: bool, message: str, **details: Any) -> dict[str, Any]: + """Create a consistently shaped validation result.""" + + return { + "ok": bool(ok), + "message": message, + **details, + } diff --git a/finetune/inflect_finetune/symbols.py b/finetune/inflect_finetune/symbols.py new file mode 100644 index 0000000..d3e9e53 --- /dev/null +++ b/finetune/inflect_finetune/symbols.py @@ -0,0 +1,174 @@ +"""Ordered Inflect symbol inventories and phoneme coverage reporting.""" + +from __future__ import annotations + +import json +import unicodedata +from collections import Counter +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Iterable, Sequence + + +_PAD = "_" +_PUNCTUATION = ';:,.!?¡¿—…"«»“” ' +_LETTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" +_LETTERS_IPA = ( + "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵ" + "ɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞" + "↓↑→↗↘'̩'ᵻ" +) +BASE_SYMBOLS: tuple[str, ...] = tuple( + [_PAD] + list(_PUNCTUATION) + list(_LETTERS) + list(_LETTERS_IPA) +) + + +class SymbolInventoryError(ValueError): + """Raised when a symbol inventory is malformed.""" + + +@dataclass(frozen=True) +class CoverageReport: + """Character coverage of phoneme strings against an ordered inventory.""" + + total_characters: int + covered_characters: int + coverage_fraction: float + unique_observed: int + unknown_counts: dict[str, int] + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-serializable representation.""" + return asdict(self) + + +@dataclass(frozen=True) +class SymbolInventory: + """An identity-preserving extension of a base symbol inventory.""" + + symbols: tuple[str, ...] + base_symbols: tuple[str, ...] + added_symbols: tuple[str, ...] + + def to_dict(self) -> dict[str, Any]: + """Return the public symbols.json contract.""" + return { + "format": "inflect_symbol_inventory_v1", + "symbols": list(self.symbols), + "base_symbols": list(self.base_symbols), + "base_size": len(self.base_symbols), + "total_size": len(self.symbols), + "added_symbols": list(self.added_symbols), + "base_identity": [ + { + "symbol": symbol, + "base_index": position, + "adapted_index": position, + } + for position, symbol in enumerate(self.base_symbols) + ], + } + + +def _validate_symbols( + symbols: Sequence[str], + *, + label: str, + allow_duplicates: bool = False, +) -> tuple[str, ...]: + result = tuple(symbols) + if not result: + raise SymbolInventoryError(f"{label} cannot be empty.") + if any(not isinstance(symbol, str) or not symbol for symbol in result): + raise SymbolInventoryError(f"{label} must contain non-empty strings.") + if not allow_duplicates and len(set(result)) != len(result): + duplicates = sorted(symbol for symbol, count in Counter(result).items() if count > 1) + raise SymbolInventoryError(f"{label} contains duplicate symbols: {duplicates!r}") + return result + + +def load_base_symbols(path: Path | None = None) -> tuple[str, ...]: + """Load a JSON symbol list, or use the exact Inflect v2 release inventory.""" + if path is None: + return BASE_SYMBOLS + path = Path(path) + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: + raise SymbolInventoryError(f"Could not read base symbols from {path}: {exc}") from exc + if isinstance(payload, dict): + payload = payload.get("symbols") + if not isinstance(payload, list): + raise SymbolInventoryError( + f"Base symbol file {path} must be a JSON list or an object with 'symbols'." + ) + return _validate_symbols( + payload, + label=f"base symbols in {path}", + allow_duplicates=True, + ) + + +def audit_symbol_coverage( + phoneme_texts: Iterable[str], + symbols: Sequence[str], +) -> CoverageReport: + """Measure character-level coverage against a supplied inventory.""" + inventory = set( + _validate_symbols(symbols, label="symbols", allow_duplicates=True) + ) + counts: Counter[str] = Counter() + for text in phoneme_texts: + if not isinstance(text, str): + raise SymbolInventoryError("Phoneme entries must be strings.") + counts.update(unicodedata.normalize("NFC", text)) + total = sum(counts.values()) + unknown = {symbol: counts[symbol] for symbol in sorted(counts) if symbol not in inventory} + unknown_total = sum(unknown.values()) + covered = total - unknown_total + return CoverageReport( + total_characters=total, + covered_characters=covered, + coverage_fraction=(covered / total) if total else 1.0, + unique_observed=len(counts), + unknown_counts=unknown, + ) + + +def build_symbol_inventory( + phoneme_texts: Iterable[str], + *, + base_symbols: Sequence[str] = BASE_SYMBOLS, + extension_symbols: Sequence[str] = (), +) -> SymbolInventory: + """Append declared then observed symbols while preserving base symbol indices.""" + base = _validate_symbols( + base_symbols, + label="base_symbols", + allow_duplicates=True, + ) + declared = _validate_symbols( + extension_symbols, + label="extension_symbols", + ) if extension_symbols else () + observed: set[str] = set() + for text in phoneme_texts: + if not isinstance(text, str): + raise SymbolInventoryError("Phoneme entries must be strings.") + observed.update(unicodedata.normalize("NFC", text)) + declared_added = tuple(symbol for symbol in declared if symbol not in base) + remaining = observed.difference(base).difference(declared_added) + added = declared_added + tuple( + sorted(remaining, key=lambda value: tuple(map(ord, value))) + ) + return SymbolInventory(symbols=base + added, base_symbols=base, added_symbols=added) + + +def write_symbol_inventory(path: Path, inventory: SymbolInventory) -> None: + """Write symbols.json deterministically.""" + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(inventory.to_dict(), ensure_ascii=False, indent=2, sort_keys=False) + "\n", + encoding="utf-8", + ) diff --git a/finetune/inflect_finetune/training.py b/finetune/inflect_finetune/training.py new file mode 100644 index 0000000..7fad881 --- /dev/null +++ b/finetune/inflect_finetune/training.py @@ -0,0 +1,869 @@ +"""Generic staged warm-start training for Inflect v2 adaptation.""" + +from __future__ import annotations + +import json +import logging +import random +import uuid +from dataclasses import MISSING, asdict, dataclass, fields, replace +from pathlib import Path +from typing import Any, Iterable + +import numpy as np +import soundfile as sf +import torch +from torch import nn +from torch.nn import functional as F + +from .checkpoint import ( + CompatibilityReport, + build_run_identity, + cpu_compatibility_report, + load_run_identity, + resume_training_checkpoint, + save_inference_checkpoint, + save_training_checkpoint, + validate_run_identity, + write_run_identity, +) +from .modeling import ( + ModelBundle, + build_training_models, + load_symbols, + optimizer_parameters, + resolve_base_model, +) +from .presets import available_presets, load_packaged_preset +from .training_data import AudioConfig, PreparedTTSDataset, create_dataloader + + +LOGGER = logging.getLogger("inflect_finetune") +STAGE_POSTERIOR = "posterior_warmup" +STAGE_ADAPT = "linguistic_adaptation" +STAGE_DECODER = "decoder_polish" +STAGES = (STAGE_POSTERIOR, STAGE_ADAPT, STAGE_DECODER) + + +@dataclass(frozen=True) +class TrainingOptions: + """Public, generic adaptation settings. + + These defaults are intentionally conservative starting points, not the + private schedule used to create an Inflect release checkpoint. + """ + + base_model: str | Path + prepared_dir: str | Path + output_dir: str | Path + preset: str | Path | None = None + resume: str | Path | None = None + device: str = "auto" + seed: int = 1234 + batch_size: int = 2 + gradient_accumulation_steps: int = 4 + num_workers: int = 2 + max_steps: int = 20_000 + learning_rate_g: float = 8.0e-5 + learning_rate_d: float = 8.0e-5 + posterior_lr_multiplier: float = 1.0 + linguistic_lr_multiplier: float = 0.5 + decoder_lr_multiplier: float = 0.1 + posterior_warmup_steps: int = 500 + decoder_unfreeze_step: int | None = 3_000 + amp: bool = True + adam_betas: tuple[float, float] = (0.8, 0.99) + adam_eps: float = 1.0e-9 + weight_decay: float = 0.0 + lr_decay: float = 0.99999 + max_grad_norm: float = 10.0 + mel_loss_weight: float = 45.0 + kl_loss_weight: float = 1.0 + feature_loss_weight: float = 1.0 + duration_loss_weight: float = 1.0 + checkpoint_interval: int = 1_000 + validation_interval: int = 500 + log_interval: int = 25 + validation_seed: int = 7 + + @classmethod + def from_preset( + cls, + preset: str | Path, + *, + base_model: str | Path, + prepared_dir: str | Path, + output_dir: str | Path, + **overrides: Any, + ) -> "TrainingOptions": + payload = load_preset(preset) + payload.update(overrides) + return cls( + base_model=base_model, + prepared_dir=prepared_dir, + output_dir=output_dir, + # The preset and caller overrides are already merged above. + preset=None, + **payload, + ) + + def resolved(self) -> "TrainingOptions": + if self.preset is None: + return self + payload = load_preset(self.preset) + defaults = { + field.name: field.default + for field in fields(self) + if field.default is not MISSING + } + explicit = asdict(self) + # Dataclass values differing from defaults are treated as explicit + # caller overrides; required paths are always preserved. + for key, value in explicit.items(): + if key in {"base_model", "prepared_dir", "output_dir", "preset", "resume"}: + continue + if key not in defaults or value != defaults[key]: + payload[key] = value + allowed = {field.name for field in fields(self)} + unknown = sorted(set(payload) - allowed) + if unknown: + raise ValueError(f"Preset contains unknown TrainingOptions: {unknown}") + return replace(self, **payload, preset=None) + + +@dataclass +class TrainingState: + step: int = 0 + epoch: int = 0 + stage: str = STAGE_POSTERIOR + + +def load_preset(preset: str | Path) -> dict[str, Any]: + source = Path(preset) + if source.is_file(): + payload = json.loads(source.read_text(encoding="utf-8")) + else: + name = str(preset) + try: + payload = load_packaged_preset(name) + except KeyError as error: + raise FileNotFoundError( + f"Unknown training preset {preset!r}. " + f"Available: {list(available_presets())}" + ) from error + if not isinstance(payload, dict): + raise ValueError(f"Training preset {preset!r} must contain a JSON object.") + return payload + + +def _validate_options(options: TrainingOptions) -> None: + positive_ints = ( + "batch_size", + "gradient_accumulation_steps", + "max_steps", + "checkpoint_interval", + "validation_interval", + "log_interval", + ) + for name in positive_ints: + if int(getattr(options, name)) <= 0: + raise ValueError(f"{name} must be positive.") + if options.posterior_warmup_steps < 0: + raise ValueError("posterior_warmup_steps must be non-negative.") + if ( + options.decoder_unfreeze_step is not None + and options.decoder_unfreeze_step < options.posterior_warmup_steps + ): + raise ValueError("decoder_unfreeze_step cannot precede posterior_warmup_steps.") + for name in ( + "learning_rate_g", + "learning_rate_d", + "posterior_lr_multiplier", + "linguistic_lr_multiplier", + "decoder_lr_multiplier", + "lr_decay", + ): + if float(getattr(options, name)) <= 0: + raise ValueError(f"{name} must be positive.") + + +def _device(name: str) -> torch.device: + if name == "auto": + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = torch.device(name) + if device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is not available.") + return device + + +def _seed_everything(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +def _stage_for_step(options: TrainingOptions, step: int) -> str: + if step < options.posterior_warmup_steps: + return STAGE_POSTERIOR + if options.decoder_unfreeze_step is not None and step >= options.decoder_unfreeze_step: + return STAGE_DECODER + return STAGE_ADAPT + + +def _generator_groups(model: nn.Module, options: TrainingOptions) -> list[dict[str, Any]]: + named = dict(model.named_parameters()) + assignments: dict[str, list[nn.Parameter]] = { + "posterior": [], + "linguistic": [], + "decoder": [], + } + for name, parameter in named.items(): + if name.startswith("enc_q."): + assignments["posterior"].append(parameter) + elif name.startswith("dec."): + assignments["decoder"].append(parameter) + elif name.startswith(("enc_p.", "dp.", "flow.")): + assignments["linguistic"].append(parameter) + else: + raise RuntimeError(f"Unclassified generator parameter {name!r}.") + if any(not values for values in assignments.values()): + empty = [name for name, values in assignments.items() if not values] + raise RuntimeError(f"Generator parameter groups are unexpectedly empty: {empty}") + multipliers = { + "posterior": options.posterior_lr_multiplier, + "linguistic": options.linguistic_lr_multiplier, + "decoder": options.decoder_lr_multiplier, + } + return [ + { + "params": parameters, + "name": name, + "lr_multiplier": multipliers[name], + "lr": 0.0, + } + for name, parameters in assignments.items() + ] + + +def _apply_stage( + model: nn.Module, + optimizer: torch.optim.Optimizer, + options: TrainingOptions, + stage: str, + *, + reset_learning_rates: bool = True, +) -> None: + if stage not in STAGES: + raise ValueError(f"Unknown training stage {stage!r}.") + enabled = { + STAGE_POSTERIOR: {"posterior"}, + STAGE_ADAPT: {"posterior", "linguistic"}, + STAGE_DECODER: {"posterior", "linguistic", "decoder"}, + }[stage] + for group in optimizer.param_groups: + name = group["name"] + active = name in enabled + if not active: + group["lr"] = 0.0 + elif reset_learning_rates: + group["lr"] = options.learning_rate_g * float(group["lr_multiplier"]) + for parameter in group["params"]: + parameter.requires_grad_(active) + if not active: + parameter.grad = None + optimizer.zero_grad(set_to_none=True) + + +def _set_requires_grad(module: nn.Module, enabled: bool) -> None: + for parameter in module.parameters(): + parameter.requires_grad_(enabled) + + +def _discriminator_loss( + real_outputs: Iterable[torch.Tensor], generated_outputs: Iterable[torch.Tensor] +) -> torch.Tensor: + return sum( + torch.mean((1.0 - real.float()) ** 2) + torch.mean(generated.float() ** 2) + for real, generated in zip(real_outputs, generated_outputs) + ) + + +def _generator_loss(outputs: Iterable[torch.Tensor]) -> torch.Tensor: + return sum(torch.mean((1.0 - output.float()) ** 2) for output in outputs) + + +def _feature_loss( + real_maps: Iterable[Iterable[torch.Tensor]], + generated_maps: Iterable[Iterable[torch.Tensor]], +) -> torch.Tensor: + return 2.0 * sum( + torch.mean(torch.abs(real.float().detach() - generated.float())) + for real_group, generated_group in zip(real_maps, generated_maps) + for real, generated in zip(real_group, generated_group) + ) + + +def _kl_loss( + z_p: torch.Tensor, + logs_q: torch.Tensor, + m_p: torch.Tensor, + logs_p: torch.Tensor, + mask: torch.Tensor, +) -> torch.Tensor: + value = logs_p.float() - logs_q.float() - 0.5 + value += 0.5 * ((z_p.float() - m_p.float()) ** 2) * torch.exp(-2.0 * logs_p.float()) + return torch.sum(value * mask.float()) / torch.sum(mask.float()).clamp_min(1.0) + + +def _hz_to_mel(value: torch.Tensor) -> torch.Tensor: + return 2595.0 * torch.log10(1.0 + value / 700.0) + + +def _mel_to_hz(value: torch.Tensor) -> torch.Tensor: + return 700.0 * (torch.pow(10.0, value / 2595.0) - 1.0) + + +def _mel_filterbank( + *, + n_fft: int, + n_mels: int, + sample_rate: int, + fmin: float, + fmax: float, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + minimum = _hz_to_mel(torch.tensor(fmin, dtype=torch.float64)) + maximum = _hz_to_mel(torch.tensor(fmax, dtype=torch.float64)) + points = _mel_to_hz(torch.linspace(minimum, maximum, n_mels + 2)) + bins = torch.floor((n_fft + 1) * points / sample_rate).long() + frequencies = n_fft // 2 + 1 + bank = torch.zeros(n_mels, frequencies, dtype=torch.float32) + for index in range(n_mels): + left, center, right = (int(value) for value in bins[index : index + 3]) + center = max(center, left + 1) + right = max(right, center + 1) + for column in range(left, min(center, frequencies)): + bank[index, column] = (column - left) / (center - left) + for column in range(center, min(right, frequencies)): + bank[index, column] = (right - column) / (right - center) + return bank.to(device=device, dtype=dtype) + + +def _mel_from_spec(spec: torch.Tensor, bundle: ModelBundle) -> torch.Tensor: + data = bundle.config["data"] + bank = _mel_filterbank( + n_fft=int(data["filter_length"]), + n_mels=int(data["n_mel_channels"]), + sample_rate=int(data["sampling_rate"]), + fmin=float(data["mel_fmin"]), + fmax=float(data["mel_fmax"]), + device=spec.device, + dtype=spec.dtype, + ) + return torch.log(torch.matmul(bank, spec).clamp_min(1.0e-5)) + + +def _mel_from_waveform(waveform: torch.Tensor, bundle: ModelBundle) -> torch.Tensor: + data = bundle.config["data"] + n_fft = int(data["filter_length"]) + hop = int(data["hop_length"]) + win = int(data["win_length"]) + padding = (n_fft - hop) // 2 + padded = F.pad(waveform.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1) + window = torch.hann_window(win, device=waveform.device, dtype=waveform.dtype) + spectrum = torch.stft( + padded, + n_fft=n_fft, + hop_length=hop, + win_length=win, + window=window, + center=False, + return_complex=True, + ).abs().clamp_min(1.0e-5) + return _mel_from_spec(spectrum, bundle) + + +def _autocast(device: torch.device, enabled: bool): + return torch.autocast(device_type=device.type, enabled=enabled and device.type == "cuda") + + +def _grad_scaler(enabled: bool): + try: + return torch.amp.GradScaler("cuda", enabled=enabled) + except (AttributeError, TypeError): + return torch.cuda.amp.GradScaler(enabled=enabled) + + +def _public_options(options: TrainingOptions) -> dict[str, Any]: + payload = asdict(options) + for key in ("base_model", "prepared_dir", "output_dir", "preset", "resume"): + payload.pop(key, None) + return json.loads(json.dumps(payload, sort_keys=True)) + + +def _optimizer_schema(options: TrainingOptions) -> dict[str, Any]: + common = { + "class": "torch.optim.AdamW", + "betas": list(options.adam_betas), + "eps": options.adam_eps, + "weight_decay": options.weight_decay, + } + return { + "generator": { + **common, + "base_learning_rate": options.learning_rate_g, + "parameter_groups": [ + { + "name": "posterior", + "lr_multiplier": options.posterior_lr_multiplier, + }, + { + "name": "linguistic", + "lr_multiplier": options.linguistic_lr_multiplier, + }, + { + "name": "decoder", + "lr_multiplier": options.decoder_lr_multiplier, + }, + ], + }, + "discriminator": { + **common, + "base_learning_rate": options.learning_rate_d, + }, + "scheduler": { + "class": "torch.optim.lr_scheduler.ExponentialLR", + "gamma": options.lr_decay, + }, + "gradient_accumulation_steps": options.gradient_accumulation_steps, + "amp_requested": options.amp, + } + + +def _validate_new_output_dir(output_dir: Path) -> None: + """Reject accidental reuse while permitting empty scaffolding.""" + + if not output_dir.exists(): + return + if not output_dir.is_dir(): + raise ValueError(f"Training output path is not a directory: {output_dir}") + allowed_directories = {"checkpoints", "exports", "validation"} + unsafe: list[str] = [] + for entry in output_dir.iterdir(): + if entry.is_file(): + if entry.name != ".gitkeep": + unsafe.append(entry.name) + continue + if not entry.is_dir() or entry.name not in allowed_directories: + unsafe.append(entry.name) + continue + for nested in entry.rglob("*"): + if nested.is_dir() or nested.name != ".gitkeep": + unsafe.append(str(nested.relative_to(output_dir))) + if unsafe: + raise ValueError( + "Refusing to start a new run in a nonempty output directory. " + f"Unexpected entries: {sorted(unsafe)}" + ) + + +def _resume_checkpoint_for_run(output_dir: Path, resume: str | Path) -> Path: + checkpoint = Path(resume).resolve() + checkpoints_dir = (output_dir / "checkpoints").resolve() + try: + checkpoint.relative_to(checkpoints_dir) + except ValueError as error: + raise ValueError( + "Resume checkpoint must be inside this run's checkpoints directory: " + f"{checkpoints_dir}" + ) from error + if not checkpoint.is_file(): + raise FileNotFoundError(f"Resume checkpoint does not exist: {checkpoint}") + return checkpoint + + +def _establish_run_identity( + *, + options: TrainingOptions, + output_dir: Path, + base_root: Path, + prepared_dir: Path, +) -> tuple[dict[str, Any], Path | None]: + marker_path = output_dir / "run-identity.json" + public_options = _public_options(options) + optimizer_schema = _optimizer_schema(options) + if options.resume is None: + _validate_new_output_dir(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + identity = build_run_identity( + run_id=uuid.uuid4().hex, + base_root=base_root, + prepared_dir=prepared_dir, + options=public_options, + optimizer_schema=optimizer_schema, + ) + write_run_identity(marker_path, identity) + return identity, None + + if not marker_path.is_file(): + raise ValueError( + "Resume requires run-identity.json in the existing output directory." + ) + checkpoint = _resume_checkpoint_for_run(output_dir, options.resume) + recorded = load_run_identity(marker_path) + expected = build_run_identity( + run_id=str(recorded.get("run_id", "")), + base_root=base_root, + prepared_dir=prepared_dir, + options=public_options, + optimizer_schema=optimizer_schema, + ) + validate_run_identity(recorded, expected, source="output directory run marker") + return expected, checkpoint + + +@torch.inference_mode() +def _validate( + bundle: ModelBundle, + loader: Any, + device: torch.device, + output_dir: Path, + step: int, + seed: int, +) -> dict[str, Any]: + model = bundle.generator + model.eval() + batch = next(iter(loader)) + x, x_lengths = batch[0].to(device), batch[1].to(device) + torch.manual_seed(seed) + output = model.infer(x[:1], x_lengths[:1], noise_scale=0.667, max_len=4000)[0] + waveform = output[0, 0].float().cpu().numpy() + sample_rate = int(bundle.config["data"]["sampling_rate"]) + sample_path = output_dir / "validation" / f"step-{step:08d}.wav" + sample_path.parent.mkdir(parents=True, exist_ok=True) + sf.write(sample_path, np.clip(waveform, -1.0, 1.0), sample_rate) + model.train() + return { + "step": step, + "sample": str(sample_path), + "samples": int(waveform.size), + "seconds": waveform.size / sample_rate, + "peak": float(np.max(np.abs(waveform))), + } + + +def train_adaptation(options: TrainingOptions) -> dict[str, Any]: + """Warm-start and train one fixed-voice, single-language checkpoint.""" + + options = options.resolved() + _validate_options(options) + _seed_everything(options.seed) + device = _device(options.device) + output_dir = Path(options.output_dir).resolve() + prepared_dir = Path(options.prepared_dir).resolve() + symbols = load_symbols(prepared_dir / "symbols.json") + base_root = resolve_base_model(options.base_model) + run_identity, resume_checkpoint = _establish_run_identity( + options=options, + output_dir=output_dir, + base_root=base_root, + prepared_dir=prepared_dir, + ) + + bundle = build_training_models(base_root, symbols, seed=options.seed) + compatibility = cpu_compatibility_report( + bundle.generator, + base_root / "model.pth", + bundle.base_symbols, + symbols, + initialization_seed=options.seed, + ) + compatibility.write(output_dir / "compatibility-report.json") + bundle.generator.to(device) + bundle.discriminator.to(device) + + generator_groups = _generator_groups(bundle.generator, options) + optimizer_g = torch.optim.AdamW( + generator_groups, + lr=options.learning_rate_g, + betas=options.adam_betas, + eps=options.adam_eps, + weight_decay=options.weight_decay, + ) + optimizer_d = torch.optim.AdamW( + optimizer_parameters(bundle.discriminator), + lr=options.learning_rate_d, + betas=options.adam_betas, + eps=options.adam_eps, + weight_decay=options.weight_decay, + ) + scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optimizer_g, gamma=options.lr_decay) + scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optimizer_d, gamma=options.lr_decay) + amp_enabled = bool(options.amp and device.type == "cuda") + scaler = _grad_scaler(amp_enabled) + + state = TrainingState(stage=_stage_for_step(options, 0)) + _apply_stage(bundle.generator, optimizer_g, options, state.stage) + if resume_checkpoint is not None: + step, epoch, saved_stage = resume_training_checkpoint( + resume_checkpoint, + generator=bundle.generator, + discriminator=bundle.discriminator, + optimizer_g=optimizer_g, + optimizer_d=optimizer_d, + scheduler_g=scheduler_g, + scheduler_d=scheduler_d, + scaler=scaler, + expected_symbols=symbols, + expected_run_identity=run_identity, + ) + expected_stage = _stage_for_step(options, step) + previous_stage = _stage_for_step(options, max(step - 1, 0)) + if saved_stage not in {previous_stage, expected_stage}: + raise ValueError( + f"Resume stage {saved_stage!r} conflicts with options-derived stage " + f"{expected_stage!r} at step {step}." + ) + # Checkpoints record the stage that produced their current step. When a + # checkpoint lands exactly on a stage boundary, resume must configure + # the model for the next step rather than reject the valid checkpoint. + state = TrainingState(step=step, epoch=epoch, stage=expected_stage) + _apply_stage( + bundle.generator, + optimizer_g, + options, + state.stage, + reset_learning_rates=False, + ) + + audio = AudioConfig( + sampling_rate=int(bundle.config["data"]["sampling_rate"]), + filter_length=int(bundle.config["data"]["filter_length"]), + hop_length=int(bundle.config["data"]["hop_length"]), + win_length=int(bundle.config["data"]["win_length"]), + add_blank=bool(bundle.config["data"].get("add_blank", True)), + ) + train_dataset = PreparedTTSDataset(prepared_dir, "train", symbols, audio) + validation_dataset = PreparedTTSDataset(prepared_dir, "validation", symbols, audio) + train_loader = create_dataloader( + train_dataset, + batch_size=options.batch_size, + shuffle=True, + num_workers=options.num_workers, + seed=options.seed, + pin_memory=device.type == "cuda", + ) + validation_loader = create_dataloader( + validation_dataset, + batch_size=1, + shuffle=False, + num_workers=min(options.num_workers, 1), + seed=options.seed, + pin_memory=device.type == "cuda", + ) + + options_path = output_dir / "training-options.json" + options_path.write_text( + json.dumps(_public_options(options), indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + log_path = output_dir / "metrics.jsonl" + bundle.generator.train() + bundle.discriminator.train() + optimizer_g.zero_grad(set_to_none=True) + optimizer_d.zero_grad(set_to_none=True) + micro_step = 0 + + while state.step < options.max_steps: + state.epoch += 1 + for batch in train_loader: + desired_stage = _stage_for_step(options, state.step) + if desired_stage != state.stage: + state.stage = desired_stage + _apply_stage(bundle.generator, optimizer_g, options, state.stage) + optimizer_d.zero_grad(set_to_none=True) + micro_step = 0 + + x, x_lengths, spec, spec_lengths, waveform, _ = ( + tensor.to(device, non_blocking=device.type == "cuda") for tensor in batch + ) + with _autocast(device, amp_enabled): + generated, duration, _, ids_slice, _, z_mask, latent = bundle.generator( + x, x_lengths, spec, spec_lengths + ) + z, z_p, m_p, logs_p, _, logs_q = latent + real = bundle.components.commons.slice_segments( + waveform, + ids_slice * audio.hop_length, + int(bundle.config["train"]["segment_size"]), + ) + + _set_requires_grad(bundle.discriminator, True) + real_scores, generated_scores, _, _ = bundle.discriminator( + real, generated.detach() + ) + loss_d = _discriminator_loss(real_scores, generated_scores) + + scaler.scale(loss_d / options.gradient_accumulation_steps).backward() + + _set_requires_grad(bundle.discriminator, False) + with _autocast(device, amp_enabled): + generated_scores_g = bundle.discriminator(real, generated) + _, fake_scores, real_maps, generated_maps = generated_scores_g + target_mel = bundle.components.commons.slice_segments( + _mel_from_spec(spec, bundle), + ids_slice, + int(bundle.config["train"]["segment_size"]) // audio.hop_length, + ) + generated_mel = _mel_from_waveform(generated.squeeze(1), bundle) + loss_mel = F.l1_loss(target_mel.float(), generated_mel.float()) + loss_duration = duration.float().sum() + loss_kl = _kl_loss(z_p, logs_q, m_p, logs_p, z_mask) + loss_feature = _feature_loss(real_maps, generated_maps) + loss_generator = _generator_loss(fake_scores) + loss_g = ( + loss_generator + + options.feature_loss_weight * loss_feature + + options.mel_loss_weight * loss_mel + + options.duration_loss_weight * loss_duration + + options.kl_loss_weight * loss_kl + ) + scaler.scale(loss_g / options.gradient_accumulation_steps).backward() + _set_requires_grad(bundle.discriminator, True) + micro_step += 1 + + if micro_step % options.gradient_accumulation_steps: + continue + scaler.unscale_(optimizer_g) + scaler.unscale_(optimizer_d) + torch.nn.utils.clip_grad_norm_( + [ + parameter + for group in optimizer_g.param_groups + for parameter in group["params"] + if parameter.requires_grad + ], + options.max_grad_norm, + ) + torch.nn.utils.clip_grad_norm_( + bundle.discriminator.parameters(), options.max_grad_norm + ) + scaler.step(optimizer_d) + scaler.step(optimizer_g) + scaler.update() + optimizer_g.zero_grad(set_to_none=True) + optimizer_d.zero_grad(set_to_none=True) + scheduler_g.step() + scheduler_d.step() + state.step += 1 + + metrics = { + "step": state.step, + "epoch": state.epoch, + "stage": state.stage, + "loss_g": float(loss_g.detach().cpu()), + "loss_d": float(loss_d.detach().cpu()), + "loss_mel": float(loss_mel.detach().cpu()), + "loss_duration": float(loss_duration.detach().cpu()), + "loss_kl": float(loss_kl.detach().cpu()), + "lr": { + group["name"]: float(group["lr"]) for group in optimizer_g.param_groups + }, + } + with log_path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(metrics, sort_keys=True) + "\n") + if state.step % options.log_interval == 0: + LOGGER.info( + "step=%d stage=%s loss_g=%.4f loss_d=%.4f", + state.step, + state.stage, + metrics["loss_g"], + metrics["loss_d"], + ) + if state.step % options.validation_interval == 0: + validation = _validate( + bundle, + validation_loader, + device, + output_dir, + state.step, + options.validation_seed, + ) + (output_dir / "validation" / f"step-{state.step:08d}.json").write_text( + json.dumps(validation, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + if state.step % options.checkpoint_interval == 0: + checkpoint_path = ( + output_dir / "checkpoints" / f"adaptation-step-{state.step:08d}.pth" + ) + save_training_checkpoint( + checkpoint_path, + generator=bundle.generator, + discriminator=bundle.discriminator, + optimizer_g=optimizer_g, + optimizer_d=optimizer_d, + scheduler_g=scheduler_g, + scheduler_d=scheduler_d, + scaler=scaler, + step=state.step, + epoch=state.epoch, + stage=state.stage, + options=_public_options(options), + symbols=symbols, + compatibility=compatibility, + run_identity=run_identity, + latest_path=output_dir / "checkpoints" / "latest.pth", + ) + save_inference_checkpoint( + output_dir / "exports" / f"model-step-{state.step:08d}.pth", + generator=bundle.generator, + iteration=state.step, + learning_rate=optimizer_g.param_groups[0]["lr"], + ) + if state.step >= options.max_steps: + break + + final_training = save_training_checkpoint( + output_dir / "checkpoints" / "adaptation-final.pth", + generator=bundle.generator, + discriminator=bundle.discriminator, + optimizer_g=optimizer_g, + optimizer_d=optimizer_d, + scheduler_g=scheduler_g, + scheduler_d=scheduler_d, + scaler=scaler, + step=state.step, + epoch=state.epoch, + stage=state.stage, + options=_public_options(options), + symbols=symbols, + compatibility=compatibility, + run_identity=run_identity, + latest_path=output_dir / "checkpoints" / "latest.pth", + ) + final_inference = save_inference_checkpoint( + output_dir / "exports" / "model.pth", + generator=bundle.generator, + iteration=state.step, + learning_rate=optimizer_g.param_groups[0]["lr"], + ) + summary = { + "step": state.step, + "epoch": state.epoch, + "stage": state.stage, + "training_checkpoint": str(final_training), + "inference_checkpoint": str(final_inference), + "compatibility_report": compatibility.to_dict(), + "run_id": run_identity["run_id"], + } + (output_dir / "training-summary.json").write_text( + json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + return summary + + +__all__ = ["TrainingOptions", "train_adaptation"] diff --git a/finetune/inflect_finetune/training_data.py b/finetune/inflect_finetune/training_data.py new file mode 100644 index 0000000..811abae --- /dev/null +++ b/finetune/inflect_finetune/training_data.py @@ -0,0 +1,184 @@ +"""Prepared-dataset loading for Inflect adaptation.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable, Sequence + +import soundfile as sf +import torch +from torch.nn import functional as F +from torch.utils.data import DataLoader, Dataset + + +@dataclass(frozen=True) +class AudioConfig: + sampling_rate: int + filter_length: int + hop_length: int + win_length: int + add_blank: bool + + +def _safe_child(root: Path, relative: str) -> Path: + candidate = (root / relative).resolve() + try: + candidate.relative_to(root.resolve()) + except ValueError as error: + raise ValueError( + f"Dataset audio path escapes the prepared directory: {relative!r}" + ) from error + if not candidate.is_file(): + raise FileNotFoundError(f"Prepared audio file does not exist: {candidate}") + return candidate + + +def read_jsonl(path: Path) -> list[dict]: + rows = [] + for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): + if not line.strip(): + continue + try: + row = json.loads(line) + except json.JSONDecodeError as error: + raise ValueError(f"{path}:{line_number}: invalid JSON: {error}") from error + if not isinstance(row, dict): + raise ValueError(f"{path}:{line_number}: each row must be a JSON object.") + rows.append(row) + if not rows: + raise ValueError(f"{path} contains no dataset rows.") + return rows + + +def spectrogram(waveform: torch.Tensor, config: AudioConfig) -> torch.Tensor: + """Match the released VITS magnitude-spectrogram convention.""" + + padding = (config.filter_length - config.hop_length) // 2 + if waveform.numel() < 2: + raise ValueError("Audio contains fewer than two samples.") + if waveform.numel() <= padding: + waveform = F.pad(waveform, (0, padding + 1 - waveform.numel())) + padded = F.pad(waveform[None, None], (padding, padding), mode="reflect")[0, 0] + window = torch.hann_window(config.win_length, dtype=waveform.dtype) + result = torch.stft( + padded, + n_fft=config.filter_length, + hop_length=config.hop_length, + win_length=config.win_length, + window=window, + center=False, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + return (result.abs().square() + 1.0e-6).sqrt() + + +class PreparedTTSDataset(Dataset): + """Load one split from the contract-defined prepared dataset.""" + + def __init__( + self, + prepared_dir: str | Path, + split: str, + symbols: Sequence[str], + audio_config: AudioConfig, + ) -> None: + if split not in {"train", "validation"}: + raise ValueError("split must be 'train' or 'validation'") + self.root = Path(prepared_dir).resolve() + self.rows = read_jsonl(self.root / f"{split}.jsonl") + self.symbol_to_id = {symbol: index for index, symbol in enumerate(symbols)} + self.audio_config = audio_config + self.records = [] + unknown: set[str] = set() + for index, row in enumerate(self.rows): + for field in ("audio", "phonemes"): + if not isinstance(row.get(field), str) or not row[field]: + raise ValueError(f"{split}.jsonl row {index} has invalid {field!r}.") + audio_path = _safe_child(self.root, row["audio"]) + missing = set(row["phonemes"]) - self.symbol_to_id.keys() + unknown.update(missing) + self.records.append((audio_path, row["phonemes"])) + if unknown: + rendered = " ".join(repr(item) for item in sorted(unknown)) + raise ValueError( + f"{split}.jsonl contains symbols absent from symbols.json: {rendered}" + ) + + def __len__(self) -> int: + return len(self.records) + + def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + audio_path, phonemes = self.records[index] + audio, sample_rate = sf.read(audio_path, dtype="float32", always_2d=True) + if sample_rate != self.audio_config.sampling_rate: + raise ValueError( + f"{audio_path} is {sample_rate} Hz; expected " + f"{self.audio_config.sampling_rate} Hz. Re-run dataset preparation." + ) + waveform = torch.from_numpy(audio.mean(axis=1).copy()) + if not torch.isfinite(waveform).all(): + raise ValueError(f"{audio_path} contains NaN or infinite samples.") + peak = float(waveform.abs().max()) + if peak > 1.05: + raise ValueError(f"{audio_path} exceeds normalized floating-point range ({peak:.3f}).") + token_ids = [self.symbol_to_id[symbol] for symbol in phonemes] + if self.audio_config.add_blank: + expanded = [0] * (len(token_ids) * 2 + 1) + expanded[1::2] = token_ids + token_ids = expanded + if not token_ids: + raise ValueError(f"{audio_path} has an empty phoneme sequence.") + tokens = torch.tensor(token_ids, dtype=torch.long) + return tokens, spectrogram(waveform, self.audio_config), waveform[None] + + +class TTSCollate: + def __call__( + self, batch: Iterable[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] + ) -> tuple[torch.Tensor, ...]: + rows = sorted(batch, key=lambda item: item[1].shape[-1], reverse=True) + batch_size = len(rows) + text_lengths = torch.tensor([row[0].numel() for row in rows], dtype=torch.long) + spec_lengths = torch.tensor([row[1].shape[-1] for row in rows], dtype=torch.long) + wav_lengths = torch.tensor([row[2].shape[-1] for row in rows], dtype=torch.long) + text = torch.zeros(batch_size, int(text_lengths.max()), dtype=torch.long) + spec = torch.zeros( + batch_size, rows[0][1].shape[0], int(spec_lengths.max()), dtype=torch.float32 + ) + wav = torch.zeros(batch_size, 1, int(wav_lengths.max()), dtype=torch.float32) + for index, (tokens, spectrum, waveform) in enumerate(rows): + text[index, : tokens.numel()] = tokens + spec[index, :, : spectrum.shape[-1]] = spectrum + wav[index, :, : waveform.shape[-1]] = waveform + return text, text_lengths, spec, spec_lengths, wav, wav_lengths + + +def create_dataloader( + dataset: Dataset, + *, + batch_size: int, + shuffle: bool, + num_workers: int, + seed: int, + pin_memory: bool, +) -> DataLoader: + generator = torch.Generator() + generator.manual_seed(seed) + kwargs = { + "dataset": dataset, + "batch_size": batch_size, + "shuffle": shuffle, + "num_workers": num_workers, + "pin_memory": pin_memory, + "collate_fn": TTSCollate(), + "drop_last": False, + "generator": generator, + } + if num_workers > 0: + kwargs.update(persistent_workers=True, prefetch_factor=2) + return DataLoader(**kwargs) diff --git a/finetune/presets/balanced.json b/finetune/presets/balanced.json new file mode 100644 index 0000000..dede24f --- /dev/null +++ b/finetune/presets/balanced.json @@ -0,0 +1,11 @@ +{ + "batch_size": 4, + "gradient_accumulation_steps": 2, + "learning_rate_g": 0.0001, + "learning_rate_d": 0.0001, + "max_steps": 20000, + "num_workers": 4, + "amp": true, + "checkpoint_interval": 1000, + "validation_interval": 500 +} diff --git a/finetune/presets/micro-12gb.json b/finetune/presets/micro-12gb.json new file mode 100644 index 0000000..23ec766 --- /dev/null +++ b/finetune/presets/micro-12gb.json @@ -0,0 +1,11 @@ +{ + "batch_size": 2, + "gradient_accumulation_steps": 4, + "learning_rate_g": 0.00008, + "learning_rate_d": 0.00008, + "max_steps": 20000, + "num_workers": 3, + "amp": true, + "checkpoint_interval": 1000, + "validation_interval": 500 +} diff --git a/finetune/presets/nano-8gb.json b/finetune/presets/nano-8gb.json new file mode 100644 index 0000000..bf31fb3 --- /dev/null +++ b/finetune/presets/nano-8gb.json @@ -0,0 +1,11 @@ +{ + "batch_size": 1, + "gradient_accumulation_steps": 8, + "learning_rate_g": 0.00008, + "learning_rate_d": 0.00008, + "max_steps": 20000, + "num_workers": 2, + "amp": true, + "checkpoint_interval": 1000, + "validation_interval": 500 +} diff --git a/finetune/pyproject.toml b/finetune/pyproject.toml new file mode 100644 index 0000000..f63a320 --- /dev/null +++ b/finetune/pyproject.toml @@ -0,0 +1,45 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "inflect-finetune" +version = "0.1.0" +description = "Generic language and fixed-voice adaptation toolkit for Inflect v2" +readme = "README.md" +requires-python = ">=3.10" +license = { text = "Apache-2.0" } +dependencies = [ + "numpy>=1.26,<3", + "scipy>=1.11,<2", + "soundfile>=0.12,<1", + "torch>=2.2", + "phonemizer>=3.2,<4", + "espeakng-loader>=0.2.4,<1", + "huggingface-hub>=0.27,<2", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8,<9", + "ruff>=0.8,<1", +] +onnx = [ + "onnx>=1.16,<2", + "onnxruntime>=1.18,<2", +] + +[project.scripts] +inflect-adapt = "inflect_finetune.cli:main" + +[tool.setuptools.packages.find] +where = ["."] +include = ["inflect_finetune*"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +addopts = "-q" + +[tool.ruff] +line-length = 100 +target-version = "py310" diff --git a/finetune/tests/test_checkpoint_identity.py b/finetune/tests/test_checkpoint_identity.py new file mode 100644 index 0000000..62e0f9b --- /dev/null +++ b/finetune/tests/test_checkpoint_identity.py @@ -0,0 +1,269 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import torch +from torch import nn + +from inflect_finetune.checkpoint import ( + CompatibilityReport, + build_run_identity, + resume_training_checkpoint, + save_training_checkpoint, + sha256_file, +) +from inflect_finetune.training import ( + TrainingOptions, + _establish_run_identity, + _public_options, + _validate_new_output_dir, +) + + +def _inputs(tmp_path: Path) -> tuple[Path, Path]: + base = tmp_path / "Inflect-Micro-v2" + base.mkdir(parents=True) + (base / "model.pth").write_bytes(b"public checkpoint") + (base / "config.json").write_text('{"model":"micro"}\n', encoding="utf-8") + prepared = tmp_path / "prepared" + prepared.mkdir() + (prepared / "dataset.json").write_text('{"language":"en-us"}\n', encoding="utf-8") + (prepared / "train.jsonl").write_text('{"audio":"a.wav"}\n', encoding="utf-8") + (prepared / "validation.jsonl").write_text( + '{"audio":"b.wav"}\n', encoding="utf-8" + ) + (prepared / "symbols.json").write_text('["a","b"]\n', encoding="utf-8") + return base, prepared + + +def _identity(tmp_path: Path, *, run_id: str = "run-a") -> dict: + base, prepared = _inputs(tmp_path) + return build_run_identity( + run_id=run_id, + base_root=base, + prepared_dir=prepared, + options={"max_steps": 10, "seed": 7}, + optimizer_schema={"generator": {"class": "torch.optim.AdamW"}}, + ) + + +def _training_parts(): + generator = nn.Linear(2, 2) + discriminator = nn.Linear(2, 1) + optimizer_g = torch.optim.AdamW(generator.parameters(), lr=1.0e-4) + optimizer_d = torch.optim.AdamW(discriminator.parameters(), lr=1.0e-4) + scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optimizer_g, gamma=0.99) + scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optimizer_d, gamma=0.99) + scaler = torch.amp.GradScaler("cuda", enabled=False) + return ( + generator, + discriminator, + optimizer_g, + optimizer_d, + scheduler_g, + scheduler_d, + scaler, + ) + + +def _report() -> CompatibilityReport: + return CompatibilityReport( + source_path="model.pth", + source_format="inflect_vits_inference_checkpoint_v1", + source_tensor_count=1, + source_parameter_count=1, + copied_tensor_count=1, + copied_parameter_count=1, + exact_tensor_count=1, + migrated_embedding_rows=0, + initialized_embedding_rows=0, + fresh_tensor_count=0, + fresh_parameter_count=0, + fresh_prefixes=("enc_q.",), + verified_equal_after_copy=True, + ) + + +def test_run_identity_hashes_every_public_input_and_is_path_independent( + tmp_path: Path, +) -> None: + identity = _identity(tmp_path) + + assert identity["toolkit_version"] == "0.1.0" + assert identity["base"]["identity"] == "Inflect-Micro-v2" + assert len(identity["base"]["checkpoint_sha256"]) == 64 + assert len(identity["base"]["config_sha256"]) == 64 + assert len(identity["prepared_dataset"]["dataset_json_sha256"]) == 64 + assert len(identity["prepared_dataset"]["train_jsonl_sha256"]) == 64 + assert len(identity["prepared_dataset"]["validation_jsonl_sha256"]) == 64 + assert len(identity["symbols_sha256"]) == 64 + rendered = json.dumps(identity) + assert str(tmp_path) not in rendered + + +def test_identity_mismatch_is_rejected_before_mutable_state_load(tmp_path: Path) -> None: + source_parts = _training_parts() + checkpoint = tmp_path / "checkpoint.pth" + latest = tmp_path / "latest.pth" + identity = _identity(tmp_path / "identity") + save_training_checkpoint( + checkpoint, + generator=source_parts[0], + discriminator=source_parts[1], + optimizer_g=source_parts[2], + optimizer_d=source_parts[3], + scheduler_g=source_parts[4], + scheduler_d=source_parts[5], + scaler=source_parts[6], + step=5, + epoch=2, + stage="posterior_warmup", + options={"max_steps": 10}, + symbols=("a", "b"), + compatibility=_report(), + run_identity=identity, + latest_path=latest, + ) + target_parts = _training_parts() + with torch.no_grad(): + target_parts[0].weight.fill_(42.0) + before = target_parts[0].weight.detach().clone() + wrong_identity = dict(identity) + wrong_identity["run_id"] = "another-run" + + with pytest.raises(ValueError, match="different adaptation run"): + resume_training_checkpoint( + checkpoint, + generator=target_parts[0], + discriminator=target_parts[1], + optimizer_g=target_parts[2], + optimizer_d=target_parts[3], + scheduler_g=target_parts[4], + scheduler_d=target_parts[5], + scaler=target_parts[6], + expected_symbols=("a", "b"), + expected_run_identity=wrong_identity, + ) + + assert torch.equal(target_parts[0].weight, before) + assert latest.read_bytes() == checkpoint.read_bytes() + assert sha256_file(latest) == sha256_file(checkpoint) + + +def test_valid_identity_resumes_stage_and_state(tmp_path: Path) -> None: + source_parts = _training_parts() + identity = _identity(tmp_path / "identity") + checkpoint = tmp_path / "checkpoint.pth" + save_training_checkpoint( + checkpoint, + generator=source_parts[0], + discriminator=source_parts[1], + optimizer_g=source_parts[2], + optimizer_d=source_parts[3], + scheduler_g=source_parts[4], + scheduler_d=source_parts[5], + scaler=source_parts[6], + step=5, + epoch=2, + stage="linguistic_adaptation", + options={"max_steps": 10}, + symbols=("a", "b"), + compatibility=_report(), + run_identity=identity, + ) + target_parts = _training_parts() + + state = resume_training_checkpoint( + checkpoint, + generator=target_parts[0], + discriminator=target_parts[1], + optimizer_g=target_parts[2], + optimizer_d=target_parts[3], + scheduler_g=target_parts[4], + scheduler_d=target_parts[5], + scaler=target_parts[6], + expected_symbols=("a", "b"), + expected_run_identity=identity, + ) + + assert state == (5, 2, "linguistic_adaptation") + assert torch.equal(target_parts[0].weight, source_parts[0].weight) + + +def test_new_run_output_rejects_unrelated_content(tmp_path: Path) -> None: + output = tmp_path / "run" + (output / "checkpoints").mkdir(parents=True) + (output / "checkpoints" / ".gitkeep").touch() + _validate_new_output_dir(output) + (output / "old-result.pth").write_bytes(b"do not overwrite") + + with pytest.raises(ValueError, match="nonempty output directory"): + _validate_new_output_dir(output) + + +def test_resume_is_bound_to_marker_and_its_checkpoints_directory(tmp_path: Path) -> None: + base, prepared = _inputs(tmp_path) + output = tmp_path / "run" + options = TrainingOptions( + base_model=base, + prepared_dir=prepared, + output_dir=output, + max_steps=10, + ) + identity, resume = _establish_run_identity( + options=options, + output_dir=output, + base_root=base, + prepared_dir=prepared, + ) + assert resume is None + inside = output / "checkpoints" / "latest.pth" + inside.parent.mkdir() + inside.write_bytes(b"checkpoint") + + resumed, checkpoint = _establish_run_identity( + options=TrainingOptions( + base_model=base, + prepared_dir=prepared, + output_dir=output, + resume=inside, + max_steps=10, + ), + output_dir=output, + base_root=base, + prepared_dir=prepared, + ) + + assert resumed == identity + assert checkpoint == inside.resolve() + outside = tmp_path / "foreign.pth" + outside.write_bytes(b"checkpoint") + with pytest.raises(ValueError, match="inside this run"): + _establish_run_identity( + options=TrainingOptions( + base_model=base, + prepared_dir=prepared, + output_dir=output, + resume=outside, + max_steps=10, + ), + output_dir=output, + base_root=base, + prepared_dir=prepared, + ) + + +def test_public_options_exclude_machine_paths() -> None: + options = TrainingOptions( + base_model=Path("private/base"), + prepared_dir=Path("private/prepared"), + output_dir=Path("private/output"), + resume=Path("private/checkpoint.pth"), + ) + payload = _public_options(options) + + assert not {"base_model", "prepared_dir", "output_dir", "preset", "resume"} & payload.keys() + assert payload["posterior_warmup_steps"] == 500 + assert payload["decoder_lr_multiplier"] == 0.1 diff --git a/finetune/tests/test_cli.py b/finetune/tests/test_cli.py new file mode 100644 index 0000000..2785005 --- /dev/null +++ b/finetune/tests/test_cli.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from inflect_finetune.cli import _optional_step, _run_train, build_parser + + +def test_cli_exposes_complete_workflow() -> None: + parser = build_parser() + subparsers = next( + action for action in parser._actions if action.dest == "command" # noqa: SLF001 + ) + assert set(subparsers.choices) == { + "prepare", + "audit", + "train", + "evaluate", + "export", + } + + +def test_prepare_cli_maps_documented_arguments() -> None: + args = build_parser().parse_args( + [ + "prepare", + "--manifest", + "metadata.jsonl", + "--audio-root", + "audio", + "--language", + "fr-fr", + "--output", + "prepared/fr", + ] + ) + assert args.command == "prepare" + assert args.manifest == Path("metadata.jsonl") + assert args.audio_root == Path("audio") + assert args.language == "fr-fr" + assert args.output == Path("prepared/fr") + + +def test_prepare_cli_exposes_custom_frontend_hook() -> None: + args = build_parser().parse_args( + [ + "prepare", + "--manifest", + "metadata.jsonl", + "--frontend", + "custom", + "--frontend-hook", + "frontend.py:create_frontend", + "--output", + "prepared/custom", + ] + ) + assert args.frontend == "custom" + assert args.frontend_hook == "frontend.py:create_frontend" + + +@pytest.mark.parametrize(("value", "expected"), [("none", None), ("off", None), ("0", 0)]) +def test_optional_decoder_step(value: str, expected: int | None) -> None: + assert _optional_step(value) == expected + + +def test_optional_decoder_step_rejects_negative_value() -> None: + with pytest.raises(Exception, match="non-negative"): + _optional_step("-1") + + +def test_train_cli_explicit_value_overrides_preset( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + captured: dict[str, object] = {} + + def fake_train(options: object) -> dict[str, object]: + captured["options"] = options + return {"ok": True} + + monkeypatch.setattr("inflect_finetune.training.train_adaptation", fake_train) + args = build_parser().parse_args( + [ + "train", + "--base", + "nano", + "--dataset", + str(tmp_path / "prepared"), + "--output", + str(tmp_path / "run"), + "--preset", + "balanced", + "--batch-size", + "2", + "--no-amp", + ] + ) + + assert _run_train(args) == {"ok": True} + options = captured["options"] + assert getattr(options, "batch_size") == 2 + assert getattr(options, "gradient_accumulation_steps") == 2 + assert getattr(options, "amp") is False + + +def test_train_cli_can_disable_decoder_unfreeze( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + captured: dict[str, object] = {} + + def fake_train(options: object) -> dict[str, object]: + captured["options"] = options + return {"ok": True} + + monkeypatch.setattr("inflect_finetune.training.train_adaptation", fake_train) + args = build_parser().parse_args( + [ + "train", + "--base", + "micro", + "--dataset", + str(tmp_path / "prepared"), + "--output", + str(tmp_path / "run"), + "--decoder-unfreeze-step", + "none", + ] + ) + + _run_train(args) + assert getattr(captured["options"], "decoder_unfreeze_step") is None + + +def test_export_cli_exposes_language_frontend_inputs() -> None: + args = build_parser().parse_args( + [ + "export", + "--checkpoint", + "run/checkpoints/adaptation-final.pth", + "--prepared-dataset", + "prepared/es", + "--frontend-hook", + "frontend.py", + "--output", + "exports/es", + ] + ) + assert args.prepared_dataset == Path("prepared/es") + assert args.frontend_hook == Path("frontend.py") diff --git a/finetune/tests/test_data_frontend_prepare.py b/finetune/tests/test_data_frontend_prepare.py new file mode 100644 index 0000000..5c23cd8 --- /dev/null +++ b/finetune/tests/test_data_frontend_prepare.py @@ -0,0 +1,313 @@ +from __future__ import annotations + +import hashlib +import importlib +import json +import sys +from pathlib import Path + +import numpy as np +import pytest +import soundfile as sf + +from inflect_finetune.audit import AuditOptions, audit_dataset +from inflect_finetune.prepare import PreparationError, PrepareOptions, prepare_dataset + + +def _write_wav(path: Path, frequency: float) -> None: + sample_rate = 24_000 + time = np.arange(sample_rate // 10, dtype=np.float32) / sample_rate + waveform = 0.1 * np.sin(2.0 * np.pi * frequency * time) + sf.write(path, waveform, sample_rate, subtype="PCM_16") + + +def _write_manifest(root: Path, rows: list[dict[str, str]]) -> Path: + path = root / "metadata.jsonl" + path.write_text( + "".join(json.dumps(row) + "\n" for row in rows), + encoding="utf-8", + ) + return path + + +def _source_dataset( + root: Path, + *, + count: int, + texts: list[str] | None = None, + groups: list[tuple[str, str] | None] | None = None, + speakers: list[str | None] | None = None, +) -> Path: + root.mkdir(parents=True) + rows: list[dict[str, str]] = [] + for index in range(count): + audio_name = f"audio-{index}.wav" + _write_wav(root / audio_name, 180.0 + index * 37.0) + row = { + "audio": audio_name, + "text": texts[index] if texts else f"Sentence {index}.", + "phonemes": f"test {index}", + } + if groups and groups[index]: + field, value = groups[index] + row[field] = value + if speakers and speakers[index]: + row["speaker"] = speakers[index] + rows.append(row) + return _write_manifest(root, rows) + + +def _prepare(manifest: Path, output: Path, **overrides: object) -> dict[str, object]: + values: dict[str, object] = { + "manifest_path": manifest, + "output_dir": output, + "frontend": "prephonemized", + "validation_fraction": 0.34, + "split_seed": 99, + } + values.update(overrides) + return prepare_dataset(PrepareOptions(**values)) + + +def _read_split(path: Path) -> list[dict[str, object]]: + return [ + json.loads(line) + for line in path.read_text(encoding="utf-8").splitlines() + if line + ] + + +def test_prepare_rejects_dataset_too_small_for_validation(tmp_path: Path) -> None: + manifest = _source_dataset(tmp_path / "source", count=1) + with pytest.raises(PreparationError, match="at least two usable rows"): + _prepare(manifest, tmp_path / "prepared") + + +def test_group_aware_split_is_deterministic_and_keeps_groups_together( + tmp_path: Path, +) -> None: + manifest = _source_dataset( + tmp_path / "source", + count=6, + groups=[ + None, + None, + ("session", "session-b"), + ("session", "session-b"), + ("group_id", "group-c"), + ("group_id", "group-d"), + ], + speakers=["alice"] * 6, + ) + first = tmp_path / "first" + second = tmp_path / "second" + first_metadata = _prepare(manifest, first) + second_metadata = _prepare(manifest, second) + + assert first_metadata["split"] == second_metadata["split"] + first_splits = { + split: _read_split(first / f"{split}.jsonl") + for split in ("train", "validation") + } + second_splits = { + split: _read_split(second / f"{split}.jsonl") + for split in ("train", "validation") + } + assert first_splits == second_splits + + placements: dict[tuple[str, str], set[str]] = {} + for split, rows in first_splits.items(): + for row in rows: + if "group_id" in row: + key = (str(row["group_field"]), str(row["group_id"])) + placements.setdefault(key, set()).add(split) + assert all(len(splits) == 1 for splits in placements.values()) + assert first_metadata["split"]["group_fields"] == [ + "group_id", + "session", + ] + assert first_metadata["speaker"] == "alice" + + +def test_same_speaker_multiple_sessions_produces_nonempty_splits( + tmp_path: Path, +) -> None: + manifest = _source_dataset( + tmp_path / "source", + count=6, + groups=[ + ("session", "session-a"), + ("session", "session-a"), + ("session", "session-b"), + ("session", "session-b"), + ("session", "session-c"), + ("session", "session-c"), + ], + speakers=["one-speaker"] * 6, + ) + output = tmp_path / "prepared" + metadata = _prepare(manifest, output) + assert metadata["speaker"] == "one-speaker" + assert metadata["row_counts"]["train"] > 0 + assert metadata["row_counts"]["validation"] > 0 + for split in ("train", "validation"): + assert {row["speaker"] for row in _read_split(output / f"{split}.jsonl")} == { + "one-speaker" + } + + +def test_prepare_rejects_multiple_speakers(tmp_path: Path) -> None: + manifest = _source_dataset( + tmp_path / "source", + count=3, + speakers=["alice", "bob", "alice"], + ) + with pytest.raises(PreparationError, match="one consistent nonempty speaker"): + _prepare(manifest, tmp_path / "prepared") + + +def test_normalized_duplicate_text_is_co_located(tmp_path: Path) -> None: + manifest = _source_dataset( + tmp_path / "source", + count=4, + texts=["Hello World", "hello world", "Independent A", "Independent B"], + ) + output = tmp_path / "prepared" + _prepare(manifest, output, validation_fraction=0.5) + placements: dict[str, set[str]] = {} + for split in ("train", "validation"): + for row in _read_split(output / f"{split}.jsonl"): + key = " ".join(str(row["normalized_text"]).casefold().split()) + placements.setdefault(key, set()).add(split) + assert placements["hello world"] in ({"train"}, {"validation"}) + + +def test_prepare_rejects_duplicate_audio_content(tmp_path: Path) -> None: + source = tmp_path / "source" + source.mkdir() + _write_wav(source / "first.wav", 220.0) + (source / "second.wav").write_bytes((source / "first.wav").read_bytes()) + _write_wav(source / "third.wav", 330.0) + manifest = _write_manifest( + source, + [ + {"audio": "first.wav", "text": "One", "phonemes": "one"}, + {"audio": "second.wav", "text": "Two", "phonemes": "two"}, + {"audio": "third.wav", "text": "Three", "phonemes": "three"}, + ], + ) + with pytest.raises(PreparationError, match="Duplicate audio content"): + _prepare(manifest, tmp_path / "prepared") + + +def _custom_hook_source() -> str: + return """ +class Frontend: + def __init__(self, language): + self.language = language + + def normalize(self, text): + return " ".join(text.lower().split()) + + def phonemize(self, normalized_text): + return "ab" + + def symbols(self): + return ["a", "b"] + + def metadata(self): + return { + "name": "test-frontend", + "version": "1", + "language": self.language, + "configuration": {"case": "lower"}, + } + +def create_frontend(*, language): + return Frontend(language) +""".lstrip() + + +def test_custom_file_hook_is_hashed_and_recorded(tmp_path: Path) -> None: + manifest = _source_dataset(tmp_path / "source", count=3) + hook_path = tmp_path / "custom_frontend.py" + hook_path.write_text(_custom_hook_source(), encoding="utf-8") + output = tmp_path / "prepared" + metadata = _prepare( + manifest, + output, + frontend="custom", + frontend_hook=f"{hook_path}:create_frontend", + ) + + hook = metadata["frontend"]["hook"] + assert hook["identity"] == "file:custom_frontend.py:create_frontend" + assert hook["source_kind"] == "file" + assert hook["source_sha256"] == hashlib.sha256(hook_path.read_bytes()).hexdigest() + assert len(hook["metadata_sha256"]) == 64 + assert hook["declared_metadata"]["language"] == "en-us" + assert {row["phonemes"] for row in _read_split(output / "train.jsonl")} == {"ab"} + + +def test_custom_module_hook_is_supported(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + manifest = _source_dataset(tmp_path / "source", count=3) + module_name = "temporary_inflect_frontend" + module_path = tmp_path / f"{module_name}.py" + module_path.write_text(_custom_hook_source(), encoding="utf-8") + monkeypatch.syspath_prepend(str(tmp_path)) + importlib.invalidate_caches() + sys.modules.pop(module_name, None) + + metadata = _prepare( + manifest, + tmp_path / "prepared", + frontend="custom", + frontend_hook=f"{module_name}:create_frontend", + ) + hook = metadata["frontend"]["hook"] + assert hook["identity"] == f"{module_name}:create_frontend" + assert hook["source_kind"] == "module" + assert hook["source_sha256"] == hashlib.sha256(module_path.read_bytes()).hexdigest() + + +def test_audit_detects_group_and_normalized_text_crossing_splits( + tmp_path: Path, +) -> None: + manifest = _source_dataset(tmp_path / "source", count=4) + output = tmp_path / "prepared" + _prepare(manifest, output, validation_fraction=0.5) + train = _read_split(output / "train.jsonl") + validation = _read_split(output / "validation.jsonl") + train[0]["group_id"] = "leaked" + train[0]["group_field"] = "session" + validation[0]["group_id"] = "leaked" + validation[0]["group_field"] = "session" + validation[0]["normalized_text"] = train[0]["normalized_text"] + (output / "train.jsonl").write_text( + "".join(json.dumps(row) + "\n" for row in train), encoding="utf-8" + ) + (output / "validation.jsonl").write_text( + "".join(json.dumps(row) + "\n" for row in validation), encoding="utf-8" + ) + + report = audit_dataset(AuditOptions(prepared_dir=output, strict=False)) + assert not report["valid"] + assert any("normalized transcript crosses" in error for error in report["errors"]) + assert any("session='leaked' crosses" in error for error in report["errors"]) + + +def test_audit_treats_empty_validation_as_error(tmp_path: Path) -> None: + manifest = _source_dataset(tmp_path / "source", count=3) + output = tmp_path / "prepared" + _prepare(manifest, output) + validation = output / "validation.jsonl" + validation.write_text("", encoding="utf-8") + dataset_path = output / "dataset.json" + dataset = json.loads(dataset_path.read_text(encoding="utf-8")) + dataset["row_counts"]["validation"] = 0 + dataset["row_counts"]["total"] = dataset["row_counts"]["train"] + dataset_path.write_text(json.dumps(dataset), encoding="utf-8") + + report = audit_dataset(AuditOptions(prepared_dir=output, strict=False)) + assert not report["valid"] + assert any("no validation rows" in error for error in report["errors"]) diff --git a/finetune/tests/test_model_resolution.py b/finetune/tests/test_model_resolution.py new file mode 100644 index 0000000..9d013de --- /dev/null +++ b/finetune/tests/test_model_resolution.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from pathlib import Path + +import huggingface_hub + +from inflect_finetune.modeling import resolve_base_model + + +def test_hugging_face_repo_id_and_revision_are_resolved( + tmp_path: Path, monkeypatch +) -> None: + downloaded = tmp_path / "snapshot" + (downloaded / "runtime").mkdir(parents=True) + (downloaded / "config.json").write_text("{}\n", encoding="utf-8") + (downloaded / "model.pth").write_bytes(b"checkpoint") + observed: dict[str, object] = {} + + def fake_snapshot_download(**kwargs): + observed.update(kwargs) + return str(downloaded) + + monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download) + + resolved = resolve_base_model("example/Inflect-Adapted@release-1") + + assert resolved == downloaded.resolve() + assert observed["repo_id"] == "example/Inflect-Adapted" + assert observed["revision"] == "release-1" + assert "model.pth" in observed["allow_patterns"] diff --git a/finetune/tests/test_monotonic_align.py b/finetune/tests/test_monotonic_align.py new file mode 100644 index 0000000..9143149 --- /dev/null +++ b/finetune/tests/test_monotonic_align.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +import pytest +import torch + +from inflect_finetune.monotonic_align import maximum_path + + +def test_maximum_path_is_monotonic_and_respects_padding() -> None: + scores = torch.randn(2, 8, 4) + mask = torch.zeros_like(scores) + mask[0, :8, :4] = 1 + mask[1, :6, :3] = 1 + + path = maximum_path(scores, mask) + + for batch_index, (audio_length, text_length) in enumerate(((8, 4), (6, 3))): + active = path[batch_index, :audio_length, :text_length] + indices = active.argmax(dim=1) + assert torch.all(active.sum(dim=1) == 1) + assert indices[0].item() == 0 + assert indices[-1].item() == text_length - 1 + assert torch.all(torch.isin(indices[1:] - indices[:-1], torch.tensor([0, 1]))) + assert torch.count_nonzero(path[batch_index] * (1 - mask[batch_index])) == 0 + + +def test_maximum_path_selects_the_highest_scoring_valid_route() -> None: + scores = torch.full((1, 5, 3), -10.0) + scores[0, 0, 0] = 4 + scores[0, 1, 0] = 4 + scores[0, 2, 1] = 4 + scores[0, 3, 1] = 4 + scores[0, 4, 2] = 4 + + path = maximum_path(scores, torch.ones_like(scores)) + + assert path[0].argmax(dim=1).tolist() == [0, 0, 1, 1, 2] + + +def test_maximum_path_rejects_more_tokens_than_frames() -> None: + with pytest.raises(ValueError, match="at least one audio frame"): + maximum_path(torch.zeros(1, 2, 3), torch.ones(1, 2, 3)) diff --git a/finetune/tests/test_packaged_presets.py b/finetune/tests/test_packaged_presets.py new file mode 100644 index 0000000..803806b --- /dev/null +++ b/finetune/tests/test_packaged_presets.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +import subprocess +import sys +import zipfile +from pathlib import Path + +import pytest + +from inflect_finetune.training import TrainingOptions, load_preset + + +TOOLKIT_ROOT = Path(__file__).resolve().parents[1] +PRESET_NAMES = ("balanced", "micro-12gb", "nano-8gb") + + +@pytest.mark.parametrize("name", PRESET_NAMES) +def test_named_preset_loads_from_the_python_package(name: str) -> None: + payload = load_preset(name) + + assert payload["max_steps"] == 20_000 + assert payload["checkpoint_interval"] == 1_000 + assert TrainingOptions.from_preset( + name, + base_model="base", + prepared_dir="prepared", + output_dir="output", + ).preset is None + + +def test_built_wheel_contains_and_can_import_packaged_presets(tmp_path: Path) -> None: + wheel_dir = tmp_path / "wheel" + wheel_dir.mkdir() + subprocess.run( + [ + sys.executable, + "-m", + "pip", + "wheel", + "--no-deps", + "--no-build-isolation", + "--wheel-dir", + str(wheel_dir), + str(TOOLKIT_ROOT), + ], + check=True, + capture_output=True, + text=True, + ) + wheels = list(wheel_dir.glob("*.whl")) + assert len(wheels) == 1 + with zipfile.ZipFile(wheels[0]) as archive: + assert "inflect_finetune/presets/__init__.py" in archive.namelist() + + script = ( + "import sys;" + f"sys.path.insert(0, {str(wheels[0])!r});" + "from inflect_finetune.presets import available_presets,load_packaged_preset;" + "assert available_presets()==('balanced','micro-12gb','nano-8gb');" + "assert load_packaged_preset('nano-8gb')['batch_size']==1" + ) + subprocess.run( + [sys.executable, "-c", script], + cwd=tmp_path, + check=True, + capture_output=True, + text=True, + ) diff --git a/finetune/tests/test_public_safety.py b/finetune/tests/test_public_safety.py new file mode 100644 index 0000000..8520710 --- /dev/null +++ b/finetune/tests/test_public_safety.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +import re +from pathlib import Path + +from inflect_finetune.exporting import _copy_public_runtime + + +TOOLKIT_ROOT = Path(__file__).resolve().parents[1] + +PROHIBITED_PATTERNS = { + "Windows user path": re.compile(r"[A-Za-z]:\\Users\\", re.IGNORECASE), + "private storage path": re.compile(r"Inflect-Storage", re.IGNORECASE), + "rental instance detail": re.compile(r"\b(vast\.ai|instance[_ -]?id|ssh_host)\b", re.IGNORECASE), + "credential assignment": re.compile( + r"\b(HF_TOKEN|HUGGING_FACE_HUB_TOKEN|WANDB_API_KEY)\s*=\s*[\"'][^\"']+", + re.IGNORECASE, + ), +} + + +def public_text_files() -> list[Path]: + extensions = { + ".py", + ".md", + ".toml", + ".json", + ".jsonl", + ".csv", + ".yaml", + ".yml", + ".txt", + } + return [ + path + for path in TOOLKIT_ROOT.rglob("*") + if path.is_file() + and path.suffix.lower() in extensions + and "__pycache__" not in path.parts + and path.name != Path(__file__).name + ] + + +def test_public_toolkit_contains_no_private_infrastructure_references() -> None: + findings: list[str] = [] + for path in public_text_files(): + text = path.read_text(encoding="utf-8") + for label, pattern in PROHIBITED_PATTERNS.items(): + if pattern.search(text): + findings.append(f"{path.relative_to(TOOLKIT_ROOT)}: {label}") + assert not findings, "\n".join(findings) + + +def test_contract_explicitly_separates_public_adaptation_from_base_recipe() -> None: + contract = (TOOLKIT_ROOT / "CONTRACT.md").read_text(encoding="utf-8") + assert "not" in contract.lower() + assert "private recipe" in contract.lower() + assert "dataset speaker becomes the checkpoint voice" in contract.lower() + + +def test_package_metadata_does_not_claim_unvalidated_universal_support() -> None: + metadata = (TOOLKIT_ROOT / "pyproject.toml").read_text(encoding="utf-8").lower() + assert "generic language and fixed-voice adaptation" in metadata + assert "universal" not in metadata + + +def test_public_runtime_export_omits_python_bytecode(tmp_path: Path) -> None: + template = tmp_path / "template" + runtime = template / "runtime" + cache = runtime / "__pycache__" + cache.mkdir(parents=True) + (runtime / "models.py").write_text("MODEL = True\n", encoding="utf-8") + (cache / "models.cpython-312.pyc").write_bytes(b"bytecode") + + destination = tmp_path / "export" + destination.mkdir() + copied = _copy_public_runtime(template, destination) + + assert (destination / "runtime" / "models.py").is_file() + assert not (destination / "runtime" / "__pycache__").exists() + assert all(path.suffix != ".pyc" for path in copied) diff --git a/finetune/tests/test_training_stages.py b/finetune/tests/test_training_stages.py new file mode 100644 index 0000000..7d0229d --- /dev/null +++ b/finetune/tests/test_training_stages.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +from inflect_finetune.training import ( + STAGE_ADAPT, + STAGE_DECODER, + STAGE_POSTERIOR, + TrainingOptions, + _stage_for_step, +) + + +def _options() -> TrainingOptions: + return TrainingOptions( + base_model="nano", + prepared_dir="prepared", + output_dir="run", + posterior_warmup_steps=10, + decoder_unfreeze_step=20, + ) + + +def test_stage_boundaries_select_the_stage_for_the_next_step() -> None: + options = _options() + + assert _stage_for_step(options, 0) == STAGE_POSTERIOR + assert _stage_for_step(options, 9) == STAGE_POSTERIOR + assert _stage_for_step(options, 10) == STAGE_ADAPT + assert _stage_for_step(options, 19) == STAGE_ADAPT + assert _stage_for_step(options, 20) == STAGE_DECODER