| 119 | |
| 120 | class TAEHV(nn.Module): |
| 121 | def __init__(self, checkpoint_path="taehv.pth", decoder_time_upscale=(True, True), decoder_space_upscale=(True, True, True), patch_size=1, latent_channels=16, model_type="wan21"): |
| 122 | super().__init__() |
| 123 | self.patch_size = patch_size |
| 124 | self.latent_channels = latent_channels |
| 125 | self.image_channels = 3 |
| 126 | self.is_cogvideox = checkpoint_path is not None and "taecvx" in checkpoint_path |
| 127 | self.model_type = model_type |
| 128 | if model_type == "wan22": |
| 129 | self.patch_size, self.latent_channels = 2, 48 |
| 130 | if model_type == "hy15": |
| 131 | act_func = nn.LeakyReLU(0.2, inplace=True) |
| 132 | else: |
| 133 | act_func = nn.ReLU(inplace=True) |
| 134 | |
| 135 | self.encoder = nn.Sequential( |
| 136 | conv(self.image_channels * self.patch_size**2, 64), |
| 137 | act_func, |
| 138 | TPool(64, 2), |
| 139 | conv(64, 64, stride=2, bias=False), |
| 140 | MemBlock(64, 64, act_func), |
| 141 | MemBlock(64, 64, act_func), |
| 142 | MemBlock(64, 64, act_func), |
| 143 | TPool(64, 2), |
| 144 | conv(64, 64, stride=2, bias=False), |
| 145 | MemBlock(64, 64, act_func), |
| 146 | MemBlock(64, 64, act_func), |
| 147 | MemBlock(64, 64, act_func), |
| 148 | TPool(64, 1), |
| 149 | conv(64, 64, stride=2, bias=False), |
| 150 | MemBlock(64, 64, act_func), |
| 151 | MemBlock(64, 64, act_func), |
| 152 | MemBlock(64, 64, act_func), |
| 153 | conv(64, self.latent_channels), |
| 154 | ) |
| 155 | n_f = [256, 128, 64, 64] |
| 156 | self.frames_to_trim = 2 ** sum(decoder_time_upscale) - 1 |
| 157 | self.decoder = nn.Sequential( |
| 158 | Clamp(), |
| 159 | conv(self.latent_channels, n_f[0]), |
| 160 | act_func, |
| 161 | MemBlock(n_f[0], n_f[0], act_func), |
| 162 | MemBlock(n_f[0], n_f[0], act_func), |
| 163 | MemBlock(n_f[0], n_f[0], act_func), |
| 164 | nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), |
| 165 | TGrow(n_f[0], 1), |
| 166 | conv(n_f[0], n_f[1], bias=False), |
| 167 | MemBlock(n_f[1], n_f[1], act_func), |
| 168 | MemBlock(n_f[1], n_f[1], act_func), |
| 169 | MemBlock(n_f[1], n_f[1], act_func), |
| 170 | nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), |
| 171 | TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), |
| 172 | conv(n_f[1], n_f[2], bias=False), |
| 173 | MemBlock(n_f[2], n_f[2], act_func), |
| 174 | MemBlock(n_f[2], n_f[2], act_func), |
| 175 | MemBlock(n_f[2], n_f[2], act_func), |
| 176 | nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1), |
| 177 | TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1), |
| 178 | conv(n_f[2], n_f[3], bias=False), |