| 168 | # ---------------------------- |
| 169 | |
| 170 | class TAEHV(nn.Module): |
| 171 | image_channels = 3 |
| 172 | def __init__( |
| 173 | self, |
| 174 | checkpoint_path="taehv.pth", |
| 175 | decoder_time_upscale=(True, True), |
| 176 | decoder_space_upscale=(True, True, True), |
| 177 | channels = [256, 128, 64, 64], |
| 178 | latent_channels = 16 |
| 179 | ): |
| 180 | """Initialize TAEHV (decoder-only) with built-in deepening after every ReLU. |
| 181 | Deepening config: how_many_each=1, k=3 (fixed as requested). |
| 182 | """ |
| 183 | super().__init__() |
| 184 | self.latent_channels = latent_channels |
| 185 | n_f = channels |
| 186 | self.frames_to_trim = 2**sum(decoder_time_upscale) - 1 |
| 187 | |
| 188 | # Build the decoder "skeleton" |
| 189 | base_decoder = nn.Sequential( |
| 190 | Clamp(), conv(self.latent_channels, n_f[0]), nn.ReLU(inplace=True), |
| 191 | |
| 192 | MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), |
| 193 | nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), |
| 194 | TGrow(n_f[0], 1), |
| 195 | conv(n_f[0], n_f[1], bias=False), |
| 196 | |
| 197 | MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), |
| 198 | nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), |
| 199 | TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), |
| 200 | conv(n_f[1], n_f[2], bias=False), |
| 201 | |
| 202 | MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), |
| 203 | nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1), |
| 204 | TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1), |
| 205 | conv(n_f[2], n_f[3], bias=False), |
| 206 | |
| 207 | nn.ReLU(inplace=True), conv(n_f[3], TAEHV.image_channels), |
| 208 | ) |
| 209 | |
| 210 | # Inline deepening: insert (IdentityConv2d(k=3) + ReLU) after every ReLU |
| 211 | self.decoder = self._apply_identity_deepen(base_decoder, how_many_each=1, k=3) |
| 212 | |
| 213 | self.pixel_shuffle = PixelShuffle3d(4, 8, 8) |
| 214 | |
| 215 | if checkpoint_path is not None: |
| 216 | missing_keys = self.load_state_dict( |
| 217 | self.patch_tgrow_layers(torch.load(checkpoint_path, map_location="cpu", weights_only=True)), |
| 218 | strict=False |
| 219 | ) |
| 220 | print('missing_keys', missing_keys) |
| 221 | |
| 222 | # Initialize decoder mem state |
| 223 | self.mem = [None] * len(self.decoder) |
| 224 | |
| 225 | @staticmethod |
| 226 | def _apply_identity_deepen(decoder: nn.Sequential, how_many_each=1, k=3) -> nn.Sequential: |
| 227 | """Return a new Sequential where every nn.ReLU is followed by how_many_each*(IdentityConv2d(k)+ReLU).""" |
no outgoing calls
no test coverage detected