spectrogram_diffusion/__init__.py

"""Shim for this checkpoint's model_index.json.

The Google weights name the note/context encoders as library
`spectrogram_diffusion`. Diffusers' from_pretrained() does:

    library = importlib.import_module("spectrogram_diffusion")
    class_obj = getattr(library, "SpectrogramContEncoder")
    class_candidates = {name: getattr(library, name, None) for name in LOADABLE_CLASSES}

so this module must export both the encoder classes *and* ModelMixin
(otherwise loading raises: no from_pretrained on the component).
"""
from __future__ import annotations

from diffusers import ModelMixin

try:
    from diffusers.pipelines.deprecated.spectrogram_diffusion.continuous_encoder import (
        SpectrogramContEncoder,
    )
    from diffusers.pipelines.deprecated.spectrogram_diffusion.notes_encoder import (
        SpectrogramNotesEncoder,
    )
except ImportError:  # older Diffusers still used the non-deprecated path
    from diffusers.pipelines.spectrogram_diffusion.continuous_encoder import (  # type: ignore
        SpectrogramContEncoder,
    )
    from diffusers.pipelines.spectrogram_diffusion.notes_encoder import (  # type: ignore
        SpectrogramNotesEncoder,
    )

__all__ = ["ModelMixin", "SpectrogramContEncoder", "SpectrogramNotesEncoder"]