oobleck ¶

@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 from `HF_TOKEN` /
    `HUGGINGFACE_HUB_TOKEN` / `HF_API_KEY` (see `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