Upload folder using huggingface_hub

train_sdxs_vae.py

@@ -26,10 +26,11 @@ 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"
-project            = "vae"
+project            = "vae2"
 batch_size         = 1
 base_learning_rate = 6e-6
 min_learning_rate  = 7e-7
@@ -52,7 +53,7 @@ clip_grad_norm     = 1.0
 mixed_precision    = "no"
 gradient_accumulation_steps = 1
 generated_folder   = "samples"
-save_as            = "vae2"
+save_as            = "vae3"
 num_workers        = 0
 device = None
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -73,9 +74,8 @@ kl_ratio           = 0.0
 loss_ratios = {
     "lpips": 0.70,#0.50,
     "fdl" :  0.10,#0.25,   
-    "edge":  0.05,
     "mse":   0.10,
-    "mae":   0.05,
+    "mae":   0.10,
     "kl":    0.00,
 }
 median_coeff_steps = 250
@@ -195,33 +195,48 @@ else:
 
 
 print(f"[INFO] Разморожено параметров: {len(unfrozen_param_names)}. Первые 200 имён:")
-for nm in unfrozen_param_names[:100]:
+for nm in unfrozen_param_names[:10]:
     print(" ", nm)
 
 # --------------------------- Датасет ---------------------------
+from torch.utils.data import Dataset
+from PIL import Image, UnidentifiedImageError
+import random
+import torchvision.transforms as transforms
+import os
+
 class PngFolderDataset(Dataset):
-    def __init__(self, root_dir, min_exts=('.png',), resolution=1024, limit=0):
-        self.root_dir = root_dir
+    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 fname in files:
-                if fname.lower().endswith(tuple(ext.lower() for ext in min_exts)):
-                    self.paths.append(os.path.join(root, fname))
+            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:
             self.paths = self.paths[:limit]
+
+        # фильтруем недопустимые картинки
         valid = []
         for p in self.paths:
             try:
-                with Image.open(p) as im:
-                    im.verify()
+                with Image.open(p) as img:
+                    img.verify()  # только метаданные
+                    w, h = img.size
+                    if w < resolution or h < resolution:
+                        continue
                 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}")
+        if not self.paths:
+            raise RuntimeError("No valid images found")
+
         random.shuffle(self.paths)
+        self.transform = transforms.ToTensor()  # конвертирует сразу [0,1] float32
 
     def __len__(self):
         return len(self.paths)
@@ -230,21 +245,10 @@ class PngFolderDataset(Dataset):
         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.BICUBIC)
+            return img
 
 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.BICUBIC)
     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))
@@ -254,11 +258,6 @@ tfm = transforms.Compose([
     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)
-print("len(dataset)",len(dataset))
-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:
@@ -266,15 +265,12 @@ def collate_fn(batch):
         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
-)
+dataset = PngFolderDataset(ds_path, min_exts=('.png',), resolution=high_resolution, limit=limit)
+print("len(dataset)",len(dataset))
+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)
 
 # --------------------------- Оптимизатор ---------------------------
 def get_param_groups(module, weight_decay=0.001):
@@ -350,15 +346,6 @@ def _get_lpips():
         _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):
@@ -532,7 +519,6 @@ for epoch in range(num_epochs):
                 "mse":   F.mse_loss(rec_f32, imgs_f32),
                 "lpips": _get_lpips()(rec_f32, imgs_f32).mean(),
                 "fdl":   fdl_loss(rec_f32, imgs_f32),
-                "edge":  F.l1_loss(sobel_edges(rec_f32), sobel_edges(imgs_f32)),
             }
 
             if full_training and not train_decoder_only: