BitDance_B_4x/autoencoder/modeling_autoencoder.py

from __future__ import annotations

from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.modeling_utils import ModelMixin


class BitDanceImageNetAutoencoder(ModelMixin, ConfigMixin):
    @register_to_config
    def __init__(self, ddconfig=None, num_codebooks: int = 4, **kwargs):
        super().__init__()
        self.ddconfig = ddconfig
        self.num_codebooks = num_codebooks

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, *args, **kwargs):
        del pretrained_model_name_or_path, args, kwargs
        return cls()