| """Адаптер на сторонний VAE (AsymmetricAutoencoderKL из sdxs-micro) под API нашего AutoEncoder. |
| |
| Что адаптер прячет от остального кода: |
| * нормировка латента своя — `latents_mean`/`latents_std` из vae_sdxs/config.json (≈ -0.04 / 1.58), |
| а не наши ≈ -4.1 / 1.03; снаружи всё так же: encode(x) → нормированный латент, |
| `to_raw()` в модели возвращает сырой; |
| * декодер асимметричный: энкодер f8, декодер f16, то есть латент 32×32 отдаёт картинку 512² |
| (вдвое больше входа) — это и есть смысл этого VAE; |
| * нет «хвоста» — поточного пиксельного декодера, поэтому пиксельный лосс при обучении не считается |
| (`--pix-weight 0`), а сэмплы декодируются один раз встроенным декодером VAE. |
| |
| API повторяет `vae/ae/model.py::AutoEncoder` в той части, которую использует train.py/sample.py: |
| `encode`, `decode_raw`, `latent_mean`, `latent_std`, `config.latent_channels`, `config.patch_size`, |
| `compute_dtype`, `.eval()`, `.to(device)`. |
| """ |
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
|
|
| PATCH = 8 |
|
|
|
|
| class SdxsAE(nn.Module): |
| def __init__(self, vae, device: str = "cuda"): |
| super().__init__() |
| self.vae = vae |
| self.device_name = device |
| self.compute_dtype = torch.float32 |
| cfg = vae.config |
| mean = torch.tensor(getattr(cfg, "latents_mean"), dtype=torch.float32) |
| std = torch.tensor(getattr(cfg, "latents_std"), dtype=torch.float32) |
| self.register_buffer("latent_mean", mean.view(1, -1, 1, 1)) |
| self.register_buffer("latent_std", std.view(1, -1, 1, 1)) |
| self.config = type("Cfg", (), {"latent_channels": int(cfg.latent_channels), |
| "patch_size": PATCH})() |
|
|
| @classmethod |
| def from_pretrained(cls, path, device: str = "cuda"): |
| from diffusers import AsymmetricAutoencoderKL |
| vae = AsymmetricAutoencoderKL.from_pretrained(str(path), torch_dtype=torch.float32) |
| vae = vae.to(device).eval() |
| vae.requires_grad_(False) |
| return cls(vae, device).to(device).eval() |
|
|
| @property |
| def device(self): |
| return self.latent_mean.device |
|
|
| def train(self, mode: bool = True): |
| super().train(mode) |
| self.vae.eval() |
| return self |
|
|
| def encode(self, x: torch.Tensor) -> torch.Tensor: |
| """x в -1..1 → нормированный латент (B, 32, H/8, W/8).""" |
| with torch.no_grad(): |
| raw = self.vae.encode(x.to(self.compute_dtype)).latent_dist.mode() |
| return ((raw - self.latent_mean) / self.latent_std).to(torch.float32) |
|
|
| def decode_raw(self, z_raw: torch.Tensor, steps: int = 0, generator=None) -> torch.Tensor: |
| """Сырой латент → картинка в -1..1. `steps`/`generator` — для совместимости с нашим API.""" |
| with torch.no_grad(): |
| return self.vae.decode(z_raw.to(self.compute_dtype), return_dict=False)[0].clamp(-1, 1) |
|
|
