### Install SoundStorm PyTorch

Source: https://github.com/lucidrains/soundstorm-pytorch/blob/main/README.md

Install the package via pip.

```bash
$ pip install soundstorm-pytorch
```

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

### Perform Generation with Acoustic Prompting

Source: https://context7.com/lucidrains/soundstorm-pytorch/llms.txt

Uses a prefix of acoustic tokens to guide the generation process for tasks like voice cloning.

```python
import torch
from soundstorm_pytorch import SoundStorm, ConformerWrapper

conformer = ConformerWrapper(
    codebook_size=1024,
    num_quantizers=12,
    conformer=dict(dim=512, depth=2),
)

model = SoundStorm(
    conformer,
    steps=18,
    schedule='cosine'
)

# Acoustic prompt tokens from reference audio (batch=1, prompt_len=128, quantizers=4)
# Using first 4 quantizers from a 128-token prompt
prompt_acoustic_token_ids = torch.randint(0, 1024, (1, 128, 4))

# Generate with acoustic prompting
generated = model.generate(
    num_latents=512,                              # Total sequence length to generate
    prompt_acoustic_token_ids=prompt_acoustic_token_ids,  # Voice prompt
    start_temperature=1.0,                        # Initial sampling temperature
    filter_thres=0.7,                             # Top-k filtering threshold
    noise_level_scale=1.0,                        # Noise scaling for token critic
    num_full_sampling_levels=1                    # Full sampling for first N quantizers
)
print(f"Generated with prompt: {generated.shape}")
```

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

### Initialize SoundStormTrainer for Distributed Training

Source: https://context7.com/lucidrains/soundstorm-pytorch/llms.txt

Sets up a complete training pipeline including dataset definition, trainer configuration, and checkpoint management.

```python
import torch
from torch.utils.data import Dataset
from soundstorm_pytorch import SoundStorm, ConformerWrapper, SoundStream
from soundstorm_pytorch import SoundStormTrainer

# Create model components
conformer = ConformerWrapper(
    codebook_size=1024,
    num_quantizers=12,
    conformer=dict(dim=512, depth=2),
)

soundstream = SoundStream(
    codebook_size=1024,
    rq_num_quantizers=12,
)

model = SoundStorm(conformer, soundstream=soundstream)

# Custom dataset returning (semantic_tokens, acoustic_tokens) tuples
class AudioDataset(Dataset):
    def __init__(self, size=1000):
        self.size = size

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        semantic_tokens = torch.randint(0, 20000, (256,))
        acoustic_tokens = torch.randint(0, 1024, (256, 12))
        return semantic_tokens, acoustic_tokens

dataset = AudioDataset()

# Initialize trainer with all configuration options
trainer = SoundStormTrainer(
    model=model,
    dataset=dataset,
    num_train_steps=100000,       # Total training steps
    num_warmup_steps=1000,        # Learning rate warmup steps
    batch_size=16,
    lr=3e-4,                      # Peak learning rate
    initial_lr=1e-5,              # Initial warmup learning rate
    grad_accum_every=4,           # Gradient accumulation steps
    max_grad_norm=0.5,            # Gradient clipping
    valid_frac=0.05,              # Validation split fraction
    save_results_every=100,       # Validation logging frequency
    save_model_every=1000,        # Checkpoint saving frequency
    results_folder='./results',
    accelerate_kwargs=dict(
        mixed_precision='fp16'    # Enable mixed precision training
    )
)

# Start training
trainer.train()

# Save and load checkpoints
trainer.save('./checkpoint.pt')
trainer.load('./checkpoint.pt', restore_optimizer=True)

# Generate samples during or after training
samples = trainer.generate(seconds=5, batch_size=4)
```

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

### Configure Conformer and ConformerWrapper

Source: https://context7.com/lucidrains/soundstorm-pytorch/llms.txt

Initializes a Conformer model with specific architectural parameters and wraps it with quantizer-aware layers for token processing.

```python
conformer = Conformer(
    dim=512,                      # Model dimension
    depth=6,                      # Number of Conformer blocks
    dim_head=64,                  # Dimension per attention head
    heads=8,                      # Number of attention heads
    ff_mult=4,                    # Feedforward expansion multiplier
    conv_expansion_factor=2,      # Convolution module expansion
    conv_kernel_size=31,          # Convolution kernel size
    attn_dropout=0.1,             # Attention dropout rate
    ff_dropout=0.1,               # Feedforward dropout rate
    conv_dropout=0.1,             # Convolution dropout rate
    attn_flash=True,              # Use Flash Attention for efficiency
    num_residual_streams=4        # Hyper-connection residual streams
)

# Wrap with quantizer-aware embedding and output layers
wrapper = ConformerWrapper(
    codebook_size=1024,
    num_quantizers=12,
    conformer=conformer,
    grouped_quantizers=1          # Group quantizers for efficiency
)

# Forward pass with token input
codes = torch.randint(0, 1024, (2, 1024 * 12))  # Flattened (batch, seq * quantizers)
logits = wrapper(codes)
print(f"Output logits shape: {logits.shape}")

# Get embeddings instead of logits
embeddings = wrapper(codes, return_embeddings=True)
print(f"Embeddings shape: {embeddings.shape}")
```

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

### ConformerWrapper Configuration

Source: https://context7.com/lucidrains/soundstorm-pytorch/llms.txt

Demonstrates the basic initialization of ConformerWrapper, the neural network backbone for SoundStorm. It handles multi-quantizer audio tokens and embedding projection.

```python
import torch
from soundstorm_pytorch import ConformerWrapper, Conformer

```

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

### Initialize SoundStorm Model and Generate Codes

Source: https://context7.com/lucidrains/soundstorm-pytorch/llms.txt

Initializes the ConformerWrapper and SoundStorm model for generating audio codes. Requires pre-encoded codebook IDs. Use for training or direct code generation.

```python
import torch
from soundstorm_pytorch import SoundStorm, ConformerWrapper

# Create the Conformer-based neural network
conformer = ConformerWrapper(
    codebook_size=1024,      # Size of the audio codec vocabulary
    num_quantizers=12,       # Number of residual vector quantizers
    conformer=dict(
        dim=512,             # Model dimension
        depth=2              # Number of Conformer blocks
    ),
)

# Initialize SoundStorm with the conformer
model = SoundStorm(
    conformer,
    steps=18,                # Number of iterative demasking steps (as in MaskGiT)
    schedule='cosine'        # Masking schedule: 'cosine' or 'linear'
)

# Pre-encoded codebook IDs from SoundStream (batch=2, seq_len=1024, quantizers=12)
codes = torch.randint(0, 1024, (2, 1024, 12))

# Forward pass returns loss and breakdown
loss, loss_breakdown = model(codes)
loss.backward()

# Generate new audio codes
generated = model.generate(1024, batch_size=2)  # Returns (2, 1024) tensor
print(f"Generated shape: {generated.shape}")
```

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

### Configure Self-Conditioning and Token Critic

Source: https://context7.com/lucidrains/soundstorm-pytorch/llms.txt

Enable self-conditioning and the token critic to improve generation quality and sampling decisions.

```python
import torch
from soundstorm_pytorch import SoundStorm, ConformerWrapper

conformer = ConformerWrapper(
    codebook_size=1024,
    num_quantizers=12,
    conformer=dict(dim=512, depth=2),
)

# Enable self-conditioning and token critic
model = SoundStorm(
    conformer,
    steps=18,
    schedule='cosine',
    self_cond=True,                    # Enable self-conditioning
    self_cond_train_prob=0.75,         # Probability of self-cond during training
    self_token_critic=True,            # Enable learned token critic
    critic_loss_weight=1.0,            # Weight for critic loss
    no_replace_prob=0.15,              # Tokens that stay unchanged (BERT-style)
    random_token_prob=0.1,             # Random token replacement probability
    can_mask_prev_unmasked=True        # Allow remasking during generation
)

# Training returns both generator and critic losses
codes = torch.randint(0, 1024, (2, 1024, 12))
total_loss, (generator_loss, critic_loss) = model(codes)
print(f"Generator loss: {generator_loss:.4f}, Critic loss: {critic_loss:.4f}")
total_loss.backward()

# Train only generator or only critic
gen_loss, _ = model(codes, only_train_generator=True)
critic_loss, _ = model(codes, only_train_critic=True)
```

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

### Train SoundStorm with Pre-encoded Codes

Source: https://github.com/lucidrains/soundstorm-pytorch/blob/main/README.md

Initialize the model with a Conformer wrapper and train using pre-encoded codebook IDs from a SoundStream model.

```python
import torch
from soundstorm_pytorch import SoundStorm, ConformerWrapper

conformer = ConformerWrapper(
    codebook_size = 1024,
    num_quantizers = 12,
    conformer = dict(
        dim = 512,
        depth = 2
    ),
)

model = SoundStorm(
    conformer,
    steps = 18,          # 18 steps, as in original maskgit paper
    schedule = 'cosine'  # currently the best schedule is cosine
)

# get your pre-encoded codebook ids from the soundstream from a lot of raw audio

codes = torch.randint(0, 1024, (2, 1024, 12)) # (batch, seq, num residual VQ)

# do the below in a loop for a ton of data

loss, _ = model(codes)
loss.backward()

# model can now generate in 18 steps. ~2 seconds sounds reasonable

generated = model.generate(1024, batch_size = 2) # (2, 1024)
```

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

### Manage Model Checkpoints

Source: https://context7.com/lucidrains/soundstorm-pytorch/llms.txt

Save and load model states using PyTorch's standard serialization methods.

```python
import torch
from soundstorm_pytorch import SoundStorm, ConformerWrapper

conformer = ConformerWrapper(
    codebook_size=1024,
    num_quantizers=12,
    conformer=dict(dim=512, depth=2),
)

model = SoundStorm(conformer, steps=18, schedule='cosine')

# Train the model...
codes = torch.randint(0, 1024, (2, 1024, 12))
for _ in range(100):
    loss, _ = model(codes)
    loss.backward()

# Save model checkpoint
checkpoint = {
    'model': model.state_dict(),
    'config': {
        'codebook_size': 1024,
        'num_quantizers': 12,
        'steps': 18
    }
}
torch.save(checkpoint, './soundstorm_checkpoint.pt')

# Load model checkpoint
loaded_pkg = model.load('./soundstorm_checkpoint.pt', strict=True)
print("Model loaded successfully")

# Generate with loaded model
model.eval()
with torch.no_grad():
    generated = model.generate(512, batch_size=1)
```

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

### Train SoundStorm on Raw Audio

Source: https://github.com/lucidrains/soundstorm-pytorch/blob/main/README.md

Pass a pretrained SoundStream model into SoundStorm to train directly on raw audio waveforms.

```python
import torch
from soundstorm_pytorch import SoundStorm, ConformerWrapper, Conformer, SoundStream

conformer = ConformerWrapper(
    codebook_size = 1024,
    num_quantizers = 12,
    conformer = dict(
        dim = 512,
        depth = 2
    ),
)

soundstream = SoundStream(
    codebook_size = 1024,
    rq_num_quantizers = 12,
    attn_window_size = 128,
    attn_depth = 2
)

model = SoundStorm(
    conformer,
    soundstream = soundstream   # pass in the soundstream
)

# find as much audio you'd like the model to learn

audio = torch.randn(2, 10080)

# course it through the model and take a gazillion tiny steps

loss, _ = model(audio)
loss.backward()

# and now you can generate state-of-the-art speech

generated_audio = model.generate(seconds = 30, batch_size = 2)  # generate 30 seconds of audio (it will calculate the length in seconds based off the sampling frequency and cumulative downsamples in the soundstream passed in above)
```

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

### Text-to-Speech with SPEAR-TTS Integration

Source: https://context7.com/lucidrains/soundstorm-pytorch/llms.txt

Enables text-to-speech by integrating SoundStorm with SPEAR-TTS for text-to-semantic token conversion. Generates speech from text input.

```python
import torch
from soundstorm_pytorch import SoundStorm, ConformerWrapper, SoundStream
from spear_tts_pytorch import TextToSemantic

# Initialize text-to-semantic model
text_to_semantic = TextToSemantic(
    dim=512,
    source_depth=12,
    target_depth=12,
    num_text_token_ids=50000,
    num_semantic_token_ids=20000,
    use_openai_tokenizer=True
)

# Load pretrained weights
text_to_semantic.load('/path/to/trained/model.pt')

# Create Conformer and SoundStream
conformer = ConformerWrapper(
    codebook_size=1024,
    num_quantizers=12,
    conformer=dict(dim=512, depth=2),
)

soundstream = SoundStream(
    codebook_size=1024,
    rq_num_quantizers=12,
    attn_window_size=128,
    attn_depth=2
)

# Create SoundStorm with TTS capability
model = SoundStorm(
    conformer,
    soundstream=soundstream,
    spear_tts_text_to_semantic=text_to_semantic
).cuda()

# Generate speech from text
generated_speech = model.generate(
    texts=[
        'the rain in spain stays mainly in the plain',
        'the quick brown fox jumps over the lazy dog'
    ]
)  # Returns (2, n) raw waveform tensor
```

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

### Integrate Text-to-Semantic Transformer

Source: https://github.com/lucidrains/soundstorm-pytorch/blob/main/README.md

Load a trained TextToSemantic model and pass it into the SoundStorm constructor for text-to-speech tasks.

```python
from spear_tts_pytorch import TextToSemantic

text_to_semantic = TextToSemantic(
    dim = 512,
    source_depth = 12,
    target_depth = 12,
    num_text_token_ids = 50000,
    num_semantic_token_ids = 20000,
    use_openai_tokenizer = True
)

# load the trained text-to-semantic transformer

text_to_semantic.load('/path/to/trained/model.pt')

# pass it into the soundstorm

model = SoundStorm(
    conformer,
    soundstream = soundstream,
    spear_tts_text_to_semantic = text_to_semantic
).cuda()
```

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

### Generate Speech with SoundStorm Model

Source: https://github.com/lucidrains/soundstorm-pytorch/blob/main/README.md

Use this snippet to generate speech from a list of text inputs. The output is a raw waveform decoded from SoundStream.

```python
generated_speech = model.generate(
    texts = [
        'the rain in spain stays mainly in the plain',
        'the quick brown fox jumps over the lazy dog'
    ]
) # (2, n) - raw waveform decoded from soundstream
```

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

### Train with Variable Length Sequences

Source: https://context7.com/lucidrains/soundstorm-pytorch/llms.txt

Use masking and a padding token ID to handle batches containing sequences of different lengths.

```python
import torch
from soundstorm_pytorch import SoundStorm, ConformerWrapper

conformer = ConformerWrapper(
    codebook_size=1024,
    num_quantizers=12,
    conformer=dict(dim=512, depth=2),
)

model = SoundStorm(
    conformer,
    steps=18,
    schedule='cosine',
    pad_id=-1  # Padding token ID for variable length
)

# Batch with different sequence lengths (padded to max length)
batch_size, max_seq_len, num_quantizers = 4, 512, 12
codes = torch.randint(0, 1024, (batch_size, max_seq_len, num_quantizers))

# Mark padded positions with pad_id
codes[0, 400:, :] = -1  # First sequence: length 400
codes[1, 350:, :] = -1  # Second sequence: length 350
codes[2, 512:, :] = -1  # Third sequence: full length 512
codes[3, 256:, :] = -1  # Fourth sequence: length 256

# Create attention mask (True for valid positions)
mask = torch.ones(batch_size, max_seq_len, dtype=torch.bool)
mask[0, 400:] = False
mask[1, 350:] = False
mask[3, 256:] = False

# Training with mask
loss, _ = model(codes, mask=mask)
loss.backward()

# Generation with mask
generated = model.generate(
    num_latents=512,
    batch_size=2,
    mask=torch.ones(2, 512, dtype=torch.bool)  # All positions valid
)
```

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

### Train SoundStorm with Raw Audio and SoundStream

Source: https://context7.com/lucidrains/soundstorm-pytorch/llms.txt

Trains SoundStorm directly on raw audio waveforms by integrating a pretrained SoundStream codec. Generates raw waveform output.

```python
import torch
from soundstorm_pytorch import SoundStorm, ConformerWrapper, SoundStream

# Create the Conformer wrapper
conformer = ConformerWrapper(
    codebook_size=1024,
    num_quantizers=12,
    conformer=dict(
        dim=512,
        depth=2
    ),
)

# Initialize SoundStream audio codec
soundstream = SoundStream(
    codebook_size=1024,
    rq_num_quantizers=12,
    attn_window_size=128,
    attn_depth=2
)

# Create SoundStorm with integrated SoundStream
model = SoundStorm(
    conformer,
    soundstream=soundstream  # Pass SoundStream for raw audio training
)

# Train on raw audio waveforms
audio = torch.randn(2, 10080)  # Batch of 2 audio samples
loss, _ = model(audio)
loss.backward()

# Generate audio directly (returns raw waveform)
generated_audio = model.generate(
    seconds=30,              # Generate 30 seconds of audio
    batch_size=2
)
print(f"Generated audio shape: {generated_audio.shape}")
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.