sa_audio

class SAAudioVAEModel(nn.Module):
    """Pipeline-glue lazy loader around `OobleckVAE`."""
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

def named_parameters(self, prefix: str = "", recurse: bool = True):
    # Hide the lazy-loaded VAE — its weights are fetched separately
    # and shouldn't appear in the host pipeline's loader sweep.
    for name, param in super().named_parameters(prefix=prefix, recurse=recurse):
        if name.startswith("_oobleck_vae.") or name == "_oobleck_vae":
            continue
        yield name, param

def _build(self, device: torch.device | None = None):
    from fastvideo.models.vaes.oobleck import OobleckVAE

    path = self.pretrained_path
    if not path:
        raise ValueError("OobleckVAEConfig.pretrained_path must be set; expected "
                         "`stabilityai/stable-audio-open-1.0` or a local path.")
    dtype = getattr(torch, self.pretrained_dtype, torch.float32)
    # If the caller already pointed us at the VAE dir directly, drop
    # the subfolder. Otherwise pass through (default "vae").
    subfolder: str | None = self.pretrained_subfolder
    if subfolder and os.path.isdir(path) and os.path.isfile(os.path.join(path, "config.json")):
        subfolder = None
    model = OobleckVAE.from_pretrained(path, subfolder=subfolder, torch_dtype=dtype)
    if device is not None:
        model = model.to(device=device)
    model.eval()
    return model

@property
def oobleck_vae(self):
    if self._oobleck_vae is None:
        self._oobleck_vae = self._build()
    return self._oobleck_vae

# Back-compat alias: callers that imported this earlier referred to
# the underlying VAE as `sa_audio_vae_model`. Both names point at the
# same object.
@property
def sa_audio_vae_model(self):
    return self.oobleck_vae

@property
def hop_length(self) -> int:
    return int(self.oobleck_vae.hop_length)

def _move_to_input_device(self, model, ref: torch.Tensor):
    if ref is None:
        return model
    first_param = next(model.parameters(), None)
    if first_param is not None and first_param.device != ref.device:
        model = model.to(device=ref.device)
        self._oobleck_vae = model
    return model

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

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()</code></pre></details>
fastvideo.models.vaes.sa_audio.SAAudioVAEModel.named_parameters
named_parameters(prefix: str = '', recurse: bool = True)
Source
def named_parameters(self, prefix: str = "", recurse: bool = True):
    # Hide the lazy-loaded VAE — its weights are fetched separately
    # and shouldn't appear in the host pipeline's loader sweep.
    for name, param in super().named_parameters(prefix=prefix, recurse=recurse):
        if name.startswith("_oobleck_vae.") or name == "_oobleck_vae":
            continue
        yield name, param
fastvideo.models.vaes.sa_audio.SAAudioVAEModel.decode
decode(latent: torch.Tensor) -> torch.Tensor
Decode an audio latent ([B, C_latent, L]) -> waveform
([B, audio_channels, samples]).
Source
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: 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).
Source
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`).
&quot;&quot;&quot;
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, &quot;latent_dist&quot;):
    out = out.latent_dist
return out.sample() if sample_posterior else out.mode()</code></pre></details>