train_sdxl_vae.py

# -*- 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            = "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         = 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            = "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("Готово!")