| 229 | |
| 230 | class WanVAE_tiny(nn.Module): |
| 231 | def __init__(self, vae_path="taew2_1.pth", dtype=torch.bfloat16, device="cuda", need_scaled=False): |
| 232 | super().__init__() |
| 233 | self.dtype = dtype |
| 234 | self.device = torch.device("cuda") |
| 235 | self.taehv = TAEHV(vae_path).to(self.dtype) |
| 236 | self.temperal_downsample = [True, True, False] |
| 237 | self.need_scaled = need_scaled |
| 238 | |
| 239 | if self.need_scaled: |
| 240 | self.latents_mean = [ |
| 241 | -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, |
| 242 | 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921, |
| 243 | ] |
| 244 | |
| 245 | self.latents_std = [ |
| 246 | 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, |
| 247 | 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160, |
| 248 | ] |
| 249 | |
| 250 | self.z_dim = 16 |
| 251 | |
| 252 | @torch.no_grad() |
| 253 | def decode(self, latents, parallel=False): |