colabcola 0.1.0__py3-none-any.whl

colabcola/__init__.py

@@ -0,0 +1,195 @@
+"""Compile trusted batch jobs into deterministic Google Colab notebooks."""
+
+from ._version import __version__
+from .bundle import (
+    ARCHIVE_WARNING_BYTES,
+    Bundle,
+    FileRecord,
+    archive_size_warnings,
+    build_bundle,
+    canonical_json_bytes,
+    verify_bundle,
+)
+from .compiler import (
+    CompilationResult,
+    DriveControllerResult,
+    StagedBundleResult,
+    compile_drive_controller,
+    compile_drive_controller_invocation,
+    compile_invocation,
+    compile_notebook,
+    prepare_drive_bundle,
+    prepare_drive_invocation,
+    write_notebook,
+)
+from .controller import (
+    DriveControllerAuth,
+    DriveControllerArtifact,
+    build_drive_controller,
+    controller_auth_from_settings,
+    verify_drive_controller,
+    write_drive_controller,
+)
+from .drive import (
+    DRIVE_ROOT_FOLDER_NAME,
+    DRIVE_STAGING_FOLDER_NAME,
+    DriveBackend,
+    DriveFile,
+    DriveFolder,
+    DriveStagingReceipt,
+    VerifiedResult,
+    authorize_abandoned_cleanup,
+    cleanup_drive_staging,
+    drive_receipt_for_observed_upload,
+    drive_receipt_json_bytes,
+    load_drive_receipt,
+    record_verified_result,
+    stage_drive_bundle,
+    verify_drive_receipt,
+    write_drive_receipt,
+)
+from .errors import BundleError, DriveTransportError, NotebookError, SpecError
+from .google_drive import (
+    DEFAULT_COLAB_SECRET_NAME,
+    GOOGLE_DRIVE_SCOPE,
+    GoogleDriveApiBackend,
+    GoogleDriveSettings,
+    GoogleOAuthClient,
+    create_google_drive_backend,
+    default_google_drive_settings_path,
+    google_drive_settings_from_clients,
+    load_google_drive_settings,
+    load_google_oauth_client,
+    verify_google_drive_settings,
+    write_google_drive_settings,
+)
+from .hifigan import (
+    HiFiGanOverlay,
+    HiFiGanProfile,
+    HiFiGanRuntimeLock,
+    HiFiGanSourceFile,
+    default_hifigan_profile_path,
+    hifigan_source_payload,
+    load_hifigan_profile,
+    verify_hifigan_checkout,
+    verify_hifigan_profile_assets,
+)
+from .invocation import JobInvocation, bind_job_arguments, load_job_invocation
+from .ljspeech import (
+    LJSpeechSelection,
+    LJSpeechSubset,
+    build_ljspeech_subset,
+    ljspeech_subset_payload,
+    select_ljspeech_ids,
+    verify_ljspeech_subset,
+)
+from .notebook import (
+    NotebookArtifact,
+    build_notebook,
+    encode_archive_chunks,
+    verify_notebook,
+)
+from .spec import JobSpec, RemoteInputSpec, RuntimeSpec, load_job_spec
+from .transport import (
+    STAGED_BUNDLE_HEADER_BYTES,
+    STAGED_BUNDLE_MAGIC,
+    STAGED_BUNDLE_SUFFIX,
+    StagedBundleArtifact,
+    build_staged_bundle,
+    verify_staged_bundle,
+    verify_staged_bundle_file,
+    write_staged_bundle,
+)
+
+__all__ = [
+    "ARCHIVE_WARNING_BYTES",
+    "Bundle",
+    "BundleError",
+    "CompilationResult",
+    "DRIVE_ROOT_FOLDER_NAME",
+    "DRIVE_STAGING_FOLDER_NAME",
+    "DEFAULT_COLAB_SECRET_NAME",
+    "DriveBackend",
+    "DriveControllerAuth",
+    "DriveControllerArtifact",
+    "DriveControllerResult",
+    "DriveFile",
+    "DriveFolder",
+    "DriveStagingReceipt",
+    "DriveTransportError",
+    "FileRecord",
+    "JobInvocation",
+    "JobSpec",
+    "LJSpeechSelection",
+    "LJSpeechSubset",
+    "HiFiGanOverlay",
+    "HiFiGanProfile",
+    "HiFiGanRuntimeLock",
+    "HiFiGanSourceFile",
+    "NotebookArtifact",
+    "NotebookError",
+    "RemoteInputSpec",
+    "STAGED_BUNDLE_HEADER_BYTES",
+    "STAGED_BUNDLE_MAGIC",
+    "STAGED_BUNDLE_SUFFIX",
+    "SpecError",
+    "StagedBundleArtifact",
+    "StagedBundleResult",
+    "VerifiedResult",
+    "authorize_abandoned_cleanup",
+    "bind_job_arguments",
+    "archive_size_warnings",
+    "build_bundle",
+    "build_drive_controller",
+    "build_notebook",
+    "build_staged_bundle",
+    "build_ljspeech_subset",
+    "canonical_json_bytes",
+    "cleanup_drive_staging",
+    "compile_drive_controller",
+    "compile_drive_controller_invocation",
+    "compile_invocation",
+    "compile_notebook",
+    "default_hifigan_profile_path",
+    "drive_receipt_for_observed_upload",
+    "drive_receipt_json_bytes",
+    "encode_archive_chunks",
+    "load_job_invocation",
+    "load_job_spec",
+    "load_drive_receipt",
+    "hifigan_source_payload",
+    "load_hifigan_profile",
+    "ljspeech_subset_payload",
+    "prepare_drive_bundle",
+    "prepare_drive_invocation",
+    "record_verified_result",
+    "select_ljspeech_ids",
+    "stage_drive_bundle",
+    "verify_bundle",
+    "verify_drive_controller",
+    "verify_drive_receipt",
+    "verify_notebook",
+    "verify_staged_bundle",
+    "verify_staged_bundle_file",
+    "verify_hifigan_checkout",
+    "verify_hifigan_profile_assets",
+    "verify_ljspeech_subset",
+    "write_drive_controller",
+    "write_drive_receipt",
+    "write_notebook",
+    "write_staged_bundle",
+    "__version__",
+    "GOOGLE_DRIVE_SCOPE",
+    "GoogleDriveApiBackend",
+    "GoogleDriveSettings",
+    "GoogleOAuthClient",
+    "RuntimeSpec",
+    "controller_auth_from_settings",
+    "create_google_drive_backend",
+    "default_google_drive_settings_path",
+    "google_drive_settings_from_clients",
+    "load_google_drive_settings",
+    "load_google_oauth_client",
+    "verify_google_drive_settings",
+    "write_google_drive_settings",
+]

colabcola/__main__.py

@@ -0,0 +1,3 @@
+from .cli import main
+
+raise SystemExit(main())

colabcola/_hifigan/__init__.py

@@ -0,0 +1 @@
+"""Checked-in assets for the pinned HiFi-GAN Phase 2-3 workload."""

colabcola/_hifigan/audio_diagnostic.py

@@ -0,0 +1,292 @@
+"""Bounded HiFi-GAN generated-audio diagnostic for a trusted checkpoint.
+
+Derived from ``inference.py`` and ``meldataset.py`` in jik876/hifi-gan commit
+4769534d45265d52a904b850da5a622601885777. The bundled upstream LICENSE
+applies to the derived portions of this file.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+import pathlib
+import shutil
+
+import numpy as np
+import torch
+from scipy.io.wavfile import read, write
+
+import meldataset as upstream_meldataset
+from env import AttrDict
+from models import Generator
+
+
+MAX_WAV_VALUE = 32768.0
+UPSTREAM_COMMIT = "4769534d45265d52a904b850da5a622601885777"
+EXPECTED_INPUTS = {
+    "LJ001-0028.wav": (
+        261_478,
+        "58831256095edbcc72346a2db9be39efd39662f442f8115f19da4d55876cfa0e",
+    ),
+    "LJ001-0068.wav": (
+        402_278,
+        "9ec4f7944fdf753a62792aef59da5f2145c0a69d8877fc07c4c77c7189dcdc0e",
+    ),
+    "LJ002-0149.wav": (
+        349_030,
+        "d71743df2a5842ae0130920afe910e53516374dae97e8f24da403efdc4cc9c60",
+    ),
+}
+
+
+def mel_spectrogram(
+    y,
+    n_fft,
+    num_mels,
+    sampling_rate,
+    hop_size,
+    win_size,
+    fmin,
+    fmax,
+    center=False,
+):
+    """Pinned upstream transform adapted to current librosa and PyTorch APIs."""
+
+    cache_key = f"{fmax}_{y.device}"
+    if cache_key not in upstream_meldataset.mel_basis:
+        mel = upstream_meldataset.librosa_mel_fn(
+            sr=sampling_rate,
+            n_fft=n_fft,
+            n_mels=num_mels,
+            fmin=fmin,
+            fmax=fmax,
+        )
+        upstream_meldataset.mel_basis[cache_key] = torch.from_numpy(mel).float().to(
+            y.device
+        )
+        upstream_meldataset.hann_window[str(y.device)] = torch.hann_window(win_size).to(
+            y.device
+        )
+
+    padded = torch.nn.functional.pad(
+        y.unsqueeze(1),
+        (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),
+        mode="reflect",
+    ).squeeze(1)
+    spectrum = torch.stft(
+        padded,
+        n_fft,
+        hop_length=hop_size,
+        win_length=win_size,
+        window=upstream_meldataset.hann_window[str(y.device)],
+        center=center,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=True,
+    ).abs()
+    spectrum = torch.matmul(upstream_meldataset.mel_basis[cache_key], spectrum)
+    return upstream_meldataset.spectral_normalize_torch(spectrum)
+
+
+def _sha256(path: pathlib.Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as stream:
+        while block := stream.read(1024 * 1024):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def _audio_stats(samples: np.ndarray) -> dict[str, object]:
+    normalized = samples.astype(np.float64) / MAX_WAV_VALUE
+    return {
+        "frames": int(samples.size),
+        "peak_absolute": float(np.max(np.abs(normalized))) if samples.size else 0.0,
+        "rms": float(np.sqrt(np.mean(np.square(normalized)))) if samples.size else 0.0,
+    }
+
+
+def _atomic_json(path: pathlib.Path, value: object) -> None:
+    content = (
+        json.dumps(
+            value,
+            allow_nan=False,
+            ensure_ascii=False,
+            separators=(",", ":"),
+            sort_keys=True,
+        ).encode("utf-8")
+        + b"\n"
+    )
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name("." + path.name + ".tmp")
+    try:
+        with temporary.open("xb") as stream:
+            stream.write(content)
+            stream.flush()
+            os.fsync(stream.fileno())
+        os.replace(temporary, path)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def run(args: argparse.Namespace) -> dict[str, object]:
+    checkpoint = pathlib.Path(args.checkpoint)
+    checkpoint_sha256 = _sha256(checkpoint)
+    if checkpoint_sha256 != args.expected_checkpoint_sha256:
+        raise RuntimeError("generator checkpoint SHA-256 does not match")
+    if not torch.cuda.is_available():
+        raise RuntimeError("audio diagnostic requires one CUDA GPU")
+
+    config_value = json.loads(pathlib.Path(args.config).read_text(encoding="utf-8"))
+    config = AttrDict(config_value)
+    device = torch.device("cuda")
+    torch.manual_seed(config.seed)
+    torch.cuda.manual_seed_all(config.seed)
+    torch.backends.cudnn.benchmark = False
+
+    generator = Generator(config).to(device)
+    checkpoint_value = torch.load(checkpoint, map_location=device, weights_only=False)
+    if not isinstance(checkpoint_value, dict) or "generator" not in checkpoint_value:
+        raise RuntimeError("generator checkpoint has an unexpected structure")
+    generator.load_state_dict(checkpoint_value["generator"], strict=True)
+    generator.eval()
+    generator.remove_weight_norm()
+
+    output_root = pathlib.Path(args.output_root)
+    reference_root = output_root / "reference"
+    generated_root = output_root / "generated"
+    reference_root.mkdir(parents=True, exist_ok=False)
+    generated_root.mkdir(parents=True, exist_ok=False)
+
+    records: list[dict[str, object]] = []
+    if [pathlib.Path(value).name for value in args.input_wav] != list(EXPECTED_INPUTS):
+        raise RuntimeError("audio diagnostic inputs do not match the pinned ordered clip set")
+    with torch.no_grad():
+        for input_value in args.input_wav:
+            source = pathlib.Path(input_value)
+            expected_size, expected_sha256 = EXPECTED_INPUTS[source.name]
+            if source.stat().st_size != expected_size or _sha256(source) != expected_sha256:
+                raise RuntimeError(f"validation WAV identity does not match for {source.name}")
+            sampling_rate, samples = read(source)
+            if sampling_rate != config.sampling_rate:
+                raise RuntimeError(f"unexpected sampling rate for {source.name}")
+            if samples.dtype != np.int16 or samples.ndim != 1 or samples.size == 0:
+                raise RuntimeError(f"unexpected WAV format for {source.name}")
+
+            reference_path = reference_root / source.name
+            generated_path = generated_root / f"{source.stem}_generated.wav"
+            shutil.copyfile(source, reference_path)
+
+            reference_tensor = (
+                torch.from_numpy(samples.astype(np.float32)).to(device) / MAX_WAV_VALUE
+            ).unsqueeze(0)
+            input_mel = mel_spectrogram(
+                reference_tensor,
+                config.n_fft,
+                config.num_mels,
+                config.sampling_rate,
+                config.hop_size,
+                config.win_size,
+                config.fmin,
+                config.fmax,
+            )
+            generated_tensor = generator(input_mel).squeeze(0).squeeze(0)
+            if not bool(torch.isfinite(generated_tensor).all().item()):
+                raise RuntimeError(f"generated non-finite audio for {source.name}")
+
+            generated_float = generated_tensor.detach().cpu().numpy()
+            clipped_samples = int(np.count_nonzero(np.abs(generated_float) > 1.0))
+            generated_samples = np.rint(
+                np.clip(generated_float, -1.0, 1.0) * 32767.0
+            ).astype(np.int16)
+            write(generated_path, config.sampling_rate, generated_samples)
+
+            generated_for_mel = torch.from_numpy(generated_float).to(device).unsqueeze(0)
+            reference_loss_mel = mel_spectrogram(
+                reference_tensor,
+                config.n_fft,
+                config.num_mels,
+                config.sampling_rate,
+                config.hop_size,
+                config.win_size,
+                config.fmin,
+                config.fmax_for_loss,
+            )
+            generated_loss_mel = mel_spectrogram(
+                generated_for_mel,
+                config.n_fft,
+                config.num_mels,
+                config.sampling_rate,
+                config.hop_size,
+                config.win_size,
+                config.fmin,
+                config.fmax_for_loss,
+            )
+            aligned_frames = min(reference_loss_mel.shape[-1], generated_loss_mel.shape[-1])
+            mel_l1 = torch.nn.functional.l1_loss(
+                reference_loss_mel[..., :aligned_frames],
+                generated_loss_mel[..., :aligned_frames],
+            ).item()
+
+            records.append(
+                {
+                    "clip_id": source.stem,
+                    "generated": {
+                        **_audio_stats(generated_samples),
+                        "clipped_samples_before_clamp": clipped_samples,
+                        "path": generated_path.as_posix(),
+                        "sha256": _sha256(generated_path),
+                        "size": generated_path.stat().st_size,
+                    },
+                    "mel_l1_aligned": float(mel_l1),
+                    "mel_l1_frames": int(aligned_frames),
+                    "reference": {
+                        **_audio_stats(samples),
+                        "path": reference_path.as_posix(),
+                        "sha256": _sha256(reference_path),
+                        "size": reference_path.stat().st_size,
+                    },
+                    "sampling_rate": int(sampling_rate),
+                }
+            )
+
+    report = {
+        "checkpoint": {
+            "path": checkpoint.as_posix(),
+            "sha256": checkpoint_sha256,
+            "size": checkpoint.stat().st_size,
+        },
+        "clips": records,
+        "diagnostic_boundary": (
+            "Operational generated-audio evidence only; not a convergence, perceptual-quality, "
+            "MOS, or paper-reproduction result."
+        ),
+        "gpu": {
+            "count": torch.cuda.device_count(),
+            "memory_bytes": torch.cuda.get_device_properties(0).total_memory,
+            "name": torch.cuda.get_device_name(0),
+        },
+        "sample_count": len(records),
+        "state": "audio_diagnostic_complete",
+        "upstream_commit": UPSTREAM_COMMIT,
+    }
+    _atomic_json(pathlib.Path(args.report), report)
+    print(json.dumps(report, allow_nan=False, separators=(",", ":"), sort_keys=True))
+    return report
+
+
+def _parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", required=True)
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--expected-checkpoint-sha256", required=True)
+    parser.add_argument("--input-wav", action="append", required=True)
+    parser.add_argument("--output-root", required=True)
+    parser.add_argument("--report", required=True)
+    return parser
+
+
+if __name__ == "__main__":
+    run(_parser().parse_args())

colabcola/_hifigan/extract_audio_diagnostic_input.py

@@ -0,0 +1,117 @@
+"""Extract the one verified HiFi-GAN checkpoint needed by the audio diagnostic."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+from pathlib import Path
+import tempfile
+import zipfile
+
+
+CHECKPOINT_MEMBER = "outputs/cp_hifigan/g_00005000"
+CHECKPOINT_SIZE = 55_825_339
+CHECKPOINT_SHA256 = "836627cc247ec688589593d2324aae2cfadc8958eb17304c9426c44de56b46e3"
+
+
+def _sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as stream:
+        for block in iter(lambda: stream.read(1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def extract_checkpoint(
+    artifacts_zip: Path,
+    output_path: Path,
+    *,
+    member: str = CHECKPOINT_MEMBER,
+    expected_size: int = CHECKPOINT_SIZE,
+    expected_sha256: str = CHECKPOINT_SHA256,
+) -> dict[str, object]:
+    artifacts_zip = artifacts_zip.resolve(strict=True)
+    if not artifacts_zip.is_file() or artifacts_zip.is_symlink():
+        raise ValueError(f"artifacts path is not a regular file: {artifacts_zip}")
+
+    output_path = output_path.resolve(strict=False)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    if output_path.exists():
+        if not output_path.is_file() or output_path.is_symlink():
+            raise ValueError(f"output path is not a regular file: {output_path}")
+        actual_size = output_path.stat().st_size
+        actual_sha256 = _sha256(output_path)
+        if actual_size != expected_size or actual_sha256 != expected_sha256:
+            raise ValueError(
+                "existing output does not match the verified checkpoint: "
+                f"size={actual_size}, sha256={actual_sha256}"
+            )
+        return {
+            "artifacts_zip": artifacts_zip.as_posix(),
+            "checkpoint_path": output_path.as_posix(),
+            "checkpoint_size": actual_size,
+            "checkpoint_sha256": actual_sha256,
+            "member": member,
+            "reused": True,
+        }
+
+    with zipfile.ZipFile(artifacts_zip) as archive:
+        matching = [info for info in archive.infolist() if info.filename == member]
+        if len(matching) != 1:
+            raise ValueError(f"expected exactly one {member!r} member; found {len(matching)}")
+        info = matching[0]
+        if info.is_dir() or info.file_size != expected_size:
+            raise ValueError(
+                f"checkpoint member size mismatch: expected {expected_size}, got {info.file_size}"
+            )
+
+        handle, temporary_name = tempfile.mkstemp(
+            dir=output_path.parent,
+            prefix=f".{output_path.name}.",
+            suffix=".tmp",
+        )
+        os.close(handle)
+        temporary_path = Path(temporary_name)
+        try:
+            digest = hashlib.sha256()
+            with archive.open(info, "r") as source, temporary_path.open("wb") as destination:
+                while block := source.read(1024 * 1024):
+                    destination.write(block)
+                    digest.update(block)
+                destination.flush()
+                os.fsync(destination.fileno())
+            actual_size = temporary_path.stat().st_size
+            actual_sha256 = digest.hexdigest()
+            if actual_size != expected_size or actual_sha256 != expected_sha256:
+                raise ValueError(
+                    "extracted checkpoint failed verification: "
+                    f"size={actual_size}, sha256={actual_sha256}"
+                )
+            os.replace(temporary_path, output_path)
+        finally:
+            temporary_path.unlink(missing_ok=True)
+
+    return {
+        "artifacts_zip": artifacts_zip.as_posix(),
+        "checkpoint_path": output_path.as_posix(),
+        "checkpoint_size": expected_size,
+        "checkpoint_sha256": expected_sha256,
+        "member": member,
+        "reused": False,
+    }
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("artifacts_zip", type=Path)
+    parser.add_argument("output_path", type=Path)
+    args = parser.parse_args()
+    result = extract_checkpoint(args.artifacts_zip, args.output_path)
+    print(json.dumps(result, sort_keys=True, separators=(",", ":")))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())