From 23148b5751c7d708345c509bc9bf5be596ff21fe Mon Sep 17 00:00:00 2001 From: Tinnci <23432137+Tinnci@users.noreply.github.com> Date: Sun, 12 Jul 2026 16:51:59 +0800 Subject: [PATCH] Add Colab MossFormer2 feasibility spike --- CHANGELOG.md | 2 + ROADMAP.md | 8 +- docs/ACCELERATORS.md | 17 ++ docs/BACKEND_CANDIDATES.md | 60 ++++-- examples/colab_mossformer2_spike.py | 300 ++++++++++++++++++++++++++++ 5 files changed, 364 insertions(+), 23 deletions(-) create mode 100644 examples/colab_mossformer2_spike.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 79054c5..d3587a5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,8 @@ This project follows semantic versioning once published to PyPI. - Added a pinned external Whisper evaluator specification and Colab-only transcript workflow for precomputed downstream WER/CER evidence without adding an ASR package dependency. - Added a Colab-only verified MUSHRA bundle export for the licensed LibriSpeech sinc/LavaSR matrix. +- Added a bounded Colab-only MossFormer2 feasibility workflow covering minimal checkpoints, safe + conversion, offline CPU inference, channel behavior, alignment, and parity-fixture constraints. ## 0.6.0 - 2026-07-12 diff --git a/ROADMAP.md b/ROADMAP.md index 9b088d4..880e6b4 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -79,7 +79,7 @@ human ratings is a follow-up study, not a prerequisite for completing the export Do not collapse quality, downstream usefulness, speed, stability, and governance into one ranking. The evaluation policy lives in [docs/EVALUATION.md](docs/EVALUATION.md). -## Priority 3: One Candidate Spike, Not Three Implementations +## Completed: One Candidate Spike, Not Three Implementations ClearerVoice `MossFormer2_SR_48K` is the first candidate to revisit because it has a clear 48 kHz speech SR contract and identifiable checkpoint files. The spike must stop before backend @@ -96,6 +96,12 @@ Resemble Enhance remains a possible external speech-enhancement backend, not an package-owned 48 kHz SR path. Detailed evidence lives in [docs/BACKEND_CANDIDATES.md](docs/BACKEND_CANDIDATES.md). +The Colab CPU spike passed with only the pointer plus `m` and `g` checkpoints. Both state dicts can +be loaded with `weights_only=True` and converted exactly to safetensors. Offline 16 kHz mono/stereo +and native 48 kHz inputs ran successfully; stereo channels were preserved. The sub-second fixture +exposed a consistent 224-sample output shortening and cross-process floating-point drift, so any +future backend needs explicit alignment plus tolerance-based parity rather than hash equality. + ## Priority 4: Optimize Only After Measurement Keep `lavasr-compat` on `torch-eager` until another provider has real-weight benchmarks and parity diff --git a/docs/ACCELERATORS.md b/docs/ACCELERATORS.md index 640d604..64775d6 100644 --- a/docs/ACCELERATORS.md +++ b/docs/ACCELERATORS.md @@ -269,6 +269,23 @@ The bundle uses MUSHRA metadata with seed `0` and four stimuli per trial: refere anchor, sinc output, and LavaSR output. Public stimuli contain only blind IDs and generic paths; backend, role, and source mapping remain in the separate answer key. +Run the bounded MossFormer2 feasibility spike on a fresh CPU session. Checkpoint download, +inspection, conversion, and real inference all remain remote: + +```sh +colab new -s asr-mossformer2 +colab exec -s asr-mossformer2 -f examples/colab_mossformer2_spike.py --timeout 3600 +colab download -s asr-mossformer2 \ + /content/asr-mossformer2-spike.tar.gz \ + runs/asr-mossformer2-spike.tar.gz +colab stop -s asr-mossformer2 +``` + +The script pins both upstream revisions, downloads only the pointer and two inference checkpoints, +forces offline mode before model construction, uses an isolated Python worker after dependency +installation, and records safetensors round-trip plus mono/stereo/alignment evidence. Do not run it +locally. + For upstream LavaSR parity, install upstream dependencies deliberately and use the environment gates in [../tests/README.md](../tests/README.md). Validate AudioSR separately because its external package owns checkpoint and accelerator behavior. diff --git a/docs/BACKEND_CANDIDATES.md b/docs/BACKEND_CANDIDATES.md index c1c6721..d9f6b18 100644 --- a/docs/BACKEND_CANDIDATES.md +++ b/docs/BACKEND_CANDIDATES.md @@ -5,14 +5,15 @@ This document owns implementation-grade evidence for model candidates. Admission ## Current Decision -Do not start another package-owned backend until v0.6.0 is released and broader LavaSR evidence is -recorded. When capacity returns, run one bounded ClearerVoice `MossFormer2_SR_48K` feasibility spike. -Do not implement MossFormer2, FlowHigh, and Resemble Enhance in parallel. +The bounded ClearerVoice `MossFormer2_SR_48K` feasibility spike is complete. It clears the checkpoint, +license, offline-loading, CPU-fallback, and basic I/O gates, but backend implementation remains a +separate decision because the model is heavyweight and its output has a fixed alignment offset. Do +not implement MossFormer2, FlowHigh, and Resemble Enhance in parallel. | Candidate | Domain | Decision | | --- | --- | --- | | `lavasr-v2-bwe` | Speech BWE | Implemented experimental baseline; broaden evidence before stable promotion. | -| ClearerVoice `MossFormer2_SR_48K` | Speech SR to 48 kHz | First future spike; defer implementation pending safe checkpoint and parity evidence. | +| ClearerVoice `MossFormer2_SR_48K` | Speech SR to 48 kHz | Feasibility passed; eligible for one bounded backend plan with tolerance-based parity. | | Resemble Enhance | Speech enhancement/BWE to 44.1 kHz | Deferred; consider only as an optional external backend. | | AudioSR `basic` / `speech` | General audio / speech | Keep the existing external-package backend; do not rewrite now. | | FlowHigh | General audio SR to 48 kHz | Deferred; CUDA-first runtime and checkpoint-license/provider blockers remain. | @@ -24,6 +25,7 @@ Source: - repository: - code license: Apache-2.0 +- checkpoint license: Apache-2.0 in the pinned Hugging Face model card and repository tag - model repository: - inspected source revision: `6b3774dc79c46ae8bed2a4fa5f706f0ac8c75c61` - inspected model revision: `39eb1f25ea84f5e0315ade9ac0070fff216fc690` @@ -38,24 +40,38 @@ Observed weight files: | `last_best_checkpoint_m.pt` | 218,471,889 | `6cbadb2b6b839e444bb65223c69eea162c8ad08f36e9d0a64144672c4095ab36` | | `do_03925000` | 1,744,458,369 | `549035d29a03928a854815b9c1b21c02d825845aa2e5f093af45694a236c1f19` | -The public SR loader appears to use the pointer plus `m` and `g` checkpoints; the larger file should -not enter a managed manifest unless the spike proves it is required. - -Blockers: - -- upstream loading can download implicitly when files are absent; -- checkpoints use PyTorch `.pt` loading and need a safe conversion/loading plan; -- package-owned module boundaries and key mapping are not reproduced; -- a short upstream CPU smoke and parity fixture have not been recorded; -- the clean-speech limitation must remain visible in catalog metadata. - -Spike exit criteria: - -1. Inspect tensor keys for the two required checkpoints from a temporary cache. -2. Run upstream CPU inference on a sub-second fixture without provider calls during inference. -3. Record input/output sample-rate, channel, alignment, and preprocessing behavior. -4. Decide whether conversion to a safer package-owned format is practical. -5. Define a parity fixture before opening backend implementation. +The Colab spike proved that the pointer plus `m` and `g` checkpoints are sufficient. The 1.74 GB +`do_03925000` file was not downloaded and must not enter a managed manifest. + +Recorded Colab CPU evidence: + +- `last_best_checkpoint_m.pt` contains a `mossformer` state with 929 tensors and 54,497,666 + parameters; `last_best_checkpoint_g.pt` contains a `generator` state with 311 tensors and + 55,150,402 parameters; +- both files load with `torch.load(..., weights_only=True)` and round-trip exactly through + safetensors after cloning shared rotary-frequency buffers; +- converted sizes are 218,118,376 and 220,633,048 bytes, with recorded deterministic conversion + hashes in the evidence bundle; +- a generated 0.25-second 16 kHz mono fixture ran on CPU with Hugging Face offline mode enabled; + model construction took 14-21 seconds and first inference took 4.1-5.2 seconds on Colab CPU; +- repeated inference in one process was bit-exact, but separate clean sessions produced small + floating-point differences, so parity must use numeric/audio tolerances rather than a WAV hash; +- 16 kHz mono input produced 48 kHz-domain shape `[1, 11776]`; 16 kHz stereo preserved two + independently processed channels with shape `[2, 11776]`; +- both a 4000-sample 16 kHz input and a 12000-sample 48 kHz input produced 11776 samples, a fixed + `-224` sample (`-4.67 ms`) alignment delta for this fixture; +- upstream preprocessing resamples inputs to 48 kHz, processes channels independently, and marks SR + output as 48 kHz rather than resampling it back to the original input rate. + +The package-owned parity fixture should generate the same seeded 440 Hz mono input, compare the +pre-alignment upstream float output, explicitly account for the 224-sample generator boundary loss, +and use bounded waveform/spectral tolerances. Do not use a cross-process file hash as the parity +oracle. The clean-speech limitation must remain visible in catalog metadata. + +Implementation is now technically admissible, but it should be accepted only as a bounded, +self-contained backend task with two managed safetensors files, explicit alignment policy, no +implicit downloads, and gated real-weight tests. Its approximately 110 million parameters and slow +CPU startup/inference make it unsuitable as the default lightweight backend. ## Resemble Enhance diff --git a/examples/colab_mossformer2_spike.py b/examples/colab_mossformer2_spike.py new file mode 100644 index 0000000..d9c4a07 --- /dev/null +++ b/examples/colab_mossformer2_spike.py @@ -0,0 +1,300 @@ +"""Run the bounded MossFormer2_SR_48K feasibility spike through Colab CLI. + +The script pins upstream source/model revisions, downloads only the three files +believed necessary for inference, then forces Hugging Face offline mode before +constructing or running the model. It writes a downloadable evidence archive. +Never run this script locally; it intentionally requires ``/content``. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import shutil +import subprocess +import sys +import tarfile +import time +from datetime import datetime, timezone +from pathlib import Path + +SOURCE_REPO = "https://github.com/modelscope/ClearerVoice-Studio.git" +SOURCE_REVISION = "6b3774dc79c46ae8bed2a4fa5f706f0ac8c75c61" +MODEL_REPO = "alibabasglab/MossFormer2_SR_48K" +MODEL_REVISION = "39eb1f25ea84f5e0315ade9ac0070fff216fc690" +EXPECTED_FILES = { + "last_best_checkpoint": (52, None), + "last_best_checkpoint_g.pt": ( + 220_712_702, + "0bdd13c21466f5963d9d1f86a9d84fc6196868318fe22c6b0a750f041805adda", + ), + "last_best_checkpoint_m.pt": ( + 218_471_889, + "6cbadb2b6b839e444bb65223c69eea162c8ad08f36e9d0a64144672c4095ab36", + ), +} +WORK_ROOT = Path(os.environ.get("ASR_COLAB_WORK_ROOT", "/content/mossformer2-spike")) +SOURCE_DIR = WORK_ROOT / "ClearerVoice-Studio" +CLEARVOICE_ROOT = SOURCE_DIR / "clearvoice" +CHECKPOINT_DIR = CLEARVOICE_ROOT / "checkpoints" / "MossFormer2_SR_48K" +EVIDENCE_DIR = WORK_ROOT / "evidence" +ARCHIVE_PATH = Path(os.environ.get("ASR_COLAB_ARCHIVE", "/content/asr-mossformer2-spike.tar.gz")) + + +def run(name: str, command: list[str], *, cwd: Path | None = None) -> None: + log = EVIDENCE_DIR / "logs" / f"{name}.log" + log.parent.mkdir(parents=True, exist_ok=True) + print(f"[mossformer2-spike] {name}: {' '.join(command)}", flush=True) + with log.open("w", encoding="utf-8") as output: + subprocess.run(command, cwd=cwd, stdout=output, stderr=subprocess.STDOUT, text=True, check=True) + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def download_checkpoints() -> list[dict[str, object]]: + from huggingface_hub import hf_hub_download + + CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True) + records = [] + for filename, (expected_size, expected_hash) in EXPECTED_FILES.items(): + cached = Path(hf_hub_download(repo_id=MODEL_REPO, filename=filename, revision=MODEL_REVISION)) + destination = CHECKPOINT_DIR / filename + shutil.copy2(cached, destination) + actual_hash = sha256(destination) + if destination.stat().st_size != expected_size: + raise RuntimeError(f"unexpected size for {filename}") + if expected_hash is not None and actual_hash != expected_hash: + raise RuntimeError(f"unexpected SHA256 for {filename}") + records.append( + { + "filename": filename, + "size": destination.stat().st_size, + "sha256": actual_hash, + } + ) + (EVIDENCE_DIR / "checkpoints.json").write_text(json.dumps(records, indent=2) + "\n") + return records + + +def inspect_and_convert() -> dict[str, object]: + import torch + from safetensors.torch import load_file, save_file + + result: dict[str, object] = {} + conversion_dir = EVIDENCE_DIR / "converted" + conversion_dir.mkdir() + state_keys = {"m": "mossformer", "g": "generator"} + for suffix, state_key in state_keys.items(): + source = CHECKPOINT_DIR / f"last_best_checkpoint_{suffix}.pt" + loaded = torch.load(source, map_location="cpu", weights_only=True) + state = loaded.get(state_key, loaded) if isinstance(loaded, dict) else loaded + if not isinstance(state, dict) or not all(isinstance(key, str) for key in state): + raise RuntimeError(f"checkpoint {suffix} is not a state dict") + # Break shared storage explicitly. MossFormer repeats one rotary-frequency + # buffer across layers, which safetensors intentionally refuses to encode + # through its generic state-dict writer. + tensors = {key: value.detach().clone().contiguous() for key, value in state.items() if torch.is_tensor(value)} + if len(tensors) != len(state): + raise RuntimeError(f"checkpoint {suffix} contains non-tensor state") + target = conversion_dir / f"mossformer2_{suffix}.safetensors" + save_file(tensors, target) + converted = load_file(target, device="cpu") + if converted.keys() != tensors.keys() or any(not torch.equal(converted[key], tensors[key]) for key in tensors): + raise RuntimeError(f"safetensors round trip failed for checkpoint {suffix}") + result[suffix] = { + "root_keys": sorted(loaded) if isinstance(loaded, dict) else None, + "selected_state_key": state_key, + "tensor_count": len(tensors), + "parameter_count": sum(value.numel() for value in tensors.values()), + "first_keys": sorted(tensors)[:25], + "safetensors_size": target.stat().st_size, + "safetensors_sha256": sha256(target), + "safetensors_round_trip_exact": True, + } + target.unlink() + (EVIDENCE_DIR / "checkpoint-inspection.json").write_text(json.dumps(result, indent=2) + "\n") + return result + + +def run_offline_inference() -> dict[str, object]: + os.environ.update( + { + "HF_HUB_OFFLINE": "1", + "TRANSFORMERS_OFFLINE": "1", + "HF_DATASETS_OFFLINE": "1", + "NO_PROXY": "*", + "no_proxy": "*", + } + ) + import numpy as np + import soundfile as sf + import torch + + sys.path.insert(0, str(SOURCE_DIR / "clearvoice")) + from clearvoice import ClearVoice + + torch.manual_seed(0) + source_rate = 16_000 + sample_count = 4_000 + timeline = np.arange(sample_count, dtype=np.float32) / source_rate + fixture = (0.08 * np.sin(2 * np.pi * 440 * timeline)).astype(np.float32) + fixture_path = WORK_ROOT / "fixture-16k.wav" + sf.write(fixture_path, fixture, source_rate, subtype="PCM_16") + stereo_path = WORK_ROOT / "fixture-16k-stereo.wav" + sf.write(stereo_path, np.column_stack((fixture, fixture * 0.5)), source_rate, subtype="PCM_16") + fixture_48k = np.repeat(fixture, 3) + fixture_48k_path = WORK_ROOT / "fixture-48k.wav" + sf.write(fixture_48k_path, fixture_48k, 48_000, subtype="PCM_16") + + original_cwd = Path.cwd() + os.chdir(CLEARVOICE_ROOT) + try: + started = time.perf_counter() + model = ClearVoice(task="speech_super_resolution", model_names=["MossFormer2_SR_48K"]) + load_seconds = time.perf_counter() - started + started = time.perf_counter() + first = np.asarray(model(input_path=str(fixture_path), online_write=False)) + inference_seconds = time.perf_counter() - started + second = np.asarray(model(input_path=str(fixture_path), online_write=False)) + stereo = np.asarray(model(input_path=str(stereo_path), online_write=False)) + native_48k = np.asarray(model(input_path=str(fixture_48k_path), online_write=False)) + finally: + os.chdir(original_cwd) + output_path = EVIDENCE_DIR / "upstream-output-48k.wav" + flat = np.squeeze(first) + sf.write(output_path, flat, 48_000, subtype="FLOAT") + expected_length = round(sample_count * 48_000 / source_rate) + record = { + "device": str(model.models[0].device), + "input_sample_rate": source_rate, + "input_channels": 1, + "input_samples": sample_count, + "output_contract_sample_rate": 48_000, + "output_shape": list(first.shape), + "output_samples": int(flat.size), + "expected_output_samples": expected_length, + "alignment_delta_samples": int(flat.size - expected_length), + "finite": bool(np.isfinite(flat).all()), + "peak": float(np.max(np.abs(flat))), + "deterministic_exact": bool(np.array_equal(first, second)), + "load_seconds": load_seconds, + "inference_seconds": inference_seconds, + "output_sha256": sha256(output_path), + "stereo_16k_output_shape": list(stereo.shape), + "stereo_channels_preserved": bool(stereo.shape[0] == 2), + "native_48k_input_samples": int(fixture_48k.size), + "native_48k_output_shape": list(native_48k.shape), + "native_48k_alignment_delta_samples": int(np.squeeze(native_48k).size - fixture_48k.size), + } + if ( + not record["finite"] + or abs(record["alignment_delta_samples"]) > 512 + or not record["stereo_channels_preserved"] + or abs(record["native_48k_alignment_delta_samples"]) > 512 + ): + raise RuntimeError("upstream inference produced an invalid or materially misaligned output") + (EVIDENCE_DIR / "inference.json").write_text(json.dumps(record, indent=2) + "\n") + return record + + +def write_summary(status: str, error: str | None, **evidence: object) -> None: + summary = { + "schema_version": 1, + "created_at": datetime.now(timezone.utc).isoformat(), + "status": status, + "source_repository": SOURCE_REPO, + "source_revision": SOURCE_REVISION, + "model_repository": MODEL_REPO, + "model_revision": MODEL_REVISION, + "offline_inference": True, + "error": error, + **evidence, + } + (EVIDENCE_DIR / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + + +def archive() -> None: + with tarfile.open(ARCHIVE_PATH, "w:gz") as bundle: + bundle.add(EVIDENCE_DIR, arcname="evidence") + print(f"[mossformer2-spike] evidence archive: {ARCHIVE_PATH}", flush=True) + + +def worker_main() -> int: + error = None + evidence: dict[str, object] = {} + try: + evidence["checkpoints"] = download_checkpoints() + evidence["inspection"] = inspect_and_convert() + evidence["inference"] = run_offline_inference() + write_summary("passed", None, **evidence) + except Exception as exc: + error = f"{type(exc).__name__}: {exc}" + write_summary("failed", error, **evidence) + if error: + print(f"[mossformer2-spike] failed: {error}", file=sys.stderr) + return 1 + return 0 + + +def main() -> int: + if "--worker" in sys.argv: + return worker_main() + if Path.cwd() != Path("/content"): + raise RuntimeError("This spike must run inside a Colab runtime rooted at /content") + if WORK_ROOT.exists(): + shutil.rmtree(WORK_ROOT) + EVIDENCE_DIR.mkdir(parents=True) + try: + runner_source = Path(__file__).read_text(encoding="utf-8") + except NameError: + runner_source = get_ipython().history_manager.input_hist_raw[-1] # noqa: F821 + runner_path = WORK_ROOT / "colab_mossformer2_spike.py" + runner_path.write_text(runner_source, encoding="utf-8") + error = None + try: + run("clone", ["git", "clone", "--filter=blob:none", SOURCE_REPO, str(SOURCE_DIR)]) + run("checkout", ["git", "checkout", SOURCE_REVISION], cwd=SOURCE_DIR) + run("install-uv", [sys.executable, "-m", "pip", "install", "uv"]) + run( + "install-dependencies", + [ + "uv", + "pip", + "install", + "--system", + "-e", + str(SOURCE_DIR / "clearvoice"), + "safetensors", + ], + ) + run( + "environment", + [ + sys.executable, + "-c", + "import platform,numpy,torch; " + "print(platform.python_version(), numpy.__version__, torch.__version__, torch.cuda.is_available())", + ], + ) + run("isolated-worker", [sys.executable, str(runner_path), "--worker"]) + except Exception as exc: + error = f"{type(exc).__name__}: {exc}" + if not (EVIDENCE_DIR / "summary.json").exists(): + write_summary("failed", error) + finally: + archive() + if error: + print(f"[mossformer2-spike] failed: {error}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())