BS-RoFormer

https://github.com/lucidrains/bs-roformer
Admin
BS-RoFormer implements the Band Split Roformer, a state-of-the-art attention network for music...
Tokens:40,987
Snippets:189
Trust Score:10
License:MIT
Update:1 month ago
Show doc for...
Tokens:
Raw
### Install BS-RoFormer with pip

Source: https://github.com/lucidrains/bs-roformer/blob/main/README.md

Install the BS-RoFormer package from PyPI using pip.

```bash
$ pip install BS-RoFormer
```

--------------------------------

### Basic manual attention example

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/attend.md

Shows how to instantiate Attend with dropout and flash=False, then run attention on random tensors. The output shape is (2, 8, 64, 64).

```python
import torch
from bs_roformer.attend import Attend

attend = Attend(dropout = 0.1, flash = False)

q = torch.randn(2, 8, 64, 64)  # b, h, i, d
k = torch.randn(2, 8, 128, 64) # b, h, j, d
v = torch.randn(2, 8, 128, 64)

out = attend(q, k, v)
print(out.shape)  # torch.Size([2, 8, 64, 64])
```

--------------------------------

### Basic training usage of MelBandRoformer

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/mel-band-roformer.md

Shows how to instantiate a MelBandRoformer model, compute loss with a target, and get output shape after training. Uses dim=32, depth=1, and disables use_pope.

```python
import torch
from bs_roformer import MelBandRoformer

model = MelBandRoformer(
    dim = 32,
    depth = 1,
    time_transformer_depth = 1,
    freq_transformer_depth = 1,
    use_pope = False # turn this on to try a successor to rotary embeddings
)

x = torch.randn(2, 352800)
target = torch.randn(2, 352800)

loss = model(x, target = target)
loss.backward()

# after much training
out = model(x)
print(out.shape)  # torch.Size([2, 1, 352800])
```

--------------------------------

### Loss breakdown with multiple stems

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/mel-band-roformer.md

Shows how to get the loss breakdown when return_loss_breakdown=True with num_stems=2. The model returns a tuple (total_loss, (loss, multi_stft_resolution_loss)).

```python
import torch
from bs_roformer import MelBandRoformer

model = MelBandRoformer(dim = 128, depth = 2, num_stems = 2)

x = torch.randn(1, 1, 352800)
target = torch.randn(1, 2, 1, 352800)

total, (l1, multi_stft) = model(x, target = target, return_loss_breakdown = True)
print(total.item(), l1.item(), multi_stft.item())
```

--------------------------------

### LinearAttention.__init__

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/mel-band-roformer.md

Initializes the LinearAttention module with configurable dimensions, heads, scale, flash attention, dropout, and value residual.

```APIDOC
## LinearAttention.__init__

### Description
Initializes the LinearAttention module with the given parameters.

### Method
Constructor

### Parameters
- **dim** (int) - Required - Model dimension (keyword-only)
- **dim_head** (int) - Optional, default `32` - Head dimension
- **heads** (int) - Optional, default `8` - Number of heads
- **scale** (float) - Optional, default `8` - Fixed attention scale passed to `Attend(scale = scale)`
- **flash** (bool) - Optional, default `False` - Flash flag passed to `Attend`
- **dropout** (float) - Optional, default `0.` - Attention dropout
- **add_value_residual** (bool) - Optional, default `False` - Builds `learned_value_residual_mix` (`Linear(dim, heads) → Rearrange('b n h -> b h 1 n') → Sigmoid`) and lerps `v` toward `value_residual`

### Example
```python
attention = LinearAttention(
    dim=512,
    dim_head=32,
    heads=8,
    scale=8,
    flash=False,
    dropout=0.1,
    add_value_residual=True
)
```
```

--------------------------------

### sample method signature with @torch.no_grad()

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/flow-bs-roformer.md

Decorated with @torch.no_grad(). Accepts raw_audio of shape (b, t) or (b, c, t) and an optional steps parameter defaulting to 10. Returns a tensor shaped (b, 1, t) for num_stems == 1 (or (b, s, t) if mono), else (b, n, s, t).

```python
@torch.no_grad()
def sample(
    self,
    raw_audio,
    steps = 10
):
```

--------------------------------

### Error handling for band size mismatch

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/bs-roformer.md

Demonstrates an AssertionError when the sum of freqs_per_bands does not match the STFT bin count (n_fft=2048 gives 1025 bins). The example catches the error and prints the message.

```python
import torch
from bs_roformer import BSRoformer

try:
    model = BSRoformer(dim = 32, depth = 1, freqs_per_bands = (4, 4))  # sums to 8, but n_fft=2048 gives 1025 bins
except AssertionError as e:
    # 'the number of freqs in the bands must equal 1025 based on the STFT settings, but got 8'
    print(e)
```

--------------------------------

### sample(raw_audio, steps = 10)

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/flow-bs-roformer.md

Generates separated audio stems from a raw audio mixture using Euler integration. The method initializes noise, iteratively denoises it using the model's forward pass, and returns the final separated audio.

```APIDOC
## sample(raw_audio, steps = 10)

### Description
Generates separated audio stems from a raw audio mixture using Euler integration. The method initializes noise, iteratively denoises it using the model's forward pass, and returns the final separated audio.

### Method
`sample`

### Signature
```python
@torch.no_grad()
def sample(
    self,
    raw_audio,
    steps = 10
)
```

### Parameters
- **raw_audio** (torch.Tensor) - Required - Input mix, shape `(b, t)` or `(b, c, t)` (channel count must match `stereo`)
- **steps** (int) - Optional, default `10` - Number of Euler integration steps

### Returns
- **torch.Tensor** - Shape `(b, 1, t)` for `num_stems == 1` (or `(b, s, t)` if mono), else `(b, n, s, t)`.

### Algorithm
1. `self.eval()` (and remains in eval mode afterwards unless the caller restores training mode).
2. Initial noise: `noise = torch.randn((batch, *stems_dim, channels, time_len), device=device) * self.noise_std + self.noise_mean`; `stems_dim = (num_stems,)` if `num_stems > 1` else `()`; `audio = noise`.
3. `times = torch.linspace(0., 1., steps + 1, device=device)`, `delta = 1. / steps`.
4. Denoising loop over `times[:-1]`: `pred_target = self.forward(raw_audio, noised_target = audio, times = time_batch)`; flow estimate `pred_flow = (pred_target - audio) / (1. - time)` (where `time` was padded via `pad_right_ndim_to(time, audio.ndim)`); Euler step `audio = audio + pred_flow * delta`.
```

--------------------------------

### forward(raw_audio, target = None, return_loss_breakdown = False)

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/bs-roformer.md

Processes raw audio through the BS-RoFormer model. If a target is provided, computes the total loss (L1 + multi-resolution STFT loss). Optionally returns a breakdown of the loss components.

```APIDOC
## forward(raw_audio, target = None, return_loss_breakdown = False)

### Description
Processes raw audio through the BS-RoFormer model. If a target is provided, computes the total loss (L1 + multi-resolution STFT loss). Optionally returns a breakdown of the loss components.

### Method
forward

### Signature
```python
def forward(
    self,
    raw_audio,
    target = None,
    return_loss_breakdown = False
):
```

### Parameters
- **raw_audio** (torch.Tensor) - Required - Raw audio. `(b, t)` mono is reshaped to `(b, 1, t)`. Channel dim must match `stereo`: 1 channel requires `stereo=False`, 2 channels require `stereo=True`.
- **target** (torch.Tensor | None) - Optional - Ground-truth separated audio for loss computation. 2D `(b, t)` is treated as `(b, 1, t)`; when `num_stems > 1` it must be `(b, n, s, t)` with `n == num_stems`.
- **return_loss_breakdown** (bool) - Optional - If `True` (and `target` given), return `(total_loss, (l1_loss, multi_stft_resolution_loss))` instead of just `total_loss`.

### Returns
- If `target is None`: `torch.Tensor` of shape `(b, s, t)` (reconstructed/separated audio) for `num_stems == 1`, or `(b, n, s, t)` for `num_stems > 1`. Spatial dim is dropped on mono: `recon_audio` is rearranged `'b n s t -> b s t'` when `num_stems == 1`.
- If `target` given: `torch.Tensor` scalar — `total_loss = l1_loss + multi_stft_resolution_loss * multi_stft_resolution_loss_weight`.
- If `target` given and `return_loss_breakdown=True`: `(total_loss, (loss, multi_stft_resolution_loss))` where `loss` is the raw L1 and `multi_stft_resolution_loss` is the unweighted multi-resolution STFT L1 sum.

### Raises
- **AssertionError** - `raw_audio` channel count / `stereo` mismatch.
- **AssertionError** - `num_stems > 1` and `target.ndim != 4` or `target.shape[1] != num_stems`.

### Forward Pipeline
1. `raw_audio` 2D → `(b, 1, t)`. Assert channel count matches `stereo`.
2. STFT: `pack_one(raw_audio, '* t')` → `torch.stft(..., return_complex=True)` → `torch.view_as_real` → unpack to `(b, s, f, t, c)`; slice `self.freq_slice` on the frequency axis; rearrange to `(b, (f s), t, c)` merging channel into frequency; reshape to `b t (f c)`.
3. `x = self.band_split(x)` → features `(b, t, n_bands, dim)`.
4. `x = self.expand_stream(x)` (hyper-connections expand; identity when `num_residual_streams == 1`).
5. Per layer group: rearrange `b t f d <-> b f t d`, run the time `Transformer` (with `time_v_residual`), then the freq `Transformer` (with `freq_v_residual`); first layer produces the value residuals used by subsequent layers (value residual learning).
6. `x = self.reduce_stream(x)`, then `self.final_norm(x)` (an `RMSNorm(dim)`).
7. Masks: `mask = torch.stack([fn(x) for fn in self.mask_estimators], dim=1)` → rearrange to `b n f t c` (c=2) → `view_as_complex` → complex-multiplication with the STFT representation → ISTFT with `freq_pad` padding, optional `zero_dc` DC bin zeroing.
8. Reconstructed audio reshaped to `(b, 1, s, t)`, then to `(b, s, t)` when `num_stems == 1`.
9. Loss path: `loss = F.l1_loss(recon_audio, target)` plus, for each `window_size` in `multi_stft_resolutions_window_sizes`, an STFT with `n_fft = max(window_size, self.multi_stft_n_fft)`, `win_length = window_size`, `hop_length = multi_stft_hop_size`, `normalized = multi_stft_normalized`, `window = self.multi_stft_window_fn(window_size, device=device)`; `multi_stft_resolution_loss` is the sum of `F.l1_loss(recon_Y, target_Y)` over resolutions. `target` is truncated with `target[..., :recon_audio.shape[-1]]` to guard against the length lost on ISTFT.
```

--------------------------------

### Catching beartype and AssertionError in BSRoformer construction

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/errors.md

Shows how to catch AssertionError for a configuration error (freqs_per_bands length <= 1) and BeartypeCallHintParamViolation for a type violation (dim as a string). Requires beartype installed and bs_roformer imported.

```python
from beartype import BeartypeCallHintParamViolation
from bs_roformer import BSRoformer

try:
    model = BSRoformer(dim = 512, depth = 1, freqs_per_bands = (128,))  # len <= 1 -> AssertionError
except AssertionError:
    pass

try:
    model = BSRoformer(dim = "not an int", depth = 1)  # beartype violation
    # raises BeartypeCallHintParamViolation (a TypeError subclass)
except BeartypeCallHintParamViolation as e:
    print("beartype rejected the argument:", e)
```

--------------------------------

### BSRoformer.__init__

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/bs-roformer.md

Constructs a BSRoformer model. The constructor accepts a feature dimension and depth as positional arguments, with all other options passed by keyword. It validates frequency range and band splits, and initializes internal components.

```APIDOC
## BSRoformer.__init__

### Description
Constructs a BSRoformer model. The constructor accepts a feature dimension and depth as positional arguments, with all other options passed by keyword. It validates frequency range and band splits, and initializes internal components.

### Method
Constructor

### Signature
```python
class BSRoformer(Module):
    @beartype
    def __init__(
        self,
        dim,
        *,
        depth,
        stereo = False,
        num_stems = 1,
        time_transformer_depth = 2,
        freq_transformer_depth = 2,
        freqs_per_bands: tuple[int, ...] = DEFAULT_FREQS_PER_BANDS,
        freq_range: tuple[int, int] | None = None,
        dim_head = 64,
        heads = 8,
        attn_dropout = 0.,
        ff_dropout = 0.,
        flash_attn = True,
        num_residual_streams = 4,
        num_residual_fracs = 1,
        mc_hyper_conn_sinkhorn_iters = 2,
        dim_freqs_in = 1025,
        stft_n_fft = 2048,
        stft_hop_length = 512,
        stft_win_length = 2048,
        stft_normalized = False,
        zero_dc = False,
        stft_window_fn: Callable | None = None,
        mask_estimator_depth = 2,
        multi_stft_resolution_loss_weight = 1.,
        multi_stft_resolutions_window_sizes: tuple[int, ...] = (4096, 2048, 1024, 512, 256),
        multi_stft_hop_size = 147,
        multi_stft_normalized = False,
        multi_stft_window_fn: Callable = torch.hann_window,
        use_pope = False
    ):
```

### Parameters
#### Positional Parameters
- **dim** (int) - Required - Feature dimension of the band embeddings
- **depth** (int) - Required - Number of layer groups; each group holds a time Transformer and a freq Transformer

#### Keyword Parameters
- **stereo** (bool) - Optional - If `True`, expects 2-channel input; `audio_channels = 2`
- **num_stems** (int) - Optional - Number of source stems to separate (one `MaskEstimator` per stem)
- **time_transformer_depth** (int) - Optional - Depth of time transformer
- **freq_transformer_depth** (int) - Optional - Depth of frequency transformer
- **freqs_per_bands** (tuple[int, ...]) - Optional - Band split sizes; `sum()` must equal the number of STFT frequency bins after `freq_range` slicing
- **freq_range** (tuple[int,int] | None) - Optional - `(min_freq, max_freq)` slice into the STFT frequency axis; `-1` means 0 (min) or `freqs` (max); `None` means `(-1, -1)` = full range
- **dim_head** (int) - Optional - Dimension of attention heads
- **heads** (int) - Optional - Number of attention heads
- **attn_dropout** (float) - Optional - Dropout for attention
- **ff_dropout** (float) - Optional - Dropout for feed-forward
- **flash_attn** (bool) - Optional - Whether to use flash attention
- **num_residual_streams** (int) - Optional - Number of residual streams; set to 1 to disable hyper connections
- **num_residual_fracs** (int) - Optional - Alternative to residual streams for memory efficiency
- **mc_hyper_conn_sinkhorn_iters** (int) - Optional - Sinkhorn iterations for hyper connections
- **dim_freqs_in** (int) - Optional - Dimension of frequency input
- **stft_n_fft** (int) - Optional - STFT FFT size
- **stft_hop_length** (int) - Optional - STFT hop length
- **stft_win_length** (int) - Optional - STFT window length
- **stft_normalized** (bool) - Optional - Whether STFT is normalized
- **zero_dc** (bool) - Optional - Whether to zero the DC component
- **stft_window_fn** (Callable | None) - Optional - STFT window function
- **mask_estimator_depth** (int) - Optional - Depth of mask estimator
- **multi_stft_resolution_loss_weight** (float) - Optional - Weight for multi-resolution STFT loss
- **multi_stft_resolutions_window_sizes** (tuple[int, ...]) - Optional - Window sizes for multi-resolution STFT
- **multi_stft_hop_size** (int) - Optional - Hop size for multi-resolution STFT
- **multi_stft_normalized** (bool) - Optional - Whether multi-resolution STFT is normalized
- **multi_stft_window_fn** (Callable) - Optional - Window function for multi-resolution STFT
- **use_pope** (bool) - Optional - Use PoPE embeddings instead of rotary (RoPE) embeddings

### Constructor Side Effects and Validation
- Computes `freqs` by running `torch.stft(torch.randn(1, 4096), **self.stft_kwargs, return_complex=True).shape[1]`.
- Resolves `freq_range`: `min_freq = 0 if min_freq == -1 else min_freq`, `max_freq = freqs if max_freq == -1 else max_freq`; asserts `min_freq >= 0 and max_freq <= freqs and min_freq < max_freq`.
- Asserts `len(freqs_per_bands) > 1` and `sum(freqs_per_bands) == freqs`.
- Computes `freqs_per_bands_with_complex = tuple(2 * f * self.audio_channels for f in freqs_per_bands)`.
- Sets `self.zero_dc = min_freq == 0 and zero_dc`.
- Registers one `BandSplit` and a `ModuleList` of `num_stems` `MaskEstimator`s.
- Initializes hyper-connection functions.
- Sets `self.stft_window_fn` and stores STFT kwargs.

### Raises
| Error type | Condition |
|-----------|-----------|
| `AssertionError` | `freq_range` values out of bounds or `min_freq >= max_freq` |
| `AssertionError` | `len(freqs_per_bands) <= 1` |
| `AssertionError` | `sum(freqs_per_bands) != freqs` |
| `beartype.TypeCheckFailure` | `dim`, `depth`, `freqs_per_bands`, `freq_range` or `multi_stft_resolutions_window_sizes` mismatch their annotations |
```

--------------------------------

### Stereo multi-stem training with loss breakdown

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/bs-roformer.md

Configures BSRoformer for stereo input/output and separates two stems (e.g., vocals + accompaniment). Uses return_loss_breakdown=True to obtain per-component losses (l1 and multi_stft) alongside the total loss.

```python
import torch
from bs_roformer import BSRoformer

model = BSRoformer(
    dim = 512,
    depth = 8,
    num_stems = 2,       # separate two stems, e.g. vocals + accompaniment
    stereo = True,       # stereo input/output
    freq_range = (0, -1) # full frequency range
)

x = torch.randn(2, 2, 352800)          # (batch, channels, samples)
target = torch.randn(2, 2, 2, 352800)  # (batch, stems, channels, samples)

total_loss, (l1, multi_stft) = model(x, target = target, return_loss_breakdown = True)

print(total_loss)       # scalar tensor
print(l1, multi_stft)   # per-component losses
```

--------------------------------

### MelBandRoformer.__init__

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/mel-band-roformer.md

Constructs a MelBandRoformer model. The constructor initializes the mel filter bank, band splitting, mask estimators, and hyper-connections based on the provided parameters.

```APIDOC
## MelBandRoformer.__init__

### Description
Constructs a MelBandRoformer model. The constructor initializes the mel filter bank, band splitting, mask estimators, and hyper-connections based on the provided parameters.

### Method
Constructor

### Signature
```python
class MelBandRoformer(Module):

    @beartype
    def __init__(
        self,
        dim,
        *,
        depth,
        stereo = False,
        num_stems = 1,
        time_transformer_depth = 2,
        freq_transformer_depth = 2,
        linear_transformer_depth = 1,
        num_bands = 60,
        dim_head = 64,
        heads = 8,
        attn_dropout = 0.1,
        ff_dropout = 0.1,
        flash_attn = True,
        linear_flash_attn = None,
        dim_freqs_in = 1025,
        sample_rate = 44100,
        stft_n_fft = 2048,
        stft_hop_length = 512,
        stft_win_length = 2048,
        stft_normalized = False,
        stft_window_fn: Callable | None = None,
        zero_dc = False,
        mask_estimator_depth = 1,
        multi_stft_resolution_loss_weight = 1.,
        multi_stft_resolutions_window_sizes: tuple[int, ...] = (4096, 2048, 1024, 512, 256),
        multi_stft_hop_size = 147,
        multi_stft_normalized = False,
        multi_stft_window_fn: Callable = torch.hann_window,
        match_input_audio_length = False,
        add_value_residual = True,
        num_residual_streams = 4,
        num_residual_fracs = 1,
        use_pope = False
    ):
```

### Parameters
#### Positional Parameters
- **dim** (int) - Required - Model dimension.
- **depth** (int) - Required - Depth of the model.

#### Keyword Parameters
- **stereo** (bool) - Optional - Default: `False` - Whether to process stereo audio.
- **num_stems** (int) - Optional - Default: `1` - Number of stems to separate.
- **time_transformer_depth** (int) - Optional - Default: `2` - Depth of the time transformer.
- **freq_transformer_depth** (int) - Optional - Default: `2` - Depth of the frequency transformer.
- **linear_transformer_depth** (int) - Optional - Default: `1` - Depth of the linear attention transformer; `0` disables it.
- **num_bands** (int) - Optional - Default: `60` - Number of mel bands.
- **dim_head** (int) - Optional - Default: `64` - Dimension of attention heads.
- **heads** (int) - Optional - Default: `8` - Number of attention heads.
- **attn_dropout** (float) - Optional - Default: `0.1` - Dropout for attention.
- **ff_dropout** (float) - Optional - Default: `0.1` - Dropout for feed-forward.
- **flash_attn** (bool) - Optional - Default: `True` - Whether to use flash attention.
- **linear_flash_attn** (bool | None) - Optional - Default: `None` - Flash flag for linear attention; defaults to `flash_attn` when `None`.
- **dim_freqs_in** (int) - Optional - Default: `1025` - Number of input frequency bins.
- **sample_rate** (int) - Optional - Default: `44100` - Sample rate for the mel filter bank.
- **stft_n_fft** (int) - Optional - Default: `2048` - FFT size for STFT.
- **stft_hop_length** (int) - Optional - Default: `512` - Hop length for STFT.
- **stft_win_length** (int) - Optional - Default: `2048` - Window length for STFT.
- **stft_normalized** (bool) - Optional - Default: `False` - Whether to normalize STFT.
- **stft_window_fn** (Callable | None) - Optional - Default: `None` - Window function for STFT.
- **zero_dc** (bool) - Optional - Default: `False` - Whether to zero the DC component.
- **mask_estimator_depth** (int) - Optional - Default: `1` - Depth of the mask estimator.
- **multi_stft_resolution_loss_weight** (float) - Optional - Default: `1.` - Weight for multi-resolution STFT loss.
- **multi_stft_resolutions_window_sizes** (tuple[int, ...]) - Optional - Default: `(4096, 2048, 1024, 512, 256)` - Window sizes for multi-resolution STFT.
- **multi_stft_hop_size** (int) - Optional - Default: `147` - Hop size for multi-resolution STFT.
- **multi_stft_normalized** (bool) - Optional - Default: `False` - Whether to normalize multi-resolution STFT.
- **multi_stft_window_fn** (Callable) - Optional - Default: `torch.hann_window` - Window function for multi-resolution STFT.
- **match_input_audio_length** (bool) - Optional - Default: `False` - If `True`, pads output to match input length.
- **add_value_residual** (bool) - Optional - Default: `True` - Enables value residual learning.
- **num_residual_streams** (int) - Optional - Default: `4` - Number of residual streams.
- **num_residual_fracs** (int) - Optional - Default: `1` - Number of residual fractions.
- **use_pope** (bool) - Optional - Default: `False` - Whether to use PoPE.

### Constructor Behavior
- No `freq_range` option; the full frequency axis is used.
- Hyper-connections are initialized with `mc_get_init_and_expand_reduce_stream_functions(num_residual_streams)`.
- Mel filter bank built with librosa; first and last band edges are adjusted.
- `freqs_per_band = mel_filter_bank > 0`; asserts all frequencies are covered.
- Computes `freq_indices` for gathering STFT rows; for stereo, interleaves channels.
- Registers non-persistent buffers: `freq_indices`, `freqs_per_band`, `num_freqs_per_band`, `num_bands_per_freq`.
- Creates one `BandSplit` and `num_stems` `MaskEstimator`s.
- Stores `multi_stft_hop_length`, `multi_stft_normalized`, and `match_input_audio_length`.

### Raises
- **AssertionError** - If some frequency is covered by no band.
- **beartype.TypeCheckFailure** - If constructor arguments mismatch their annotations.

```

--------------------------------

### Attend.__init__

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/attend.md

Constructor for the Attend module. Initializes dropout, flash attention flag, and optional scale. Validates PyTorch version when flash is enabled and sets up CPU/CUDA flash attention configurations.

```APIDOC
## Attend.__init__

### Description
Initializes the Attend module with dropout, flash attention flag, and optional scale. Validates PyTorch version when flash is enabled and sets up CPU/CUDA flash attention configurations.

### Method
Constructor

### Signature
```python
class Attend(nn.Module):
    def __init__(
        self,
        dropout = 0.,
        flash = False,
        scale = None
    ):
```

### Parameters
- **dropout** (float) - Optional - Dropout probability applied to attention weights (via `nn.Dropout`); only active in training mode. Default: `0.`
- **flash** (bool) - Optional - If `True`, use `F.scaled_dot_product_attention` (flash/math/mem-efficient) instead of the manual einsum attention path. Default: `False`
- **scale** (float | None) - Optional - Explicit attention scale; if set, it overrides the default `q.shape[-1] ** -0.5` head-dimension scaling. Default: `None`

### Behavior / Side Effects
- When `flash=True` and `torch.__version__` < 2.0.0, raises `AssertionError` ("in order to use flash attention, you must be using pytorch 2.0 or above").
- `self.cpu_config = FlashAttentionConfig(True, True, True)` unconditionally.
- If CUDA is unavailable or `flash` is `False`, `self.cuda_config` stays `None` (a guard `if not torch.cuda.is_available() or not flash: return`).
- On an A100 (`device_properties.major == 8 and device_properties.minor == 0`), `cuda_config = FlashAttentionConfig(True, False, False)` and a one-time message "A100 GPU detected, using flash attention if input tensor is on cuda" is printed.
- Otherwise `cuda_config = FlashAttentionConfig(False, True, True)` with one-time message "Non-A100 GPU detected, using math or mem efficient attention if input tensor is on cuda".

### Instance Attributes
- `scale`, `dropout`, `attn_dropout` (`nn.Dropout`), `flash` (bool), `cpu_config`, `cuda_config`.

### Raises
- **AssertionError** - If `flash=True` and PyTorch version is < 2.0.
```

--------------------------------

### Import MelBandRoformer from top-level and submodule

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/mel-band-roformer.md

Shows two equivalent ways to import the `MelBandRoformer` class: from the package top level and from the submodule.

```python
from bs_roformer import MelBandRoformer                       # top-level export
from bs_roformer.mel_band_roformer import MelBandRoformer      # same class, submodule import
```

--------------------------------

### Training loop usage (mono, single stem)

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/bs-roformer.md

Basic usage of BSRoformer for mono audio with a single stem. Initializes the model with specified dimensions and transformer depths, performs a forward pass to compute loss, and demonstrates inference output shape.

```python
import torch
from bs_roformer import BSRoformer

model = BSRoformer(
    dim = 512,
    depth = 12,
    time_transformer_depth = 1,
    freq_transformer_depth = 1,
    use_pope = False # turn this on to try a successor to rotary embeddings
)

x = torch.randn(2, 352800)
target = torch.randn(2, 352800)

loss = model(x, target = target)
loss.backward()

# after much training
out = model(x)
print(out.shape)  # torch.Size([2, 1, 352800])
```

--------------------------------

### Import BS-RoFormer classes from package top-level

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/index.md

Shows the three classes re-exported from the package's __init__.py: BSRoformer, MelBandRoformer, and FlowBSRoformer. Use these imports to access the main model classes.

```python
from bs_roformer.bs_roformer import BSRoformer
from bs_roformer.mel_band_roformer import MelBandRoformer
from bs_roformer.flow_bs_roformer import FlowBSRoformer
```

--------------------------------

### Training with flow matching

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/flow-bs-roformer.md

Basic training loop for FlowBSRoformer. Creates a model with specified dimensions, generates random input and target tensors, computes the loss, and performs a backward pass. After training, samples the model and prints the output shape.

```python
import torch
from bs_roformer import FlowBSRoformer

model = FlowBSRoformer(
    dim = 512,
    depth = 12,
    time_transformer_depth = 1,
    freq_transformer_depth = 1
)

x = torch.randn(2, 352800)
target = torch.randn(2, 352800)

loss = model(x, target = target)
loss.backward()

# after much training
out = model.sample(x)
print(out.shape)  # torch.Size([2, 1, 352800])
```

--------------------------------

### forward method definition

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/bs-roformer.md

Defines the forward method of the BS-RoFormer model. It accepts raw audio, optional target, and a flag to return loss breakdown. The docstring explains the einops notation used in the implementation.

```python
def forward(
    self,
    raw_audio,
    target = None,
    return_loss_breakdown = False
):
    """
    einops

    b - batch
    f - freq
    t - time
    s - audio channel (1 for mono, 2 for stereo)
    n - number of 'stems'
    c - complex (2)
    d - feature dimension
    """
```

--------------------------------

### forward method signature

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/attend.md

Method signature for forward, with docstring explaining einstein notation for batch, heads, sequence lengths, and feature dimension.

```python
def forward(self, q, k, v):
    """
    einstein notation
    b - batch
    h - heads
    n, i, j - sequence length (base sequence length, source, target)
    d - feature dimension
    """
```

--------------------------------

### FlowBSRoformer.__init__

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/flow-bs-roformer.md

Constructor for the FlowBSRoformer model. Initializes the model with specified dimensions, depth, and flow-specific options.

```APIDOC
## FlowBSRoformer.__init__

### Description
Initializes a FlowBSRoformer instance. This constructor sets up the model architecture, including time conditioning, band splitting, and per-stem STFT estimators.

### Signature
```python
class FlowBSRoformer(Module):
    @beartype
    def __init__(
        self,
        dim,
        *,
        depth,
        stereo = False,
        num_stems = 1,
        time_transformer_depth = 2,
        freq_transformer_depth = 2,
        freqs_per_bands: tuple[int, ...] = DEFAULT_FREQS_PER_BANDS,
        freq_range: tuple[int, int] | None = None,
        dim_head = 64,
        heads = 8,
        attn_dropout = 0.,
        ff_dropout = 0.,
        flash_attn = True,
        num_residual_streams = 4,
        num_residual_fracs = 1,
        mc_hyper_conn_sinkhorn_iters = 2,
        stft_n_fft = 2048,
        stft_hop_length = 512,
        stft_win_length = 2048,
        stft_normalized = False,
        zero_dc = False,
        stft_window_fn: Callable | None = None,
        stft_estimator_depth = 2,
        multi_stft_resolution_loss_weight = 1.,
        multi_stft_resolutions_window_sizes: tuple[int, ...] = (4096, 2048, 1024, 512, 256),
        multi_stft_hop_size = 147,
        multi_stft_normalized = False,
        multi_stft_window_fn: Callable = torch.hann_window,
        use_pope = False,
        noise_std = 0.1,
        noise_mean = 0.0,
        time_to_loss_weight_fn: Callable | None = None
    ):
```

### Parameters

#### Positional Parameters
- **dim** (int) - Required - The model dimension.
- **depth** (int) - Required - The depth of the model.

#### Keyword Parameters
- **stereo** (bool) - Optional - Whether to use stereo audio. Default: `False`.
- **num_stems** (int) - Optional - Number of stems to separate. Default: `1`.
- **time_transformer_depth** (int) - Optional - Depth of the time transformer. Default: `2`.
- **freq_transformer_depth** (int) - Optional - Depth of the frequency transformer. Default: `2`.
- **freqs_per_bands** (tuple[int, ...]) - Optional - Frequencies per band. Default: `DEFAULT_FREQS_PER_BANDS`.
- **freq_range** (tuple[int, int] | None) - Optional - Frequency range as (min, max). `-1` implies 0 and inf. Default: `None`.
- **dim_head** (int) - Optional - Dimension per attention head. Default: `64`.
- **heads** (int) - Optional - Number of attention heads. Default: `8`.
- **attn_dropout** (float) - Optional - Attention dropout rate. Default: `0.`.
- **ff_dropout** (float) - Optional - Feed-forward dropout rate. Default: `0.`.
- **flash_attn** (bool) - Optional - Whether to use flash attention. Default: `True`.
- **num_residual_streams** (int) - Optional - Number of residual streams. Set to 1 to disable hyper connections. Default: `4`.
- **num_residual_fracs** (int) - Optional - Alternative to residual streams for memory efficiency. Default: `1`.
- **mc_hyper_conn_sinkhorn_iters** (int) - Optional - Number of Sinkhorn iterations for hyper connections. Default: `2`.
- **stft_n_fft** (int) - Optional - STFT FFT size. Default: `2048`.
- **stft_hop_length** (int) - Optional - STFT hop length. Default: `512`.
- **stft_win_length** (int) - Optional - STFT window length. Default: `2048`.
- **stft_normalized** (bool) - Optional - Whether to normalize STFT. Default: `False`.
- **zero_dc** (bool) - Optional - Whether to zero the DC component. Default: `False`.
- **stft_window_fn** (Callable | None) - Optional - STFT window function. Default: `None`.
- **stft_estimator_depth** (int) - Optional - Depth of each per-stem STFTEstimator MLP. Default: `2`.
- **multi_stft_resolution_loss_weight** (float) - Optional - Weight for multi-resolution STFT loss. Default: `1.`.
- **multi_stft_resolutions_window_sizes** (tuple[int, ...]) - Optional - Window sizes for multi-resolution STFT. Default: `(4096, 2048, 1024, 512, 256)`.
- **multi_stft_hop_size** (int) - Optional - Hop size for multi-resolution STFT. Default: `147`.
- **multi_stft_normalized** (bool) - Optional - Whether to normalize multi-resolution STFT. Default: `False`.
- **multi_stft_window_fn** (Callable) - Optional - Window function for multi-resolution STFT. Default: `torch.hann_window`.
- **use_pope** (bool) - Optional - Whether to use Pope. Default: `False`.
- **noise_std** (float) - Optional - Standard deviation of Gaussian noise used in flow-matching interpolation. Default: `0.1`.
- **noise_mean** (float) - Optional - Mean of the noise distribution. Default: `0.0`.
- **time_to_loss_weight_fn** (Callable | None) - Optional - Callable taking per-example sampled times and returning per-example loss weights. Default: `None`.

### Behavior
- Sets `self.time_cond_dim = dim * 4` and builds `self.time_mlp` as a sequential MLP with sinusoidal positional embedding.
- Sets `self.stft_channels = self.audio_channels * (1 + self.num_stems)`.
- Initializes hyper-connection streams using `mc_get_init_and_expand_reduce_stream_functions`.
- Validates `freq_range` and `freqs_per_bands` with assertions.
- Builds band split input widths using `stft_channels` and STFTEstimator output widths using `audio_channels`.
- Creates a `ModuleList` of `STFTEstimator` per stem.
- Sets `self.zero_dc` based on `min_freq == 0 and zero_dc`.
- Registers a non-persistent buffer `zero` for multi-resolution loss accumulator.
- Sets `self.final_norm = RMSNorm(dim)`.

### Raises
- `AssertionError` if `freq_range` or `freqs_per_bands` are invalid (same as `BSRoformer`).
- Beartype validation errors for type mismatches.
```

--------------------------------

### Import BSRoformer from top-level and submodule

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/bs-roformer.md

Shows the two equivalent import paths for the BSRoformer class: the top-level package export and the direct submodule import. Use either to access the model class.

```python
from bs_roformer import BSRoformer            # top-level export
from bs_roformer.bs_roformer import BSRoformer # same class, submodule import
```

--------------------------------

### Attend.forward(q, k, v)

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/attend.md

Main forward pass of the Attend module. Computes attention either via flash attention or manual einsum path, returning the aggregated values.

```APIDOC
## Attend.forward(q, k, v)

### Description
Computes attention output. If `flash` is enabled, delegates to `flash_attn`; otherwise uses manual einsum with softmax and dropout.

### Method
forward

### Signature
```python
def forward(self, q, k, v):
    """
    einstein notation
    b - batch
    h - heads
    n, i, j - sequence length (base sequence length, source, target)
    d - feature dimension
    """
```

### Parameters
- **q** (torch.Tensor) - Required - Queries tensor of shape `(b, h, n, d)`.
- **k** (torch.Tensor) - Required - Keys tensor of shape `(b, h, j, d)`.
- **v** (torch.Tensor) - Required - Values tensor of shape `(b, h, j, d)`.

### Returns
- **torch.Tensor** - Attention output of shape `(b, h, i, d)` — aggregated values.

### Logic
1. `scale = default(self.scale, q.shape[-1] ** -0.5)`.
2. If `self.flash`, delegate to `flash_attn`.
3. Otherwise compute `sim = einsum("b h i d, b h j d -> b h i j", q, k) * scale`, then `attn = sim.softmax(dim=-1)`, apply `self.attn_dropout`, and aggregate `out = einsum("b h i j, b h j d -> b h i d", attn, v)`.

### Raises
- **AssertionError** - If `flash=True` with PyTorch < 2.0.
```

--------------------------------

### Attend.flash_attn(q, k, v)

Source: https://github.com/lucidrains/bs-roformer/blob/main/_autodocs/api-reference/attend.md

Internal method that applies flash attention using PyTorch's scaled_dot_product_attention. It pre-multiplies queries by scale if set, selects the appropriate backend config, and calls the attention function.

```APIDOC
## Attend.flash_attn(q, k, v)

### Description
Applies flash attention using `F.scaled_dot_product_attention`. This method is not normally called directly; `forward` routes here when `self.flash` is `True`.

### Method
flash_attn

### Signature
```python
def flash_attn(self, q, k, v):
```

### Parameters
- **q** (torch.Tensor) - Required - Queries tensor of shape `(b, h, n, d)`.
- **k** (torch.Tensor) - Required - Keys tensor of shape `(b, h, j, d)`.
- **v** (torch.Tensor) - Required - Values tensor of shape `(b, h, j, d)`.

### Behavior
- Applies `self.scale` by pre-multiplying `q` with `self.scale / (q.shape[-1] ** -0.5)` if a scale is set.
- Selects `self.cuda_config` when `q.is_cuda` else `self.cpu_config`.
- Calls `F.scaled_dot_product_attention(q, k, v, dropout_p = self.dropout if self.training else 0.)` inside `torch.backends.cuda.sdp_kernel(**config._asdict())`.

### Returns
- **torch.Tensor** - Attention output of shape `(b, h, i, d)`.
```
Tokens:
Raw