| from ..configuration_videobase import VideoBaseConfiguration | |
| from typing import Union, Tuple | |
| class VQVAEConfiguration(VideoBaseConfiguration): | |
| def __init__( | |
| self, | |
| embedding_dim: int = 256, | |
| n_codes: int = 2048, | |
| n_hiddens: int = 240, | |
| n_res_layers: int = 4, | |
| resolution: int = 128, | |
| sequence_length: int = 16, | |
| downsample: Union[Tuple[int, int, int], str] = (4, 4, 4), | |
| no_pos_embd: bool = True, | |
| **kwargs, | |
| ): | |
| super().__init__(**kwargs) | |
| self.embedding_dim = embedding_dim | |
| self.n_codes = n_codes | |
| self.n_hiddens = n_hiddens | |
| self.n_res_layers = n_res_layers | |
| self.resolution = resolution | |
| self.sequence_length = sequence_length | |
| if isinstance(downsample, str): | |
| self.downsample = tuple(map(int, downsample.split(","))) | |
| else: | |
| self.downsample = downsample | |
| self.no_pos_embd = no_pos_embd | |
| self.hidden_size = n_hiddens | |