asym

README.md

@@ -9,7 +9,7 @@ library_name: diffusers
 ---
 # SDXL-VAE finetuned
 
-
+Imagenet eval (256px)
 |  Model                     | MSE         | PSNR      | LPIPS      |
 |----------------------------|-------------|-----------|------------|
 | madebyollin/sdxl-vae-fp16-fix | 3.680e-03   | 25.2100   | 0.1314     |
@@ -17,6 +17,16 @@ library_name: diffusers
 | **AiArtLab/sdxl_vae**       | **3.321e-03** | **25.6389** | **0.1251** |
 
 
+Alchemist eval (512px)
+
+| Model                         |        MSE |       PSNR |      LPIPS |
+|--------------------------------|------------|------------|------------|
+| madebyollin/sdxl-vae-fp16      |       100% |       100% |       100% |
+| KBlueLeaf/EQ-SDXL-VAE          |     107.8% |     100.1% |      95.5% |
+| AiArtLab/sdxl_vae              |     112.3% |     101.8% |     106.6% |
+| AiArtLab/sdxl_vae_asym         |     111.7% |     101.1% |      89.4% |
+| FLUX.1-schnell-vae             |     324.0% |     119.8% |     292.0% |
+
 [![Click it](vae.png)](https://imgsli.com/NDA3OTgz)
 
 

asymmetric_vae/config.json

@@ -3,12 +3,6 @@
   "_diffusers_version": "0.34.0",
   "_name_or_path": "asymmetric_vae_empty",
   "act_fn": "silu",
-  "block_out_channels": [
-    128,
-    256,
-    512,
-    512
-  ],
   "down_block_out_channels": [
     128,
     256,
@@ -21,7 +15,6 @@
     "DownEncoderBlock2D",
     "DownEncoderBlock2D"
   ],
-  "force_upcast": false,
   "in_channels": 3,
   "latent_channels": 4,
   "layers_per_down_block": 2,

asymmetric_vae/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8de44e4f21835eb457785a63f7e96c7ddba34b9b812bdeee79012d8bd0dae199
+oid sha256:ded3c30322578e3371f32a58423b6a3be3a2c3b81d3eb5d35433772be796a1ba
 size 421473052

asymmetric_vae_new/config.json

@@ -1,6 +1,6 @@
 {
   "_class_name": "AsymmetricAutoencoderKL",
-  "_diffusers_version": "0.34.0",
+  "_diffusers_version": "0.35.0.dev0",
   "_name_or_path": "asymmetric_vae",
   "act_fn": "silu",
   "block_out_channels": [

asymmetric_vae_new/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ded3c30322578e3371f32a58423b6a3be3a2c3b81d3eb5d35433772be796a1ba
+oid sha256:df9380b1e8d8b1a36b3d0f9501a854717a911ae9b8d2aebe18809a6eefa9318b
 size 421473052

down.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+
+TARGET_DIR="/workspace/d23"
+mkdir -p "$TARGET_DIR"
+
+BASE_URL="https://huggingface.co/datasets/AI-Art-Collab/dtasettar23/resolve/main/d23.tar."
+
+(
+  # Устанавливаем `set -e` внутри subshell, чтобы он завершился при первой ошибке curl
+  set -e
+  # Попробуем от 'a' до 'z' для первого символа суффикса
+  for c1 in {a..z}; do
+    # Попробуем от 'a' до 'z' для второго символа суффикса
+    for c2 in {a..z}; do
+      suffix="${c1}${c2}"
+      url="${BASE_URL}${suffix}"
+      echo "Fetching: $url" >&2
+      # Качаем часть архива. --fail заставит curl завершиться с ошибкой, если файла нет.
+      curl -LsS --fail "$url"
+    done
+  done
+) 2>/dev/null | tar -xv -C "$TARGET_DIR" --wildcards '*.png'
+#    └─ 1 ─┘   └────────── 2 ──────────┘ └─────────── 3 ───────────┘
+
+echo "Extraction of PNG files finished. Check $TARGET_DIR"

eval_alchemist.py

@@ -0,0 +1,186 @@
+import os
+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
+from diffusers import AutoencoderKL, AsymmetricAutoencoderKL
+import random
+
+# --------------------------- Параметры ---------------------------
+DEVICE = "cuda"
+DTYPE = torch.float16
+IMAGE_FOLDER = "/workspace/alchemist"
+MIN_SIZE = 1280
+CROP_SIZE = 512
+BATCH_SIZE = 4  # можно увеличить для ускорения
+MAX_IMAGES = None
+NUM_WORKERS = 4  # параллельная загрузка
+
+# Список VAE для тестирования
+VAE_LIST = [
+    ("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/sdxl_vae_asym", AsymmetricAutoencoderKL, "AiArtLab/sdxl_vae", "asymmetric_vae"),
+    ("FLUX.1-schnell-vae", AutoencoderKL, "black-forest-labs/FLUX.1-schnell", "vae"),
+]
+
+# --------------------------- Dataset ---------------------------
+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("Сканирование папки...")
+        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))
+        
+        # Ограничение количества
+        if limit:
+            self.paths = self.paths[:limit]
+        
+        # Быстрая проверка валидности (опционально, можно убрать для скорости)
+        print("Проверка изображений...")
+        valid = []
+        for p in tqdm(self.paths, desc="Проверка"):
+            try:
+                with Image.open(p) as im:
+                    im.verify()
+                valid.append(p)
+            except:
+                continue
+        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),
+            CenterCrop(crop_size),
+            ToTensor(),
+        ])
+    
+    def __len__(self):
+        return len(self.paths)
+    
+    def __getitem__(self, idx):
+        path = self.paths[idx]
+        with Image.open(path) 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
+
+# --------------------------- Основной код ---------------------------
+if __name__ == "__main__":
+    # Создаем датасет и загрузчик
+    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
+    lpips_net = lpips.LPIPS(net="vgg").eval().to(DEVICE).requires_grad_(False)
+    
+    # Загрузка VAE моделей
+    print("\nЗагрузка VAE моделей...")
+    vaes = []
+    names = []
+    
+    for name, vae_class, model_path, subfolder in VAE_LIST:
+        try:
+            print(f"  Загружаю {name}...")
+            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, "count": 0} for name in names}
+    
+    with torch.no_grad():
+        for batch in tqdm(dataloader, desc="Обработка батчей"):
+            batch = batch.to(DEVICE)
+            test_inp = process(batch).to(DTYPE)
+            
+            for vae, name in zip(vaes, names):
+                # Encode/decode
+                latent = vae.encode(test_inp).latent_dist.mode()
+                recon = deprocess(vae.decode(latent).sample.float())
+                
+                # Метрики для батча
+                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()
+                    
+                    results[name]["mse"] += mse
+                    results[name]["psnr"] += psnr
+                    results[name]["lpips"] += lpips_val
+                    results[name]["count"] += 1
+    
+    # Усреднение результатов
+    for name in names:
+        count = results[name]["count"]
+        results[name]["mse"] /= count
+        results[name]["psnr"] /= count
+        results[name]["lpips"] /= count
+    
+    # Вывод абсолютных значений
+    print("\n=== Абсолютные значения ===")
+    for name in names:
+        print(f"{name:30s}: MSE: {results[name]['mse']:.3e}, PSNR: {results[name]['psnr']:.4f}, LPIPS: {results[name]['lpips']:.4f}")
+    
+    # Вывод таблицы с процентами
+    print("\n=== Сравнение с первой моделью (%) ===")
+    print(f"| {'Модель':30s} | {'MSE':>10s} | {'PSNR':>10s} | {'LPIPS':>10s} |")
+    print(f"|{'-'*32}|{'-'*12}|{'-'*12}|{'-'*12}|")
+    
+    baseline = names[0]
+    for name in names:
+        mse_pct = (results[baseline]["mse"] / results[name]["mse"]) * 100
+        psnr_pct = (results[name]["psnr"] / results[baseline]["psnr"]) * 100
+        lpips_pct = (results[baseline]["lpips"] / results[name]["lpips"]) * 100
+        
+        if name == baseline:
+            print(f"| {name:30s} | {'100%':>10s} | {'100%':>10s} | {'100%':>10s} |")
+        else:
+            print(f"| {name:30s} | {f'{mse_pct:.1f}%':>10s} | {f'{psnr_pct:.1f}%':>10s} | {f'{lpips_pct:.1f}%':>10s} |")
+    
+    print("\n✅ Готово!")

eval_asym.py

@@ -1,159 +0,0 @@
-import warnings
-import logging
-import torch
-import torch.nn.functional as F
-import torch.utils.data as data
-import lpips
-from tqdm import tqdm
-from torchvision.transforms import (
-    Compose,
-    Resize,
-    ToTensor,
-    CenterCrop,
-)
-from diffusers import AutoencoderKL,AsymmetricAutoencoderKL
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-warnings.filterwarnings(
-    "ignore",
-    ".*Found keys that are not in the model state dict but in the checkpoint.*",
-)
-
-DEVICE = "cuda"
-DTYPE = torch.float16
-SHORT_AXIS_SIZE = 256
-batch_size = 1
-
-NAMES = [
-#    "asymmetric_vae",
-    "asymmetric_vae_new",
-#    "madebyollin/sdxl-vae-fp16-fix",
-#    "KBlueLeaf/EQ-SDXL-VAE        ",
-#    "AiArtLab/simplevae           ",
-]
-BASE_MODELS = [
-#    "./asymmetric_vae",
-    "./asymmetric_vae_new",    
-#    "madebyollin/sdxl-vae-fp16-fix",
-#    "KBlueLeaf/EQ-SDXL-VAE",
-#    "AiArtLab/simplevae",
-]
-SUB_FOLDERS = [
-#    "sdxs_vae",
-    None, 
-#    None, 
-#    "sdxl_vae"
-]
-
-def process(x):
-    return x * 2 - 1
-
-def deprocess(x):
-    return x * 0.5 + 0.5
-
-import torch.utils.data as data
-from datasets import load_dataset
-
-class ImageNetDataset(data.IterableDataset):
-    def __init__(self, split, transform=None, max_len=10, streaming=True):
-        self.split = split
-        self.transform = transform
-        self.dataset = load_dataset("evanarlian/imagenet_1k_resized_256", split=split, streaming=streaming)
-        self.max_len = max_len
-        self.iterator = iter(self.dataset)
-
-    def __iter__(self):
-        for i, entry in enumerate(self.iterator):
-            if self.max_len and i >= self.max_len:
-                break
-            img = entry["image"]
-            target = entry["label"]
-            if self.transform is not None:
-                img = self.transform(img)
-            yield img, target
-
-if __name__ == "__main__":
-    lpips_loss = torch.compile(
-        lpips.LPIPS(net="vgg").eval().to(DEVICE).requires_grad_(False)
-    )
-
-    @torch.compile
-    def metrics(inp, recon):
-        mse = F.mse_loss(inp, recon)
-        psnr = 10 * torch.log10(1 / mse)
-        return (
-            mse.cpu(),
-            psnr.cpu(),
-            lpips_loss(inp, recon, normalize=True).mean().cpu(),
-        )
-
-    transform = Compose(
-        [
-            Resize(SHORT_AXIS_SIZE),
-            CenterCrop(SHORT_AXIS_SIZE),
-            ToTensor(),
-        ]
-    )
-    valid_dataset = ImageNetDataset("val", transform=transform, max_len=50000, streaming=True)
-    valid_loader = data.DataLoader(
-        valid_dataset,
-        batch_size=batch_size,
-        shuffle=False,
-        num_workers=2,
-        pin_memory=True,
-        pin_memory_device=DEVICE,
-    )
-
-    # Проверяем, что данные грузятся
-    for batch in valid_loader:
-        print("Batch shape:", batch[0].shape)
-        break
-
-    logger.info("Loading models...")
-    vaes = []
-    for base_model, sub_folder in zip(
-        BASE_MODELS, SUB_FOLDERS
-    ):
-        vae = AsymmetricAutoencoderKL.from_pretrained(base_model, subfolder=sub_folder)
-        vae = vae.to(DTYPE).eval().requires_grad_(False).to(DEVICE)
-        vae.encoder = torch.compile(vae.encoder)
-        vae.decoder = torch.compile(vae.decoder)
-        vaes.append(torch.compile(vae))
-
-    logger.info("Running Validation")
-    total = 0
-    all_latents = [[] for _ in range(len(vaes))]
-    all_mse = [[] for _ in range(len(vaes))]
-    all_psnr = [[] for _ in range(len(vaes))]
-    all_lpips = [[] for _ in range(len(vaes))]
-
-    for idx, batch in enumerate(tqdm(valid_loader)):
-        image = batch[0].to(DEVICE)
-        test_inp = process(image).to(DTYPE)
-        batch_size = test_inp.size(0)
-
-        for i, vae in enumerate(vaes):
-            latent = vae.encode(test_inp).latent_dist.mode()
-            recon = deprocess(vae.decode(latent).sample.float())
-            all_latents[i].append(latent.cpu().float())
-            mse, psnr, lpips_ = metrics(image, recon)
-            all_mse[i].append(mse.cpu() * batch_size)
-            all_psnr[i].append(psnr.cpu() * batch_size)
-            all_lpips[i].append(lpips_.cpu() * batch_size)
-
-        total += batch_size
-
-    for i in range(len(vaes)):
-        all_latents[i] = torch.cat(all_latents[i], dim=0)
-        all_mse[i] = torch.stack(all_mse[i]).sum() / total
-        all_psnr[i] = torch.stack(all_psnr[i]).sum() / total
-        all_lpips[i] = torch.stack(all_lpips[i]).sum() / total
-
-        logger.info(
-            f"  - {NAMES[i]}: MSE: {all_mse[i]:.3e}, PSNR: {all_psnr[i]:.4f}, "
-            f"LPIPS: {all_lpips[i]:.4f}"
-        )
-
-    logger.info("End")

train_sdxl_vae_gpt5.py

@@ -0,0 +1,547 @@
+# -*- 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            = "asymmetric_vae"
+batch_size         = 3
+base_learning_rate = 6e-6
+min_learning_rate  = 1e-6
+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    = 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            = "asymmetric_vae_new"
+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 — ресайза не будет; рекомендовано 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)
+
+# 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.py → train_sdxl_vae_my.py

@@ -38,7 +38,7 @@ dtype              = torch.float32
 # model_resolution — то, что подавается в VAE (низкое разрешение)
 model_resolution   = 512   # бывший `resolution`
 # high_resolution — настоящий «высокий» кроп, на котором считаем метрики и сохраняем сэмплы
-high_resolution    = 512  # >>> CHANGED: обучаемся на входах 1024 -> даунсемплим до 512 для модели
+high_resolution    = 1024 
 limit              = 0
 save_barrier       = 1.03
 warmup_percent     = 0.01