init

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,13 @@
+# Jupyter Notebook
+__pycache__/
+*.pyc
+.ipynb_checkpoints/
+*.ipynb_checkpoints/*
+.ipynb_checkpoints/*
+src/samples
+# cache
+cache
+datasets
+test
+wandb
+nohup.out

README.md

@@ -10,4 +10,48 @@ library_name: diffusers
 # SDXL-VAE finetuned
  - madebyollin/sdxl-vae-fp16-fix: MSE: 3.680e-03, PSNR: 25.2100, LPIPS: 0.1314
  - KBlueLeaf/EQ-SDXL-VAE        : MSE: 3.530e-03, PSNR: 25.2827, LPIPS: 0.1298
- - AiArtLab/sdxl_vae            : MSE: <span style="color:red">3.321e-03</span>, PSNR: <span style="color:red">25.6389</span>, LPIPS: <span style="color:red">0.1251</span>
+ - AiArtLab/sdxl_vae            : MSE: <span style="color:red">3.321e-03</span>, PSNR: <span style="color:red">25.6389</span>, LPIPS: <span style="color:red">0.1251</span>
+
+### Train status, in progress:
+
+![result](result.png)
+
+## VAE Training Process
+
+Dataset: 100,000 PNG images
+Training Time: 4 days
+Hardware: Single RTX 4090
+Resolution: 512px
+Precision: FP32
+Effective Batch Size: 16 (batch size 2 + gradient accumulation 8)
+Optimizer: AdamW (8-bit)
+
+## Implementation
+
+Base Code: Used a simple diffusion model training script.
+Encoder: Frozen (to avoid retraining SDXL for the new VAE).
+Training Target: Only the decoder, focusing on image reconstruction.
+
+## Loss Functions
+
+Initially used LPIPS and MSE.
+Noticed FID score improving, but images becoming blurry (FID overfits to blurry images—improving FID is not always good).
+Switched to MAE (Mean Absolute Error) instead of MSE (not sure is MSE bad).
+Balanced LPIPS and MAE at 90/10 ratio.
+Used median perceptual_loss_weight for better balance.
+
+## Results
+
+https://imgsli.com/NDA3NTEy/1/2
+
+## Donations
+
+Please contact with us if you may provide some GPU's or money on training
+
+DOGE: DEw2DR8C7BnF8GgcrfTzUjSnGkuMeJhg83
+
+BTC: 3JHv9Hb8kEW8zMAccdgCdZGfrHeMhH1rpN
+
+## Contacts
+
+[recoilme](https://t.me/recoilme)

config.json

@@ -0,0 +1,38 @@
+{
+  "_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

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

eval.py

@@ -0,0 +1,167 @@
+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
+
+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
+
+NAMES = [
+    "madebyollin/sdxl-vae-fp16-fix",
+    "KBlueLeaf/EQ-SDXL-VAE        ",
+    "AiArtLab/simplevae           ",
+]
+BASE_MODELS = [
+    "madebyollin/sdxl-vae-fp16-fix",
+    "KBlueLeaf/EQ-SDXL-VAE",
+    "AiArtLab/simplevae",
+]
+SUB_FOLDERS = [None, None, "sdxl_vae"]
+CKPT_PATHS = [
+    None,
+    None,
+    None,
+]
+USE_APPROXS = [False, False, False]
+
+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=4,
+        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, ckpt_path, use_approx in zip(
+        BASE_MODELS, SUB_FOLDERS, CKPT_PATHS, USE_APPROXS
+    ):
+        vae = AutoencoderKL.from_pretrained(base_model, subfolder=sub_folder)
+        if use_approx:
+            vae.decoder = LatentApproxDecoder(
+                latent_dim=vae.config.latent_channels,
+                out_channels=3,
+                shuffle=2,
+            )
+            vae.decode = lambda x: vae.decoder(x)
+            vae.get_last_layer = lambda: vae.decoder.conv_out.weight
+        if ckpt_path:
+            LatentTrainer.load_from_checkpoint(
+                ckpt_path, vae=vae, map_location="cpu", strict=False
+            )
+        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.py

@@ -0,0 +1,504 @@
+# -*- 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            = "sdxl_vae"
+batch_size         = 1
+base_learning_rate = 1e-6
+min_learning_rate  = 8e-7
+num_epochs         = 8
+sample_interval_share = 20
+use_wandb          = True
+save_model         = True
+use_decay          = True
+optimizer_type     = "adam8bit"
+dtype              = torch.float32
+# model_resolution — то, что подавается в VAE (низкое разрешение)
+model_resolution   = 768   # бывший `resolution`
+# high_resolution — настоящий «высокий» кроп, на котором считаем метрики и сохраняем сэмплы
+high_resolution    = 768  # >>> CHANGED: обучаемся на входах 1024 -> даунсемплим до 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 = 16
+generated_folder   = "samples"
+save_as            = "sdxl_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 ---------------------------
+vae = AutoencoderKL.from_pretrained(project).to(dtype)
+#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)
+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("Готово!")