sdxlvaev2

config.json

@@ -1,38 +0,0 @@
-{
-  "_class_name": "AutoencoderKL",
-  "_diffusers_version": "0.34.0",
-  "_name_or_path": "sdxl_vae",
-  "act_fn": "silu",
-  "block_out_channels": [
-    128,
-    256,
-    512,
-    512
-  ],
-  "down_block_types": [
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D"
-  ],
-  "force_upcast": false,
-  "in_channels": 3,
-  "latent_channels": 4,
-  "latents_mean": null,
-  "latents_std": null,
-  "layers_per_block": 2,
-  "mid_block_add_attention": true,
-  "norm_num_groups": 32,
-  "out_channels": 3,
-  "sample_size": 512,
-  "scaling_factor": 0.13025,
-  "shift_factor": null,
-  "up_block_types": [
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D"
-  ],
-  "use_post_quant_conv": true,
-  "use_quant_conv": true
-}

diffusion_pytorch_model.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:03f2412467f6bedce9efeddba5860b5ec0d3267931d14c500d4bd7a878e14cbd
-size 334643268

eval_alchemist.py

@@ -1,330 +1,520 @@
 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
-import lpips
-from PIL import Image, UnidentifiedImageError
-from tqdm import tqdm
 from torch.utils.data import Dataset, DataLoader
-from torchvision.transforms import Compose, Resize, ToTensor, CenterCrop,ToPILImage 
-from diffusers import AutoencoderKL, AsymmetricAutoencoderKL, AutoencoderKLWan,AutoencoderKLLTXVideo
-import random
+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 = "/workspace/alchemist" #wget https://huggingface.co/datasets/AiArtLab/alchemist/resolve/main/alchemist.zip
+IMAGE_FOLDER = "/home/recoilme/dataset/alchemist"
 MIN_SIZE = 1280
 CROP_SIZE = 512
 BATCH_SIZE = 10
 MAX_IMAGES = 0
 NUM_WORKERS = 4
-NUM_SAMPLES_TO_SAVE = 2  # Сколько примеров сохранить (0 - не сохранять)
-SAMPLES_FOLDER = "vaetest"
+SAMPLES_DIR = "vaetest"
 
-# Список VAE для тестирования
 VAE_LIST = [
-#    ("stable-diffusion-v1-5/stable-diffusion-v1-5", AutoencoderKL, "stable-diffusion-v1-5/stable-diffusion-v1-5", "vae"),
-#    ("cross-attention/asymmetric-autoencoder-kl-x-1-5", AsymmetricAutoencoderKL, "cross-attention/asymmetric-autoencoder-kl-x-1-5", None),
-#    ("madebyollin/sdxl-vae-fp16", AutoencoderKL, "madebyollin/sdxl-vae-fp16-fix", None),
-#    ("KBlueLeaf/EQ-SDXL-VAE", AutoencoderKL, "KBlueLeaf/EQ-SDXL-VAE", None),
-#    ("AiArtLab/sdxl_vae", AutoencoderKL, "AiArtLab/sdxl_vae", None),
-#    ("AiArtLab/sdxlvae_nightly", AutoencoderKL, "AiArtLab/sdxl_vae", "vae_nightly"),
-#    ("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, "AiArtLab/simplevae", "vae"),
-    ("simple_vae2", AutoencoderKL, "AiArtLab/simplevae", None),
-    ("simple_vae_nightly", AutoencoderKL, "/workspace/sdxl_vae/simple_vae_nightly", None),
-    
+    # ("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"),    
+    #("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"),
+    #("wan16x_vae_nightly", AutoencoderKLWan, "AiArtLab/simplevae","wan16x_vae_nightly"),
+    #("wan16x_vae_nightly2", AutoencoderKLWan, "AiArtLab/simplevae","wan16x_vae_nightly2"),
+    #("SimpleVAE        ", AutoencoderKL, "AiArtLab/simplevae", None),
+    #("AuraDiffusion/16ch-vae", AutoencoderKL, "AuraDiffusion/16ch-vae", None),
+    #("SimpleVAE nightly", AutoencoderKL, "AiArtLab/simplevae", "simple_vae_nightly"),
+    #("SimpleVAE nightly2", AutoencoderKL, "AiArtLab/simplevae", "simple_vae_nightly2"),    
+    ("aiartlab/SDXLVAE", AutoencoderKL, "/home/recoilme/vae", "sdxlvae"),
 ]
 
-# --------------------------- Sobel Edge Detection ---------------------------
-# Определяем фильтры Собеля глобально
-_sobel_kx = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).view(1, 1, 3, 3)
-_sobel_ky = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32).view(1, 1, 3, 3)
 
-def sobel_edges(x: torch.Tensor) -> torch.Tensor:
-    """
-    Вычисляет карту границ с помощью оператора Собеля
-    x: [B,C,H,W] в диапазоне [-1,1]
-    Возвращает: [B,C,H,W] - магнитуда градиента
-    """
-    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)
-
-def compute_edge_loss(real: torch.Tensor, fake: torch.Tensor) -> float:
+# ========================== Утилиты ==========================
+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):
     """
-    Вычисляет Edge Loss между реальным и сгенерированным изображением
-    real, fake: [B,C,H,W] в диапазоне [0,1]
-    Возвращает: скалярное значение loss
+    Нормализация латентов: глобальная и поканальная.
+    Применение: сначала глобальная (scalar), затем поканальная (vector).
+    Если в конфиге есть несколько ключей — аккумулируем.
     """
-    # Конвертируем в [-1,1] для sobel_edges
-    real_norm = real * 2 - 1
-    fake_norm = fake * 2 - 1
-    
-    # Получаем карты границ
-    edges_real = sobel_edges(real_norm)
-    edges_fake = sobel_edges(fake_norm)
+    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
     
-    # L1 loss между картами границ
-    return F.l1_loss(edges_fake, edges_real).item()
+    scale_channel = np.ones(C, dtype=np.float32)
+    shift_channel = np.zeros(C, dtype=np.float32)
 
-# --------------------------- Dataset ---------------------------
+    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, extensions=('.png',), min_size=1024, crop_size=512, limit=None):
-        self.root_dir = root_dir
-        self.min_size = min_size
-        self.crop_size = crop_size
-        self.paths = []
-        
-        print("Сканирование папки...")
+    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):
-                    self.paths.append(os.path.join(root, fname))
-        
+                    paths.append(os.path.join(root, fname))
         if limit:
-            self.paths = self.paths[:limit]
-        
-        print("Проверка изображений...")
+            paths = paths[:limit]
+
         valid = []
-        for p in tqdm(self.paths, desc="Проверка"):
+        for p in tqdm(paths, desc="Проверяем файлы"):
             try:
                 with Image.open(p) as im:
                     im.verify()
                 valid.append(p)
-            except:
-                continue
+            except Exception:
+                pass
+        if not valid:
+            raise RuntimeError(f"Нет валидных изображений в {root_dir}")
+        random.shuffle(valid)
         self.paths = valid
-        
-        if len(self.paths) == 0:
-            raise RuntimeError(f"Не найдено валидных изображений в {root_dir}")
-        
-        random.shuffle(self.paths)
         print(f"Найдено {len(self.paths)} изображений")
-        
+
         self.transform = Compose([
-            Resize(min_size, interpolation=Image.LANCZOS),
+            Resize(min_size),
             CenterCrop(crop_size),
-            ToTensor(),
+            ToTensor(),  # 0..1, float32
         ])
-    
+
     def __len__(self):
         return len(self.paths)
-    
+
     def __getitem__(self, idx):
-        path = self.paths[idx]
-        with Image.open(path) as img:
+        with Image.open(self.paths[idx]) as img:
             img = img.convert("RGB")
             return self.transform(img)
 
-# --------------------------- Функции ---------------------------
-def process(x):
-    return x * 2 - 1
 
-def deprocess(x):
-    return x * 0.5 + 0.5
+# ========================== Основное ==========================
+def main():
+    torch.set_grad_enabled(False)
+    os.makedirs(SAMPLES_DIR, exist_ok=True)
 
-def _sanitize_name(name: str) -> str:
-    return name.replace('/', '_').replace('-', '_')
+    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)
 
-# --------------------------- Анализ VAE ---------------------------
-@torch.no_grad()
-def tensor_stats(name, x: torch.Tensor):
-    finite = torch.isfinite(x)
-    fin_ratio = finite.float().mean().item()
-    x_f = x[finite]
-    minv = x_f.min().item() if x_f.numel() else float('nan')
-    maxv = x_f.max().item() if x_f.numel() else float('nan')
-    mean = x_f.mean().item() if x_f.numel() else float('nan')
-    std  = x_f.std().item()  if x_f.numel() else float('nan')
-    big  = (x_f.abs() > 20).float().mean().item() if x_f.numel() else float('nan')
-    print(f"[{name}] shape={tuple(x.shape)} dtype={x.dtype} "
-          f"finite={fin_ratio:.6f} min={minv:.4g} max={maxv:.4g} mean={mean:.4g} std={std:.4g} |x|>20={big:.6f}")
+    lpips_net = lpips.LPIPS(net="vgg").to(DEVICE).eval()
 
-@torch.no_grad()
-def analyze_vae_latents(vae, name, images):
-    """
-    images: [B,3,H,W] в [-1,1]
-    """
-    try:
-        enc = vae.encode(images)
-        if hasattr(enc, "latent_dist"):
-            mu, logvar = enc.latent_dist.mean, enc.latent_dist.logvar
-            z = enc.latent_dist.sample()
-        else:
-            mu, logvar = enc[0], enc[1]
-            z = mu
-        tensor_stats(f"{name}.mu", mu)
-        tensor_stats(f"{name}.logvar", logvar)
-        tensor_stats(f"{name}.z_raw", z)
+    # Загрузка 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}")
 
-        sf = getattr(vae.config, "scaling_factor", 1.0)
-        z_scaled = z * sf
-        tensor_stats(f"{name}.z_scaled(x{sf})", z_scaled)
-    except Exception as e:
-        print(f"⚠️ Ошибка анализа VAE {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__ == "__main__":
-    if NUM_SAMPLES_TO_SAVE > 0:
-        os.makedirs(SAMPLES_FOLDER, exist_ok=True)
-        
-    dataset = ImageFolderDataset(
-        IMAGE_FOLDER, 
-        extensions=('.png',), 
-        min_size=MIN_SIZE,
-        crop_size=CROP_SIZE,
-        limit=MAX_IMAGES
-    )
-    
-    dataloader = DataLoader(
-        dataset,
-        batch_size=BATCH_SIZE,
-        shuffle=False,
-        num_workers=NUM_WORKERS,
-        pin_memory=True,
-        drop_last=False
-    )
-    
-    lpips_net = lpips.LPIPS(net="vgg").eval().to(DEVICE).requires_grad_(False)
-    
-    print("\nЗагрузка VAE моделей...")
-    vaes = []
-    names = []
-    
-    for name, vae_class, model_path, subfolder in VAE_LIST:
-        try:
-            print(f"  Загружаю {name}...")
-            # Исправлена загрузка для variant
-            if "sdxs" in model_path:
-                vae = vae_class.from_pretrained(model_path, subfolder=subfolder, variant="fp16")
-            else:
-                vae = vae_class.from_pretrained(model_path, subfolder=subfolder)
-            vae = vae.to(DEVICE, DTYPE).eval()
-            vaes.append(vae)
-            names.append(name)
-        except Exception as e:
-            print(f"  ❌ Ошибка загрузки {name}: {e}")
-    
-    print("\nОценка метрик...")
-    results = {name: {"mse": 0.0, "psnr": 0.0, "lpips": 0.0, "edge": 0.0, "count": 0} for name in names}
-
-    to_pil = ToPILImage()
-
-    # >>>>>>>> ОСНОВНЫЕ ИЗМЕНЕНИЯ ЗДЕСЬ (KISS) <<<<<<<<
-    with torch.no_grad():
-        images_saved = 0  # считаем именно КОЛ-ВО ИЗОБРАЖЕНИЙ, а не сохранённых файлов
-        for batch in tqdm(dataloader, desc="Обработка батчей"):
-            batch = batch.to(DEVICE)                                # [B,3,H,W] в [0,1]
-            test_inp = process(batch).to(DTYPE)                     # [-1,1] для энкодера
-            # >>> Анализируем латенты каждой VAE на первой итерации
-            if images_saved == 0:   # только для первого батча, чтобы не засорять лог
-                for vae, name in zip(vaes, names):
-                    analyze_vae_latents(vae, name, test_inp)
-
-            # 1) считаем реконструкции для всех VAE на весь батч
-            recon_list = []
-            for vae, name in zip(vaes, names):
-                test_inp_vae = test_inp  # локальная копия
-                #if name == "Wan2.2-T2V-A14B-Diffusers" and test_inp_vae.ndim == 4:
-                if (isinstance(vae, AutoencoderKLWan) or isinstance(vae, AutoencoderKLLTXVideo)) and test_inp_vae.ndim == 4:
-                    test_inp_vae = test_inp_vae.unsqueeze(2)   # только для Wan
-                latent = vae.encode(test_inp_vae).latent_dist.mode()
-                dec = vae.decode(latent).sample.float()
-                if dec.ndim == 5:
-                    dec = dec.squeeze(2)
-                recon = deprocess(dec).clamp(0.0, 1.0)
-                recon_list.append(recon)
-
-            # 2) обновляем метрики (по каждой VAE)
-            for recon, name in zip(recon_list, names):
-                for i in range(batch.shape[0]):
-                    img_orig = batch[i:i+1]
-                    img_recon = recon[i:i+1]
-                    mse = F.mse_loss(img_orig, img_recon).item()
-                    psnr = 10 * torch.log10(1 / torch.tensor(mse)).item()
-                    lpips_val = lpips_net(img_orig, img_recon, normalize=True).mean().item()
-                    edge_loss = compute_edge_loss(img_orig, img_recon)
-                    results[name]["mse"] += mse
-                    results[name]["psnr"] += psnr
-                    results[name]["lpips"] += lpips_val
-                    results[name]["edge"] += edge_loss
-                    results[name]["count"] += 1
-
-            # 3) сохраняем ровно NUM_SAMPLES_TO_SAVE изображений (orig + все VAE + общий коллаж)
-            if NUM_SAMPLES_TO_SAVE > 0:
-                for i in range(batch.shape[0]):
-                    if images_saved >= NUM_SAMPLES_TO_SAVE:
-                        break
-                    idx_str = f"{images_saved + 1:03d}"
-
-                    # original
-                    orig_pil = to_pil(batch[i].detach().float().cpu())
-                    orig_pil.save(os.path.join(SAMPLES_FOLDER, f"{idx_str}_orig.png"))
-
-                    # per-VAE decodes
-                    tiles = [orig_pil]
-                    for recon, name in zip(recon_list, names):
-                        recon_pil = to_pil(recon[i].detach().cpu())
-                        recon_pil.save(os.path.join(
-                            SAMPLES_FOLDER, f"{idx_str}_decoded_{_sanitize_name(name)}.png"
-                        ))
-                        tiles.append(recon_pil)
-
-                    # общий коллаж: [orig | vae1 | vae2 | ...]
-                    collage_w = CROP_SIZE * len(tiles)
-                    collage_h = CROP_SIZE
-                    collage = Image.new("RGB", (collage_w, collage_h))
-                    x = 0
-                    for tile in tiles:
-                        collage.paste(tile, (x, 0))
-                        x += CROP_SIZE
-                    collage.save(os.path.join(SAMPLES_FOLDER, f"{idx_str}_all.png"))
-
-                    images_saved += 1
-
-
-    # Усреднение результатов
-    for name in names:
-        count = results[name]["count"]
-        results[name]["mse"] /= count
-        results[name]["psnr"] /= count
-        results[name]["lpips"] /= count
-        results[name]["edge"] /= count
-    
-    # Вывод абсолютных значений
-    print("\n=== Абсолютные значения ===")
-    for name in names:
-        print(f"{name:30s}: MSE: {results[name]['mse']:.3e}, PSNR: {results[name]['psnr']:.4f}, "
-              f"LPIPS: {results[name]['lpips']:.4f}, Edge: {results[name]['edge']:.4f}")
-    
-    # Вывод таблицы с процентами
-    print("\n=== Сравнение с первой моделью (%) ===")
-    print(f"| {'Модель':30s} | {'MSE':>10s} | {'PSNR':>10s} | {'LPIPS':>10s} | {'Edge':>10s} |")
-    print(f"|{'-'*32}|{'-'*12}|{'-'*12}|{'-'*12}|{'-'*12}|")
-    
-    baseline = names[0]
-    for name in names:
-        # Для MSE, LPIPS и Edge: меньше = лучше, поэтому инвертируем
-        mse_pct = (results[baseline]["mse"] / results[name]["mse"]) * 100
-        # Для PSNR: больше = лучше
-        psnr_pct = (results[name]["psnr"] / results[baseline]["psnr"]) * 100
-        # Для LPIPS и Edge: меньше = лучше
-        lpips_pct = (results[baseline]["lpips"] / results[name]["lpips"]) * 100
-        edge_pct = (results[baseline]["edge"] / results[name]["edge"]) * 100
-        
         if name == baseline:
-            print(f"| {name:30s} | {'100%':>10s} | {'100%':>10s} | {'100%':>10s} | {'100%':>10s} |")
+            print(f"| {name:26s} | {'100%':>9s} | {'100%':>9s} | {'100%':>9s} | {'100%':>9s} | {'100%':>9s} | {'100%':>9s} |")
         else:
-            print(f"| {name:30s} | {f'{mse_pct:.1f}%':>10s} | {f'{psnr_pct:.1f}%':>10s} | "
-                  f"{f'{lpips_pct:.1f}%':>10s} | {f'{edge_pct:.1f}%':>10s} |")
-    
-    print("\n✅ Готово!")
+            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()

eval_alchemist2.py

@@ -1,516 +0,0 @@
-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 = "/workspace/alchemist"
-MIN_SIZE = 1280
-CROP_SIZE = 512
-BATCH_SIZE = 10
-MAX_IMAGES = 0
-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 ", AutoencoderKL, "AiArtLab/simplevae", None),
-    #("SimpleVAE nightly", AutoencoderKL, "AiArtLab/simplevae", "simple_vae_nightly"),
-    ("FLUX.1-schnell VAE", AutoencoderKL, "black-forest-labs/FLUX.1-schnell", "vae"),
-    ("SimpleVAE nightly", AutoencoderKL, "AiArtLab/simplevae", "simple_vae_nightly"),
-    # ("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()

sdxl_vae_a1111.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ebe642d26e14851e98eb3d06575009e0d1a669704a1c9c8dcf06573d82233a21
+oid sha256:8f8e696f579f70d185f4b944a0d821ab5578a0915ac079fe44c148ce5102cc5b
 size 334640988

simple_vae/config.json

@@ -1,38 +0,0 @@
-{
-  "_class_name": "AutoencoderKL",
-  "_diffusers_version": "0.35.0.dev0",
-  "_name_or_path": "simple_vae",
-  "act_fn": "silu",
-  "block_out_channels": [
-    128,
-    256,
-    512,
-    512
-  ],
-  "down_block_types": [
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D"
-  ],
-  "force_upcast": false,
-  "in_channels": 3,
-  "latent_channels": 16,
-  "latents_mean": null,
-  "latents_std": null,
-  "layers_per_block": 2,
-  "mid_block_add_attention": true,
-  "norm_num_groups": 32,
-  "out_channels": 3,
-  "sample_size": 1024,
-  "scaling_factor": 1.0,
-  "shift_factor": 0,
-  "up_block_types": [
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D"
-  ],
-  "use_post_quant_conv": true,
-  "use_quant_conv": true
-}

simple_vae/diffusion_pytorch_model.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8ba1d500c4bd376a7c8662a35fa1857c7e577da0635414b524180852143ef2f6
-size 335311892

simple_vae_nightly/config.json

@@ -1,38 +0,0 @@
-{
-  "_class_name": "AutoencoderKL",
-  "_diffusers_version": "0.35.0.dev0",
-  "_name_or_path": "simple_vae",
-  "act_fn": "silu",
-  "block_out_channels": [
-    128,
-    256,
-    512,
-    512
-  ],
-  "down_block_types": [
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D"
-  ],
-  "force_upcast": false,
-  "in_channels": 3,
-  "latent_channels": 16,
-  "latents_mean": null,
-  "latents_std": null,
-  "layers_per_block": 2,
-  "mid_block_add_attention": true,
-  "norm_num_groups": 32,
-  "out_channels": 3,
-  "sample_size": 1024,
-  "scaling_factor": 1.0,
-  "shift_factor": 0,
-  "up_block_types": [
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D"
-  ],
-  "use_post_quant_conv": true,
-  "use_quant_conv": true
-}

simple_vae_nightly/diffusion_pytorch_model.safetensors

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

train_sdxl_vae.py

@@ -1,547 +0,0 @@
-# -*- 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
-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            = "simple_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
-asymmetric         = False
-optimizer_type     = "adam8bit"
-dtype              = torch.float32
-# model_resolution — то, что подавается в VAE (низкое разрешение)
-model_resolution   = 512   # бывший `resolution`
-# high_resolution — настоящий «высокий» кроп, на котором считаем метрики и сохраняем сэмплы
-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"   # или "fp16"/"bf16" при поддержке
-gradient_accumulation_steps = 5
-generated_folder   = "samples"
-save_as            = "simple_vae_nightly"
-num_workers        = 0
-device = None  # accelerator задаст устройство
-
-# --- Пропорции лоссов и окно медианного нормирования (КОЭФ., не значения) ---
-# Итоговые доли в total loss (сумма = 1.0)
-loss_ratios = {
-    "lpips": 0.85,
-    "edge":  0.05,
-    "mse":   0.05,
-    "mae":   0.05,
-}
-median_coeff_steps = 256  # за сколько шагов считать медианные коэффициенты
-
-# --------------------------- параметры препроцессинга ---------------------------
-resize_long_side = 1280  # если None или 0 — ресайза не будет; рекомендовано 1280
-
-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,
-    })
-
-# --------------------------- VAE ---------------------------
-if model_resolution==high_resolution and not asymmetric:
-    vae = AutoencoderKL.from_pretrained(project).to(dtype)
-else:
-    vae = AsymmetricAutoencoderKL.from_pretrained(project).to(dtype)
-
-# torch.compile (если доступно) — просто и без лишней логики
-if hasattr(torch, "compile"):
-    try:
-        vae = torch.compile(vae)
-    except Exception as e:
-        print(f"[WARN] torch.compile failed: {e}")
-
-# >>> Заморозка всех параметров, затем выборочная разморозка
-for p in vae.parameters():
-    p.requires_grad = False
-
-decoder = getattr(vae, "decoder", None)
-if decoder is None:
-    raise RuntimeError("vae.decoder not found — не могу применить стратегию разморозки. Проверь структуру модели.")
-
-unfrozen_param_names = []
-
-if not hasattr(decoder, "up_blocks"):
-    raise RuntimeError("decoder.up_blocks не найдены — ожидается список блоков декодера.")
-
-# >>> Размораживаем все up_blocks и mid_block (как было в твоём варианте start_idx=0)
-n_up = len(decoder.up_blocks)
-start_idx = 0
-for idx in range(start_idx, n_up):
-    block = decoder.up_blocks[idx]
-    for name, p in block.named_parameters():
-        p.requires_grad = True
-        unfrozen_param_names.append(f"decoder.up_blocks.{idx}.{name}")
-
-if hasattr(decoder, "mid_block"):
-    for name, p in decoder.mid_block.named_parameters():
-        p.requires_grad = True
-        unfrozen_param_names.append(f"decoder.mid_block.{name}")
-else:
-    print("[WARN] decoder.mid_block не найден — mid_block не разморожен.")
-
-print(f"[INFO] Разморожено параметров: {len(unfrozen_param_names)}. Первые 200 имён:")
-for nm in unfrozen_param_names[:200]:
-    print("  ", nm)
-
-# сохраняем trainable_module (get_param_groups будет учитывать p.requires_grad)
-trainable_module = vae.decoder
-
-# --------------------------- Custom PNG Dataset (only .png, skip corrupted) -----------
-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 = []
-        # collect png files recursively
-        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))
-        # optional limit
-        if limit:
-            self.paths = self.paths[:limit]
-        # verify images and keep only valid ones
-        valid = []
-        for p in self.paths:
-            try:
-                with Image.open(p) as im:
-                    im.verify()  # fast check for truncated/corrupted images
-                valid.append(p)
-            except (OSError, UnidentifiedImageError):
-                # skip corrupted image
-                continue
-        self.paths = valid
-        if len(self.paths) == 0:
-            raise RuntimeError(f"No valid PNG images found under {root_dir}")
-        # final shuffle for randomness
-        random.shuffle(self.paths)
-
-    def __len__(self):
-        return len(self.paths)
-
-    def __getitem__(self, idx):
-        p = self.paths[idx % len(self.paths)]
-        # open and convert to RGB; ensure file is closed promptly
-        with Image.open(p) as img:
-            img = img.convert("RGB")
-            # пережимаем длинную сторону до resize_long_side (Lanczos)
-            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])
-])
-
-# build dataset using high_resolution crops
-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}")
-
-# collate_fn кропит до high_resolution
-
-def collate_fn(batch):
-    imgs = []
-    for img in batch:  # img is PIL.Image
-        img = random_crop(img, high_resolution)   # кропим high-res
-        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 module.named_parameters():
-        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)
-
-# --------------------------- Подготовка Accelerate (вместе) ---------------------------
-
-batches_per_epoch = len(dataloader)  # число микро-батчей (dataloader steps)
-steps_per_epoch = int(math.ceil(batches_per_epoch / float(gradient_accumulation_steps)))  # чис��о optimizer.step() за эпоху
-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.decoder.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
-
-# Собель для edge loss
-_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:
-    # x: [B,C,H,W] в [-1,1]
-    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):
-        # нормируем доли на случай, если сумма != 1
-        s = sum(desired_ratios.values())
-        self.ratios = {k: (v / s) 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().cpu()))
-        # Медианы (устойчивые к выбросам)
-        meds = {k: (np.median(self.buffers[k]) if len(self.buffers[k]) > 0 else 1.0) for k in self.buffers}
-        # Вычисляем КОЭФФИЦИЕНТЫ как ratio_k / median_k — т.е. именно коэффициенты, а не значения
-        coeffs = {k: (self.ratios[k] / max(meds[k], 1e-12)) for k in self.ratios}
-        # Важно: при таких коэффициентах сумма (coeff_k * median_k) = сумма(ratio_k) = 1, т.е. масштаб стабилен
-        total = sum(coeffs[k] * abs_losses[k] for k in coeffs)
-        return total, coeffs, meds
-
-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]  # dataset returns PIL.Image
-    tensors = []
-    for img in pil_imgs:
-        img = random_crop(img, high_resolution)  # high-res fixed samples
-        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:
-    # img_tensor: [C,H,W] in [-1,1]
-    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():
-            # Готовим low-res вход для кодера ВСЕГДА под model_resolution
-            orig_high = fixed_samples  # [B,C,H,W] в [-1,1]
-            orig_low = F.interpolate(orig_high, size=(model_resolution, model_resolution), mode="bilinear", align_corners=False)
-            # dtype как у модели
-            model_dtype = next(temp_vae.parameters()).dtype
-            orig_low = orig_low.to(dtype=model_dtype)
-            # encode/decode
-            latents = temp_vae.encode(orig_low).latent_dist.mean
-            rec = temp_vae.decode(latents).sample
-
-        # Приводим spatial размер рекона к high-res (downsample для асимметричных VAE)
-        if rec.shape[-2:] != orig_high.shape[-2:]:
-            rec = F.interpolate(rec, size=orig_high.shape[-2:], mode="bilinear", align_corners=False)
-
-        # Сохраняем ПЕРВЫЙ семпл: real и decoded без номера шага в имени
-        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 на полном изображении (high-res) — для лога
-        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: high-res tensor from dataloader ([-1,1]), move to device
-            imgs = imgs.to(accelerator.device)
-
-            # ВСЕГДА даунсемплим вход под model_resolution для кодера
-            # Тупая железяка норовит все по своему сделать
-            if high_resolution != model_resolution:
-                imgs_low = F.interpolate(imgs, size=(model_resolution, model_resolution), mode="bilinear", align_corners=False)
-            else:
-                imgs_low = imgs 
-
-            # ensure dtype matches model params to avoid float/half mismatch
-            model_dtype = next(vae.parameters()).dtype
-            if imgs_low.dtype != model_dtype:
-                imgs_low_model = imgs_low.to(dtype=model_dtype)
-            else:
-                imgs_low_model = imgs_low
-
-            # Encode/decode
-            latents = vae.encode(imgs_low_model).latent_dist.mean
-            rec = vae.decode(latents).sample  # rec может быть увеличенным (асимметричный VAE)
-
-            # Приводим размер к high-res
-            if rec.shape[-2:] != imgs.shape[-2:]:
-                rec = F.interpolate(rec, size=imgs.shape[-2:], mode="bilinear", align_corners=False)
-
-            # Лоссы считаем на high-res
-            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)),
-            }
-
-            # Total с медианными КОЭФФИЦИЕНТАМИ
-            # Не надо так орать когда у тебя получилось понять мою идею
-            total_loss, coeffs, meds = normalizer.update_and_total(abs_losses)
-
-            if torch.isnan(total_loss) or torch.isinf(total_loss):
-                print("NaN/Inf loss – stopping")
-                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 if isinstance(grad_norm, (float, int)) else grad_norm.cpu().item()))
-                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],
-                    }
-                    # добавляем отдельные лоссы
-                    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
-                if len(batch_losses) >= n_micro:
-                    avg_loss = float(np.mean(batch_losses[-n_micro:]))
-                else:
-                    avg_loss = 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("Готово!")

train_sdxl_vae_full.py

@@ -1,594 +0,0 @@
-# -*- 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
-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            = "simple_vae"
-batch_size         = 3
-base_learning_rate = 2e-6
-min_learning_rate  = 8e-7
-num_epochs         = 8
-sample_interval_share = 10
-use_wandb          = True
-save_model         = True
-use_decay          = True
-asymmetric         = False
-optimizer_type     = "adam8bit"
-dtype              = torch.float32
-# model_resolution — то, что подавается в VAE (низкое разрешение)
-model_resolution   = 512   # бывший `resolution`
-# high_resolution — настоящий «высокий» кроп, на котором считаем метрики и сохраняем сэмплы
-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"   # или "fp16"/"bf16" при поддержке
-gradient_accumulation_steps = 5
-generated_folder   = "samples"
-save_as            = "simple_vae_nightly"
-num_workers        = 0
-device = None  # accelerator задаст устройство
-
-# --------------------------- Тренировочные режимы ---------------------------
-# CHANGED: добавлен параметр для полного обучения VAE (а не только декодера).
-# Если False — поведение прежнее: учим только decoder.* (up_blocks + mid_block).
-# Если True  — размораживаем ВСЮ модель и добавляем KL-loss для энкодера.
-full_training      = False
-
-# CHANGED: добавлен вес (через долю в нормализаторе) для KL, используется только при full_training=True.
-kl_ratio           = 0.00  # простая доля для KL в общей смеси (KISS). Игнорируется, если full_training=False.
-
-# --- Пропорции лоссов и окно медианного нормирования (КОЭФ., не значения) ---
-# Итоговые доли в total loss (сумма = 1.0 после нормализации).
-loss_ratios = {
-    "lpips": 0.60,
-    "edge":  0.10,
-    "mse":   0.15,
-    "mae":   0.15,
-    # CHANGED: заранее добавлен ключ "kl" (по умолчанию 0.0). Если включаем full_training — активируем ниже.
-    "kl":    0.00,
-}
-median_coeff_steps = 256  # за сколько шагов считать медианные коэффициенты
-
-# --------------------------- параметры препроцессинга ---------------------------
-resize_long_side = 1280  # если None или 0 — ресайза не будет; рекомендовано 1280
-
-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,
-        "full_training": full_training,  # CHANGED: логируем режим
-        "kl_ratio": kl_ratio,            # CHANGED: логируем долю KL
-    })
-
-# --------------------------- VAE ---------------------------
-if model_resolution==high_resolution and not asymmetric:
-    vae = AutoencoderKL.from_pretrained(project).to(dtype)
-else:
-    vae = AsymmetricAutoencoderKL.from_pretrained(project).to(dtype)
-
-# torch.compile (если доступно) — просто и без лишней логики
-if hasattr(torch, "compile"):
-    try:
-        vae = torch.compile(vae)
-    except Exception as e:
-        print(f"[WARN] torch.compile failed: {e}")
-
-# >>> Стратегия заморозки / разморозки
-for p in vae.parameters():
-    p.requires_grad = False
-
-decoder = getattr(vae, "decoder", None)
-if decoder is None:
-    raise RuntimeError("vae.decoder not found — не могу применить стратегию разморозки. Проверь структуру модели.")
-
-unfrozen_param_names = []
-
-if not full_training:
-    # === Прежнее поведение: обучаем только decoder.up_blocks и decoder.mid_block ===
-    if not hasattr(decoder, "up_blocks"):
-        raise RuntimeError("decoder.up_blocks не найдены — ожидается список блоков декодера.")
-
-    n_up = len(decoder.up_blocks)
-    start_idx = 0
-    for idx in range(start_idx, n_up):
-        block = decoder.up_blocks[idx]
-        for name, p in block.named_parameters():
-            p.requires_grad = True
-            unfrozen_param_names.append(f"decoder.up_blocks.{idx}.{name}")
-
-    if hasattr(decoder, "mid_block"):
-        for name, p in decoder.mid_block.named_parameters():
-            p.requires_grad = True
-            unfrozen_param_names.append(f"decoder.mid_block.{name}")
-    else:
-        print("[WARN] decoder.mid_block не найден — mid_block не разморожен.")
-
-    # Обучаем только декодер
-    trainable_module = vae.decoder
-else:
-    # === CHANGED: Полное обучение — размораживаем ВСЕ слои VAE (и энкодер, и декодер, и пост-проекцию) ===
-    for name, p in vae.named_parameters():
-        p.requires_grad = True
-        unfrozen_param_names.append(name)
-    trainable_module = vae  # CHANGED: учим всю модель
-
-    # CHANGED: активируем KL-долю в нормализаторе
-    loss_ratios["kl"] = float(kl_ratio)
-
-print(f"[INFO] Разморожено параметров: {len(unfrozen_param_names)}. Первые 200 имён:")
-for nm in unfrozen_param_names[:200]:
-    print("  ", nm)
-
-# --------------------------- Custom PNG Dataset (only .png, skip corrupted) -----------
-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 = []
-        # collect png files recursively
-        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))
-        # optional limit
-        if limit:
-            self.paths = self.paths[:limit]
-        # verify images and keep only valid ones
-        valid = []
-        for p in self.paths:
-            try:
-                with Image.open(p) as im:
-                    im.verify()  # fast check for truncated/corrupted images
-                valid.append(p)
-            except (OSError, UnidentifiedImageError):
-                # skip corrupted image
-                continue
-        self.paths = valid
-        if len(self.paths) == 0:
-            raise RuntimeError(f"No valid PNG images found under {root_dir}")
-        # final shuffle for randomness
-        random.shuffle(self.paths)
-
-    def __len__(self):
-        return len(self.paths)
-
-    def __getitem__(self, idx):
-        p = self.paths[idx % len(self.paths)]
-        # open and convert to RGB; ensure file is closed promptly
-        with Image.open(p) as img:
-            img = img.convert("RGB")
-            # пережимаем длинную сторону до resize_long_side (Lanczos)
-            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])
-])
-
-# build dataset using high_resolution crops
-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}")
-
-# collate_fn кропит до high_resolution
-def collate_fn(batch):
-    imgs = []
-    for img in batch:  # img is PIL.Image
-        img = random_crop(img, high_resolution)   # кропим high-res
-        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 module.named_parameters():
-        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 ---------------------------
-
-batches_per_epoch = len(dataloader)  # число микро-батчей (dataloader steps)
-steps_per_epoch = int(math.ceil(batches_per_epoch / float(gradient_accumulation_steps)))  # число optimizer.step() за эпоху
-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)
-
-# CHANGED: формируем список trainable_params исходя из выбранного trainable_module
-trainable_params = [p for p in (trainable_module.parameters() if hasattr(trainable_module, "parameters") else []) 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
-
-# Собель для edge loss
-_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:
-    # x: [B,C,H,W] в [-1,1]
-    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):
-        # нормируем доли на случай, если сумма != 1
-        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}
-        # Вычисляем КОЭФФИЦИЕНТЫ как ratio_k / median_k — т.е. именно коэффициенты, а не значения
-        coeffs = {k: (self.ratios[k] / max(meds[k], 1e-12)) for k in self.ratios}
-        # Итоговый total — сумма по ключам, присутствующим в abs_losses
-        total = sum(coeffs[k] * abs_losses[k] for k in abs_losses if k in coeffs)
-        return total, coeffs, meds
-
-# CHANGED: создаём нормализатор ПОСЛЕ возможной активации 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]  # dataset returns PIL.Image
-    tensors = []
-    for img in pil_imgs:
-        img = random_crop(img, high_resolution)  # high-res fixed samples
-        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:
-    # img_tensor: [C,H,W] in [-1,1]
-    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():
-            # Готовим low-res вход для кодера ВСЕГДА под model_resolution
-            orig_high = fixed_samples  # [B,C,H,W] в [-1,1]
-            orig_low = F.interpolate(orig_high, size=(model_resolution, model_resolution), mode="bilinear", align_corners=False)
-            # dtype как у модели
-            model_dtype = next(temp_vae.parameters()).dtype
-            orig_low = orig_low.to(dtype=model_dtype)
-            # encode/decode
-            # CHANGED: при валидации/сэмплах всегда используем mean (стабильно и детерминированно)
-            enc = temp_vae.encode(orig_low)
-            latents_mean = enc.latent_dist.mean
-            rec = temp_vae.decode(latents_mean).sample
-
-        # Приводим spatial размер рекона к high-res (downsample для асимметричных VAE)
-        if rec.shape[-2:] != orig_high.shape[-2:]:
-            rec = F.interpolate(rec, size=orig_high.shape[-2:], mode="bilinear", align_corners=False)
-
-        # Сохраняем ПЕРВЫЙ семпл: real и decoded без номера шага в имени
-        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 на полном изображении (high-res) — для лога
-        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: high-res tensor from dataloader ([-1,1]), move to device
-            imgs = imgs.to(accelerator.device)
-
-            # ВСЕГДА даунсемплим вход под model_resolution для кодера
-            if high_resolution != model_resolution:
-                imgs_low = F.interpolate(imgs, size=(model_resolution, model_resolution), mode="bilinear", align_corners=False)
-            else:
-                imgs_low = imgs 
-
-            # ensure dtype matches model params to avoid float/half mismatch
-            model_dtype = next(vae.parameters()).dtype
-            imgs_low_model = imgs_low.to(dtype=model_dtype) if imgs_low.dtype != model_dtype else imgs_low
-
-            # Encode/decode
-            enc = vae.encode(imgs_low_model)
-
-            # CHANGED: если тренируем всю модель — используем reparameterization sample()
-            # это важно для стохастичности и согласованности с KL.
-            latents = enc.latent_dist.sample() if full_training else enc.latent_dist.mean
-
-            rec = vae.decode(latents).sample  # rec может быть увеличенным (асимметричный VAE)
-
-            # Приводим размер к high-res
-            if rec.shape[-2:] != imgs.shape[-2:]:
-                rec = F.interpolate(rec, size=imgs.shape[-2:], mode="bilinear", align_corners=False)
-
-            # Лоссы считаем на high-res
-            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)),
-            }
-
-            # CHANGED: KL-loss добавляется ТОЛЬКО при полном обучении.
-            # KL(q(z|x) || N(0,1)) = -0.5 * sum(1 + logσ^2 - μ^2 - σ^2).
-            if full_training:
-                mean   = enc.latent_dist.mean
-                logvar = enc.latent_dist.logvar
-                # стабильное усреднение по батчу и пространству
-                # СТАРОЕ (неправильное):
-                #kl = -0.5 * torch.mean(1 + logvar - mean.pow(2) - logvar.exp())
-                # НОВОЕ (правильное):
-                kl_per_sample = -0.5 * torch.sum(1 + logvar - mean.pow(2) - logvar.exp(), dim=[1, 2, 3])
-                kl = torch.mean(kl_per_sample)
-                abs_losses["kl"] = kl
-            else:
-                # ключ присутствует в ratios, но при partial-training его доля = 0 и он не влияет
-                abs_losses["kl"] = torch.tensor(0.0, device=accelerator.device, dtype=torch.float32)
-
-            # Total с медианными КОЭФФИЦИЕНТАМИ
-            total_loss, coeffs, meds = normalizer.update_and_total(abs_losses)
-
-            if torch.isnan(total_loss) or torch.isinf(total_loss):
-                print("NaN/Inf loss – stopping")
-                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())
-                # CHANGED: корректно извлекаем scalar из разн. типов
-                if isinstance(grad_norm, torch.Tensor):
-                    batch_grads.append(float(grad_norm.detach().cpu().item()))
-                else:
-                    batch_grads.append(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/full_training": int(full_training),  # CHANGED: для наглядности в логах
-                    }
-                    # добавляем отдельные лоссы
-                    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
-                if len(batch_losses) >= n_micro:
-                    avg_loss = float(np.mean(batch_losses[-n_micro:]))
-                else:
-                    avg_loss = 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("Готово!")

train_sdxl_vae_my.py

@@ -1,507 +0,0 @@
-# -*- 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
-from accelerate import Accelerator
-from PIL import Image, UnidentifiedImageError
-from tqdm import tqdm
-import bitsandbytes as bnb
-import wandb
-import lpips   # pip install lpips
-
-# --------------------------- Параметры ---------------------------
-ds_path            = "/workspace/png"
-project            = "asymmetric_vae"
-batch_size         = 2
-base_learning_rate = 1e-6
-min_learning_rate  = 8e-7
-num_epochs         = 8
-sample_interval_share = 10
-use_wandb          = True
-save_model         = True
-use_decay          = True
-asymmetric         = True
-optimizer_type     = "adam8bit"
-dtype              = torch.float32
-# model_resolution — то, что подавается в VAE (низкое разрешение)
-model_resolution   = 512   # бывший `resolution`
-# high_resolution — настоящий «высокий» кроп, на котором считаем метрики и сохраняем сэмплы
-high_resolution    = 1024 
-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"   # или "fp16"/"bf16" при поддержке
-gradient_accumulation_steps = 8
-generated_folder   = "samples"
-save_as            = "asymmetric_vae_new"
-perceptual_loss_weight = 0.03  # начальное значение веса (будет перезаписываться каждый шаг)
-num_workers        = 0
-device = None  # accelerator задаст устройство
-
-# --- Параметры динамической нормализации LPIPS 
-lpips_ratio = 0.9 #percent lpips in loss
-
-min_perceptual_weight = 0.1    # минимальный предел веса
-max_perceptual_weight = 99     # максимальный предел веса (защита от взрывов)
-
-# --------------------------- параметры препроцессинга ---------------------------
-resize_long_side = 1280  # если None или 0 — ресайза не будет; рекомендовано 1024
-
-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 = True
-
-# --------------------------- 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,
-    })
-
-# --------------------------- VAE ---------------------------
-if model_resolution==high_resolution and not asymmetric:
-    vae = AutoencoderKL.from_pretrained(project).to(dtype)
-else:
-    vae = AsymmetricAutoencoderKL.from_pretrained(project).to(dtype)
-
-# >>> CHANGED: заморозка всех параметров, затем разморозка mid_block + up_blocks[-2:]
-for p in vae.parameters():
-    p.requires_grad = False
-
-decoder = getattr(vae, "decoder", None)
-if decoder is None:
-    raise RuntimeError("vae.decoder not found — не могу применить стратегию разморозки. Проверь структуру модели.")
-
-unfrozen_param_names = []
-
-if not hasattr(decoder, "up_blocks"):
-    raise RuntimeError("decoder.up_blocks не найдены — ожидается список блоков декодера.")
-
-# >>> CHANGED: размораживаем последние 2 up_blocks (как просил) и mid_block
-n_up = len(decoder.up_blocks)
-start_idx = 0 #max(0, n_up - 2) # all
-for idx in range(start_idx, n_up):
-    block = decoder.up_blocks[idx]
-    for name, p in block.named_parameters():
-        p.requires_grad = True
-        unfrozen_param_names.append(f"decoder.up_blocks.{idx}.{name}")
-
-if hasattr(decoder, "mid_block"):
-    for name, p in decoder.mid_block.named_parameters():
-        p.requires_grad = True
-        unfrozen_param_names.append(f"decoder.mid_block.{name}")
-else:
-    print("[WARN] decoder.mid_block не найден — mid_block не разморожен.")
-
-print(f"[INFO] Разморожено параметров: {len(unfrozen_param_names)}. Первые 200 имён:")
-for nm in unfrozen_param_names[:200]:
-    print("  ", nm)
-
-# сохраняем trainable_module (get_param_groups будет учитывать p.requires_grad)
-trainable_module = vae.decoder
-
-# --------------------------- Custom PNG Dataset (only .png, skip corrupted) -----------
-class PngFolderDataset(Dataset):
-    def __init__(self, root_dir, min_exts=('.png',), resolution=1024, limit=0):
-        # >>> CHANGED: default resolution argument is high-resolution (1024)
-        self.root_dir = root_dir
-        self.resolution = resolution
-        self.paths = []
-        # collect png files recursively
-        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))
-        # optional limit
-        if limit:
-            self.paths = self.paths[:limit]
-        # verify images and keep only valid ones
-        valid = []
-        for p in self.paths:
-            try:
-                with Image.open(p) as im:
-                    im.verify()  # fast check for truncated/corrupted images
-                valid.append(p)
-            except (OSError, UnidentifiedImageError):
-                # skip corrupted image
-                continue
-        self.paths = valid
-        if len(self.paths) == 0:
-            raise RuntimeError(f"No valid PNG images found under {root_dir}")
-        # final shuffle for randomness
-        random.shuffle(self.paths)
-
-    def __len__(self):
-        return len(self.paths)
-
-    def __getitem__(self, idx):
-        p = self.paths[idx % len(self.paths)]
-        # open and convert to RGB; ensure file is closed promptly
-        with Image.open(p) as img:
-            img = img.convert("RGB")
-            # return PIL image (collate will transform)
-            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])
-])
-
-# build dataset using high_resolution crops
-dataset = PngFolderDataset(ds_path, min_exts=('.png',), resolution=high_resolution, limit=limit)  # >>> CHANGED
-if len(dataset) < batch_size:
-    raise RuntimeError(f"Not enough valid images ({len(dataset)}) to form a batch of size {batch_size}")
-
-# collate_fn кропит до high_resolution
-def collate_fn(batch):
-    imgs = []
-    for img in batch:  # img is PIL.Image
-        img = random_crop(img, high_resolution)   # >>> CHANGED: crop high-res
-        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 module.named_parameters():
-        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)
-
-# --------------------------- Подготовка Accelerate (вместе) ---------------------------
-batches_per_epoch = len(dataloader)  # число микро-батчей (dataloader steps)
-steps_per_epoch = int(math.ceil(batches_per_epoch / float(gradient_accumulation_steps)))  # число optimizer.step() за эпоху
-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.decoder.parameters() if p.requires_grad]
-
-# --------------------------- Сэмплы и LPIPS helper ---------------------------
-@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]  # dataset returns PIL.Image
-    tensors = []
-    for img in pil_imgs:
-        img = random_crop(img, high_resolution)  # >>> CHANGED: high-res fixed samples
-        tensors.append(tfm(img))
-    return torch.stack(tensors).to(accelerator.device, dtype)
-
-fixed_samples = get_fixed_samples()
-
-_lpips_net = None
-def _get_lpips():
-    global _lpips_net
-    if _lpips_net is None:
-        # lpips uses its internal vgg, but we use it as-is.
-        _lpips_net = lpips.LPIPS(net='vgg', verbose=False).eval().to(accelerator.device).eval()
-    return _lpips_net
-
-@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():
-            # >>> CHANGED: use high-res fixed_samples, downsample to model_res for encoding
-            orig_high = fixed_samples  # already on device
-            # make low-res input for model
-            if model_resolution==high_resolution:
-                orig_low = F.interpolate(orig_high, size=(model_resolution, model_resolution), mode="bilinear", align_corners=False)
-            else:
-                orig_low =orig_high 
-
-            # ensure dtype matches model params to avoid dtype mismatch
-            model_dtype = next(temp_vae.parameters()).dtype
-            orig_low = orig_low.to(dtype=model_dtype)
-
-            latent_dist = temp_vae.encode(orig_low).latent_dist
-            latents = latent_dist.mean
-            rec = temp_vae.decode(latents).sample  # expected to be upscaled to high_res
-
-        # make sure rec is float32 in range [0,1] for saving
-        # if rec spatial size differs from orig_high, resize rec to orig_high
-        if rec.shape[-2:] != orig_high.shape[-2:]:
-            rec = F.interpolate(rec, size=orig_high.shape[-2:], mode="bilinear", align_corners=False)
-
-        rec_img = ((rec.float() / 2.0 + 0.5).clamp(0, 1) * 255).cpu().numpy()
-        for i in range(rec_img.shape[0]):
-            arr = rec_img[i].transpose(1, 2, 0).astype(np.uint8)
-            Image.fromarray(arr).save(f"{generated_folder}/sample_{step if step is not None else 'init'}_{i}.jpg", quality=95)
-
-        # LPIPS на полном изображении (high-res)
-        lpips_scores = []
-        for i in range(rec.shape[0]):
-            orig_full = orig_high[i:i+1]  # [B, C, H, W], in [-1,1]
-            rec_full  = rec[i:i+1]
-            # ensure same spatial size/dtype
-            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)
-            rec_full = rec_full.to(torch.float32)
-            orig_full = orig_full.to(torch.float32)
-            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({
-                "generated_images": [wandb.Image(Image.fromarray(rec_img[i].transpose(1,2,0).astype(np.uint8))) for i in range(rec_img.shape[0])],
-                "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_losses_mae = []
-    batch_losses_lpips = []
-    batch_losses_perc = []
-    batch_grads = []
-    for imgs in dataloader:
-        with accelerator.accumulate(vae):
-            # imgs: high-res tensor from dataloader ([-1,1]), move to device
-            imgs = imgs.to(accelerator.device)
-
-            # >>> CHANGED: create low-res input for model by downsampling high-res crop
-            if model_resolution==high_resolution:
-                imgs_low = F.interpolate(imgs, size=(model_resolution, model_resolution), mode="bilinear", align_corners=False)
-            else:
-                imgs_low = imgs
-
-            # ensure dtype matches model params to avoid float/half mismatch
-            model_dtype = next(vae.parameters()).dtype
-            if imgs_low.dtype != model_dtype:
-                imgs_low_model = imgs_low.to(dtype=model_dtype)
-            else:
-                imgs_low_model = imgs_low
-
-            # Encode/decode on low-res input
-            latent_dist = vae.encode(imgs_low_model).latent_dist
-            latents = latent_dist.mean
-            rec = vae.decode(latents).sample  # rec is expected to be high-res (upscaled)
-
-            # If rec isn't the same spatial size as original high-res input, resize to high-res
-            if rec.shape[-2:] != imgs.shape[-2:]:
-                rec = F.interpolate(rec, size=imgs.shape[-2:], mode="bilinear", align_corners=False)
-
-            # Now compute losses **on high-res** (rec vs imgs)
-            rec_f32 = rec.to(torch.float32)
-            imgs_f32 = imgs.to(torch.float32)
-
-            # MAE
-            mae_loss = F.l1_loss(rec_f32, imgs_f32)
-
-            # LPIPS (ensure float32)
-            lpips_loss = _get_lpips()(rec_f32, imgs_f32).mean()
-            
-            # dynamic perceptual weighting (same as before)
-            if float(mae_loss.detach().cpu().item()) > 1e-12:
-                desired_multiplier = lpips_ratio / max(1.0 - lpips_ratio, 1e-12)
-                new_weight = (mae_loss.item() / float(lpips_loss.detach().cpu().item())) * desired_multiplier
-            else:
-                new_weight = perceptual_loss_weight
-            
-            perceptual_loss_weight = float(np.clip(new_weight, min_perceptual_weight, max_perceptual_weight))
-            batch_losses_perc.append(perceptual_loss_weight)
-            if len(batch_losses_perc) >= sample_interval:
-                avg_perc = float(np.mean(batch_losses_perc[-sample_interval:]))
-            else:
-                avg_perc = float(np.mean(batch_losses_perc[-sample_interval:]))
-
-            total_loss = mae_loss + avg_perc * lpips_loss
-
-            if torch.isnan(total_loss) or torch.isinf(total_loss):
-                print("NaN/Inf loss – stopping")
-                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_losses_mae.append(mae_loss.detach().item())
-                batch_losses_lpips.append(lpips_loss.detach().item())
-                batch_grads.append(float(grad_norm if isinstance(grad_norm, (float, int)) else grad_norm.cpu().item()))
-
-                if use_wandb and accelerator.sync_gradients:
-                    wandb.log({
-                        "mae_loss": mae_loss.detach().item(),
-                        "lpips_loss": lpips_loss.detach().item(),
-                        "perceptual_loss_weight": avg_perc,
-                        "total_loss": total_loss.detach().item(),
-                        "learning_rate": current_lr,
-                        "epoch": epoch,
-                        "grad_norm": batch_grads[-1],
-                    }, step=global_step)
-
-            # периодические сэмплы и чекпоинты
-            if global_step > 0 and global_step % sample_interval == 0:
-                # делаем генерацию и лог только в main process (генерация использует fixed_samples high-res)
-                if accelerator.is_main_process:
-                    generate_and_save_samples(global_step)
-
-                accelerator.wait_for_everyone()
-                
-                # сколько микро-батчей нужно взять для усреднения
-                n_micro = sample_interval * gradient_accumulation_steps
-                # защищаем от выхода за пределы
-                if len(batch_losses) >= n_micro:
-                    avg_loss = float(np.mean(batch_losses[-n_micro:]))
-                    avg_loss_mae = float(np.mean(batch_losses_mae[-n_micro:]))
-                    avg_loss_lpips = float(np.mean(batch_losses_lpips[-n_micro:]))
-                else:
-                    avg_loss = float(np.mean(batch_losses)) if batch_losses else float("nan")
-                    avg_loss_mae = float(np.mean(batch_losses_mae)) if batch_losses_mae else float("nan")
-                    avg_loss_lpips = float(np.mean(batch_losses_lpips)) if batch_losses_lpips 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_loss_mae": avg_loss_mae,"interm_loss_lpips": avg_loss_lpips, "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("Готово!")

train_sdxl_vae_qwen.py

@@ -1,526 +0,0 @@
-# -*- 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("Готово!")

train_sdxl_vae_simple.py

@@ -1,547 +0,0 @@
-# -*- 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
-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            = "simple_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
-asymmetric         = False
-optimizer_type     = "adam8bit"
-dtype              = torch.float32
-# model_resolution — то, что подавается в VAE (низкое разрешение)
-model_resolution   = 512   # бывший `resolution`
-# high_resolution — настоящий «высокий» кроп, на котором считаем метрики и сохраняем сэмплы
-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"   # или "fp16"/"bf16" при поддержке
-gradient_accumulation_steps = 5
-generated_folder   = "samples"
-save_as            = "simple_vae_nightly"
-num_workers        = 0
-device = None  # accelerator задаст устройство
-
-# --- Пропорции лоссов и окно медианного нормирования (КОЭФ., не значения) ---
-# Итоговые доли в total loss (сумма = 1.0)
-loss_ratios = {
-    "lpips": 0.85,
-    "edge":  0.05,
-    "mse":   0.05,
-    "mae":   0.05,
-}
-median_coeff_steps = 256  # за сколько шагов считать медианные коэффициенты
-
-# --------------------------- параметры препроцессинга ---------------------------
-resize_long_side = 1280  # если None или 0 — ресайза не будет; рекомендовано 1280
-
-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 = True
-
-# --------------------------- 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,
-    })
-
-# --------------------------- VAE ---------------------------
-if model_resolution==high_resolution and not asymmetric:
-    vae = AutoencoderKL.from_pretrained(project).to(dtype)
-else:
-    vae = AsymmetricAutoencoderKL.from_pretrained(project).to(dtype)
-
-# torch.compile (если доступно) — просто и без лишней логики
-if hasattr(torch, "compile"):
-    try:
-        vae = torch.compile(vae)
-    except Exception as e:
-        print(f"[WARN] torch.compile failed: {e}")
-
-# >>> Заморозка всех параметров, затем выборочная разморозка
-for p in vae.parameters():
-    p.requires_grad = False
-
-decoder = getattr(vae, "decoder", None)
-if decoder is None:
-    raise RuntimeError("vae.decoder not found — не могу применить стратегию разморозки. Проверь структуру модели.")
-
-unfrozen_param_names = []
-
-if not hasattr(decoder, "up_blocks"):
-    raise RuntimeError("decoder.up_blocks не найдены — ожидается список блоков декодера.")
-
-# >>> Размораживаем все up_blocks и mid_block (как было в твоём варианте start_idx=0)
-n_up = len(decoder.up_blocks)
-start_idx = 0
-for idx in range(start_idx, n_up):
-    block = decoder.up_blocks[idx]
-    for name, p in block.named_parameters():
-        p.requires_grad = True
-        unfrozen_param_names.append(f"decoder.up_blocks.{idx}.{name}")
-
-if hasattr(decoder, "mid_block"):
-    for name, p in decoder.mid_block.named_parameters():
-        p.requires_grad = True
-        unfrozen_param_names.append(f"decoder.mid_block.{name}")
-else:
-    print("[WARN] decoder.mid_block не найден — mid_block не разморожен.")
-
-print(f"[INFO] Разморожено параметров: {len(unfrozen_param_names)}. Первые 200 имён:")
-for nm in unfrozen_param_names[:200]:
-    print("  ", nm)
-
-# сохраняем trainable_module (get_param_groups будет учитывать p.requires_grad)
-trainable_module = vae.decoder
-
-# --------------------------- Custom PNG Dataset (only .png, skip corrupted) -----------
-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 = []
-        # collect png files recursively
-        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))
-        # optional limit
-        if limit:
-            self.paths = self.paths[:limit]
-        # verify images and keep only valid ones
-        valid = []
-        for p in self.paths:
-            try:
-                with Image.open(p) as im:
-                    im.verify()  # fast check for truncated/corrupted images
-                valid.append(p)
-            except (OSError, UnidentifiedImageError):
-                # skip corrupted image
-                continue
-        self.paths = valid
-        if len(self.paths) == 0:
-            raise RuntimeError(f"No valid PNG images found under {root_dir}")
-        # final shuffle for randomness
-        random.shuffle(self.paths)
-
-    def __len__(self):
-        return len(self.paths)
-
-    def __getitem__(self, idx):
-        p = self.paths[idx % len(self.paths)]
-        # open and convert to RGB; ensure file is closed promptly
-        with Image.open(p) as img:
-            img = img.convert("RGB")
-            # пережимаем длинную сторону до resize_long_side (Lanczos)
-            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])
-])
-
-# build dataset using high_resolution crops
-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}")
-
-# collate_fn кропит до high_resolution
-
-def collate_fn(batch):
-    imgs = []
-    for img in batch:  # img is PIL.Image
-        img = random_crop(img, high_resolution)   # кропим high-res
-        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 module.named_parameters():
-        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)
-
-# --------------------------- Подготовка Accelerate (вместе) ---------------------------
-
-batches_per_epoch = len(dataloader)  # число микро-батчей (dataloader steps)
-steps_per_epoch = int(math.ceil(batches_per_epoch / float(gradient_accumulation_steps)))  # число optimizer.step() за эпоху
-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.decoder.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
-
-# Собель для edge loss
-_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:
-    # x: [B,C,H,W] в [-1,1]
-    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):
-        # нормируем доли на случай, если сумма != 1
-        s = sum(desired_ratios.values())
-        self.ratios = {k: (v / s) 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().cpu()))
-        # Медианы (устойчивые к выбросам)
-        meds = {k: (np.median(self.buffers[k]) if len(self.buffers[k]) > 0 else 1.0) for k in self.buffers}
-        # Вычисляем КОЭФФИЦИЕНТЫ как ratio_k / median_k — т.е. именно коэффициенты, а не значения
-        coeffs = {k: (self.ratios[k] / max(meds[k], 1e-12)) for k in self.ratios}
-        # Важно: при таких коэффициентах сумма (coeff_k * median_k) = сумма(ratio_k) = 1, т.е. масштаб стабилен
-        total = sum(coeffs[k] * abs_losses[k] for k in coeffs)
-        return total, coeffs, meds
-
-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]  # dataset returns PIL.Image
-    tensors = []
-    for img in pil_imgs:
-        img = random_crop(img, high_resolution)  # high-res fixed samples
-        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:
-    # img_tensor: [C,H,W] in [-1,1]
-    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():
-            # Готовим low-res вход для кодера ВСЕГДА под model_resolution
-            orig_high = fixed_samples  # [B,C,H,W] в [-1,1]
-            orig_low = F.interpolate(orig_high, size=(model_resolution, model_resolution), mode="bilinear", align_corners=False)
-            # dtype как у модели
-            model_dtype = next(temp_vae.parameters()).dtype
-            orig_low = orig_low.to(dtype=model_dtype)
-            # encode/decode
-            latents = temp_vae.encode(orig_low).latent_dist.mean
-            rec = temp_vae.decode(latents).sample
-
-        # Приводим spatial размер рекона к high-res (downsample для асимметричных VAE)
-        if rec.shape[-2:] != orig_high.shape[-2:]:
-            rec = F.interpolate(rec, size=orig_high.shape[-2:], mode="bilinear", align_corners=False)
-
-        # Сохраняем ПЕРВЫЙ семпл: real и decoded без номера шага в имени
-        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 на полном изображении (high-res) — для лога
-        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: high-res tensor from dataloader ([-1,1]), move to device
-            imgs = imgs.to(accelerator.device)
-
-            # ВСЕГДА даунсемплим вход под model_resolution для кодера
-            # Тупая железяка норовит все по своему сделать
-            if high_resolution != model_resolution:
-                imgs_low = F.interpolate(imgs, size=(model_resolution, model_resolution), mode="bilinear", align_corners=False)
-            else:
-                imgs_low = imgs 
-
-            # ensure dtype matches model params to avoid float/half mismatch
-            model_dtype = next(vae.parameters()).dtype
-            if imgs_low.dtype != model_dtype:
-                imgs_low_model = imgs_low.to(dtype=model_dtype)
-            else:
-                imgs_low_model = imgs_low
-
-            # Encode/decode
-            latents = vae.encode(imgs_low_model).latent_dist.mean
-            rec = vae.decode(latents).sample  # rec может быть увеличенным (асимметричный VAE)
-
-            # Приводим размер к high-res
-            if rec.shape[-2:] != imgs.shape[-2:]:
-                rec = F.interpolate(rec, size=imgs.shape[-2:], mode="bilinear", align_corners=False)
-
-            # Лоссы считаем на high-res
-            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)),
-            }
-
-            # Total с медианными КОЭФФИЦИЕНТАМИ
-            # Не надо так орать когда у тебя получилось понять мою идею
-            total_loss, coeffs, meds = normalizer.update_and_total(abs_losses)
-
-            if torch.isnan(total_loss) or torch.isinf(total_loss):
-                print("NaN/Inf loss – stopping")
-                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 if isinstance(grad_norm, (float, int)) else grad_norm.cpu().item()))
-                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],
-                    }
-                    # добавляем отдельные лоссы
-                    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
-                if len(batch_losses) >= n_micro:
-                    avg_loss = float(np.mean(batch_losses[-n_micro:]))
-                else:
-                    avg_loss = 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("Готово!")

vae/config.json

@@ -1,7 +1,7 @@
 {
   "_class_name": "AutoencoderKL",
-  "_diffusers_version": "0.34.0",
-  "_name_or_path": "sdxl_vae",
+  "_diffusers_version": "0.35.1",
+  "_name_or_path": "AiArtLab/sdxl_vae",
   "act_fn": "silu",
   "block_out_channels": [
     128,

vae/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:03f2412467f6bedce9efeddba5860b5ec0d3267931d14c500d4bd7a878e14cbd
-size 334643268
+oid sha256:f9f3bf86e95df913a45a4a238709c47f24530c07d10e0f923b0dae2f679799ea
+size 167335342

vae_nightly/config.json

@@ -1,38 +0,0 @@
-{
-  "_class_name": "AutoencoderKL",
-  "_diffusers_version": "0.35.0.dev0",
-  "_name_or_path": "vae",
-  "act_fn": "silu",
-  "block_out_channels": [
-    128,
-    256,
-    512,
-    512
-  ],
-  "down_block_types": [
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D"
-  ],
-  "force_upcast": false,
-  "in_channels": 3,
-  "latent_channels": 4,
-  "latents_mean": null,
-  "latents_std": null,
-  "layers_per_block": 2,
-  "mid_block_add_attention": true,
-  "norm_num_groups": 32,
-  "out_channels": 3,
-  "sample_size": 512,
-  "scaling_factor": 0.13025,
-  "shift_factor": null,
-  "up_block_types": [
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D"
-  ],
-  "use_post_quant_conv": true,
-  "use_quant_conv": true
-}

vae_nightly/diffusion_pytorch_model.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:588db8438a9dea0c4c68dfd4cbdc7747b1ed3601f2a71f46d1608fae9bdb96a3
-size 334643268