kugelaudio_open/models/__init__.py

"""KugelAudio model components."""

from .kugelaudio_model import (
    KugelAudioModel,
    KugelAudioPreTrainedModel,
    KugelAudioForConditionalGeneration,
)
from .kugelaudio_inference import (
    KugelAudioForConditionalGenerationInference,
    KugelAudioCausalLMOutputWithPast,
    KugelAudioGenerationOutput,
)
from .tokenizer import (
    KugelAudioAcousticTokenizerModel,
    KugelAudioSemanticTokenizerModel,
    KugelAudioTokenizerEncoderOutput,
)
from .diffusion_head import KugelAudioDiffusionHead
from .conv_layers import (
    RMSNorm,
    ConvRMSNorm,
    ConvLayerNorm,
    SConv1d,
    SConvTranspose1d,
)

__all__ = [
    # Main models
    "KugelAudioModel",
    "KugelAudioPreTrainedModel",
    "KugelAudioForConditionalGeneration",
    "KugelAudioForConditionalGenerationInference",
    # Outputs
    "KugelAudioCausalLMOutputWithPast",
    "KugelAudioGenerationOutput",
    # Tokenizers
    "KugelAudioAcousticTokenizerModel",
    "KugelAudioSemanticTokenizerModel",
    "KugelAudioTokenizerEncoderOutput",
    # Components
    "KugelAudioDiffusionHead",
    "RMSNorm",
    "ConvRMSNorm",
    "ConvLayerNorm",
    "SConv1d",
    "SConvTranspose1d",
]