Skip to content

Add FlowUNet1D generative bridge flow-matching denoiser - #61

Merged
will-rice merged 2 commits into
mainfrom
claude/dependency-updates-7x8xno
Jul 23, 2026
Merged

Add FlowUNet1D generative bridge flow-matching denoiser#61
will-rice merged 2 commits into
mainfrom
claude/dependency-updates-7x8xno

Conversation

@will-rice

Copy link
Copy Markdown
Owner

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 as docs/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 for num_inference_steps iterations — 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"), sinusoidal TimeEmbedding, zero-initialized FiLM1D conditioning, FlowUNet1D backbone reusing the existing unet1d Down/Mid/Up blocks with a 2-channel input (bridge state ‖ noisy), and FlowUNet1DModel with predict_clean() (training) and a deterministic-by-default bridge sampler in forward() (optional stochastic sampling).
  • FlowMatchingLightningModule — subclass of DenoisersLightningModule inheriting metrics, DNSMOS, EMA, and W&B logging; overrides train/val steps for the bridge objective. Parent now reads autoencoder via getattr(..., True) so configs without the field work.
  • Wiring — registered in MODELS/CONFIGS, exported from the package root, and train.py selects the flow module for the flowunet1d key (train flowunet1d <name> <data_root>).
  • Tests — config round-trip, sampler output shape, bridge-sigma endpoint behavior, single-step ≡ predictive-call equivalence, and Lightning train/val steps.
  • Docs — cited survey + design proposal, including 48 kHz risk analysis and fallbacks (StoRM-style regeneration seeded by the existing pretrained unet1d-vctk-48khz, complex-STFT front-end).

Test plan

  • pytest — new model + lightning tests pass (29 passed locally; the 3 test_pretrained HF-download tests require network and pass in CI)
  • ruff check . clean
  • mypy . clean (40 files)

🤖 Generated with Claude Code

https://claude.ai/code/session_01CnieBtBKfRzKAYpc7N7hWP


Generated by Claude Code

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
@cursor

cursor Bot commented Jul 23, 2026

Copy link
Copy Markdown

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
@will-rice
will-rice merged commit 6fb2fa7 into main Jul 23, 2026
6 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants