| """2D VAE encoder — wraps PixArt SDXL AutoencoderKL encoder half. |
| |
| Encodes images or mel spectrograms to [B, 4, H/8, W/8] latents. |
| Same encoder used for images AND audio spectrograms (via MelSpectrogram3Band). |
| Frozen float32 sidecar (no gradients). |
| """ |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| def load_vae2d(device="cuda", quantize=None): |
| from diffusers import AutoencoderKL |
|
|
| vae = AutoencoderKL.from_pretrained( |
| "PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers", |
| subfolder="vae", |
| torch_dtype=torch.float32, |
| ).to(device) |
| vae.eval() |
| for p in vae.parameters(): |
| p.requires_grad = False |
|
|
| return VAE2DEncoder(vae) |
|
|
|
|
| class VAE2DEncoder(nn.Module): |
| def __init__(self, vae): |
| super().__init__() |
| self.encoder = vae.encoder |
| self.quant_conv = vae.quant_conv |
| self.latent_channels = 4 |
| self.input_scale = 0.18215 |
|
|
| def forward(self, x): |
| H, W = x.shape[-2], x.shape[-1] |
| pad_h = (8 - H % 8) % 8 |
| pad_w = (8 - W % 8) % 8 |
| if pad_h > 0 or pad_w > 0: |
| x = F.pad(x, (0, pad_w, 0, pad_h)) |
|
|
| h = self.encoder(x) |
| moments = self.quant_conv(h) |
| posterior = torch.distributions.Normal( |
| moments[:, :self.latent_channels], |
| torch.nn.functional.softplus(moments[:, self.latent_channels:]) |
| ) |
| latent = posterior.rsample() |
| latent = latent * self.input_scale |
|
|
| if pad_h > 0 or pad_w > 0: |
| out_h = H // 8 if H >= 8 else 1 |
| out_w = W // 8 if W >= 8 else 1 |
| latent = latent[:, :, :out_h, :out_w] |
|
|
| return latent |
|
|
