feat: finish #28 (safetensors, config.json, latent_length) + implement #38 (per-sample conditioning) - #39
Merged
Merged
Conversation
… + implement #38 (per-sample conditioning) Issue #28 remaining items: - #3: Added safetensors dependency + export after torch.save in both training scripts (decoder.safetensors, dit.safetensors). Prior is excluded (full pickled object, not a state_dict). - #4: --latent-length now defaults to None (sentinel); when not explicitly set, derived as audio_length // 320 (EnCodec stride). Applied to train_audio_diffusion.py and sample_audio.py. train_audio_decoder.py doesn't need --latent-length (decoder operates on any length). - #5: Both training scripts now write config.json alongside checkpoints with full architecture hyperparams (dim, depth, num_heads, patch_size, latent_channels, latent_length, spec_dim, q, k, upsample_strides, sample_rate, etc.) so the Rust generation worker can reconstruct the exact architecture without hard-coding. Issue #38: - Added --num-samples (default 1) and --per-sample-conditioning (flag) to sample_audio.py. - When --per-sample-conditioning, c_spec is resampled per iteration with a deterministic per-sample seed (base_seed + idx + 1000 offset) so each sample is steered toward a different spectral region. - When off (default), shared c_spec produces a coherent bank. - Multi-sample outputs saved as sample_00.wav, sample_01.wav, etc. (single sample keeps original filename for backward compat). - Added --audio-length to sample_audio.py for latent_length derivation. Tests (TDD, 13 new): - test_config_export_28.py: config.json export from both scripts, safetensors export + loadability, latent_length derivation logic. - test_per_sample_conditioning_38.py: CLI flags present, per-sample c_spec diversity > shared (with guidance_scale=50, perturbed AdaLN), shared c_spec condenses vs unconditional, per-sample is deterministic. Full suite: 238 passed, ruff clean. Bumps 0.7.0 -> 0.8.0.
6 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Finishes the remaining #28 items and implements #38 (per-sample conditioning). Single branch, separate commits per issue.
Issue #28 remaining items
#3 — safetensors export
Added safetensors dependency. Both training scripts now export .safetensors alongside .pt files after training:
Prior excluded (full pickled object, not a state_dict).
#4 — derive latent_length from audio_length
--latent-length now defaults to None (sentinel). When not explicitly set, derived as audio_length // 320 (EnCodec stride). Applied to train_audio_diffusion.py and sample_audio.py. Backward compatible: explicit --latent-length overrides derivation.
#5 — config.json alongside checkpoints
Both training scripts now write config.json with full architecture hyperparams:
The Rust generation worker can now reconstruct exact architecture without hard-coding.
Issue #38 — per-sample conditioning
sample_audio.py now supports:
When active, each sample gets a fresh c_spec probe with a deterministic per-sample seed (base_seed + idx + 1000), steering toward different spectral regions. When off (default), shared c_spec produces a coherent bank.
Multi-sample outputs: sample_00.wav, sample_01.wav, etc. Single sample keeps original filename.
Also added --audio-length to sample_audio.py for latent_length derivation.
Tests (TDD, 13 new)
Full suite: 238 passed, ruff clean, pre-commit hook passed.
Version
0.7.0 to 0.8.0
Closes #28, closes #38.