train_sdxs_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
# Import standard and asymmetric VAEs only
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 FDL_pytorch import FDL_loss # pip install fdl-pytorch
from collections import deque

# --- Configuration ---
DATASET_PATH            = "/workspace/laion-coco-1024-1520-10000-data-10000"
PROJECT_NAME            = "vae4"
BATCH_SIZE              = 1
BASE_LEARNING_RATE      = 4e-6
MIN_LEARNING_RATE       = 4e-7
NUM_EPOCHS              = 6
SAMPLE_INTERVAL_SHARE   = 10
USE_WANDB               = False
SAVE_MODEL              = True
USE_DECAY               = True
OPTIMIZER_TYPE          = "adam8bit"
DTYPE                   = torch.float32

MODEL_RESOLUTION        = 512
HIGH_RESOLUTION         = 1024
DATA_LIMIT              = 0  # Limit dataset size (0 for no limit)
SAVE_BARRIER            = 1.3
WARMUP_PERCENT          = 0.005
BETA2                   = 0.997
EPSILON                 = 1e-8
CLIP_GRAD_NORM          = 1.0
MIXED_PRECISION         = "no"
GRADIENT_ACCUMULATION_STEPS = 1
GENERATED_FOLDER        = "samples"
SAVE_AS                 = "vae5"
NUM_WORKERS             = 0

# Enable deterministic training and optimizations
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(True)
torch.backends.cuda.enable_math_sdp(False)

# --- Training Modes ---
TRAIN_DECODER_ONLY = True
TRAIN_UP_ONLY      = False
FULL_TRAINING      = False
KL_RATIO           = 0.0

# --- Loss Ratios ---
LOSS_RATIOS = {
    "lpips": 0.60,
    "fdl" :  0.10,
    "mse":   0.05,
    "mae":   0.15,
    "dssim": 0.05,
    "kl":    0.00,
    "edge":  0.05,
}
MEDIAN_COEFF_STEPS = 250

# --- VAE Type ---
# 'kl' for standard AutoencoderKL, 'asymmetric' for AsymmetricAutoencoderKL
VAE_TYPE = "asymmetric"

Path(GENERATED_FOLDER).mkdir(parents=True, exist_ok=True)

# Initialize Accelerator
accelerator = Accelerator(
    mixed_precision=MIXED_PRECISION,
    gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS
)
device = accelerator.device

# Set seeds for reproducibility
seed = int(datetime.now().strftime("%Y%m%d")) + 42
torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
torch.backends.cudnn.benchmark = False

# --------------------------- WandB Logging ---------------------------
if USE_WANDB and accelerator.is_main_process:
    wandb.init(project=PROJECT_NAME, 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_type": VAE_TYPE,
    })

# --------------------------- VAE Model Loading ---------------------------
def get_core_model(model):
    """Unwraps a model potentially wrapped by torch.compile."""
    if hasattr(model, "_orig_mod"):
        model = model._orig_mod
    return model

# Load the appropriate VAE model (Video VAEs completely removed)
if VAE_TYPE == "asymmetric":
    vae = AsymmetricAutoencoderKL.from_pretrained(PROJECT_NAME)
elif VAE_TYPE == "kl":
    vae = AutoencoderKL.from_pretrained(PROJECT_NAME)
else:
    raise ValueError(f"Unsupported VAE_TYPE: {VAE_TYPE}")

vae = vae.to(DTYPE)

# Apply torch.compile
if hasattr(torch, "compile"):
    try:
        vae = torch.compile(vae)
        print("[INFO] torch.compile applied successfully.")
    except Exception as e:
        print(f"[WARN] torch.compile failed: {e}")

# --------------------------- Freeze/Unfreeze Parameters ---------------------------
core = get_core_model(vae)

for p in core.parameters():
    p.requires_grad = False

unfrozen_param_names = []

if FULL_TRAINING and not TRAIN_DECODER_ONLY:
    for name, p in core.named_parameters():
        p.requires_grad = True
        unfrozen_param_names.append(name)
    LOSS_RATIOS["kl"] = float(KL_RATIO)
    trainable_module = core
else:
    if hasattr(core, "decoder"):
        if TRAIN_UP_ONLY and hasattr(core.decoder, "up_blocks") and len(core.decoder.up_blocks) > 0:
            for name, p in core.decoder.up_blocks[0].named_parameters():
                p.requires_grad = True
                unfrozen_param_names.append(f"decoder.up_blocks[0].{name}")
        else:
            print("[INFO] Decoder: Falling back to training the full decoder.")
            for name, p in core.decoder.named_parameters():
                p.requires_grad = True
                unfrozen_param_names.append(f"decoder.{name}")
    if hasattr(core, "post_quant_conv"):
        for name, p in core.post_quant_conv.named_parameters():
            p.requires_grad = True
            unfrozen_param_names.append(f"post_quant_conv.{name}")
    trainable_module = core.decoder if hasattr(core, "decoder") else core

print(f"[INFO] Unfrozen parameters: {len(unfrozen_param_names)}. First 10 names:")
for nm in unfrozen_param_names[:10]:
    print(f"  {nm}")

# --------------------------- Dataset Preparation ---------------------------
class PngFolderDataset(Dataset):
    def __init__(self, root_dir, resolution=1024, min_exts=('.png',), limit=0):
        self.resolution = resolution
        self.paths = []
        for root, _, files in os.walk(root_dir):
            for f in files:
                if f.lower().endswith(tuple(ext.lower() for ext in min_exts)):
                    self.paths.append(os.path.join(root, f))
        if limit > 0:
            self.paths = self.paths[:limit]

        valid_paths = []
        for p in self.paths:
            try:
                with Image.open(p) as img:
                    img.verify()
                    w, h = img.size
                    if w < resolution or h < resolution:
                        continue
                valid_paths.append(p)
            except (OSError, UnidentifiedImageError) as e:
                print(f"[WARN] Skipping invalid image file {p}: {e}")
        self.paths = valid_paths
        if not self.paths:
            raise RuntimeError(f"No valid images found in {root_dir}")
        random.shuffle(self.paths)
        self.transform = transforms.ToTensor()

    def __len__(self): return len(self.paths)

    def __getitem__(self, idx):
        p = self.paths[idx % len(self.paths)]
        try:
            with Image.open(p) as img:
                return img.convert("RGB")
        except Exception as e:
            print(f"[ERROR] Failed to load image {p}: {e}")
            return Image.new("RGB", (self.resolution, self.resolution), 'red')

def random_crop(img, sz):
    w, h = img.size
    crop_w = min(sz, w)
    crop_h = min(sz, h)
    x = random.randint(0, max(0, w - crop_w))
    y = random.randint(0, max(0, h - crop_h))
    return img.crop((x, y, x + crop_w, y + crop_h))

input_tfm = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])

def collate_fn(batch):
    imgs = []
    for img in batch:
        img = random_crop(img, HIGH_RESOLUTION)
        imgs.append(input_tfm(img))
    return torch.stack(imgs)

try:
    dataset = PngFolderDataset(DATASET_PATH, min_exts=('.png', '.PNG'), resolution=HIGH_RESOLUTION, limit=DATA_LIMIT)
    print(f"[INFO] Dataset loaded: {len(dataset)} images.")
    if len(dataset) < BATCH_SIZE:
        raise RuntimeError(f"Not enough valid images ({len(dataset)}) to form a batch of size {BATCH_SIZE}")

    dataloader = DataLoader(
        dataset, batch_size=BATCH_SIZE, shuffle=True,
        collate_fn=collate_fn, num_workers=NUM_WORKERS,
        pin_memory=True, drop_last=True
    )
except RuntimeError as e:
    print(f"[ERROR] Failed to initialize dataloader: {e}")
    exit()

# --------------------------- Optimizer Setup ---------------------------
def get_param_groups(module, weight_decay=0.001):
    no_decay_tokens = ("bias", "norm", "rms", "layernorm")
    decay_params, no_decay_params = [], []
    for name, param in module.named_parameters():
        if not param.requires_grad: continue
        name_lower = name.lower()
        if any(token in name_lower for token in no_decay_tokens):
            no_decay_params.append(param)
        else:
            decay_params.append(param)
    return [
        {"params": decay_params, "weight_decay": weight_decay},
        {"params": no_decay_params, "weight_decay": 0.0},
    ]

param_groups = get_param_groups(get_core_model(vae), weight_decay=0.001)
optimizer = bnb.optim.AdamW8bit(param_groups, lr=BASE_LEARNING_RATE, betas=(0.9, BETA2), eps=EPSILON)

# --------------------------- Learning Rate Scheduler ---------------------------
batches_per_epoch = len(dataloader)
steps_per_epoch = 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
    current_step_fraction = float(step) / float(max(1, total_steps))
    warmup_fraction = float(WARMUP_PERCENT)
    min_lr_ratio = float(MIN_LEARNING_RATE) / float(BASE_LEARNING_RATE)
    if current_step_fraction < warmup_fraction:
        return min_lr_ratio + (1.0 - min_lr_ratio) * (current_step_fraction / warmup_fraction)
    else:
        decay_fraction = (current_step_fraction - warmup_fraction) / (1.0 - warmup_fraction)
        return min_lr_ratio + 0.5 * (1.0 - min_lr_ratio) * (1.0 + math.cos(math.pi * decay_fraction))

scheduler = LambdaLR(optimizer, lr_lambda)

# --------------------------- Prepare for Training ---------------------------
(dataloader, vae, optimizer, scheduler) = accelerator.prepare(dataloader, vae, optimizer, scheduler)
trainable_params = [p for p in vae.parameters() if p.requires_grad]

fdl_loss_fn = FDL_loss().to(accelerator.device)

_lpips_net = None
def get_lpips_loss():
    global _lpips_net
    if _lpips_net is None:
        _lpips_net = lpips.LPIPS(net='vgg', verbose=False).eval().to(accelerator.device)
    return _lpips_net

def _gaussian_kernel(window_size, sigma, device, dtype):
    coords = torch.arange(window_size, dtype=dtype, device=device) - (window_size - 1) / 2
    k = torch.exp(-coords**2 / (2 * sigma**2))
    return k / k.sum()

def _ssim(img1, img2, window_size=11, sigma=1.5):
    channels = img1.shape[1]
    kernel = _gaussian_kernel(window_size, sigma, img1.device, img1.dtype)
    win = (kernel.view(1, 1, -1, 1) * kernel.view(1, 1, 1, -1)).expand(channels, 1, window_size, window_size).contiguous()
    mu1 = F.conv2d(img1, win, padding=window_size//2, groups=channels)
    mu2 = F.conv2d(img2, win, padding=window_size//2, groups=channels)
    mu1_sq, mu2_sq, mu1_mu2 = mu1.pow(2), mu2.pow(2), mu1 * mu2
    sigma1_sq = F.conv2d(img1*img1, win, padding=window_size//2, groups=channels) - mu1_sq
    sigma2_sq = F.conv2d(img2*img2, win, padding=window_size//2, groups=channels) - mu2_sq
    sigma12 = F.conv2d(img1*img2, win, padding=window_size//2, groups=channels) - mu1_mu2
    
    # ИСПРАВЛЕНО: Разделено присваивание, чтобы избежать UnboundLocalError
    L = 2.0
    C1 = (0.01 * L) ** 2
    C2 = (0.03 * L) ** 2
    
    num = (2 * mu1_mu2 + C1) * (2 * sigma12 + C2)
    den = (mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)
    return (num / den).mean()

def edge_loss(img1, img2):
    def get_edges(img):
        C = img.shape[1]
        # Sobel x kernel (horizontal edges)
        sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=img.dtype, device=img.device).view(1, 1, 3, 3)
        # Sobel y kernel (vertical edges)
        sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=img.dtype, device=img.device).view(1, 1, 3, 3)
        # Repeat for each channel
        sobel_x_c = sobel_x.repeat(C, 1, 1, 1)
        sobel_y_c = sobel_y.repeat(C, 1, 1, 1)
        # Apply convolution per channel
        grad_x = F.conv2d(img, sobel_x_c, padding=1, groups=C)
        grad_y = F.conv2d(img, sobel_y_c, padding=1, groups=C)
        # Gradient magnitude
        return torch.sqrt(grad_x**2 + grad_y**2 + 1e-12)
    return F.l1_loss(get_edges(img1), get_edges(img2))

def dssim_loss(img1, img2):
    return 1.0 - _ssim(img1, img2)

class MedianLossNormalizer:
    def __init__(self, desired_ratios: dict, window_steps: int):
        total_ratio = sum(desired_ratios.values())
        self.ratios = {k: (v / total_ratio) if total_ratio > 0 else 0.0 for k, v in desired_ratios.items()}
        self.buffers = {k: deque(maxlen=window_steps) for k in self.ratios.keys()}

    def update_and_total(self, absolute_losses: dict):
        for k, v in absolute_losses.items():
            if k in self.buffers:
                self.buffers[k].append(float(v.detach().abs().cpu()))
        medians = {k: (np.median(self.buffers[k]) if len(self.buffers[k]) > 0 else 1.0) for k in self.buffers}
        coefficients = {k: (self.ratios[k] / max(medians[k], 1e-12)) for k in self.ratios}
        total_loss = sum(coefficients[k] * absolute_losses[k] for k in absolute_losses if k in coefficients)
        return total_loss, coefficients, medians

loss_normalizer = MedianLossNormalizer(LOSS_RATIOS, MEDIAN_COEFF_STEPS)

# --------------------------- Sample Generation ---------------------------
@torch.no_grad()
def get_fixed_samples(n=3):
    indices = random.sample(range(len(dataset)), min(n, len(dataset)))
    tensors = [input_tfm(random_crop(dataset[i], HIGH_RESOLUTION)) for i in indices]
    return torch.stack(tensors).to(accelerator.device, DTYPE)

fixed_samples = get_fixed_samples()

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:
        unwrapped_vae = accelerator.unwrap_model(vae)
        temp_vae = get_core_model(unwrapped_vae).eval()
        lpips_net = get_lpips_loss()
        
        original_high_res = fixed_samples
        input_tensor = original_high_res.to(dtype=next(temp_vae.parameters()).dtype) if MODEL_RESOLUTION == HIGH_RESOLUTION else F.interpolate(original_high_res, size=(MODEL_RESOLUTION, MODEL_RESOLUTION), mode="area")

        encoder_output = temp_vae.encode(input_tensor)
        latents = encoder_output.latent_dist.mean if TRAIN_DECODER_ONLY else encoder_output.latent_dist.sample()
        reconstructed_images = temp_vae.decode(latents).sample

        if reconstructed_images.shape[-2:] != original_high_res.shape[-2:]:
             reconstructed_images = F.interpolate(reconstructed_images, size=original_high_res.shape[-2:], mode="bilinear", align_corners=False)

        for i in range(reconstructed_images.shape[0]):
            _to_pil_uint8(original_high_res[i]).save(os.path.join(GENERATED_FOLDER, f"sample_real_{i}.png"))
            _to_pil_uint8(reconstructed_images[i]).save(os.path.join(GENERATED_FOLDER, f"sample_decoded_{i}.png"))

        if USE_WANDB and accelerator.is_main_process:
            log_data = {"lpips_mean": float(np.mean([lpips_net(original_high_res[i:i+1], reconstructed_images[i:i+1]).item() for i in range(len(original_high_res))]))}
            for i in range(len(original_high_res)):
                log_data[f"sample/real_{i}"] = wandb.Image(os.path.join(GENERATED_FOLDER, f"sample_real_{i}.png"))
                log_data[f"sample/decoded_{i}"] = wandb.Image(os.path.join(GENERATED_FOLDER, f"sample_decoded_{i}.png"))
            wandb.log(log_data, step=step)
    finally:
        gc.collect()
        torch.cuda.empty_cache()

if accelerator.is_main_process and SAVE_MODEL:
    print("[INFO] Generating initial samples before training...")
    generate_and_save_samples(step=0)

accelerator.wait_for_everyone()

# --------------------------- Training Loop ---------------------------
progress_bar = tqdm(total=total_steps, desc="Training", disable=not accelerator.is_local_main_process)
global_step = 0
min_loss = float("inf")
num_samples_per_epoch = max(1, int(total_steps / max(1, SAMPLE_INTERVAL_SHARE * NUM_EPOCHS)))
sample_interval = max(1, int(round(num_samples_per_epoch / GRADIENT_ACCUMULATION_STEPS)))

for epoch in range(NUM_EPOCHS):
    vae.train()
    batch_losses_history, batch_grads_history = [], []
    tracked_losses = {k: [] for k in LOSS_RATIOS.keys()}

    for batch_idx, imgs in enumerate(dataloader):
        with accelerator.accumulate(vae):
            imgs = imgs.to(accelerator.device)
            imgs_low = imgs if MODEL_RESOLUTION == HIGH_RESOLUTION else F.interpolate(imgs, size=(MODEL_RESOLUTION, MODEL_RESOLUTION), mode="area")
            model_dtype = next(vae.parameters()).dtype
            input_images = imgs_low.to(dtype=model_dtype) if imgs_low.dtype != model_dtype else imgs_low

            current_vae_model = get_core_model(accelerator.unwrap_model(vae))
            encoder_output = current_vae_model.encode(input_images)
            latents = encoder_output.latent_dist.mean if TRAIN_DECODER_ONLY else encoder_output.latent_dist.sample()
            rec_f32 = current_vae_model.decode(latents).sample.to(torch.float32)
            imgs_f32 = imgs.to(torch.float32)

            mae_loss = F.l1_loss(rec_f32, imgs_f32)
            mse_loss = F.mse_loss(rec_f32, imgs_f32)
            lpips_loss_val = get_lpips_loss()(rec_f32, imgs_f32).mean()
            fdl_loss_val = fdl_loss_fn(rec_f32, imgs_f32)
            dssim_loss_val = dssim_loss(rec_f32, imgs_f32)
            edge_loss_val = edge_loss(rec_f32, imgs_f32)

            kl_loss = torch.tensor(0.0, device=accelerator.device, dtype=torch.float32)
            if FULL_TRAINING and not TRAIN_DECODER_ONLY:
                mean = encoder_output.latent_dist.mean
                logvar = encoder_output.latent_dist.logvar
                kl_loss = -0.5 * torch.mean(1 + logvar - mean.pow(2) - logvar.exp())
            
            absolute_losses = {
                "mae": mae_loss, "mse": mse_loss, "lpips": lpips_loss_val,
                "fdl": fdl_loss_val, "dssim": dssim_loss_val, "kl": kl_loss,
                "edge": edge_loss_val,
            }

            total_loss, coeffs, medians = loss_normalizer.update_and_total(absolute_losses)
            if torch.isnan(total_loss) or torch.isinf(total_loss):
                raise RuntimeError("NaN/Inf loss encountered during training.")

            accelerator.backward(total_loss)

            current_grad_norm = torch.tensor(0.0, device=accelerator.device)
            if accelerator.sync_gradients:
                current_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_bar.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_history.append(total_loss.detach().item())
                batch_grads_history.append(float(current_grad_norm.detach().cpu().item()))
                for k, v in absolute_losses.items():
                    tracked_losses[k].append(float(v.detach().item()))

                if USE_WANDB and accelerator.sync_gradients:
                    log_dict = {"total_loss": batch_losses_history[-1], "learning_rate": current_lr, "epoch": epoch, "grad_norm": batch_grads_history[-1]}
                    for k, v in absolute_losses.items(): log_dict[f"loss_{k}"] = float(v.detach().item())
                    for k in coeffs: log_dict[f"coeff_{k}"] = float(coeffs[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(step=global_step)
                accelerator.wait_for_everyone()

                n_logs = min(len(batch_losses_history), sample_interval)
                avg_total = float(np.mean(batch_losses_history[-n_logs:]))
                avg_grad = float(np.mean(batch_grads_history[-n_logs:]))
                
                # ЯВНОЕ ЛОГИРОВАНИЕ КОМПОНЕНТ ПОТЕРЬ
                loss_avgs = {k: float(np.mean(tracked_losses[k][-n_logs:])) for k in tracked_losses if len(tracked_losses[k]) >= n_logs}

                print(f"Epoch {epoch} | Step {global_step} | "
                      f"Total: {avg_total:.5f} | "
                      f"LPIPS: {loss_avgs.get('lpips', 0):.5f} | "
                      f"DSSIM: {loss_avgs.get('dssim', 0):.5f} | "
                      f"MAE: {loss_avgs.get('mae', 0):.5f} | "
                      f"FDL: {loss_avgs.get('fdl', 0):.5f} | "
                      f"EDGE: {loss_avgs.get('edge', 0):.5f} | "
                      f"MSE: {loss_avgs.get('mse', 0):.5f} | "
                      f"Grad: {avg_grad:.5f} | LR: {current_lr:.9f}")

                if SAVE_MODEL and avg_total < min_loss * SAVE_BARRIER:
                    min_loss = avg_total
                    print(f"[INFO] Saving model with improved loss: {min_loss:.6f}")
                    get_core_model(accelerator.unwrap_model(vae)).save_pretrained(SAVE_AS)

    if accelerator.is_main_process:
        print(f"Epoch {epoch} completed. Average Loss: {float(np.mean(batch_losses_history)):.6f}")

if accelerator.is_main_process:
    print("Training finished – saving final model.")
    if SAVE_MODEL:
        get_core_model(accelerator.unwrap_model(vae)).save_pretrained(SAVE_AS)

accelerator.free_memory()
if torch.distributed.is_initialized():
    torch.distributed.destroy_process_group()

print("Training complete. Done!")