{
 "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"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The config attributes {'block_out_channels': [128, 256, 512, 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, 142199.23it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "✅ Перенос завершён.\n",
      "Статистика:\n",
      "  перенесено: 142\n",
      "  дублировано: 26\n",
      "  сдвинуто: 106\n",
      "  пропущено: 0\n",
      "\n",
      "Неперенесённые ключи (первые 20):\n",
      "   decoder.condition_encoder.layers.0.weight\n",
      "   decoder.condition_encoder.layers.0.bias\n",
      "   decoder.condition_encoder.layers.1.weight\n",
      "   decoder.condition_encoder.layers.1.bias\n",
      "   decoder.condition_encoder.layers.2.weight\n",
      "   decoder.condition_encoder.layers.2.bias\n",
      "   decoder.condition_encoder.layers.3.weight\n",
      "   decoder.condition_encoder.layers.3.bias\n",
      "   decoder.condition_encoder.layers.4.weight\n",
      "   decoder.condition_encoder.layers.4.bias\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",
    "    \"up_block_out_channels\": [64, 128, 256, 512, 512],  # +1 блок\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",
    "# ---- Функция переноса весов ----\n",
    "def transfer_weights_with_duplication(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, \"пропущено\": 0}\n",
    "\n",
    "    print(\"\\n--- Перенос весов ---\")\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",
    "            parts = k.split(\".\")\n",
    "            idx = int(parts[2])\n",
    "\n",
    "            # сдвигаем индекс на +1 (так как добавлен новый блок в начало)\n",
    "            new_idx = idx + 1\n",
    "            new_k = \".\".join([parts[0], parts[1], str(new_idx)] + parts[3:])\n",
    "            if new_k in new_sd and 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",
    "            continue\n",
    "\n",
    "        # === Перенос прочих совпадающих ключей ===\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",
    "\n",
    "    # === Дублирование весов старого 512→512 блока в новый ===\n",
    "    ref_prefix = \"decoder.up_blocks.1\"  # старый первый up-блок (512→512)\n",
    "    new_prefix = \"decoder.up_blocks.0\"  # новый добавленный блок\n",
    "    for k, v in old_sd.items():\n",
    "        if k.startswith(ref_prefix):\n",
    "            new_k = k.replace(ref_prefix, new_prefix)\n",
    "            if new_k in new_sd and 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",
    "    missing = [k for k in new_sd.keys() if k not in transferred_keys]\n",
    "    if missing:\n",
    "        print(\"\\nНеперенесённые ключи (первые 20):\")\n",
    "        for k in missing[:20]:\n",
    "            print(\"  \", k)\n",
    "\n",
    "# ---- Запуск ----\n",
    "transfer_weights_with_duplication(\"AiArtLab/simplevae\", \"asymmetric_vae_empty\", save_path=\"vae\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "65653a65-e7c2-4b67-bc17-62c21cfd1db8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting hf_transfer\n",
      "  Downloading hf_transfer-0.1.9-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (1.7 kB)\n",
      "Downloading hf_transfer-0.1.9-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.6 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.6/3.6 MB\u001b[0m \u001b[31m34.5 MB/s\u001b[0m  \u001b[33m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: hf_transfer\n",
      "Successfully installed hf_transfer-0.1.9\n"
     ]
    }
   ],
   "source": [
    "!pip install hf_transfer"
   ]
  },
  {
   "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
}