PlayDiffusion

During inference, PlayDiffusion uses an iterative decoding process starting from a fully masked token sequence. The process follows these steps:

1. Preliminary Prediction: The model generates an initial prediction ($\hat{X}_0$) conditioned on the current masked audio and textual input.
2. Confidence Scoring: Tokens are assigned confidence scores. Newly predicted (previously masked) tokens get a score based on their predicted probabilities. Tokens that were previously determined retain a confidence score of $1$.
3. Adaptive Remasking: Using a progressively decreasing schedule ($\gamma$), the lowest-confidence tokens are selected for remasking in the next iteration. The number of tokens to remask decreases over time, focusing refinement on high-uncertainty areas.

This iterative refinement continues until the decoding steps are complete, resulting in coherent audio.

You can containerize PlayDiffusion using Docker or Podman. Ensure you mount your Hugging Face and Whisper caches to avoid re-downloading models. The demo runs on port 7860.

Using Podman

podman build -t playdiffusion-py:latest .

podman run -it --rm \
  --device nvidia.com/gpu=all \
  --network=host \
  -v $HOME/.cache/huggingface:/app/.cache/huggingface \
  -v $HOME/.cache/whisper:$HOME/.cache/whisper \
  -p 7860:7860 \
  playdiffusion-py:latest

Using Docker

docker build -t playdiffusion-py:latest .

docker run -it --rm \
  --gpus all \
  --network=host \
  -v $HOME/.cache/huggingface:/app/.cache/huggingface \
  -v $HOME/.cache/whisper:$HOME/.cache/whisper \
  -p 7860:7860 \
  playdiffusion-py:latest

To use PlayDiffusion locally, you must set the OPENAI_API_KEY environment variable for Automatic Speech Recognition (ASR) and word timings. Follow these steps to set up a virtual environment and install the package with the demo extra:

1. Create a virtual environment: python3.11 -m venv .venv or uv venv
2. Activate the environment: source .venv/bin/activate
3. Install the package and demo dependencies: pip install '.[demo]' or uv sync --extra demo
4. Run the Gradio demo: python demo/gradio-demo.py or uv run demo/gradio-demo.py

# Install with pip
pip install '.[demo']

# Or with uv
uv sync --extra demo

# Run the demo
python demo/gradio-demo.py
# or
uv run demo/gradio-demo.py

When initializing SpeechTokenizer, you can specify the following:

• checkpoint (str): Path to the XLS-R encoder checkpoint. Defaults to "data/checkpoints/xlsr2_1b_v2_custom.pt".
• kmeans_layer_checkpoint (str): Path to the pre-trained Kmeans .npy file. Defaults to "data/checkpoints/kmeans_10k.npy".
• dtype (DataType): The precision for the encoder and kmeans layers. Defaults to torch.float16.
• device (Device): The target device.

When initializing SpeechEncoder, you can control the model depth and precision:

• checkpoint (str): Path to the XLS-R checkpoint. Defaults to "data/checkpoints/xlsr2_1b_v2_custom.pt".
• max_layer (int): The number of layers to load. If set (e.g., 35), the model will only load up to that layer and disable the final layer_norm to preserve intermediate representation quality.
• device (Device): The target device (e.g., torch.device('cuda')).
• dtype (DataType): The precision (e.g., torch.float32 or torch.float16).
• strict (bool): If max_layer is used, this is automatically set to False to allow loading partial state dicts.
• eval (bool): If True, sets the model to evaluation mode.

Use load_xlsr_encoder() to initialize a Wav2Vec2Model along with its corresponding Wav2Vec2Config and Wav2Vec2EncoderConfig. This function is designed to build the correct configurations for the XLS-R 1B v2 model architecture.

To optimize memory usage and avoid loading unnecessary weights, you can specify a max_layer parameter to truncate the number of transformer encoder layers.

from src.playdiffusion.models.speech_tokenizer.xlsr_encoder import load_xlsr_encoder
from fairseq2.typing import Device, DataType
import torch

# Define device and dtype
device = Device("cuda") if torch.cuda.is_available() else Device("cpu")
dtype = torch.float16

# Load the model, potentially limiting the number of layers to 35 (default)
model, config, encoder_config = load_xlsr_encoder(
    device=device, 
    dtype=dtype, 
    max_layer=35
)

The SpeechEncoder is a wrapper for the XLS-R 1B model designed to extract intermediate representations (latents) from a waveform. It allows loading only a subset of the model's layers (up to max_layer) to save memory and compute.

When using max_layer, the model automatically disables the final layer_norm to ensure the intermediate representations are not incorrectly normalized for the final layer's distribution.

from playdiffusion.models.speech_tokenizer.speech_tokenizer import SpeechEncoder
import torch

# Initialize encoder with a specific checkpoint and layer limit
encoder = SpeechEncoder(
    checkpoint="data/checkpoints/xlsr2_1b_v2_custom.pt",
    max_layer=35,
    device="cuda",
    dtype=torch.float32
)

# Input batch must be a SequenceBatch (from fairseq2)
# encoder(batch) returns (encoder_output, padding_mask)
latents, padding_mask = encoder(batch)

The Timer class provides a simple way to measure and print the elapsed time between consecutive calls or from the moment of initialization.

Usage:

1. Instantiate Timer().
2. Call the instance as a function with a description string: timer("task_name").
3. The timer prints the elapsed time in milliseconds (ms) and stores it in an internal dictionary.
4. Use get_times() to retrieve the dictionary of all recorded durations.

from playdiffusion.utils.audio_utils import Timer

timer = Timer()

# Perform some operation
dosomething()
timer("operation_1") # Prints: operation_1 time: X.X ms

# Perform another operation
dosomething_else()
timer("operation_2") # Prints: operation_2 time: X.X ms

# Retrieve all recorded times
all_times = timer.get_times()