qwen

eval_alchemist.py

@@ -15,8 +15,8 @@ DTYPE = torch.float16
 IMAGE_FOLDER = "/workspace/alchemist" #wget https://huggingface.co/datasets/AiArtLab/alchemist/resolve/main/alchemist.zip
 MIN_SIZE = 1280
 CROP_SIZE = 512
-BATCH_SIZE = 1
-MAX_IMAGES = 100
+BATCH_SIZE = 10
+MAX_IMAGES = 0
 NUM_WORKERS = 4
 NUM_SAMPLES_TO_SAVE = 2  # Сколько примеров сохранить (0 - не сохранять)
 SAMPLES_FOLDER = "vaetest"
@@ -32,9 +32,10 @@ VAE_LIST = [
 #    ("Lightricks/LTX-Video", AutoencoderKLLTXVideo, "Lightricks/LTX-Video", "vae"),
 #    ("Wan2.2-TI2V-5B-Diffusers", AutoencoderKLWan, "Wan-AI/Wan2.2-TI2V-5B-Diffusers", "vae"),
 #    ("Wan2.2-T2V-A14B-Diffusers", AutoencoderKLWan, "Wan-AI/Wan2.2-T2V-A14B-Diffusers", "vae"),
-    ("AiArtLab/sdxs", AutoencoderKL, "AiArtLab/sdxs", "vae"),
-#    ("FLUX.1-schnell-vae", AutoencoderKL, "black-forest-labs/FLUX.1-schnell", "vae"),
-#    ("simple_vae", AutoencoderKL, "/workspace/sdxl_vae/simple_vae", None),
+#    ("AiArtLab/sdxs", AutoencoderKL, "AiArtLab/sdxs", "vae"),
+    ("FLUX.1-schnell-vae", AutoencoderKL, "black-forest-labs/FLUX.1-schnell", "vae"),
+    ("simple_vae", AutoencoderKL, "AiArtLab/simplevae", "vae"),
+    ("simple_vae2", AutoencoderKL, "AiArtLab/simplevae", None),
     ("simple_vae_nightly", AutoencoderKL, "/workspace/sdxl_vae/simple_vae_nightly", None),
     
 ]

eval_alchemist2.py

@@ -0,0 +1,514 @@
+import os
+import json
+import random
+from typing import Dict, List, Tuple, Optional, Any
+
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from torchvision.transforms import Compose, Resize, ToTensor, CenterCrop
+from torchvision.utils import save_image
+import lpips
+
+from diffusers import (
+    AutoencoderKL,
+    AutoencoderKLWan,
+    AutoencoderKLLTXVideo,
+    AutoencoderKLQwenImage
+)
+
+from scipy.stats import skew, kurtosis
+
+
+# ========================== Конфиг ==========================
+DEVICE = "cuda"
+DTYPE = torch.float16
+IMAGE_FOLDER = "/home/recoilme/dataset/alchemist"
+MIN_SIZE = 1280
+CROP_SIZE = 512
+BATCH_SIZE = 10
+MAX_IMAGES = 500
+NUM_WORKERS = 4
+SAMPLES_DIR = "vaetest"
+
+VAE_LIST = [
+    # ("SD15 VAE", AutoencoderKL, "stable-diffusion-v1-5/stable-diffusion-v1-5", "vae"),
+    # ("SDXL VAE fp16 fix", AutoencoderKL, "madebyollin/sdxl-vae-fp16-fix", None),
+    ("Wan2.2-TI2V-5B", AutoencoderKLWan, "Wan-AI/Wan2.2-TI2V-5B-Diffusers", "vae"),
+    ("Wan2.2-T2V-A14B", AutoencoderKLWan, "Wan-AI/Wan2.2-T2V-A14B-Diffusers", "vae"),
+    #("SimpleVAE1", AutoencoderKL, "/home/recoilme/simplevae/simplevae", "simple_vae_nightly"),
+    #("SimpleVAE2", AutoencoderKL, "/home/recoilme/simplevae/simplevae", "simple_vae_nightly2"),    
+    #("SimpleVAE nightly", AutoencoderKL, "AiArtLab/simplevae", "simple_vae_nightly"),
+    #("FLUX.1-schnell VAE", AutoencoderKL, "black-forest-labs/FLUX.1-schnell", "vae"),
+    # ("LTX-Video VAE", AutoencoderKLLTXVideo, "Lightricks/LTX-Video", "vae"),
+    ("QwenImage", AutoencoderKLQwenImage, "Qwen/Qwen-Image", "vae"),
+]
+
+
+# ========================== Утилиты ==========================
+def to_neg1_1(x: torch.Tensor) -> torch.Tensor:
+    return x * 2 - 1
+
+
+def to_0_1(x: torch.Tensor) -> torch.Tensor:
+    return (x + 1) * 0.5
+
+
+def safe_psnr(mse: float) -> float:
+    if mse <= 1e-12:
+        return float("inf")
+    return 10.0 * float(np.log10(1.0 / mse))
+
+
+def is_video_like_vae(vae) -> bool:
+    # Wan и LTX-Video ждут [B, C, T, H, W]
+    return isinstance(vae, (AutoencoderKLWan, AutoencoderKLLTXVideo,AutoencoderKLQwenImage))
+
+
+def add_time_dim_if_needed(x: torch.Tensor, vae) -> torch.Tensor:
+    if is_video_like_vae(vae) and x.ndim == 4:
+        return x.unsqueeze(2)  # -> [B, C, 1, H, W]
+    return x
+
+
+def strip_time_dim_if_possible(x: torch.Tensor, vae) -> torch.Tensor:
+    if is_video_like_vae(vae) and x.ndim == 5 and x.shape[2] == 1:
+        return x.squeeze(2)  # -> [B, C, H, W]
+    return x
+
+
+@torch.no_grad()
+def sobel_edge_l1(real_0_1: torch.Tensor, fake_0_1: torch.Tensor) -> float:
+    real = to_neg1_1(real_0_1)
+    fake = to_neg1_1(fake_0_1)
+    kx = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32, device=real.device).view(1, 1, 3, 3)
+    ky = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32, device=real.device).view(1, 1, 3, 3)
+    C = real.shape[1]
+    kx = kx.to(real.dtype).repeat(C, 1, 1, 1)
+    ky = ky.to(real.dtype).repeat(C, 1, 1, 1)
+
+    def grad_mag(x):
+        gx = F.conv2d(x, kx, padding=1, groups=C)
+        gy = F.conv2d(x, ky, padding=1, groups=C)
+        return torch.sqrt(gx * gx + gy * gy + 1e-12)
+
+    return F.l1_loss(grad_mag(fake), grad_mag(real)).item()
+
+
+def flatten_channels(x: torch.Tensor) -> torch.Tensor:
+    # -> [C, N*H*W] или [C, N*T*H*W]
+    if x.ndim == 4:
+        return x.permute(1, 0, 2, 3).reshape(x.shape[1], -1)
+    elif x.ndim == 5:
+        return x.permute(1, 0, 2, 3, 4).reshape(x.shape[1], -1)
+    else:
+        raise ValueError(f"Unexpected tensor ndim={x.ndim}")
+
+
+def _to_numpy_1d(x: Any) -> Optional[np.ndarray]:
+    if x is None:
+        return None
+    if isinstance(x, (int, float)):
+        return None
+    if isinstance(x, torch.Tensor):
+        x = x.detach().cpu().float().numpy()
+    elif isinstance(x, (list, tuple)):
+        x = np.array(x, dtype=np.float32)
+    elif isinstance(x, np.ndarray):
+        x = x.astype(np.float32, copy=False)
+    else:
+        return None
+    x = x.reshape(-1)
+    return x
+
+
+def _to_float(x: Any) -> Optional[float]:
+    if x is None:
+        return None
+    if isinstance(x, (int, float)):
+        return float(x)
+    if isinstance(x, np.ndarray) and x.size == 1:
+        return float(x.item())
+    if isinstance(x, torch.Tensor) and x.numel() == 1:
+        return float(x.item())
+    return None
+
+
+def get_norm_tensors_and_summary(vae, latent_like: torch.Tensor):
+    """
+    Нормализация латентов: глобальная и поканальная.
+    Применение: сначала глобальная (scalar), затем поканальная (vector).
+    Если в конфиге есть несколько ключей — аккумулируем.
+    """
+    cfg = getattr(vae, "config", vae)
+
+    scale_keys = [
+        "latents_std"
+    ]
+    shift_keys = [
+        "latents_mean"
+    ]
+
+    C = latent_like.shape[1]
+    nd = latent_like.ndim  # 4 или 5
+    dev = latent_like.device
+    dt = latent_like.dtype
+
+    scale_global = getattr(vae.config, "scaling_factor", 1.0)
+    shift_global = getattr(vae.config, "shift_factor", 0.0)
+    if scale_global is None:
+        scale_global = 1.0
+    if shift_global is None:
+        shift_global = 0.0
+    
+    scale_channel = np.ones(C, dtype=np.float32)
+    shift_channel = np.zeros(C, dtype=np.float32)
+
+    for k in scale_keys:
+        v = getattr(cfg, k, None)
+        if v is None:
+            continue
+        vec = _to_numpy_1d(v)
+        if vec is not None and vec.size == C:
+            scale_channel *= vec
+        else:
+            s = _to_float(v)
+            if s is not None:
+                scale_global *= s
+
+    for k in shift_keys:
+        v = getattr(cfg, k, None)
+        if v is None:
+            continue
+        vec = _to_numpy_1d(v)
+        if vec is not None and vec.size == C:
+            shift_channel += vec
+        else:
+            s = _to_float(v)
+            if s is not None:
+                shift_global += s
+
+    g_shape = [1] * nd
+    c_shape = [1] * nd
+    c_shape[1] = C
+
+    t_scale_g = torch.tensor(scale_global, dtype=dt, device=dev).view(*g_shape)
+    t_shift_g = torch.tensor(shift_global, dtype=dt, device=dev).view(*g_shape)
+    t_scale_c = torch.from_numpy(scale_channel).to(device=dev, dtype=dt).view(*c_shape)
+    t_shift_c = torch.from_numpy(shift_channel).to(device=dev, dtype=dt).view(*c_shape)
+
+    summary = {
+        "scale_global": float(scale_global),
+        "shift_global": float(shift_global),
+        "scale_channel_min": float(scale_channel.min()),
+        "scale_channel_mean": float(scale_channel.mean()),
+        "scale_channel_max": float(scale_channel.max()),
+        "shift_channel_min": float(shift_channel.min()),
+        "shift_channel_mean": float(shift_channel.mean()),
+        "shift_channel_max": float(shift_channel.max()),
+    }
+    return t_shift_g, t_scale_g, t_shift_c, t_scale_c, summary
+
+
+@torch.no_grad()
+def kl_divergence_per_image(mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
+    kl_map = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())  # [B, ...]
+    return kl_map.float().view(kl_map.shape[0], -1).mean(dim=1)  # [B]
+
+
+def sanitize_filename(name: str) -> str:
+    name = name.replace("/", "_").replace("\\", "_").replace(" ", "_")
+    return "".join(ch if (ch.isalnum() or ch in "._-") else "_" for ch in name)
+
+
+# ========================== Датасет ==========================
+class ImageFolderDataset(Dataset):
+    def __init__(self, root_dir: str, extensions=(".png", ".jpg", ".jpeg", ".webp"), min_size=1024, crop_size=512, limit=None):
+        paths = []
+        for root, _, files in os.walk(root_dir):
+            for fname in files:
+                if fname.lower().endswith(extensions):
+                    paths.append(os.path.join(root, fname))
+        if limit:
+            paths = paths[:limit]
+
+        valid = []
+        for p in tqdm(paths, desc="Проверяем файлы"):
+            try:
+                with Image.open(p) as im:
+                    im.verify()
+                valid.append(p)
+            except Exception:
+                pass
+        if not valid:
+            raise RuntimeError(f"Нет валидных изображений в {root_dir}")
+        random.shuffle(valid)
+        self.paths = valid
+        print(f"Найдено {len(self.paths)} изображений")
+
+        self.transform = Compose([
+            Resize(min_size),
+            CenterCrop(crop_size),
+            ToTensor(),  # 0..1, float32
+        ])
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, idx):
+        with Image.open(self.paths[idx]) as img:
+            img = img.convert("RGB")
+            return self.transform(img)
+
+
+# ========================== Основное ==========================
+def main():
+    torch.set_grad_enabled(False)
+    os.makedirs(SAMPLES_DIR, exist_ok=True)
+
+    dataset = ImageFolderDataset(IMAGE_FOLDER, min_size=MIN_SIZE, crop_size=CROP_SIZE, limit=MAX_IMAGES)
+    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, pin_memory=True)
+
+    lpips_net = lpips.LPIPS(net="vgg").to(DEVICE).eval()
+
+    # Загрузка VAE
+    vaes: List[Tuple[str, object]] = []
+    print("\nЗагрузка VAE...")
+    for human_name, vae_class, model_path, subfolder in VAE_LIST:
+        try:
+            vae = vae_class.from_pretrained(model_path, subfolder=subfolder, torch_dtype=DTYPE)
+            vae = vae.to(DEVICE).eval()
+            vaes.append((human_name, vae))
+            print(f"  ✅ {human_name}")
+        except Exception as e:
+            print(f"  ❌ {human_name}: {e}")
+
+    if not vaes:
+        print("Нет успешно загруженных VAE. Выходим.")
+        return
+
+    # Агрегаторы
+    per_model_metrics: Dict[str, Dict[str, float]] = {
+        name: {"mse": 0.0, "psnr": 0.0, "lpips": 0.0, "edge": 0.0, "kl": 0.0, "count": 0.0}
+        for name, _ in vaes
+    }
+
+    buffers_zmodel: Dict[str, List[torch.Tensor]] = {name: [] for name, _ in vaes}
+    norm_summaries: Dict[str, Dict[str, float]] = {}
+
+    # Флаг для сохранения первой картинки
+    saved_first_for: Dict[str, bool] = {name: False for name, _ in vaes}
+
+    for batch_0_1 in tqdm(loader, desc="Батчи"):
+        batch_0_1 = batch_0_1.to(DEVICE, torch.float32)
+        batch_neg1_1 = to_neg1_1(batch_0_1).to(DTYPE)
+
+        for model_name, vae in vaes:
+            x_in = add_time_dim_if_needed(batch_neg1_1, vae)
+
+            posterior = vae.encode(x_in).latent_dist
+            mu, logvar = posterior.mean, posterior.logvar
+
+            # Реконструкция (детерминированно)
+            z_raw_mode = posterior.mode()
+            x_dec = vae.decode(z_raw_mode).sample  # [-1, 1]
+            x_dec = strip_time_dim_if_possible(x_dec, vae)
+            x_rec_0_1 = to_0_1(x_dec.float()).clamp(0, 1)
+
+            # Латенты для UNet: global -> channelwise
+            z_raw_sample = posterior.sample()
+            t_shift_g, t_scale_g, t_shift_c, t_scale_c, summary = get_norm_tensors_and_summary(vae, z_raw_sample)
+ 
+            if model_name not in norm_summaries:
+                norm_summaries[model_name] = summary
+
+            z_tmp = (z_raw_sample - t_shift_g) * t_scale_g
+            z_model = (z_tmp - t_shift_c) * t_scale_c
+            z_model = strip_time_dim_if_possible(z_model, vae)
+
+            buffers_zmodel[model_name].append(z_model.detach().to("cpu", torch.float32))
+
+            # Сохранить первую картинку (оригинал и реконструкцию) для каждого VAE
+            if not saved_first_for[model_name]:
+                safe = sanitize_filename(model_name)
+                orig_path = os.path.join(SAMPLES_DIR, f"{safe}_original.png")
+                dec_path  = os.path.join(SAMPLES_DIR, f"{safe}_decoded.png")
+                save_image(batch_0_1[0:1].cpu(), orig_path)
+                save_image(x_rec_0_1[0:1].cpu(),  dec_path)
+                saved_first_for[model_name] = True
+
+            # Метрики по картинкам
+            B = batch_0_1.shape[0]
+            for i in range(B):
+                gt = batch_0_1[i:i+1]
+                rec = x_rec_0_1[i:i+1]
+
+                mse = F.mse_loss(gt, rec).item()
+                psnr = safe_psnr(mse)
+                lp = float(lpips_net(gt, rec, normalize=True).mean().item())
+                edge = sobel_edge_l1(gt, rec)
+
+                per_model_metrics[model_name]["mse"] += mse
+                per_model_metrics[model_name]["psnr"] += psnr
+                per_model_metrics[model_name]["lpips"] += lp
+                per_model_metrics[model_name]["edge"] += edge
+
+            # KL per-image
+            kl_pi = kl_divergence_per_image(mu, logvar)  # [B]
+            per_model_metrics[model_name]["kl"] += float(kl_pi.sum().item())
+            per_model_metrics[model_name]["count"] += B
+
+    # Усреднение метрик
+    for name in per_model_metrics:
+        c = max(1.0, per_model_metrics[name]["count"])
+        for k in ["mse", "psnr", "lpips", "edge", "kl"]:
+            per_model_metrics[name][k] /= c
+
+    # Подсчёт статистик латентов и нормальности
+    per_model_latent_stats = {}
+    for name, _ in vaes:
+        if not buffers_zmodel[name]:
+            continue
+        Z = torch.cat(buffers_zmodel[name], dim=0)   # [N, C, H, W]
+
+        # Глобальные
+        z_min = float(Z.min().item())
+        z_mean = float(Z.mean().item())
+        z_max = float(Z.max().item())
+        z_std = float(Z.std(unbiased=True).item())
+
+        # Пер-канально: skew/kurtosis
+        Z_ch = flatten_channels(Z).numpy()  # [C, *]
+        C = Z_ch.shape[0]
+        sk = np.zeros(C, dtype=np.float64)
+        ku = np.zeros(C, dtype=np.float64)
+        for c in range(C):
+            v = Z_ch[c]
+            sk[c] = float(skew(v, bias=False))
+            ku[c] = float(kurtosis(v, fisher=True, bias=False))
+
+        skew_min, skew_mean, skew_max = float(sk.min()), float(sk.mean()), float(sk.max())
+        kurt_min, kurt_mean, kurt_max = float(ku.min()), float(ku.mean()), float(ku.max())
+        mean_abs_skew = float(np.mean(np.abs(sk)))
+        mean_abs_kurt = float(np.mean(np.abs(ku)))
+
+        per_model_latent_stats[name] = {
+            "Z_min": z_min, "Z_mean": z_mean, "Z_max": z_max, "Z_std": z_std,
+            "skew_min": skew_min, "skew_mean": skew_mean, "skew_max": skew_max,
+            "kurt_min": kurt_min, "kurt_mean": kurt_mean, "kurt_max": kurt_max,
+            "mean_abs_skew": mean_abs_skew, "mean_abs_kurt": mean_abs_kurt,
+        }
+
+    # Печать параметров нормализации (shift/scale)
+    print("\n=== Параметры нормализации ��атентов (как применялись) ===")
+    for name, _ in vaes:
+        if name not in norm_summaries:
+            continue
+        s = norm_summaries[name]
+        print(
+            f"{name:26s} | "
+            f"shift_g={s['shift_global']:.6g} scale_g={s['scale_global']:.6g} | "
+            f"shift_c[min/mean/max]=[{s['shift_channel_min']:.6g}, {s['shift_channel_mean']:.6g}, {s['shift_channel_max']:.6g}] | "
+            f"scale_c[min/mean/max]=[{s['scale_channel_min']:.6g}, {s['scale_channel_mean']:.6g}, {s['scale_channel_max']:.6g}]"
+        )
+
+    # Абсолютные метрики
+    print("\n=== Абсолютные метрики реконструкции и латентов ===")
+    for name, _ in vaes:
+        if name not in per_model_latent_stats:
+            continue
+        m = per_model_metrics[name]
+        s = per_model_latent_stats[name]
+        print(
+            f"{name:26s} | "
+            f"MSE={m['mse']:.3e} PSNR={m['psnr']:.2f} LPIPS={m['lpips']:.3f} Edge={m['edge']:.3f} KL={m['kl']:.3f} | "
+            f"Z[min/mean/max/std]=[{s['Z_min']:.3f}, {s['Z_mean']:.3f}, {s['Z_max']:.3f}, {s['Z_std']:.3f}] | "
+            f"Skew[min/mean/max]=[{s['skew_min']:.3f}, {s['skew_mean']:.3f}, {s['skew_max']:.3f}] | "
+            f"Kurt[min/mean/max]=[{s['kurt_min']:.3f}, {s['kurt_mean']:.3f}, {s['kurt_max']:.3f}]"
+        )
+
+    # Сравнение с первой моделью
+    baseline = vaes[0][0]
+    print("\n=== Сравнение с первой моделью (проценты) ===")
+    print(f"| {'Модель':26s} | {'MSE':>9s} | {'PSNR':>9s} | {'LPIPS':>9s} | {'Edge':>9s} | {'Skew|0':>9s} | {'Kurt|0':>9s} |")
+    print(f"|{'-'*28}|{'-'*11}|{'-'*11}|{'-'*11}|{'-'*11}|{'-'*11}|{'-'*11}|")
+
+    b_m = per_model_metrics[baseline]
+    b_s = per_model_latent_stats[baseline]
+
+    for name, _ in vaes:
+        m = per_model_metrics[name]
+        s = per_model_latent_stats[name]
+
+        mse_pct  = (b_m["mse"] / max(1e-12, m["mse"])) * 100.0               # меньше лучше
+        psnr_pct = (m["psnr"] / max(1e-12, b_m["psnr"])) * 100.0             # больше лучше
+        lpips_pct= (b_m["lpips"] / max(1e-12, m["lpips"])) * 100.0           # меньше лучше
+        edge_pct = (b_m["edge"] / max(1e-12, m["edge"])) * 100.0             # меньше лучше
+
+        skew0_pct = (b_s["mean_abs_skew"] / max(1e-12, s["mean_abs_skew"])) * 100.0
+        kurt0_pct = (b_s["mean_abs_kurt"] / max(1e-12, s["mean_abs_kurt"])) * 100.0
+
+        if name == baseline:
+            print(f"| {name:26s} | {'100%':>9s} | {'100%':>9s} | {'100%':>9s} | {'100%':>9s} | {'100%':>9s} | {'100%':>9s} |")
+        else:
+            print(f"| {name:26s} | {mse_pct:8.1f}% | {psnr_pct:8.1f}% | {lpips_pct:8.1f}% | {edge_pct:8.1f}% | {skew0_pct:8.1f}% | {kurt0_pct:8.1f}% |")
+
+    # ========================== Коррекции для последнего VAE + сохранение в JSON ==========================
+    last_name = vaes[-1][0]
+    if buffers_zmodel[last_name]:
+        Z = torch.cat(buffers_zmodel[last_name], dim=0)  # [N, C, H, W]
+
+        # Глобальная коррекция (по всем каналам/пикселям)
+        z_mean = float(Z.mean().item())
+        z_std  = float(Z.std(unbiased=True).item())
+        correction_global = {
+            "shift": -z_mean,
+            "scale": (1.0 / z_std) if z_std > 1e-12 else 1.0
+        }
+
+        # Поканальная коррекция
+        Z_ch = flatten_channels(Z)  # [C, M]
+        ch_means_t = Z_ch.mean(dim=1)                       # [C]
+        ch_stds_t  = Z_ch.std(dim=1, unbiased=True) + 1e-12 # [C]
+        ch_means = [float(x) for x in ch_means_t.tolist()]
+        ch_stds  = [float(x) for x in ch_stds_t.tolist()]
+
+        correction_per_channel = [
+            {"shift": float(-m), "scale": float(1.0 / s)}
+            for m, s in zip(ch_means, ch_stds)
+        ]
+
+        print(f"\n=== Доп. коррекция для {last_name} (поверх VAE-нормализации) ===")
+        print(f"global_correction = {correction_global}")
+        print(f"channelwise_means = {ch_means}")
+        print(f"channelwise_stds  = {ch_stds}")
+        print(f"channelwise_correction = {correction_per_channel}")
+
+        # Сохранение в JSON
+        json_path = os.path.join(SAMPLES_DIR, f"{sanitize_filename(last_name)}_correction.json")
+        to_save = {
+            "model_name": last_name,
+            "vae_normalization_summary": norm_summaries.get(last_name, {}),
+            "global_correction": correction_global,
+            "per_channel_means": ch_means,
+            "per_channel_stds": ch_stds,
+            "per_channel_correction": correction_per_channel,
+            "apply_order": {
+                "forward": "z_model -> (z - global_shift)*global_scale -> (per-channel: (z - mean_c)/std_c)",
+                "inverse":  "z_corr -> (per-channel: z*std_c + mean_c) -> (z/global_scale + global_shift)"
+            },
+            "note": "Эти коэффициенты рассчитаны по z_model (после встроенных VAE shift/scale), чтобы привести распределение к N(0,1)."
+        }
+        with open(json_path, "w", encoding="utf-8") as f:
+            json.dump(to_save, f, ensure_ascii=False, indent=2)
+        print("Corrections JSON saved to:", os.path.abspath(json_path))
+
+    print("\n✅ Готово. Сэмплы сохранены в:", os.path.abspath(SAMPLES_DIR))
+
+
+if __name__ == "__main__":
+    main()

simple_vae/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:010d2cb8824a347425be4e41d662b22492965ffb61393621eb1253be8b7fa0ce
+oid sha256:f5f0a20e403669e880b510514ee575a2a9cb74a1b36ab0e31fc68ef66c2173d7
 size 335311892

simple_vae_nightly/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ccd57f2cd9455d6c66ed2fee9396dbb53cbeb675fa0c1fbee87a9b0f94c3de79
+oid sha256:7b705da7f401289eefa22570514d7c1b9b2f9fd32a71159e2d3d5888f74e41cd
 size 335311892

train_sdxl_vae_full.py

@@ -58,7 +58,7 @@ device = None  # accelerator задаст устройство
 # CHANGED: добавлен параметр для полного обучения VAE (а не только декодера).
 # Если False — поведение прежнее: учим только decoder.* (up_blocks + mid_block).
 # Если True  — размораживаем ВСЮ модель и добавляем KL-loss для энкодера.
-full_training      = False
+full_training      = True
 
 # CHANGED: добавлен вес (через долю в нормализаторе) для KL, используется только при full_training=True.
 kl_ratio           = 0.05  # простая доля для KL в общей смеси (KISS). Игнорируется, если full_training=False.
@@ -66,12 +66,12 @@ kl_ratio           = 0.05  # простая доля для KL в общей с
 # --- Пропорции лоссов и окно медианного нормирования (КОЭФ., не значения) ---
 # Итоговые доли в total loss (сумма = 1.0 после нормализации).
 loss_ratios = {
-    "lpips": 0.85,
+    "lpips": 0.80,
     "edge":  0.05,
     "mse":   0.05,
     "mae":   0.05,
     # CHANGED: заранее добавлен ключ "kl" (по умолчанию 0.0). Если включаем full_training — активируем ниже.
-    "kl":    0.00,
+    "kl":    0.05,
 }
 median_coeff_steps = 256  # за сколько шагов считать медианные коэффициенты
 

train_sdxl_vae_qwen.py

@@ -0,0 +1,526 @@
+# -*- coding: utf-8 -*-
+import os
+import math
+import re
+import torch
+import numpy as np
+import random
+import gc
+from datetime import datetime
+from pathlib import Path
+
+import torchvision.transforms as transforms
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+from torch.optim.lr_scheduler import LambdaLR
+from diffusers import AutoencoderKL, AsymmetricAutoencoderKL
+# QWEN: импорт класса
+from diffusers import AutoencoderKLQwenImage
+
+from accelerate import Accelerator
+from PIL import Image, UnidentifiedImageError
+from tqdm import tqdm
+import bitsandbytes as bnb
+import wandb
+import lpips   # pip install lpips
+from collections import deque
+
+# --------------------------- Параметры ---------------------------
+ds_path            = "/workspace/png"
+project            = "qwen_vae"
+batch_size         = 3
+base_learning_rate = 5e-5
+min_learning_rate  = 9e-7
+num_epochs         = 16
+sample_interval_share = 10
+use_wandb          = True
+save_model         = True
+use_decay          = True
+optimizer_type     = "adam8bit"
+dtype              = torch.float32
+
+model_resolution   = 512
+high_resolution    = 512 
+limit              = 0
+save_barrier       = 1.03
+warmup_percent     = 0.01
+percentile_clipping = 95
+beta2              = 0.97
+eps                = 1e-6
+clip_grad_norm     = 1.0
+mixed_precision    = "no"
+gradient_accumulation_steps = 5
+generated_folder   = "samples"
+save_as            = "wen_vae_nightly"
+num_workers        = 0
+device = None
+
+# --- Режимы обучения ---
+# QWEN: учим только декодер
+train_decoder_only = True
+full_training      = False  # если True — учим весь VAE и добавляем KL (ниже)
+kl_ratio           = 0.05
+
+# Доли лоссов
+loss_ratios = {
+    "lpips": 0.80,
+    "edge":  0.05,
+    "mse":   0.10,
+    "mae":   0.05,
+    "kl":    0.00,  # активируем при full_training=True
+}
+median_coeff_steps = 256
+
+resize_long_side = 1280  # ресайз длинной стороны исходных картинок
+
+# QWEN: конфиг загрузки модели
+vae_kind      = "qwen"  # "qwen" или "kl" (обычный)
+vae_model_id  = "Qwen/Qwen-Image"
+vae_subfolder = "vae"
+
+Path(generated_folder).mkdir(parents=True, exist_ok=True)
+
+accelerator = Accelerator(
+    mixed_precision=mixed_precision,
+    gradient_accumulation_steps=gradient_accumulation_steps
+)
+device = accelerator.device
+
+# reproducibility
+seed = int(datetime.now().strftime("%Y%m%d"))
+torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
+torch.backends.cudnn.benchmark = False
+
+# --------------------------- WandB ---------------------------
+if use_wandb and accelerator.is_main_process:
+    wandb.init(project=project, config={
+        "batch_size": batch_size,
+        "base_learning_rate": base_learning_rate,
+        "num_epochs": num_epochs,
+        "optimizer_type": optimizer_type,
+        "model_resolution": model_resolution,
+        "high_resolution": high_resolution,
+        "gradient_accumulation_steps": gradient_accumulation_steps,
+        "train_decoder_only": train_decoder_only,
+        "full_training": full_training,
+        "kl_ratio": kl_ratio,
+        "vae_kind": vae_kind,
+        "vae_model_id": vae_model_id,
+    })
+
+# --------------------------- VAE ---------------------------
+def is_qwen_vae(vae) -> bool:
+    return isinstance(vae, AutoencoderKLQwenImage) or ("Qwen" in vae.__class__.__name__)
+
+# загрузка
+if vae_kind == "qwen":
+    vae = AutoencoderKLQwenImage.from_pretrained(vae_model_id, subfolder=vae_subfolder)
+else:
+    # старое поведение (пример)
+    if model_resolution==high_resolution:
+        vae = AutoencoderKL.from_pretrained(project)
+    else:
+        vae = AsymmetricAutoencoderKL.from_pretrained(project)
+
+vae = vae.to(dtype)
+
+# torch.compile (опционально)
+if hasattr(torch, "compile"):
+    try:
+        vae = torch.compile(vae)
+    except Exception as e:
+        print(f"[WARN] torch.compile failed: {e}")
+
+# --------------------------- Freeze/Unfreeze ---------------------------
+for p in vae.parameters():
+    p.requires_grad = False
+
+unfrozen_param_names = []
+
+if full_training and not train_decoder_only:
+    # учим всю модель
+    for name, p in vae.named_parameters():
+        p.requires_grad = True
+        unfrozen_param_names.append(name)
+    loss_ratios["kl"] = float(kl_ratio)
+    trainable_module = vae
+else:
+    # QWEN: учим только декодер (и post_quant_conv — часть декодерного тракта)
+    # универсально: всё, что начинается с "decoder." или "post_quant_conv"
+    for name, p in vae.named_parameters():
+        if name.startswith("decoder.") or name.startswith("post_quant_conv"):
+            p.requires_grad = True
+            unfrozen_param_names.append(name)
+    trainable_module = vae.decoder if hasattr(vae, "decoder") else vae
+
+print(f"[INFO] Разморожено параметров: {len(unfrozen_param_names)}. Первые 200 имён:")
+for nm in unfrozen_param_names[:200]:
+    print("  ", nm)
+
+# --------------------------- Датасет ---------------------------
+class PngFolderDataset(Dataset):
+    def __init__(self, root_dir, min_exts=('.png',), resolution=1024, limit=0):
+        self.root_dir = root_dir
+        self.resolution = resolution
+        self.paths = []
+        for root, _, files in os.walk(root_dir):
+            for fname in files:
+                if fname.lower().endswith(tuple(ext.lower() for ext in min_exts)):
+                    self.paths.append(os.path.join(root, fname))
+        if limit:
+            self.paths = self.paths[:limit]
+        valid = []
+        for p in self.paths:
+            try:
+                with Image.open(p) as im:
+                    im.verify()
+                valid.append(p)
+            except (OSError, UnidentifiedImageError):
+                continue
+        self.paths = valid
+        if len(self.paths) == 0:
+            raise RuntimeError(f"No valid PNG images found under {root_dir}")
+        random.shuffle(self.paths)
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, idx):
+        p = self.paths[idx % len(self.paths)]
+        with Image.open(p) as img:
+            img = img.convert("RGB")
+            if not resize_long_side or resize_long_side <= 0:
+                return img
+            w, h = img.size
+            long = max(w, h)
+            if long <= resize_long_side:
+                return img
+            scale = resize_long_side / float(long)
+            new_w = int(round(w * scale))
+            new_h = int(round(h * scale))
+            return img.resize((new_w, new_h), Image.LANCZOS)
+
+def random_crop(img, sz):
+    w, h = img.size
+    if w < sz or h < sz:
+        img = img.resize((max(sz, w), max(sz, h)), Image.LANCZOS)
+    x = random.randint(0, max(1, img.width - sz))
+    y = random.randint(0, max(1, img.height - sz))
+    return img.crop((x, y, x + sz, y + sz))
+
+tfm = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+])
+
+dataset = PngFolderDataset(ds_path, min_exts=('.png',), resolution=high_resolution, limit=limit)
+if len(dataset) < batch_size:
+    raise RuntimeError(f"Not enough valid images ({len(dataset)}) to form a batch of size {batch_size}")
+
+def collate_fn(batch):
+    imgs = []
+    for img in batch:
+        img = random_crop(img, high_resolution)
+        imgs.append(tfm(img))
+    return torch.stack(imgs)
+
+dataloader = DataLoader(
+    dataset,
+    batch_size=batch_size,
+    shuffle=True,
+    collate_fn=collate_fn,
+    num_workers=num_workers,
+    pin_memory=True,
+    drop_last=True
+)
+
+# --------------------------- Оптимизатор ---------------------------
+def get_param_groups(module, weight_decay=0.001):
+    no_decay = ["bias", "LayerNorm.weight", "layer_norm.weight", "ln_1.weight", "ln_f.weight"]
+    decay_params, no_decay_params = [], []
+    for n, p in vae.named_parameters():  # глобально по vae, с фильтром requires_grad
+        if not p.requires_grad:
+            continue
+        if any(nd in n for nd in no_decay):
+            no_decay_params.append(p)
+        else:
+            decay_params.append(p)
+    return [
+        {"params": decay_params, "weight_decay": weight_decay},
+        {"params": no_decay_params, "weight_decay": 0.0},
+    ]
+
+def create_optimizer(name, param_groups):
+    if name == "adam8bit":
+        return bnb.optim.AdamW8bit(param_groups, lr=base_learning_rate, betas=(0.9, beta2), eps=eps)
+    raise ValueError(name)
+
+param_groups = get_param_groups(trainable_module, weight_decay=0.001)
+optimizer = create_optimizer(optimizer_type, param_groups)
+
+# --------------------------- LR schedule ---------------------------
+batches_per_epoch = len(dataloader)
+steps_per_epoch = int(math.ceil(batches_per_epoch / float(gradient_accumulation_steps)))
+total_steps = steps_per_epoch * num_epochs
+
+def lr_lambda(step):
+    if not use_decay:
+        return 1.0
+    x = float(step) / float(max(1, total_steps))
+    warmup = float(warmup_percent)
+    min_ratio = float(min_learning_rate) / float(base_learning_rate)
+    if x < warmup:
+        return min_ratio + (1.0 - min_ratio) * (x / warmup)
+    decay_ratio = (x - warmup) / (1.0 - warmup)
+    return min_ratio + 0.5 * (1.0 - min_ratio) * (1.0 + math.cos(math.pi * decay_ratio))
+
+scheduler = LambdaLR(optimizer, lr_lambda)
+
+# Подготовка
+dataloader, vae, optimizer, scheduler = accelerator.prepare(dataloader, vae, optimizer, scheduler)
+trainable_params = [p for p in vae.parameters() if p.requires_grad]
+
+# --------------------------- LPIPS и вспомогательные ---------------------------
+_lpips_net = None
+def _get_lpips():
+    global _lpips_net
+    if _lpips_net is None:
+        _lpips_net = lpips.LPIPS(net='vgg', verbose=False).eval().to(accelerator.device).eval()
+    return _lpips_net
+
+_sobel_kx = torch.tensor([[[[-1,0,1],[-2,0,2],[-1,0,1]]]], dtype=torch.float32)
+_sobel_ky = torch.tensor([[[[-1,-2,-1],[0,0,0],[1,2,1]]]], dtype=torch.float32)
+def sobel_edges(x: torch.Tensor) -> torch.Tensor:
+    C = x.shape[1]
+    kx = _sobel_kx.to(x.device, x.dtype).repeat(C, 1, 1, 1)
+    ky = _sobel_ky.to(x.device, x.dtype).repeat(C, 1, 1, 1)
+    gx = F.conv2d(x, kx, padding=1, groups=C)
+    gy = F.conv2d(x, ky, padding=1, groups=C)
+    return torch.sqrt(gx * gx + gy * gy + 1e-12)
+
+class MedianLossNormalizer:
+    def __init__(self, desired_ratios: dict, window_steps: int):
+        s = sum(desired_ratios.values())
+        self.ratios = {k: (v / s) if s > 0 else 0.0 for k, v in desired_ratios.items()}
+        self.buffers = {k: deque(maxlen=window_steps) for k in self.ratios.keys()}
+        self.window = window_steps
+
+    def update_and_total(self, abs_losses: dict):
+        for k, v in abs_losses.items():
+            if k in self.buffers:
+                self.buffers[k].append(float(v.detach().abs().cpu()))
+        meds = {k: (np.median(self.buffers[k]) if len(self.buffers[k]) > 0 else 1.0) for k in self.buffers}
+        coeffs = {k: (self.ratios[k] / max(meds[k], 1e-12)) for k in self.ratios}
+        total = sum(coeffs[k] * abs_losses[k] for k in abs_losses if k in coeffs)
+        return total, coeffs, meds
+
+if full_training and not train_decoder_only:
+    loss_ratios["kl"] = float(kl_ratio)
+normalizer = MedianLossNormalizer(loss_ratios, median_coeff_steps)
+
+# --------------------------- Сэмплы ---------------------------
+@torch.no_grad()
+def get_fixed_samples(n=3):
+    idx = random.sample(range(len(dataset)), min(n, len(dataset)))
+    pil_imgs = [dataset[i] for i in idx]
+    tensors = []
+    for img in pil_imgs:
+        img = random_crop(img, high_resolution)
+        tensors.append(tfm(img))
+    return torch.stack(tensors).to(accelerator.device, dtype)
+
+fixed_samples = get_fixed_samples()
+
+@torch.no_grad()
+def _to_pil_uint8(img_tensor: torch.Tensor) -> Image.Image:
+    arr = ((img_tensor.float().clamp(-1, 1) + 1.0) * 127.5).clamp(0, 255).byte().cpu().numpy().transpose(1, 2, 0)
+    return Image.fromarray(arr)
+
+@torch.no_grad()
+def generate_and_save_samples(step=None):
+    try:
+        temp_vae = accelerator.unwrap_model(vae).eval()
+        lpips_net = _get_lpips()
+        with torch.no_grad():
+            orig_high = fixed_samples
+            orig_low = F.interpolate(orig_high, size=(model_resolution, model_resolution), mode="bilinear", align_corners=False)
+            model_dtype = next(temp_vae.parameters()).dtype
+            orig_low = orig_low.to(dtype=model_dtype)
+
+            # QWEN: добавляем T=1 на encode/decode и снимаем при сравнении
+            if is_qwen_vae(temp_vae):
+                x_in = orig_low.unsqueeze(2)           # [B,3,1,H,W]
+                enc = temp_vae.encode(x_in)
+                latents_mean = enc.latent_dist.mean
+                dec = temp_vae.decode(latents_mean).sample  # [B,3,1,H,W]
+                rec = dec.squeeze(2)                   # [B,3,H,W]
+            else:
+                enc = temp_vae.encode(orig_low)
+                latents_mean = enc.latent_dist.mean
+                rec = temp_vae.decode(latents_mean).sample
+
+        if rec.shape[-2:] != orig_high.shape[-2:]:
+            rec = F.interpolate(rec, size=orig_high.shape[-2:], mode="bilinear", align_corners=False)
+
+        first_real = _to_pil_uint8(orig_high[0])
+        first_dec  = _to_pil_uint8(rec[0])
+        first_real.save(f"{generated_folder}/sample_real.jpg", quality=95)
+        first_dec.save(f"{generated_folder}/sample_decoded.jpg", quality=95)
+
+        for i in range(rec.shape[0]):
+            _to_pil_uint8(rec[i]).save(f"{generated_folder}/sample_{i}.jpg", quality=95)
+
+        lpips_scores = []
+        for i in range(rec.shape[0]):
+            orig_full = orig_high[i:i+1].to(torch.float32)
+            rec_full  = rec[i:i+1].to(torch.float32)
+            if rec_full.shape[-2:] != orig_full.shape[-2:]:
+                rec_full = F.interpolate(rec_full, size=orig_full.shape[-2:], mode="bilinear", align_corners=False)
+            lpips_val = lpips_net(orig_full, rec_full).item()
+            lpips_scores.append(lpips_val)
+        avg_lpips = float(np.mean(lpips_scores))
+
+        if use_wandb and accelerator.is_main_process:
+            wandb.log({"lpips_mean": avg_lpips}, step=step)
+    finally:
+        gc.collect()
+        torch.cuda.empty_cache()
+
+if accelerator.is_main_process and save_model:
+    print("Генерация сэмплов до старта обучения...")
+    generate_and_save_samples(0)
+
+accelerator.wait_for_everyone()
+
+# --------------------------- Тренировка ---------------------------
+progress = tqdm(total=total_steps, disable=not accelerator.is_local_main_process)
+global_step = 0
+min_loss = float("inf")
+sample_interval = max(1, total_steps // max(1, sample_interval_share * num_epochs))
+
+for epoch in range(num_epochs):
+    vae.train()
+    batch_losses, batch_grads = [], []
+    track_losses = {k: [] for k in loss_ratios.keys()}
+
+    for imgs in dataloader:
+        with accelerator.accumulate(vae):
+            imgs = imgs.to(accelerator.device)
+
+            if high_resolution != model_resolution:
+                imgs_low = F.interpolate(imgs, size=(model_resolution, model_resolution), mode="bilinear", align_corners=False)
+            else:
+                imgs_low = imgs
+
+            model_dtype = next(vae.parameters()).dtype
+            imgs_low_model = imgs_low.to(dtype=model_dtype) if imgs_low.dtype != model_dtype else imgs_low
+
+            # QWEN: encode/decode с T=1
+            if is_qwen_vae(vae):
+                x_in = imgs_low_model.unsqueeze(2)             # [B,3,1,H,W]
+                enc = vae.encode(x_in)
+                latents = enc.latent_dist.mean if train_decoder_only else enc.latent_dist.sample()
+                dec = vae.decode(latents).sample               # [B,3,1,H,W]
+                rec = dec.squeeze(2)                           # [B,3,H,W]
+            else:
+                enc = vae.encode(imgs_low_model)
+                latents = enc.latent_dist.mean if train_decoder_only else enc.latent_dist.sample()
+                rec = vae.decode(latents).sample
+
+            if rec.shape[-2:] != imgs.shape[-2:]:
+                rec = F.interpolate(rec, size=imgs.shape[-2:], mode="bilinear", align_corners=False)
+
+            rec_f32 = rec.to(torch.float32)
+            imgs_f32 = imgs.to(torch.float32)
+
+            abs_losses = {
+                "mae":   F.l1_loss(rec_f32, imgs_f32),
+                "mse":   F.mse_loss(rec_f32, imgs_f32),
+                "lpips": _get_lpips()(rec_f32, imgs_f32).mean(),
+                "edge":  F.l1_loss(sobel_edges(rec_f32), sobel_edges(imgs_f32)),
+            }
+
+            if full_training and not train_decoder_only:
+                mean   = enc.latent_dist.mean
+                logvar = enc.latent_dist.logvar
+                kl = -0.5 * torch.mean(1 + logvar - mean.pow(2) - logvar.exp())
+                abs_losses["kl"] = kl
+            else:
+                abs_losses["kl"] = torch.tensor(0.0, device=accelerator.device, dtype=torch.float32)
+
+            total_loss, coeffs, meds = normalizer.update_and_total(abs_losses)
+
+            if torch.isnan(total_loss) or torch.isinf(total_loss):
+                raise RuntimeError("NaN/Inf loss")
+
+            accelerator.backward(total_loss)
+
+            grad_norm = torch.tensor(0.0, device=accelerator.device)
+            if accelerator.sync_gradients:
+                grad_norm = accelerator.clip_grad_norm_(trainable_params, clip_grad_norm)
+                optimizer.step()
+                scheduler.step()
+                optimizer.zero_grad(set_to_none=True)
+                global_step += 1
+                progress.update(1)
+
+            if accelerator.is_main_process:
+                try:
+                    current_lr = optimizer.param_groups[0]["lr"]
+                except Exception:
+                    current_lr = scheduler.get_last_lr()[0]
+
+                batch_losses.append(total_loss.detach().item())
+                batch_grads.append(float(grad_norm.detach().cpu().item()) if isinstance(grad_norm, torch.Tensor) else float(grad_norm))
+                for k, v in abs_losses.items():
+                    track_losses[k].append(float(v.detach().item()))
+
+                if use_wandb and accelerator.sync_gradients:
+                    log_dict = {
+                        "total_loss": float(total_loss.detach().item()),
+                        "learning_rate": current_lr,
+                        "epoch": epoch,
+                        "grad_norm": batch_grads[-1],
+                        "mode/train_decoder_only": int(train_decoder_only),
+                        "mode/full_training": int(full_training),
+                    }
+                    for k, v in abs_losses.items():
+                        log_dict[f"loss_{k}"] = float(v.detach().item())
+                    for k in coeffs:
+                        log_dict[f"coeff_{k}"] = float(coeffs[k])
+                        log_dict[f"median_{k}"] = float(meds[k])
+                    wandb.log(log_dict, step=global_step)
+
+            if global_step > 0 and global_step % sample_interval == 0:
+                if accelerator.is_main_process:
+                    generate_and_save_samples(global_step)
+                accelerator.wait_for_everyone()
+
+                n_micro = sample_interval * gradient_accumulation_steps
+                avg_loss = float(np.mean(batch_losses[-n_micro:])) if len(batch_losses) >= n_micro else float(np.mean(batch_losses)) if batch_losses else float("nan")
+                avg_grad = float(np.mean(batch_grads[-n_micro:])) if len(batch_grads) >= 1 else float(np.mean(batch_grads)) if batch_grads else 0.0
+
+                if accelerator.is_main_process:
+                    print(f"Epoch {epoch} step {global_step} loss: {avg_loss:.6f}, grad_norm: {avg_grad:.6f}, lr: {current_lr:.9f}")
+                    if save_model and avg_loss < min_loss * save_barrier:
+                        min_loss = avg_loss
+                        accelerator.unwrap_model(vae).save_pretrained(save_as)
+                    if use_wandb:
+                        wandb.log({"interm_loss": avg_loss, "interm_grad": avg_grad}, step=global_step)
+
+    if accelerator.is_main_process:
+        epoch_avg = float(np.mean(batch_losses)) if batch_losses else float("nan")
+        print(f"Epoch {epoch} done, avg loss {epoch_avg:.6f}")
+        if use_wandb:
+            wandb.log({"epoch_loss": epoch_avg, "epoch": epoch + 1}, step=global_step)
+
+# --------------------------- Финальное сохранение ---------------------------
+if accelerator.is_main_process:
+    print("Training finished – saving final model")
+    if save_model:
+        accelerator.unwrap_model(vae).save_pretrained(save_as)
+
+accelerator.free_memory()
+if torch.distributed.is_initialized():
+    torch.distributed.destroy_process_group()
+print("Готово!")