Upload folder using huggingface_hub

train_sdxs_vae.py

@@ -26,7 +26,6 @@ import wandb
 import lpips   # pip install lpips
 from FDL_pytorch import FDL_loss # pip install fdl-pytorch
 from collections import deque
-import torch.nn.functional as F
 
 # --------------------------- Параметры ---------------------------
 ds_path            = "/workspace/d23"
@@ -49,7 +48,7 @@ save_barrier       = 1.3
 warmup_percent     = 0.005
 beta2              = 0.997
 eps                = 1e-8
-clip_grad_norm     = 1.0
+clip_grad_norm     = 2.0  # Увеличили с 1.0 для лучшей прорисовки деталей
 mixed_precision    = "no"
 gradient_accumulation_steps = 1
 generated_folder   = "samples"
@@ -72,10 +71,11 @@ kl_ratio           = 0.0
 
 # Доли лоссов
 loss_ratios = {
-    "lpips": 0.70,#0.50,
-    "fdl" :  0.10,#0.25,   
-    "mse":   0.10,
-    "mae":   0.10,
+    "lpips": 0.60,
+    "fdl" :  0.15,   
+    "mse":   0.00,
+    "mae":   0.20,
+    "dssim": 0.05,  # <-- Заменён edge на DSSIM
     "kl":    0.00,
 }
 median_coeff_steps = 250
@@ -273,21 +273,6 @@ if len(dataset) < 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)
 
 # --------------------------- Оптимизатор ---------------------------
-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 get_param_groups(module, weight_decay=0.001):
     no_decay_tokens = ("bias", "norm", "rms", "layernorm")
     decay_params, no_decay_params = [], []
@@ -338,7 +323,7 @@ trainable_params = [p for p in vae.parameters() if p.requires_grad]
 fdl_loss = FDL_loss()
 fdl_loss = fdl_loss.to(accelerator.device)
 
-# --------------------------- LPIPS и вспомогательные ---------------------------
+# --------------------------- LPIPS и DSSIM ---------------------------
 _lpips_net = None
 def _get_lpips():
     global _lpips_net
@@ -346,7 +331,41 @@ def _get_lpips():
         _lpips_net = lpips.LPIPS(net='vgg', verbose=False).eval().to(accelerator.device).eval()
     return _lpips_net
 
-
+# --- DSSIM Implementation (замена Sobel) ---
+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(x, y, window_size=11, sigma=1.5):
+    """Вычисляет SSIM для тензоров в диапазоне [-1, 1]"""
+    C = x.shape[1]
+    k = _gaussian_kernel(window_size, sigma, x.device, x.dtype)
+    win = (k.view(1, 1, -1, 1) * k.view(1, 1, 1, -1)).expand(C, 1, window_size, window_size).contiguous()
+    
+    mu1 = F.conv2d(x, win, padding=window_size//2, groups=C)
+    mu2 = F.conv2d(y, win, padding=window_size//2, groups=C)
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1 * mu2
+    
+    sigma1_sq = F.conv2d(x*x, win, padding=window_size//2, groups=C) - mu1_sq
+    sigma2_sq = F.conv2d(y*y, win, padding=window_size//2, groups=C) - mu2_sq
+    sigma12 = F.conv2d(x*y, win, padding=window_size//2, groups=C) - mu1_mu2
+    
+    L = 2.0  # диапазон для [-1, 1] -> 2
+    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)
+    ssim_map = num / den
+    return ssim_map.mean()
+
+def dssim_loss(x, y):
+    """DSSIM = 1 - SSIM. Отлично сохраняет локальную структуру штрихов."""
+    return 1.0 - _ssim(x, y)
+    
 class MedianLossNormalizer:
     def __init__(self, desired_ratios: dict, window_steps: int):
         s = sum(desired_ratios.values())
@@ -508,17 +527,16 @@ for epoch in range(num_epochs):
                 latents = enc.latent_dist.mean if train_decoder_only else enc.latent_dist.sample()
                 rec = current_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),
+                "mse":   torch.tensor(0.0, device=accelerator.device, dtype=torch.float32), # Отключён, экономим вычисления
                 "lpips": _get_lpips()(rec_f32, imgs_f32).mean(),
                 "fdl":   fdl_loss(rec_f32, imgs_f32),
+                "dssim": dssim_loss(rec_f32, imgs_f32), # <-- DSSIM вместо Edge
             }
 
             if full_training and not train_decoder_only: