| 72 | |
| 73 | |
| 74 | class WanVAE: |
| 75 | def __init__(self, |
| 76 | z_dim=16, |
| 77 | vae_pth=None, |
| 78 | dtype=torch.float, |
| 79 | device="cpu"): |
| 80 | self.dtype = dtype |
| 81 | self.device = device |
| 82 | |
| 83 | mean = [ |
| 84 | -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, |
| 85 | 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 |
| 86 | ] |
| 87 | std = [ |
| 88 | 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, |
| 89 | 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160 |
| 90 | ] |
| 91 | self.mean = torch.tensor(mean, dtype=dtype, device=device) |
| 92 | self.std = torch.tensor(std, dtype=dtype, device=device) |
| 93 | self.scale = [self.mean, 1.0 / self.std] |
| 94 | |
| 95 | # init model |
| 96 | self.model = _video_vae( |
| 97 | pretrained_path=vae_pth, |
| 98 | z_dim=z_dim, |
| 99 | ).eval().requires_grad_(False).to(device) |
| 100 | |
| 101 | |
| 102 | def vae_encode(tensor, vae): |
no outgoing calls
no test coverage detected