| 914 | ) |
| 915 | |
| 916 | def _build_modules(self, module_configs): |
| 917 | model_channels = self.hidden_size |
| 918 | time_embed_dim = self.time_embed_dim |
| 919 | self.time_embed = nn.Sequential( |
| 920 | linear(model_channels, time_embed_dim), |
| 921 | nn.SiLU(), |
| 922 | linear(time_embed_dim, time_embed_dim), |
| 923 | ) |
| 924 | |
| 925 | if self.num_classes is not None: |
| 926 | if isinstance(self.num_classes, int): |
| 927 | self.label_emb = nn.Embedding(self.num_classes, time_embed_dim) |
| 928 | elif self.num_classes == "continuous": |
| 929 | print("setting up linear c_adm embedding layer") |
| 930 | self.label_emb = nn.Linear(1, time_embed_dim) |
| 931 | elif self.num_classes == "timestep": |
| 932 | self.label_emb = nn.Sequential( |
| 933 | Timestep(model_channels), |
| 934 | nn.Sequential( |
| 935 | linear(model_channels, time_embed_dim), |
| 936 | nn.SiLU(), |
| 937 | linear(time_embed_dim, time_embed_dim), |
| 938 | ), |
| 939 | ) |
| 940 | elif self.num_classes == "sequential": |
| 941 | assert self.adm_in_channels is not None |
| 942 | self.label_emb = nn.Sequential( |
| 943 | nn.Sequential( |
| 944 | linear(self.adm_in_channels, time_embed_dim), |
| 945 | nn.SiLU(), |
| 946 | linear(time_embed_dim, time_embed_dim), |
| 947 | ) |
| 948 | ) |
| 949 | if self.zero_init_y_embed: |
| 950 | nn.init.constant_(self.label_emb[0][2].weight, 0) |
| 951 | nn.init.constant_(self.label_emb[0][2].bias, 0) |
| 952 | else: |
| 953 | raise ValueError() |
| 954 | |
| 955 | pos_embed_config = module_configs["pos_embed_config"] |
| 956 | self.add_mixin( |
| 957 | "pos_embed", |
| 958 | instantiate_from_config( |
| 959 | pos_embed_config, |
| 960 | height=self.latent_height // self.patch_size, |
| 961 | width=self.latent_width // self.patch_size, |
| 962 | compressed_num_frames=(self.num_frames - 1) // self.time_compressed_rate + 1, |
| 963 | hidden_size=self.hidden_size, |
| 964 | ), |
| 965 | reinit=True, |
| 966 | ) |
| 967 | |
| 968 | patch_embed_config = module_configs["patch_embed_config"] |
| 969 | self.add_mixin( |
| 970 | "patch_embed", |
| 971 | instantiate_from_config( |
| 972 | patch_embed_config, |
| 973 | patch_size=self.patch_size, |