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