Skip to content

sa_audio

Lazy-loading pipeline wrapper around OobleckVAE.

Two reasons this exists rather than using OobleckVAE directly:

1. The underlying VAE is fetched on first encode/decode call, not at construction — lets pipelines build the module tree on CPU before knowing the target device. 2. The lazy VAE's params are hidden from named_parameters() so the FastVideo pipeline-component loader doesn't try to match Oobleck's safetensors against the host pipeline's converted-repo state dict.

For standalone use prefer OobleckVAE.from_pretrained(...) directly.

Classes

fastvideo.models.vaes.sa_audio.SAAudioVAEModel 

SAAudioVAEModel(config: OobleckVAEConfig)

Bases: Module

Pipeline-glue lazy loader around OobleckVAE.

Source code in fastvideo/models/vaes/sa_audio.py 
28
29
30
31
32
33
34
35
36
37
38
def __init__(self, config: OobleckVAEConfig) -> None:
    super().__init__()
    self.config = config
    arch = config.arch_config
    self.pretrained_path: str = config.pretrained_path
    self.pretrained_subfolder: str | None = config.pretrained_subfolder
    self.pretrained_dtype: str = config.pretrained_dtype
    self.sampling_rate: int = arch.sampling_rate
    self.audio_channels: int = arch.audio_channels
    self.decoder_input_channels: int = arch.decoder_input_channels
    self._oobleck_vae = None

Methods:

fastvideo.models.vaes.sa_audio.SAAudioVAEModel.decode 
decode(latent: Tensor) -> Tensor

Decode an audio latent ([B, C_latent, L]) -> waveform ([B, audio_channels, samples]).

Source code in fastvideo/models/vaes/sa_audio.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
def decode(self, latent: torch.Tensor) -> torch.Tensor:
    """Decode an audio latent (`[B, C_latent, L]`) -> waveform
    (`[B, audio_channels, samples]`).
    """
    model = self.oobleck_vae
    model = self._move_to_input_device(model, latent)
    with torch.no_grad():
        out = model.decode(latent.to(next(model.parameters()).dtype))
    if hasattr(out, "sample"):
        return out.sample
    return out
fastvideo.models.vaes.sa_audio.SAAudioVAEModel.encode 
encode(waveform: Tensor, sample_posterior: bool = False) -> Tensor

Encode [B, C_audio, samples] -> latent [B, C_latent, L].

sample_posterior=False (default): deterministic mean. sample_posterior=True: stochastic sample (mean + softplus(scale) * randn).

Source code in fastvideo/models/vaes/sa_audio.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
def encode(self, waveform: torch.Tensor, sample_posterior: bool = False) -> torch.Tensor:
    """Encode `[B, C_audio, samples]` -> latent `[B, C_latent, L]`.

    `sample_posterior=False` (default): deterministic mean.
    `sample_posterior=True`: stochastic sample (`mean + softplus(scale) * randn`).
    """
    model = self.oobleck_vae
    model = self._move_to_input_device(model, waveform)
    with torch.no_grad():
        out = model.encode(waveform.to(next(model.parameters()).dtype))
    if hasattr(out, "latent_dist"):
        out = out.latent_dist
    return out.sample() if sample_posterior else out.mode()