Add FlowUNet1D generative bridge flow-matching denoiser - #61
Merged
Conversation
Introduce the repo's first generative denoiser: a waveform-domain Schrödinger-bridge-style flow-matching model with a data-prediction objective, per the research survey in docs/generative_denoiser_proposal.md. - FlowUNet1D backbone reuses the UNet1D Down/Mid/Up blocks, adding a sinusoidal time embedding and zero-initialized FiLM conditioning after each block, with a 2-channel input (bridge state + noisy waveform). - Training corrupts clean audio along a Brownian bridge between clean and noisy endpoints (sigma_max * sqrt(t(1-t)) noise, exact endpoints) and regresses the clean waveform with the existing L1 + multi-resolution STFT losses. - Sampling starts at the noisy waveform and iteratively re-projects onto the bridge; num_inference_steps=1 reduces to a single predictive forward pass, more steps buy generative refinement. Deterministic by default with optional stochastic sampling. - FlowMatchingLightningModule subclasses DenoisersLightningModule, reusing metrics, EMA, and logging; train.py selects it for the flowunet1d model key. - Register flowunet1d in MODELS/CONFIGS, export from package root, add model and lightning tests, and commit the research survey + design proposal under docs/. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01CnieBtBKfRzKAYpc7N7hWP
|
Bugbot is not enabled for your account, so this pull request was not reviewed. Enable Bugbot in the Cursor dashboard to get automatic reviews on future PRs. |
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01CnieBtBKfRzKAYpc7N7hWP
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
Adds the repo's first generative denoiser:
FlowUNet1D, a waveform-domain Schrödinger-bridge-style flow-matching model with a data-prediction objective. The full research survey and design rationale (diffusion vs flow/bridge vs GAN vs latent/codec, with citations) is committed asdocs/generative_denoiser_proposal.md.Key idea: training corrupts clean audio along a Brownian bridge between the clean and noisy waveforms (
x_t = (1−t)·x₀ + t·y + σ_max√(t(1−t))·ε) and the network regresses the clean endpoint using the existing L1 + multi-resolution STFT losses. Sampling starts at the noisy waveform and re-projects onto the bridge fornum_inference_stepsiterations — at 1 step it reduces exactly to a single predictive forward pass (verified by test), and more steps buy generative refinement.Changes
src/denoisers/modeling/flowunet1d/—FlowUNet1DConfig(model_type="flowunet1d"), sinusoidalTimeEmbedding, zero-initializedFiLM1Dconditioning,FlowUNet1Dbackbone reusing the existingunet1dDown/Mid/Up blocks with a 2-channel input (bridge state ‖ noisy), andFlowUNet1DModelwithpredict_clean()(training) and a deterministic-by-default bridge sampler inforward()(optional stochastic sampling).FlowMatchingLightningModule— subclass ofDenoisersLightningModuleinheriting metrics, DNSMOS, EMA, and W&B logging; overrides train/val steps for the bridge objective. Parent now readsautoencoderviagetattr(..., True)so configs without the field work.MODELS/CONFIGS, exported from the package root, andtrain.pyselects the flow module for theflowunet1dkey (train flowunet1d <name> <data_root>).unet1d-vctk-48khz, complex-STFT front-end).Test plan
pytest— new model + lightning tests pass (29 passed locally; the 3test_pretrainedHF-download tests require network and pass in CI)ruff check .cleanmypy .clean (40 files)🤖 Generated with Claude Code
https://claude.ai/code/session_01CnieBtBKfRzKAYpc7N7hWP
Generated by Claude Code