remove asymmetric

README.md

@@ -24,7 +24,6 @@ Alchemist eval (512px)
 | madebyollin/sdxl-vae-fp16      |       100% |       100% |       100% |
 | KBlueLeaf/EQ-SDXL-VAE          |     107.8% |     100.1% |      95.5% |
 | AiArtLab/sdxl_vae              |     112.3% |     101.8% |     106.6% |
-| AiArtLab/sdxl_vae_asym         |     111.7% |     101.1% |      89.4% |
 | FLUX.1-schnell-vae             |     324.0% |     119.8% |     292.0% |
 
 [![Click it](vae.png)](https://imgsli.com/NDA3OTgz)

asymmetric_vae/config.json

@@ -1,38 +0,0 @@
-{
-  "_class_name": "AsymmetricAutoencoderKL",
-  "_diffusers_version": "0.34.0",
-  "_name_or_path": "asymmetric_vae_empty",
-  "act_fn": "silu",
-  "down_block_out_channels": [
-    128,
-    256,
-    512,
-    512
-  ],
-  "down_block_types": [
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D"
-  ],
-  "in_channels": 3,
-  "latent_channels": 4,
-  "layers_per_down_block": 2,
-  "layers_per_up_block": 3,
-  "norm_num_groups": 32,
-  "out_channels": 3,
-  "sample_size": 1024,
-  "scaling_factor": 0.13025,
-  "up_block_out_channels": [
-    128,
-    256,
-    512,
-    512
-  ],
-  "up_block_types": [
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D"
-  ]
-}

asymmetric_vae/diffusion_pytorch_model.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ded3c30322578e3371f32a58423b6a3be3a2c3b81d3eb5d35433772be796a1ba
-size 421473052

asymmetric_vae_new/config.json

@@ -1,45 +0,0 @@
-{
-  "_class_name": "AsymmetricAutoencoderKL",
-  "_diffusers_version": "0.35.0.dev0",
-  "_name_or_path": "asymmetric_vae",
-  "act_fn": "silu",
-  "block_out_channels": [
-    128,
-    256,
-    512,
-    512
-  ],
-  "down_block_out_channels": [
-    128,
-    256,
-    512,
-    512
-  ],
-  "down_block_types": [
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D"
-  ],
-  "force_upcast": false,
-  "in_channels": 3,
-  "latent_channels": 4,
-  "layers_per_down_block": 2,
-  "layers_per_up_block": 3,
-  "norm_num_groups": 32,
-  "out_channels": 3,
-  "sample_size": 1024,
-  "scaling_factor": 0.13025,
-  "up_block_out_channels": [
-    128,
-    256,
-    512,
-    512
-  ],
-  "up_block_types": [
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D"
-  ]
-}

asymmetric_vae_new/diffusion_pytorch_model.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:df9380b1e8d8b1a36b3d0f9501a854717a911ae9b8d2aebe18809a6eefa9318b
-size 421473052

eval_alchemist.py

@@ -5,29 +5,71 @@ import lpips
 from PIL import Image, UnidentifiedImageError
 from tqdm import tqdm
 from torch.utils.data import Dataset, DataLoader
-from torchvision.transforms import Compose, Resize, ToTensor, CenterCrop
+from torchvision.transforms import Compose, Resize, ToTensor, CenterCrop,ToPILImage 
 from diffusers import AutoencoderKL, AsymmetricAutoencoderKL
 import random
 
 # --------------------------- Параметры ---------------------------
 DEVICE = "cuda"
 DTYPE = torch.float16
-IMAGE_FOLDER = "/workspace/alchemist"
+IMAGE_FOLDER = "/workspace/alchemist" #wget https://huggingface.co/datasets/AiArtLab/alchemist/resolve/main/alchemist.zip
 MIN_SIZE = 1280
 CROP_SIZE = 512
-BATCH_SIZE = 4  # можно увеличить для ускорения
-MAX_IMAGES = None
-NUM_WORKERS = 4  # параллельная загрузка
+BATCH_SIZE = 5
+MAX_IMAGES = 100
+NUM_WORKERS = 4
+NUM_SAMPLES_TO_SAVE = 10  # Сколько примеров сохранить (0 - не сохранять)
+SAMPLES_FOLDER = "vaetest"
 
 # Список VAE для тестирования
 VAE_LIST = [
+
+#    ("stable-diffusion-v1-5/stable-diffusion-v1-5", AutoencoderKL, "stable-diffusion-v1-5/stable-diffusion-v1-5", "vae"),
+#    ("cross-attention/asymmetric-autoencoder-kl-x-1-5", AsymmetricAutoencoderKL, "cross-attention/asymmetric-autoencoder-kl-x-1-5", None),
     ("madebyollin/sdxl-vae-fp16", AutoencoderKL, "madebyollin/sdxl-vae-fp16-fix", None),
-    ("KBlueLeaf/EQ-SDXL-VAE", AutoencoderKL, "KBlueLeaf/EQ-SDXL-VAE", None),
+#    ("AiArtLab/sdxs", AutoencoderKL, "AiArtLab/sdxs", "vae"),
     ("AiArtLab/sdxl_vae", AutoencoderKL, "AiArtLab/sdxl_vae", None),
-    ("AiArtLab/sdxl_vae_asym", AsymmetricAutoencoderKL, "AiArtLab/sdxl_vae", "asymmetric_vae"),
-    ("FLUX.1-schnell-vae", AutoencoderKL, "black-forest-labs/FLUX.1-schnell", "vae"),
+#    ("AiArtLab/sdxl_vae_asym", AsymmetricAutoencoderKL, "AiArtLab/sdxl_vae", "asymmetric_vae"),
+    ("AiArtLab/sdxl_vae_asym_new", AsymmetricAutoencoderKL, "AiArtLab/sdxl_vae", "asymmetric_vae_new"),
+#    ("KBlueLeaf/EQ-SDXL-VAE", AutoencoderKL, "KBlueLeaf/EQ-SDXL-VAE", None),
+#    ("FLUX.1-schnell-vae", AutoencoderKL, "black-forest-labs/FLUX.1-schnell", "vae"),
 ]
 
+# --------------------------- Sobel Edge Detection ---------------------------
+# Определяем фильтры Собеля глобально
+_sobel_kx = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).view(1, 1, 3, 3)
+_sobel_ky = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32).view(1, 1, 3, 3)
+
+def sobel_edges(x: torch.Tensor) -> torch.Tensor:
+    """
+    Вычисляет карту границ с помощью оператора Собеля
+    x: [B,C,H,W] в диапазоне [-1,1]
+    Возвращает: [B,C,H,W] - магнитуда градиента
+    """
+    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)
+
+def compute_edge_loss(real: torch.Tensor, fake: torch.Tensor) -> float:
+    """
+    Вычисляет Edge Loss между реальным и сгенерированным изображением
+    real, fake: [B,C,H,W] в диапазоне [0,1]
+    Возвращает: скалярное значение loss
+    """
+    # Конвертируем в [-1,1] для sobel_edges
+    real_norm = real * 2 - 1
+    fake_norm = fake * 2 - 1
+    
+    # Получаем карты границ
+    edges_real = sobel_edges(real_norm)
+    edges_fake = sobel_edges(fake_norm)
+    
+    # L1 loss между картами границ
+    return F.l1_loss(edges_fake, edges_real).item()
+
 # --------------------------- Dataset ---------------------------
 class ImageFolderDataset(Dataset):
     def __init__(self, root_dir, extensions=('.png',), min_size=1024, crop_size=512, limit=None):
@@ -36,18 +78,15 @@ class ImageFolderDataset(Dataset):
         self.crop_size = crop_size
         self.paths = []
         
-        # Собираем пути к файлам
         print("Сканирование папки...")
         for root, _, files in os.walk(root_dir):
             for fname in files:
                 if fname.lower().endswith(extensions):
                     self.paths.append(os.path.join(root, fname))
         
-        # Ограничение количества
         if limit:
             self.paths = self.paths[:limit]
         
-        # Быстрая проверка валидности (опционально, можно убрать для скорости)
         print("Проверка изображений...")
         valid = []
         for p in tqdm(self.paths, desc="Проверка"):
@@ -62,11 +101,9 @@ class ImageFolderDataset(Dataset):
         if len(self.paths) == 0:
             raise RuntimeError(f"Не найдено валидных изображений в {root_dir}")
         
-        # Перемешиваем для случайности
         random.shuffle(self.paths)
         print(f"Найдено {len(self.paths)} изображений")
         
-        # Трансформации
         self.transform = Compose([
             Resize(min_size, interpolation=Image.LANCZOS),
             CenterCrop(crop_size),
@@ -89,9 +126,14 @@ def process(x):
 def deprocess(x):
     return x * 0.5 + 0.5
 
+def _sanitize_name(name: str) -> str:
+    return name.replace('/', '_').replace('-', '_')
+
 # --------------------------- Основной код ---------------------------
 if __name__ == "__main__":
-    # Создаем датасет и загрузчик
+    if NUM_SAMPLES_TO_SAVE > 0:
+        os.makedirs(SAMPLES_FOLDER, exist_ok=True)
+        
     dataset = ImageFolderDataset(
         IMAGE_FOLDER, 
         extensions=('.png',), 
@@ -103,16 +145,14 @@ if __name__ == "__main__":
     dataloader = DataLoader(
         dataset,
         batch_size=BATCH_SIZE,
-        shuffle=False,  # уже перемешали в датасете
+        shuffle=False,
         num_workers=NUM_WORKERS,
         pin_memory=True,
         drop_last=False
     )
     
-    # Инициализация LPIPS
     lpips_net = lpips.LPIPS(net="vgg").eval().to(DEVICE).requires_grad_(False)
     
-    # Загрузка VAE моделей
     print("\nЗагрузка VAE моделей...")
     vaes = []
     names = []
@@ -120,67 +160,118 @@ if __name__ == "__main__":
     for name, vae_class, model_path, subfolder in VAE_LIST:
         try:
             print(f"  Загружаю {name}...")
-            vae = vae_class.from_pretrained(model_path, subfolder=subfolder)
+            # Исправлена загрузка для variant
+            if "sdxs" in model_path:
+                vae = vae_class.from_pretrained(model_path, subfolder=subfolder, variant="fp16")
+            else:
+                vae = vae_class.from_pretrained(model_path, subfolder=subfolder)
             vae = vae.to(DEVICE, DTYPE).eval()
             vaes.append(vae)
             names.append(name)
         except Exception as e:
             print(f"  ❌ Ошибка загрузки {name}: {e}")
     
-    # Оценка метрик
     print("\nОценка метрик...")
-    results = {name: {"mse": 0.0, "psnr": 0.0, "lpips": 0.0, "count": 0} for name in names}
-    
+    results = {name: {"mse": 0.0, "psnr": 0.0, "lpips": 0.0, "edge": 0.0, "count": 0} for name in names}
+
+    to_pil = ToPILImage()
+
+    # >>>>>>>> ОСНОВНЫЕ ИЗМЕНЕНИЯ ЗДЕСЬ (KISS) <<<<<<<<
     with torch.no_grad():
+        images_saved = 0  # считаем именно КОЛ-ВО ИЗОБРАЖЕНИЙ, а не сохранённых файлов
         for batch in tqdm(dataloader, desc="Обработка батчей"):
-            batch = batch.to(DEVICE)
-            test_inp = process(batch).to(DTYPE)
-            
-            for vae, name in zip(vaes, names):
-                # Encode/decode
+            batch = batch.to(DEVICE)                                # [B,3,H,W] в [0,1]
+            test_inp = process(batch).to(DTYPE)                     # [-1,1] для энкодера
+
+            # 1) считаем реконструкции для всех VAE на весь батч
+            recon_list = []
+            for vae in vaes:
                 latent = vae.encode(test_inp).latent_dist.mode()
-                recon = deprocess(vae.decode(latent).sample.float())
-                
-                # Метрики для батча
+                dec = vae.decode(latent).sample.float()             # [-1,1] (как правило)
+                recon = deprocess(dec).clamp(0.0, 1.0)              # -> [0,1], clamp убирает артефакты
+                recon_list.append(recon)
+
+            # 2) обновляем метрики (по каждой VAE)
+            for recon, name in zip(recon_list, names):
                 for i in range(batch.shape[0]):
                     img_orig = batch[i:i+1]
                     img_recon = recon[i:i+1]
-                    
                     mse = F.mse_loss(img_orig, img_recon).item()
                     psnr = 10 * torch.log10(1 / torch.tensor(mse)).item()
                     lpips_val = lpips_net(img_orig, img_recon, normalize=True).mean().item()
-                    
+                    edge_loss = compute_edge_loss(img_orig, img_recon)
                     results[name]["mse"] += mse
                     results[name]["psnr"] += psnr
                     results[name]["lpips"] += lpips_val
+                    results[name]["edge"] += edge_loss
                     results[name]["count"] += 1
-    
+
+            # 3) сохраняем ровно NUM_SAMPLES_TO_SAVE изображений (orig + все VAE + общий коллаж)
+            if NUM_SAMPLES_TO_SAVE > 0:
+                for i in range(batch.shape[0]):
+                    if images_saved >= NUM_SAMPLES_TO_SAVE:
+                        break
+                    idx_str = f"{images_saved + 1:03d}"
+
+                    # original
+                    orig_pil = to_pil(batch[i].detach().float().cpu())
+                    orig_pil.save(os.path.join(SAMPLES_FOLDER, f"{idx_str}_orig.png"))
+
+                    # per-VAE decodes
+                    tiles = [orig_pil]
+                    for recon, name in zip(recon_list, names):
+                        recon_pil = to_pil(recon[i].detach().cpu())
+                        recon_pil.save(os.path.join(
+                            SAMPLES_FOLDER, f"{idx_str}_decoded_{_sanitize_name(name)}.png"
+                        ))
+                        tiles.append(recon_pil)
+
+                    # общий коллаж: [orig | vae1 | vae2 | ...]
+                    collage_w = CROP_SIZE * len(tiles)
+                    collage_h = CROP_SIZE
+                    collage = Image.new("RGB", (collage_w, collage_h))
+                    x = 0
+                    for tile in tiles:
+                        collage.paste(tile, (x, 0))
+                        x += CROP_SIZE
+                    collage.save(os.path.join(SAMPLES_FOLDER, f"{idx_str}_all.png"))
+
+                    images_saved += 1
+
+
     # Усреднение результатов
     for name in names:
         count = results[name]["count"]
         results[name]["mse"] /= count
         results[name]["psnr"] /= count
         results[name]["lpips"] /= count
+        results[name]["edge"] /= count
     
     # Вывод абсолютных значений
     print("\n=== Абсолютные значения ===")
     for name in names:
-        print(f"{name:30s}: MSE: {results[name]['mse']:.3e}, PSNR: {results[name]['psnr']:.4f}, LPIPS: {results[name]['lpips']:.4f}")
+        print(f"{name:30s}: MSE: {results[name]['mse']:.3e}, PSNR: {results[name]['psnr']:.4f}, "
+              f"LPIPS: {results[name]['lpips']:.4f}, Edge: {results[name]['edge']:.4f}")
     
     # Вывод таблицы с процентами
     print("\n=== Сравнение с первой моделью (%) ===")
-    print(f"| {'Модель':30s} | {'MSE':>10s} | {'PSNR':>10s} | {'LPIPS':>10s} |")
-    print(f"|{'-'*32}|{'-'*12}|{'-'*12}|{'-'*12}|")
+    print(f"| {'Модель':30s} | {'MSE':>10s} | {'PSNR':>10s} | {'LPIPS':>10s} | {'Edge':>10s} |")
+    print(f"|{'-'*32}|{'-'*12}|{'-'*12}|{'-'*12}|{'-'*12}|")
     
     baseline = names[0]
     for name in names:
+        # Для MSE, LPIPS и Edge: меньше = лучше, поэтому инвертируем
         mse_pct = (results[baseline]["mse"] / results[name]["mse"]) * 100
+        # Для PSNR: больше = лучше
         psnr_pct = (results[name]["psnr"] / results[baseline]["psnr"]) * 100
+        # Для LPIPS и Edge: меньше = лучше
         lpips_pct = (results[baseline]["lpips"] / results[name]["lpips"]) * 100
+        edge_pct = (results[baseline]["edge"] / results[name]["edge"]) * 100
         
         if name == baseline:
-            print(f"| {name:30s} | {'100%':>10s} | {'100%':>10s} | {'100%':>10s} |")
+            print(f"| {name:30s} | {'100%':>10s} | {'100%':>10s} | {'100%':>10s} | {'100%':>10s} |")
         else:
-            print(f"| {name:30s} | {f'{mse_pct:.1f}%':>10s} | {f'{psnr_pct:.1f}%':>10s} | {f'{lpips_pct:.1f}%':>10s} |")
+            print(f"| {name:30s} | {f'{mse_pct:.1f}%':>10s} | {f'{psnr_pct:.1f}%':>10s} | "
+                  f"{f'{lpips_pct:.1f}%':>10s} | {f'{edge_pct:.1f}%':>10s} |")
     
-    print("\n✅ Готово!")
+    print("\n✅ Готово!")

train_sdxl_vae_gpt5.py → train_sdxl_vae.py

@@ -24,7 +24,7 @@ from collections import deque
 
 # --------------------------- Параметры ---------------------------
 ds_path            = "/workspace/png"
-project            = "asymmetric_vae"
+project            = "vae"
 batch_size         = 3
 base_learning_rate = 6e-6
 min_learning_rate  = 1e-6
@@ -50,7 +50,7 @@ clip_grad_norm     = 1.0
 mixed_precision    = "no"   # или "fp16"/"bf16" при поддержке
 gradient_accumulation_steps = 5
 generated_folder   = "samples"
-save_as            = "asymmetric_vae_new"
+save_as            = "vae_nightly"
 num_workers        = 0
 device = None  # accelerator задаст устройство
 
@@ -65,7 +65,7 @@ loss_ratios = {
 median_coeff_steps = 256  # за сколько шагов считать медианные коэффициенты
 
 # --------------------------- параметры препроцессинга ---------------------------
-resize_long_side = 1280  # если None или 0 — ресайза не будет; рекомендовано 1024
+resize_long_side = 1280  # если None или 0 — ресайза не будет; рекомендовано 1280
 
 Path(generated_folder).mkdir(parents=True, exist_ok=True)
 

vae/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f17d5c9503862b25a273b8874851a99de817dbfae6094432f51381bb1cdd60c8
+oid sha256:03f2412467f6bedce9efeddba5860b5ec0d3267931d14c500d4bd7a878e14cbd
 size 334643268