File size: 3,543 Bytes
8096951 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | """Адаптер на сторонний 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 # энкодер сжимает в 8 раз; декодер потом растягивает в 16 → 2× на выходе
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() # VAE заморожен всегда
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)
|