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