--- title: oobleck ---

Stable Audio Open 1.0 "Oobleck" VAE.

5-stage Conv1d autoencoder with Snake activations + diagonal-Gaussian bottleneck. Loads stabilityai/stable-audio-open-1.0/vae/ weights directly via OobleckVAE.from_pretrained(...).

vae = OobleckVAE.from_pretrained("stabilityai/stable-audio-open-1.0", subfolder="vae")
waveform = vae.decode(latent)            # (B, audio_channels, samples)
latent   = vae.encode(waveform).sample()  # or .mode()

fastvideo.models.vaes.oobleck.init_logger

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

fastvideo.models.vaes.oobleck.Snake1d

Snake1d(hidden_dim: int, logscale: bool = True)

A 1D Snake activation with learnable per-channel alpha/beta.

Source
class Snake1d(nn.Module):
    """A 1D Snake activation with learnable per-channel alpha/beta."""

    def __init__(self, hidden_dim: int, logscale: bool = True):
        super().__init__()
        self.alpha = nn.Parameter(torch.zeros(1, hidden_dim, 1))
        self.beta = nn.Parameter(torch.zeros(1, hidden_dim, 1))
        self.alpha.requires_grad = True
        self.beta.requires_grad = True
        self.logscale = logscale

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shape = x.shape
        alpha = self.alpha if not self.logscale else torch.exp(self.alpha)
        beta = self.beta if not self.logscale else torch.exp(self.beta)
        x = x.reshape(shape[0], shape[1], -1)
        x = x + (beta + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
        return x.reshape(shape)

fastvideo.models.vaes.oobleck.Snake1d.forward

forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    shape = x.shape
    alpha = self.alpha if not self.logscale else torch.exp(self.alpha)
    beta = self.beta if not self.logscale else torch.exp(self.beta)
    x = x.reshape(shape[0], shape[1], -1)
    x = x + (beta + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
    return x.reshape(shape)

fastvideo.models.vaes.oobleck.OobleckResidualUnit

OobleckResidualUnit(dimension: int = 16, dilation: int = 1)
Source
class OobleckResidualUnit(nn.Module):

    def __init__(self, dimension: int = 16, dilation: int = 1):
        super().__init__()
        pad = ((7 - 1) * dilation) // 2
        self.snake1 = Snake1d(dimension)
        self.conv1 = weight_norm(nn.Conv1d(
            dimension,
            dimension,
            kernel_size=7,
            dilation=dilation,
            padding=pad,
        ))
        self.snake2 = Snake1d(dimension)
        self.conv2 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=1))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.conv1(self.snake1(x))
        out = self.conv2(self.snake2(out))
        pad = (x.shape[-1] - out.shape[-1]) // 2
        if pad > 0:
            x = x[..., pad:-pad]
        return x + out

fastvideo.models.vaes.oobleck.OobleckResidualUnit.forward

forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    out = self.conv1(self.snake1(x))
    out = self.conv2(self.snake2(out))
    pad = (x.shape[-1] - out.shape[-1]) // 2
    if pad > 0:
        x = x[..., pad:-pad]
    return x + out

fastvideo.models.vaes.oobleck.OobleckEncoderBlock

OobleckEncoderBlock(input_dim: int, output_dim: int, stride: int = 1)
Source
class OobleckEncoderBlock(nn.Module):

    def __init__(self, input_dim: int, output_dim: int, stride: int = 1):
        super().__init__()
        self.res_unit1 = OobleckResidualUnit(input_dim, dilation=1)
        self.res_unit2 = OobleckResidualUnit(input_dim, dilation=3)
        self.res_unit3 = OobleckResidualUnit(input_dim, dilation=9)
        self.snake1 = Snake1d(input_dim)
        self.conv1 = weight_norm(
            nn.Conv1d(
                input_dim,
                output_dim,
                kernel_size=2 * stride,
                stride=stride,
                padding=math.ceil(stride / 2),
            ))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.res_unit1(x)
        x = self.res_unit2(x)
        x = self.snake1(self.res_unit3(x))
        return self.conv1(x)

fastvideo.models.vaes.oobleck.OobleckEncoderBlock.forward

forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    x = self.res_unit1(x)
    x = self.res_unit2(x)
    x = self.snake1(self.res_unit3(x))
    return self.conv1(x)

fastvideo.models.vaes.oobleck.OobleckDecoderBlock

OobleckDecoderBlock(input_dim: int, output_dim: int, stride: int = 1)
Source
class OobleckDecoderBlock(nn.Module):

    def __init__(self, input_dim: int, output_dim: int, stride: int = 1):
        super().__init__()
        self.snake1 = Snake1d(input_dim)
        self.conv_t1 = weight_norm(
            nn.ConvTranspose1d(
                input_dim,
                output_dim,
                kernel_size=2 * stride,
                stride=stride,
                padding=math.ceil(stride / 2),
            ))
        self.res_unit1 = OobleckResidualUnit(output_dim, dilation=1)
        self.res_unit2 = OobleckResidualUnit(output_dim, dilation=3)
        self.res_unit3 = OobleckResidualUnit(output_dim, dilation=9)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.snake1(x)
        x = self.conv_t1(x)
        x = self.res_unit1(x)
        x = self.res_unit2(x)
        return self.res_unit3(x)

fastvideo.models.vaes.oobleck.OobleckDecoderBlock.forward

forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    x = self.snake1(x)
    x = self.conv_t1(x)
    x = self.res_unit1(x)
    x = self.res_unit2(x)
    return self.res_unit3(x)

fastvideo.models.vaes.oobleck.OobleckDiagonalGaussianDistribution

OobleckDiagonalGaussianDistribution(parameters: torch.Tensor, deterministic: bool = False)

Diagonal-Gaussian VAE posterior with softplus(scale) + 1e-4 std.

Source
class OobleckDiagonalGaussianDistribution:
    """Diagonal-Gaussian VAE posterior with `softplus(scale) + 1e-4` std."""

    def __init__(self, parameters: torch.Tensor, deterministic: bool = False):
        self.parameters = parameters
        self.mean, self.scale = parameters.chunk(2, dim=1)
        self.std = nn.functional.softplus(self.scale) + 1e-4
        self.var = self.std * self.std
        self.logvar = torch.log(self.var)
        self.deterministic = deterministic

    def sample(self, generator: torch.Generator | None = None) -> torch.Tensor:
        noise = torch.randn(
            self.mean.shape,
            generator=generator,
            device=self.parameters.device,
            dtype=self.parameters.dtype,
        )
        return self.mean + self.std * noise

    def mode(self) -> torch.Tensor:
        return self.mean

fastvideo.models.vaes.oobleck.OobleckDiagonalGaussianDistribution.sample

sample(generator: torch.Generator | None = None) -> torch.Tensor
Source
def sample(self, generator: torch.Generator | None = None) -> torch.Tensor:
    noise = torch.randn(
        self.mean.shape,
        generator=generator,
        device=self.parameters.device,
        dtype=self.parameters.dtype,
    )
    return self.mean + self.std * noise

fastvideo.models.vaes.oobleck.OobleckDiagonalGaussianDistribution.mode

mode() -> torch.Tensor
Source
def mode(self) -> torch.Tensor:
    return self.mean

fastvideo.models.vaes.oobleck.OobleckDecoderOutput

OobleckDecoderOutput(sample: torch.Tensor)
Source
@dataclass
class OobleckDecoderOutput:
    sample: torch.Tensor

fastvideo.models.vaes.oobleck.OobleckEncoder

OobleckEncoder(encoder_hidden_size: int, audio_channels: int, downsampling_ratios: list[int], channel_multiples: list[int])
Source
class OobleckEncoder(nn.Module):

    def __init__(
        self,
        encoder_hidden_size: int,
        audio_channels: int,
        downsampling_ratios: list[int],
        channel_multiples: list[int],
    ):
        super().__init__()
        strides = downsampling_ratios
        channel_multiples = [1] + list(channel_multiples)
        self.conv1 = weight_norm(nn.Conv1d(
            audio_channels,
            encoder_hidden_size,
            kernel_size=7,
            padding=3,
        ))
        self.block = nn.ModuleList([
            OobleckEncoderBlock(
                input_dim=encoder_hidden_size * channel_multiples[i],
                output_dim=encoder_hidden_size * channel_multiples[i + 1],
                stride=s,
            ) for i, s in enumerate(strides)
        ])
        d_model = encoder_hidden_size * channel_multiples[-1]
        self.snake1 = Snake1d(d_model)
        self.conv2 = weight_norm(nn.Conv1d(
            d_model,
            encoder_hidden_size,
            kernel_size=3,
            padding=1,
        ))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv1(x)
        for m in self.block:
            x = m(x)
        x = self.snake1(x)
        return self.conv2(x)

fastvideo.models.vaes.oobleck.OobleckEncoder.forward

forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    x = self.conv1(x)
    for m in self.block:
        x = m(x)
    x = self.snake1(x)
    return self.conv2(x)

fastvideo.models.vaes.oobleck.OobleckDecoder

OobleckDecoder(channels: int, input_channels: int, audio_channels: int, upsampling_ratios: list[int], channel_multiples: list[int])
Source
class OobleckDecoder(nn.Module):

    def __init__(
        self,
        channels: int,
        input_channels: int,
        audio_channels: int,
        upsampling_ratios: list[int],
        channel_multiples: list[int],
    ):
        super().__init__()
        strides = upsampling_ratios
        channel_multiples = [1] + list(channel_multiples)
        self.conv1 = weight_norm(nn.Conv1d(
            input_channels,
            channels * channel_multiples[-1],
            kernel_size=7,
            padding=3,
        ))
        self.block = nn.ModuleList([
            OobleckDecoderBlock(
                input_dim=channels * channel_multiples[len(strides) - i],
                output_dim=channels * channel_multiples[len(strides) - i - 1],
                stride=s,
            ) for i, s in enumerate(strides)
        ])
        self.snake1 = Snake1d(channels)
        self.conv2 = weight_norm(nn.Conv1d(
            channels,
            audio_channels,
            kernel_size=7,
            padding=3,
            bias=False,
        ))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv1(x)
        for layer in self.block:
            x = layer(x)
        x = self.snake1(x)
        return self.conv2(x)

fastvideo.models.vaes.oobleck.OobleckDecoder.forward

forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    x = self.conv1(x)
    for layer in self.block:
        x = layer(x)
    x = self.snake1(x)
    return self.conv2(x)

fastvideo.models.vaes.oobleck.OobleckVAE

OobleckVAE(config = None, *, encoder_hidden_size: int = 128, downsampling_ratios: list[int] | None = None, channel_multiples: list[int] | None = None, decoder_channels: int = 128, decoder_input_channels: int = 64, audio_channels: int = 2, sampling_rate: int = 44100)

Stable Audio Open 1.0 VAE.

Constructed either from an OobleckVAEConfig (the standard VAELoader path) or from explicit kwargs (back-compat for tests and from_pretrained callers).

Source
class OobleckVAE(nn.Module):
    """Stable Audio Open 1.0 VAE.

    Constructed either from an `OobleckVAEConfig` (the standard
    `VAELoader` path) or from explicit kwargs (back-compat for tests
    and `from_pretrained` callers).
    """

    def __init__(
        self,
        config=None,  # type: OobleckVAEConfig | None
        *,
        encoder_hidden_size: int = 128,
        downsampling_ratios: list[int] | None = None,
        channel_multiples: list[int] | None = None,
        decoder_channels: int = 128,
        decoder_input_channels: int = 64,
        audio_channels: int = 2,
        sampling_rate: int = 44100,
    ):
        super().__init__()
        if config is not None:
            arch = config.arch_config
            encoder_hidden_size = arch.encoder_hidden_size
            downsampling_ratios = list(arch.downsampling_ratios)
            channel_multiples = list(arch.channel_multiples)
            decoder_channels = arch.decoder_channels
            decoder_input_channels = arch.decoder_input_channels
            audio_channels = arch.audio_channels
            sampling_rate = arch.sampling_rate
        if downsampling_ratios is None:
            downsampling_ratios = [2, 4, 4, 8, 8]
        if channel_multiples is None:
            channel_multiples = [1, 2, 4, 8, 16]
        self.encoder_hidden_size = encoder_hidden_size
        self.downsampling_ratios = downsampling_ratios
        self.decoder_channels = decoder_channels
        self.upsampling_ratios = list(reversed(downsampling_ratios))
        self.hop_length = int(np.prod(downsampling_ratios))
        self.sampling_rate = sampling_rate
        self.audio_channels = audio_channels
        self.decoder_input_channels = decoder_input_channels

        self.encoder = OobleckEncoder(
            encoder_hidden_size=encoder_hidden_size,
            audio_channels=audio_channels,
            downsampling_ratios=downsampling_ratios,
            channel_multiples=channel_multiples,
        )
        self.decoder = OobleckDecoder(
            channels=decoder_channels,
            input_channels=decoder_input_channels,
            audio_channels=audio_channels,
            upsampling_ratios=self.upsampling_ratios,
            channel_multiples=channel_multiples,
        )

    def encode(
        self,
        x: torch.Tensor,
    ) -> OobleckDiagonalGaussianDistribution:
        return OobleckDiagonalGaussianDistribution(self.encoder(x))

    def decode(self, z: torch.Tensor) -> OobleckDecoderOutput:
        return OobleckDecoderOutput(sample=self.decoder(z))

    def forward(
        self,
        sample: torch.Tensor,
        sample_posterior: bool = False,
    ) -> OobleckDecoderOutput:
        posterior = self.encode(sample)
        z = posterior.sample() if sample_posterior else posterior.mode()
        return self.decode(z)

    # -------------------------------------------------------------------
    # Loader
    # -------------------------------------------------------------------
    @classmethod
    def from_pretrained(
        cls,
        model_path: str,
        *,
        subfolder: str | None = None,
        torch_dtype: torch.dtype | None = None,
    ) -> "OobleckVAE":
        """Instantiate and load weights from a Stable Audio VAE dir.

        `model_path` may be:
          * a HF repo id (e.g. `stabilityai/stable-audio-open-1.0`),
          * a local directory containing `config.json` + safetensors,
          * a local directory whose `subfolder="vae"` holds those files.

        For gated repos, the HF token is read by `resolve_hf_token`.
        """
        import inspect
        from safetensors.torch import load_file

        from fastvideo.utils import resolve_hf_token

        # Resolve to a local directory.
        if os.path.isdir(model_path):
            root = model_path
        else:
            from huggingface_hub import snapshot_download
            allow = ["vae/*"] if subfolder else ["*"]
            root = snapshot_download(
                repo_id=model_path,
                token=resolve_hf_token(),
                allow_patterns=allow,
            )
        if subfolder:
            root = os.path.join(root, subfolder)
        if not os.path.isdir(root):
            raise FileNotFoundError(f"Not a directory: {root}")

        cfg_path = os.path.join(root, "config.json")
        if not os.path.isfile(cfg_path):
            raise FileNotFoundError(f"Expected config.json at {cfg_path}. If using a HF repo, "
                                    f"pass subfolder='vae'.")
        with open(cfg_path) as f:
            cfg = json.load(f)
        cfg_fields = {k: v for k, v in cfg.items() if not k.startswith("_")}
        # Diffusers configs commonly carry extra fields (`scaling_factor`,
        # `_diffusers_version`, ...) the bare `OobleckVAE` ctor doesn't accept.
        init_params = inspect.signature(cls.__init__).parameters
        cfg_fields = {k: v for k, v in cfg_fields.items() if k in init_params}

        model = cls(**cfg_fields)

        weights_path = os.path.join(root, "diffusion_pytorch_model.safetensors")
        if not os.path.isfile(weights_path):
            # Allow `model.safetensors` as a fallback.
            alt = os.path.join(root, "model.safetensors")
            if os.path.isfile(alt):
                weights_path = alt
            else:
                raise FileNotFoundError(f"No safetensors weights under {root}. Expected "
                                        f"diffusion_pytorch_model.safetensors.")
        state = load_file(weights_path)
        missing, unexpected = model.load_state_dict(state, strict=False)
        if missing:
            raise RuntimeError(f"OobleckVAE missing {len(missing)} keys from {weights_path}: "
                               f"{missing[:5]}")
        if unexpected:
            # Non-critical: some checkpoints embed the VAE inside a larger
            # container (e.g. `pretransform.model.*`). Log the count so
            # genuine loader regressions don't go unnoticed.
            logger.debug(
                "OobleckVAE: ignored %d unexpected keys from %s "
                "(first 3: %s)",
                len(unexpected),
                weights_path,
                unexpected[:3],
            )

        if torch_dtype is not None:
            model = model.to(dtype=torch_dtype)
        model.eval()
        return model

fastvideo.models.vaes.oobleck.OobleckVAE.encode

encode(x: torch.Tensor) -> OobleckDiagonalGaussianDistribution
Source
def encode(
    self,
    x: torch.Tensor,
) -> OobleckDiagonalGaussianDistribution:
    return OobleckDiagonalGaussianDistribution(self.encoder(x))

fastvideo.models.vaes.oobleck.OobleckVAE.decode

decode(z: torch.Tensor) -> OobleckDecoderOutput
Source
def decode(self, z: torch.Tensor) -> OobleckDecoderOutput:
    return OobleckDecoderOutput(sample=self.decoder(z))

fastvideo.models.vaes.oobleck.OobleckVAE.forward

forward(sample: torch.Tensor, sample_posterior: bool = False) -> OobleckDecoderOutput
Source
def forward(
    self,
    sample: torch.Tensor,
    sample_posterior: bool = False,
) -> OobleckDecoderOutput:
    posterior = self.encode(sample)
    z = posterior.sample() if sample_posterior else posterior.mode()
    return self.decode(z)

fastvideo.models.vaes.oobleck.OobleckVAE.from_pretrained

from_pretrained(model_path: str, *, subfolder: str | None = None, torch_dtype: torch.dtype | None = None) -> 'OobleckVAE'

Instantiate and load weights from a Stable Audio VAE dir.

model_path may be: * a HF repo id (e.g. stabilityai/stable-audio-open-1.0), * a local directory containing config.json + safetensors, * a local directory whose subfolder="vae" holds those files.

For gated repos, the HF token is read by resolve_hf_token.

Source
@classmethod
def from_pretrained(
    cls,
    model_path: str,
    *,
    subfolder: str | None = None,
    torch_dtype: torch.dtype | None = None,
) -> "OobleckVAE":
    """Instantiate and load weights from a Stable Audio VAE dir.

    `model_path` may be:
      * a HF repo id (e.g. `stabilityai/stable-audio-open-1.0`),
      * a local directory containing `config.json` + safetensors,
      * a local directory whose `subfolder="vae"` holds those files.

    For gated repos, the HF token is read by `resolve_hf_token`.
    """
    import inspect
    from safetensors.torch import load_file

    from fastvideo.utils import resolve_hf_token

    # Resolve to a local directory.
    if os.path.isdir(model_path):
        root = model_path
    else:
        from huggingface_hub import snapshot_download
        allow = ["vae/*"] if subfolder else ["*"]
        root = snapshot_download(
            repo_id=model_path,
            token=resolve_hf_token(),
            allow_patterns=allow,
        )
    if subfolder:
        root = os.path.join(root, subfolder)
    if not os.path.isdir(root):
        raise FileNotFoundError(f"Not a directory: {root}")

    cfg_path = os.path.join(root, "config.json")
    if not os.path.isfile(cfg_path):
        raise FileNotFoundError(f"Expected config.json at {cfg_path}. If using a HF repo, "
                                f"pass subfolder='vae'.")
    with open(cfg_path) as f:
        cfg = json.load(f)
    cfg_fields = {k: v for k, v in cfg.items() if not k.startswith("_")}
    # Diffusers configs commonly carry extra fields (`scaling_factor`,
    # `_diffusers_version`, ...) the bare `OobleckVAE` ctor doesn't accept.
    init_params = inspect.signature(cls.__init__).parameters
    cfg_fields = {k: v for k, v in cfg_fields.items() if k in init_params}

    model = cls(**cfg_fields)

    weights_path = os.path.join(root, "diffusion_pytorch_model.safetensors")
    if not os.path.isfile(weights_path):
        # Allow `model.safetensors` as a fallback.
        alt = os.path.join(root, "model.safetensors")
        if os.path.isfile(alt):
            weights_path = alt
        else:
            raise FileNotFoundError(f"No safetensors weights under {root}. Expected "
                                    f"diffusion_pytorch_model.safetensors.")
    state = load_file(weights_path)
    missing, unexpected = model.load_state_dict(state, strict=False)
    if missing:
        raise RuntimeError(f"OobleckVAE missing {len(missing)} keys from {weights_path}: "
                           f"{missing[:5]}")
    if unexpected:
        # Non-critical: some checkpoints embed the VAE inside a larger
        # container (e.g. `pretransform.model.*`). Log the count so
        # genuine loader regressions don't go unnoticed.
        logger.debug(
            "OobleckVAE: ignored %d unexpected keys from %s "
            "(first 3: %s)",
            len(unexpected),
            weights_path,
            unexpected[:3],
        )

    if torch_dtype is not None:
        model = model.to(dtype=torch_dtype)
    model.eval()
    return model