audio-transition-inpainting-checkpoints

Craftic/audio-transition-inpainting-checkpoints

Research models trained for musical audio transition inpainting as part of the final NIR project.

The task is to reconstruct a missing musical transition T from left and right contexts:

A | masked transition t | B  →  reconstructed transition T

Audio is represented with WavTokenizer tokens. The canonical evaluation protocol uses MusicNet, fixed left/right contexts of 96 tokens, and transition lengths of 24, 48, 96, 144, and 192 tokens.

This repository contains:

• 24 PyTorch checkpoints from six controlled training lines
• four model families: Transformer, diffusion denoiser, discrete diffusion Transformer, and AIDD
• checkpoint-selection and evaluation tables
• the architecture and loading code used by the project
• SHA-256 checksums for every project checkpoint

Model summary

• Task: musical audio transition inpainting
• Input: left context A, masked transition t, and right context B
• Output: reconstructed transition T
• Audio representation: WavTokenizer discrete tokens
• Sample rate: 24 kHz
• Context: 96 tokens on each side
• Transition lengths: 24, 48, 96, 144, and 192 tokens
• Model families: Transformer, continuous diffusion denoiser, discrete diffusion Transformer, AIDD
• Training objectives: DWDSE and cross-entropy
• Training corruption: full-gap and span-corruption

Configuration

The repository contains 24 project-owned checkpoints covering six training lines and four model families. Its models/ tree mirrors the canonical project layout.

Each training line includes:

• Transformer;
• diffusion denoiser;
• discrete diffusion Transformer;
• AIDD.

Training data

• Dataset: MusicNet
• Domain: classical music recordings
• Split unit: complete tracks
• Evaluation: 75 fixed cases covering five gap lengths and three positions
• Tokenizer: WavTokenizer at 24 kHz

The tokenizer weights and raw MusicNet data are not included in this repository.

Training setup

• Base runs: up to 100,000 optimization steps
• Checkpoint selection: best validation loss
• Objectives: DWDSE or token cross-entropy
• Corruption protocols: full-gap or span-corruption
• Primary evaluation: strict full-gap reconstruction
• Diagnostic evaluation: span-corruption with 50% masked transition tokens

Exact selected steps, validation values, and run metadata are stored under metadata/<training-line>/.

Results

The strongest project-owned direction is Base Cross-Entropy AIDD. In strict full-gap evaluation it substantially closes the gap to the DSP reference, while the AIDD family is more robust than the other learned architectures as transition length increases.

Strict full-gap results for the Base CE line:

Span-corruption is reported separately and is not directly comparable to strict full-gap evaluation. In that diagnostic protocol, Base CE AIDD reaches transition LSD 2.906.

The metadata/ directories contain the exact checkpoint-selection records, training summaries, and full-gap/span-corruption evaluation summaries used in the project report.

Usage

Files are PyTorch training checkpoints rather than standalone Transformers-format models.

From the root of the canonical project, download all 24 project checkpoints directly into the expected models/<training-line>/ directories with:

hf download Craftic/audio-transition-inpainting-checkpoints models/ --local-dir .

The Hub repository uses the same relative paths as the project, so no additional copying or renaming is required. Existing third-party files under models/pretrained/ are not modified. File integrity can be checked against CHECKSUMS.csv.

Base checkpoints generally contain:

{
    "step": int,
    "model_name": str,
    "model": state_dict,
    "optimizer": optimizer_state,
    "config": dict,
    "last_loss": float,
    "params": int,
    "ema_model": state_dict,
}

Small checkpoints may contain only the selected model state and selection metadata. Inspect the keys before loading:

import torch

checkpoint = torch.load("path/to/checkpoint.pt", map_location="cpu", weights_only=False)
print(checkpoint.keys())
state_dict = checkpoint["model"]

The corresponding architecture and loader implementation is preserved under code/. These are custom research models and are not loadable with transformers.AutoModel.

Repository structure

models/
├── small_dwdse_full_gap/
├── base_dwdse_full_gap/
├── base_dwdse_span_corruption/
├── small_ce_full_gap/
├── base_ce_full_gap/
└── base_ce_span_corruption/
metadata/
└── <training-line>/*.csv|*.json
code/
├── models/*.py
├── training_runs.py
├── model_training.py
├── experiment_protocol.py
└── samplers.py
CHECKSUMS.csv

External weights intentionally excluded

The following locally used third-party checkpoints are not redistributed here:

• pretrained AIDD MusicNet and MAESTRO checkpoints;
• WavTokenizer checkpoint;
• DAC checkpoint;
• EnCodec checkpoint.

They must be obtained from their original authors/distribution channels under the applicable licenses.

Notes

• Full-gap and span-corruption results are different protocols and must not be averaged together.
• FAD-style and ODG-style are reproducible local proxy metrics, not official external FAD/PEAQ implementations.
• Model selection was performed using validation loss; evaluation tables contain the fixed MusicNet cases used in the project.
• The checkpoints include optimizer/EMA state where available and are therefore significantly larger than inference-only exports.

Integrity

CHECKSUMS.csv records SHA-256, byte size, training line, architecture, checkpoint step, and selection metadata for every uploaded checkpoint.

Limitations

• Custom PyTorch checkpoint format
• Requires the original WavTokenizer setup
• Trained and evaluated only on the project MusicNet protocol
• Objective metrics do not replace controlled listening tests
• Full training checkpoints are large because optimizer and EMA states are preserved
• No packaged production inference API

Intended use

• Research reproducibility
• Audio inpainting experiments
• Comparison of architecture, loss, corruption, and scale
• Continued work on Base CE AIDD
• Development of boundary-aware and efficient inference strategies

Not intended use

• General-purpose music generation
• Production audio restoration without additional validation
• Safety-critical or forensic audio processing
• Redistribution of excluded third-party checkpoints

Access

This repository is private and intended for research preservation, reproducibility, and future continuation of the project. Redistribution and publication decisions should be made separately from the licenses of external datasets and tokenizers.