vae3

requirements.txt

@@ -0,0 +1,10 @@
+diffusers>=0.32.2
+accelerate>=1.5.2
+datasets>=3.5.0
+matplotlib>=3.10.1
+wandb>=0.19.8
+huggingface_hub>=0.29.3
+bitsandbytes>=0.45.4
+transformers
+hf_transfer
+lpips

train_vae.py

@@ -28,20 +28,20 @@ from collections import deque
 
 # --------------------------- Параметры ---------------------------
 ds_path            = "/workspace/d23"
-project            = "vae2"
-batch_size         = 1
-base_learning_rate = 6e-6
-min_learning_rate  = 8e-8
-num_epochs         = 20
-sample_interval_share = 10
+project            = "vae3"
+batch_size         = 5
+base_learning_rate = 5e-5
+min_learning_rate  = 1e-5
+num_epochs         = 50
+sample_interval_share = 2
 use_wandb          = True
 save_model         = True
 use_decay          = True
 optimizer_type     = "adam8bit"
 dtype              = torch.float32
 
-model_resolution   = 512
-high_resolution    = 1024 
+model_resolution   = 256
+high_resolution    = 512 
 limit              = 0
 save_barrier       = 1.3
 warmup_percent     = 0.001
@@ -50,9 +50,9 @@ beta2              = 0.997
 eps                = 1e-8
 clip_grad_norm     = 1.0
 mixed_precision    = "no"
-gradient_accumulation_steps = 16
+gradient_accumulation_steps = 2
 generated_folder   = "samples"
-save_as            = "vae2"
+save_as            = "vae3"
 num_workers        = 0
 device = None
 
@@ -70,7 +70,7 @@ loss_ratios = {
     "mae":   0.10,
     "kl":    0.00,  # активируем при full_training=True
 }
-median_coeff_steps = 256
+median_coeff_steps = 1000
 
 resize_long_side = 1280  # ресайз длинной стороны исходных картинок
 
@@ -385,8 +385,72 @@ 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:
+        temp_vae = accelerator.unwrap_model(vae).eval()
+        lpips_net = _get_lpips()
+        with torch.no_grad():
+            orig_high = fixed_samples
+            orig_low = F.interpolate(
+                orig_high,
+                size=(model_resolution, model_resolution),
+                mode="bilinear",
+                align_corners=False
+            )
+            model_dtype = next(temp_vae.parameters()).dtype
+            orig_low = orig_low.to(dtype=model_dtype)
+
+            # Encode/decode с учётом видео-режима
+            if is_video_vae(temp_vae):
+                x_in = orig_low.unsqueeze(2)  # [B,3,1,H,W]
+                enc = temp_vae.encode(x_in)
+                latents_mean = enc.latent_dist.mean
+                dec = temp_vae.decode(latents_mean).sample  # [B,3,1,H,W]
+                rec = dec.squeeze(2)  # [B,3,H,W]
+            else:
+                enc = temp_vae.encode(orig_low)
+                latents_mean = enc.latent_dist.mean
+                rec = temp_vae.decode(latents_mean).sample
+
+        # Подгон размеров, если надо
+        if rec.shape[-2:] != orig_high.shape[-2:]:
+            rec = F.interpolate(rec, size=orig_high.shape[-2:], mode="bilinear", align_corners=False)
+
+        # Сохраняем все real/decoded
+        for i in range(rec.shape[0]):
+            real_img = _to_pil_uint8(orig_high[i])
+            dec_img  = _to_pil_uint8(rec[i])
+            real_img.save(f"{generated_folder}/sample_real_{i}.jpg", quality=95)
+            dec_img.save(f"{generated_folder}/sample_decoded_{i}.jpg", quality=95)
+
+        # LPIPS
+        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))
+
+        # W&B логирование
+        if use_wandb and accelerator.is_main_process:
+            log_data = {"lpips_mean": avg_lpips}
+            for i in range(rec.shape[0]):
+                log_data[f"sample/real_{i}"] = wandb.Image(f"{generated_folder}/sample_real_{i}.jpg", caption=f"real_{i}")
+                log_data[f"sample/decoded_{i}"] = wandb.Image(f"{generated_folder}/sample_decoded_{i}.jpg", caption=f"decoded_{i}")
+            wandb.log(log_data, step=step)
+
+    finally:
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+
+def generate_and_save_samples2(step=None):
     try:
         temp_vae = accelerator.unwrap_model(vae).eval()
         lpips_net = _get_lpips()

transfer_simplevae3.ipynb

@@ -0,0 +1,221 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "c15deb04-94a0-4073-a174-adcd22af10b8",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "✅ Создана новая модель: <class 'diffusers.models.autoencoders.autoencoder_asym_kl.AsymmetricAutoencoderKL'>\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e2063f203ab844489f3c02cb9c2ae70b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "config.json:   0%|          | 0.00/801 [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "d33d67a744ee43b3b9eaeba9228ba976",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "vae/diffusion_pytorch_model.safetensors:   0%|          | 0.00/168M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "The config attributes {'block_out_channels': [128, 128, 256, 512, 512], 'force_upcast': False} were passed to AsymmetricAutoencoderKL, but are not expected and will be ignored. Please verify your config.json configuration file.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "--- Перенос весов ---\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 248/248 [00:00<00:00, 87271.36it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "✅ Перенос завершён.\n",
+      "Статистика:\n",
+      "  перенесено: 216\n",
+      "  дублировано: 26\n",
+      "  пропущено: 0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from diffusers.models import AsymmetricAutoencoderKL, AutoencoderKL\n",
+    "import torch\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "# ---- Конфиг новой модели ----\n",
+    "config = {\n",
+    "    \"_class_name\": \"AsymmetricAutoencoderKL\",\n",
+    "    \"act_fn\": \"silu\",\n",
+    "    \"in_channels\": 3,\n",
+    "    \"out_channels\": 3,\n",
+    "    \"scaling_factor\": 1.0,\n",
+    "    \"norm_num_groups\": 32,\n",
+    "    \"down_block_out_channels\": [128, 256, 512, 512],\n",
+    "    \"down_block_types\": [\n",
+    "        \"DownEncoderBlock2D\",\n",
+    "        \"DownEncoderBlock2D\",\n",
+    "        \"DownEncoderBlock2D\",\n",
+    "        \"DownEncoderBlock2D\",\n",
+    "    ],\n",
+    "    \"latent_channels\": 16,\n",
+    "    # Новый UpDecoderBlock добавлен в начало\n",
+    "    \"up_block_out_channels\": [128, 128, 256, 512, 512],\n",
+    "    \"up_block_types\": [\n",
+    "        \"UpDecoderBlock2D\",\n",
+    "        \"UpDecoderBlock2D\",\n",
+    "        \"UpDecoderBlock2D\",\n",
+    "        \"UpDecoderBlock2D\",\n",
+    "        \"UpDecoderBlock2D\",\n",
+    "    ],\n",
+    "}\n",
+    "\n",
+    "# ---- Создание пустой асимметричной модели ----\n",
+    "vae = AsymmetricAutoencoderKL(\n",
+    "    act_fn=config[\"act_fn\"],\n",
+    "    down_block_out_channels=config[\"down_block_out_channels\"],\n",
+    "    down_block_types=config[\"down_block_types\"],\n",
+    "    latent_channels=config[\"latent_channels\"],\n",
+    "    up_block_out_channels=config[\"up_block_out_channels\"],\n",
+    "    up_block_types=config[\"up_block_types\"],\n",
+    "    in_channels=config[\"in_channels\"],\n",
+    "    out_channels=config[\"out_channels\"],\n",
+    "    scaling_factor=config[\"scaling_factor\"],\n",
+    "    norm_num_groups=config[\"norm_num_groups\"],\n",
+    "    layers_per_down_block=2,\n",
+    "    layers_per_up_block=2,\n",
+    "    sample_size=1024\n",
+    ")\n",
+    "\n",
+    "vae.save_pretrained(\"asymmetric_vae_empty\")\n",
+    "print(\"✅ Создана новая модель:\", type(vae))\n",
+    "\n",
+    "# ---- Функция переноса весов старого VAE ----\n",
+    "def transfer_weights(old_path, new_path, save_path=\"asymmetric_vae\", device=\"cuda\", dtype=torch.float16):\n",
+    "    old_vae = AutoencoderKL.from_pretrained(old_path, subfolder=\"vae\").to(device, dtype=dtype)\n",
+    "    new_vae = AsymmetricAutoencoderKL.from_pretrained(new_path).to(device, dtype=dtype)\n",
+    "\n",
+    "    old_sd = old_vae.state_dict()\n",
+    "    new_sd = new_vae.state_dict()\n",
+    "\n",
+    "    transferred_keys = set()\n",
+    "    transfer_stats = {\"перенесено\": 0, \"дублировано\": 0, \"пропущено\": 0}\n",
+    "\n",
+    "    print(\"\\n--- Перенос весов ---\")\n",
+    "    for k, v in tqdm(old_sd.items()):\n",
+    "        # Копирование энкодера и прочих совпадающих ключей\n",
+    "        if (\"encoder\" in k) or (\"quant_conv\" in k) or (\"post_quant_conv\" in k):\n",
+    "            if k in new_sd and new_sd[k].shape == v.shape:\n",
+    "                new_sd[k] = v.clone()\n",
+    "                transferred_keys.add(k)\n",
+    "                transfer_stats[\"перенесено\"] += 1\n",
+    "            continue\n",
+    "\n",
+    "        # Копирование декодера (без сдвига)\n",
+    "        if \"decoder.up_blocks\" in k:\n",
+    "            if k in new_sd and new_sd[k].shape == v.shape:\n",
+    "                new_sd[k] = v.clone()\n",
+    "                transferred_keys.add(k)\n",
+    "                transfer_stats[\"перенесено\"] += 1\n",
+    "            continue\n",
+    "\n",
+    "    # Дублирование весов старого первого 512→512 блока в новый блок 64→128 для апскейла\n",
+    "    ref_prefix = \"decoder.up_blocks.1\"\n",
+    "    new_prefix = \"decoder.up_blocks.0\"\n",
+    "    for k, v in old_sd.items():\n",
+    "        if k.startswith(ref_prefix) and new_prefix + k[len(ref_prefix):] in new_sd:\n",
+    "            new_k = k.replace(ref_prefix, new_prefix)\n",
+    "            if new_sd[new_k].shape == v.shape:\n",
+    "                new_sd[new_k] = v.clone()\n",
+    "                transferred_keys.add(new_k)\n",
+    "                transfer_stats[\"дублировано\"] += 1\n",
+    "\n",
+    "    # Загрузка и сохранение\n",
+    "    new_vae.load_state_dict(new_sd, strict=False)\n",
+    "    new_vae.save_pretrained(save_path)\n",
+    "\n",
+    "    print(\"\\n✅ Перенос завершён.\")\n",
+    "    print(\"Статистика:\")\n",
+    "    for k, v in transfer_stats.items():\n",
+    "        print(f\"  {k}: {v}\")\n",
+    "\n",
+    "# ---- Запуск переноса ----\n",
+    "transfer_weights(\"AiArtLab/simplevae\", \"asymmetric_vae_empty\", save_path=\"vae3\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "59fcafb9-6d89-49b4-8362-b4891f591687",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

vae3/config.json

@@ -0,0 +1,48 @@
+{
+  "_class_name": "AsymmetricAutoencoderKL",
+  "_diffusers_version": "0.35.2",
+  "_name_or_path": "vae3",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    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": 16,
+  "layers_per_down_block": 2,
+  "layers_per_up_block": 2,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 1024,
+  "scaling_factor": 1.0,
+  "up_block_out_channels": [
+    128,
+    128,
+    256,
+    512,
+    512
+  ],
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ]
+}

vae3/diffusion_pytorch_model.safetensors

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