DOCSEARCHbyAlgolia
Submit Docs

BS-RoFormer Documentation

repository·main·Indexed 21 days ago

https://github.com/lucidrains/bs-roformer

A Band-Split Rotary Transformer implementation for music source separation. It features axial attention across frequency and time bands, supports stereo training and multiple stem outputs, and includes the BSRoformer, MelBandRoformer, and FlowBSRoformer (a Flow-Matching variant) model architectures.

Tokens
962
Snippets
4
Records
4
Agent score
25%

What's inside BS-RoFormer

  2. Use the MelBandRoformer model

    main

    The MelBandRoformer class implements the Mel-Band Roformer architecture. It follows the same API pattern as BSRoformer for training and inference.

    Arguments:

    • dim: Dimension of the model.
    • depth: Number of layers.
    • time_transformer_depth: Depth of the time transformer.
    • freq_transformer_depth: Depth of the frequency transformer.
    • use_pope: Boolean flag to enable a successor to rotary embeddings.
    import torch
from bs_roformer import MelBandRoformer

model = MelBandRoformer(
    dim = 32,
    depth = 1,
    time_transformer_depth = 1,
    freq_transformer_depth = 1,
    use_pope = False
)

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

# Training mode
loss = model(x, target = target)
loss.backward()

# Inference mode (after training)
out = model(x)

  3. Use the BSRoformer model

    main

    The BSRoformer class implements the Band Split Roformer architecture. It can be used for training by passing both input x and target audio to the model to compute a loss, or for inference by passing only x after training.

    Arguments:

    • dim: Dimension of the model.
    • depth: Number of layers.
    • time_transformer_depth: Depth of the time transformer.
    • freq_transformer_depth: Depth of the frequency transformer.
    • use_pope: Boolean flag to enable a successor to rotary embeddings.
    import torch
from bs_roformer import BSRoformer

model = BSRoformer(
    dim = 512,
    depth = 12,
    time_transformer_depth = 1,
    freq_transformer_depth = 1,
    use_pope = False
)

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

# Training mode
loss = model(x, target = target)
loss.backward()

# Inference mode (after training)
out = model(x)

  4. Use the FlowBSRoformer model

    main

    The FlowBSRoformer is a Flow-Matching variant of the BS-Roformer. Instead of masking, it predicts the flow between pure noise and the target audio.

    Note on Inference: Unlike the standard Roformer models which use a direct forward pass for inference, FlowBSRoformer uses a .sample(x) method to generate output from noise.

    Arguments:

    • dim: Dimension of the model.
    • depth: Number of layers.
    • time_transformer_depth: Depth of the time transformer.
    • freq_transformer_depth: Depth of the frequency transformer.
    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)

# Training mode
loss = model(x, target = target)
loss.backward()

# Inference mode (after training)
out = model.sample(x)