Upload folder using huggingface_hub

.ipynb_checkpoints/create_asymmetric-checkpoint.ipynb

@@ -0,0 +1,620 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "407171be-ab46-442b-a0bd-83ca75173eba",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "AsymmetricAutoencoderKL(\n",
+      "  (encoder): Encoder(\n",
+      "    (conv_in): Conv2d(3, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "    (down_blocks): ModuleList(\n",
+      "      (0): DownEncoderBlock2D(\n",
+      "        (resnets): ModuleList(\n",
+      "          (0-1): 2 x ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "          )\n",
+      "        )\n",
+      "        (downsamplers): ModuleList(\n",
+      "          (0): Downsample2D(\n",
+      "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2))\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "      (1): DownEncoderBlock2D(\n",
+      "        (resnets): ModuleList(\n",
+      "          (0): ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "            (conv_shortcut): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1))\n",
+      "          )\n",
+      "          (1): ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "          )\n",
+      "        )\n",
+      "        (downsamplers): ModuleList(\n",
+      "          (0): Downsample2D(\n",
+      "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2))\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "      (2): DownEncoderBlock2D(\n",
+      "        (resnets): ModuleList(\n",
+      "          (0): ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "            (conv_shortcut): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1))\n",
+      "          )\n",
+      "          (1): ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "          )\n",
+      "        )\n",
+      "        (downsamplers): ModuleList(\n",
+      "          (0): Downsample2D(\n",
+      "            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2))\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "      (3): DownEncoderBlock2D(\n",
+      "        (resnets): ModuleList(\n",
+      "          (0-1): 2 x ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "    )\n",
+      "    (mid_block): UNetMidBlock2D(\n",
+      "      (attentions): ModuleList(\n",
+      "        (0): Attention(\n",
+      "          (group_norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "          (to_q): Linear(in_features=512, out_features=512, bias=True)\n",
+      "          (to_k): Linear(in_features=512, out_features=512, bias=True)\n",
+      "          (to_v): Linear(in_features=512, out_features=512, bias=True)\n",
+      "          (to_out): ModuleList(\n",
+      "            (0): Linear(in_features=512, out_features=512, bias=True)\n",
+      "            (1): Dropout(p=0.0, inplace=False)\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "      (resnets): ModuleList(\n",
+      "        (0-1): 2 x ResnetBlock2D(\n",
+      "          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "          (dropout): Dropout(p=0.0, inplace=False)\n",
+      "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "          (nonlinearity): SiLU()\n",
+      "        )\n",
+      "      )\n",
+      "    )\n",
+      "    (conv_norm_out): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "    (conv_act): SiLU()\n",
+      "    (conv_out): Conv2d(512, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "  )\n",
+      "  (decoder): MaskConditionDecoder(\n",
+      "    (conv_in): Conv2d(32, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "    (up_blocks): ModuleList(\n",
+      "      (0-1): 2 x UpDecoderBlock2D(\n",
+      "        (resnets): ModuleList(\n",
+      "          (0-2): 3 x ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "          )\n",
+      "        )\n",
+      "        (upsamplers): ModuleList(\n",
+      "          (0): Upsample2D(\n",
+      "            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "      (2): UpDecoderBlock2D(\n",
+      "        (resnets): ModuleList(\n",
+      "          (0): ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "            (conv_shortcut): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n",
+      "          )\n",
+      "          (1-2): 2 x ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "          )\n",
+      "        )\n",
+      "        (upsamplers): ModuleList(\n",
+      "          (0): Upsample2D(\n",
+      "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "      (3): UpDecoderBlock2D(\n",
+      "        (resnets): ModuleList(\n",
+      "          (0): ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "            (conv_shortcut): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1))\n",
+      "          )\n",
+      "          (1-2): 2 x ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "          )\n",
+      "        )\n",
+      "        (upsamplers): ModuleList(\n",
+      "          (0): Upsample2D(\n",
+      "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "      (4): UpDecoderBlock2D(\n",
+      "        (resnets): ModuleList(\n",
+      "          (0-2): 3 x ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "    )\n",
+      "    (mid_block): UNetMidBlock2D(\n",
+      "      (attentions): ModuleList(\n",
+      "        (0): Attention(\n",
+      "          (group_norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "          (to_q): Linear(in_features=512, out_features=512, bias=True)\n",
+      "          (to_k): Linear(in_features=512, out_features=512, bias=True)\n",
+      "          (to_v): Linear(in_features=512, out_features=512, bias=True)\n",
+      "          (to_out): ModuleList(\n",
+      "            (0): Linear(in_features=512, out_features=512, bias=True)\n",
+      "            (1): Dropout(p=0.0, inplace=False)\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "      (resnets): ModuleList(\n",
+      "        (0-1): 2 x ResnetBlock2D(\n",
+      "          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "          (dropout): Dropout(p=0.0, inplace=False)\n",
+      "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "          (nonlinearity): SiLU()\n",
+      "        )\n",
+      "      )\n",
+      "    )\n",
+      "    (condition_encoder): MaskConditionEncoder(\n",
+      "      (layers): Sequential(\n",
+      "        (0): Conv2d(3, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "        (1): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "        (2): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
+      "        (3): Conv2d(512, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
+      "        (4): Conv2d(512, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
+      "      )\n",
+      "    )\n",
+      "    (conv_norm_out): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "    (conv_act): SiLU()\n",
+      "    (conv_out): Conv2d(128, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "  )\n",
+      "  (quant_conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n",
+      "  (post_quant_conv): Conv2d(32, 32, kernel_size=(1, 1), stride=(1, 1))\n",
+      ")\n"
+     ]
+    }
+   ],
+   "source": [
+    "from diffusers.models import AsymmetricAutoencoderKL\n",
+    "import torch\n",
+    "\n",
+    "config = {\n",
+    "  \"_class_name\": \"AutoencoderKL\",\n",
+    "  \"_diffusers_version\": \"0.36.0\",\n",
+    "  \"act_fn\": \"silu\",\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",
+    "  \"down_block_out_channels\": [128, 256, 512, 512],\n",
+    "  \"down_block_types\": [\n",
+    "    \"DownEncoderBlock2D\",\n",
+    "    \"DownEncoderBlock2D\",\n",
+    "    \"DownEncoderBlock2D\",\n",
+    "    \"DownEncoderBlock2D\"\n",
+    "  ],\n",
+    "  \"force_upcast\": False,\n",
+    "  \"in_channels\": 3,\n",
+    "  \"latent_channels\": 32,\n",
+    "  \"latents_mean\": [\n",
+    "    -0.03542253375053406,\n",
+    "    0.20086465775966644,\n",
+    "    -0.016413161531090736,\n",
+    "    -0.0956302210688591,\n",
+    "    -0.2672063112258911,\n",
+    "    0.2609933018684387,\n",
+    "    -0.07806991040706635,\n",
+    "    -0.48407721519470215,\n",
+    "    0.21844269335269928,\n",
+    "    -0.1122383326292038,\n",
+    "    0.27197545766830444,\n",
+    "    -0.18958772718906403,\n",
+    "    0.18776826560497284,\n",
+    "    0.0987580344080925,\n",
+    "    0.2837068736553192,\n",
+    "    -0.4486690163612366,\n",
+    "    0.4816776514053345,\n",
+    "    0.02947971224784851,\n",
+    "    -0.1337375044822693,\n",
+    "    -0.39750921726226807,\n",
+    "    -0.08513020724058151,\n",
+    "    -0.054023586213588715,\n",
+    "    -0.3943594992160797,\n",
+    "    0.23918119072914124,\n",
+    "    -0.12466679513454437,\n",
+    "    0.09935147315263748,\n",
+    "    0.31858691573143005,\n",
+    "    0.48585832118988037,\n",
+    "    -0.6416525840759277,\n",
+    "    -0.15164820849895477,\n",
+    "    -0.4693508744239807,\n",
+    "    -0.13071806728839874\n",
+    "  ],\n",
+    "  \"latents_std\": [\n",
+    "    1.5792087316513062,\n",
+    "    1.5769503116607666,\n",
+    "    1.5864241123199463,\n",
+    "    1.6454921960830688,\n",
+    "    1.5336694717407227,\n",
+    "    1.5587652921676636,\n",
+    "    1.5838669538497925,\n",
+    "    1.5659377574920654,\n",
+    "    1.6860467195510864,\n",
+    "    1.5192310810089111,\n",
+    "    1.573639988899231,\n",
+    "    1.5953549146652222,\n",
+    "    1.5271092653274536,\n",
+    "    1.6246271133422852,\n",
+    "    1.7054023742675781,\n",
+    "    1.607722282409668,\n",
+    "    1.558642864227295,\n",
+    "    1.5824549198150635,\n",
+    "    1.6202995777130127,\n",
+    "    1.6206320524215698,\n",
+    "    1.6379750967025757,\n",
+    "    1.6527063846588135,\n",
+    "    1.498811960220337,\n",
+    "    1.5706247091293335,\n",
+    "    1.5854856967926025,\n",
+    "    1.4828169345855713,\n",
+    "    1.5693111419677734,\n",
+    "    1.692481517791748,\n",
+    "    1.6409776210784912,\n",
+    "    1.6216280460357666,\n",
+    "    1.6087706089019775,\n",
+    "    1.5776633024215698\n",
+    "  ],\n",
+    "  \"layers_per_block\": 2,\n",
+    "  \"mid_block_add_attention\": True,\n",
+    "  \"norm_num_groups\": 32,\n",
+    "  \"out_channels\": 3,\n",
+    "  \"sample_size\": 32,\n",
+    "  \"scaling_factor\": 1.0,\n",
+    "  \"shift_factor\": 0.0,\n",
+    "  \"use_post_quant_conv\": True,\n",
+    "  \"use_quant_conv\": True\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",
+    "    up_block_out_channels=config[\"up_block_out_channels\"],\n",
+    "    up_block_types=config[\"up_block_types\"],\n",
+    "    in_channels=config[\"in_channels\"],\n",
+    "    latent_channels=config[\"latent_channels\"],\n",
+    "    norm_num_groups=config[\"norm_num_groups\"],\n",
+    "    out_channels=config[\"out_channels\"],\n",
+    "    sample_size=config[\"sample_size\"],\n",
+    "    scaling_factor=config[\"latents_mean\"],\n",
+    "    #shift_factor=config[\"shift_factor\"],\n",
+    "    layers_per_down_block=config[\"layers_per_block\"],\n",
+    "    layers_per_up_block=config[\"layers_per_block\"],\n",
+    "    #latents_mean = config[\"latents_mean\"],\n",
+    "    #force_upcast=config[\"force_upcast\"],\n",
+    "    #use_post_quant_conv=config[\"use_post_quant_conv\"],\n",
+    "    #use_quant_conv=config[\"use_quant_conv\"]\n",
+    ")\n",
+    "#if hasattr(vae, 'latents_mean'):\n",
+    "#vae.latents_mean = torch.tensor(config[\"latents_mean\"])\n",
+    "#if hasattr(vae, 'latents_std'):\n",
+    "#vae.latents_std = torch.tensor(config[\"latents_std\"])\n",
+    "\n",
+    "# 👇 Добавляем новые поля в config diffusers\n",
+    "vae.register_to_config(\n",
+    "    latents_mean=config[\"latents_mean\"],\n",
+    "    latents_std=config[\"latents_std\"],\n",
+    ")\n",
+    "\n",
+    "# 👇 создаём tensor версии (удобно для инференса)\n",
+    "vae.latents_mean = torch.tensor(vae.config.latents_mean)\n",
+    "vae.latents_std = torch.tensor(vae.config.latents_std)\n",
+    "\n",
+    "vae.save_pretrained(\"asymmetric_vae_empty\")\n",
+    "print(vae)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "a2950158-5203-42b9-8791-e231ddbf1063",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "The config attributes {'block_out_channels': [128, 128, 256, 512, 512], 'force_upcast': False, 'latents_mean': [-0.03542253375053406, 0.20086465775966644, -0.016413161531090736, -0.0956302210688591, -0.2672063112258911, 0.2609933018684387, -0.07806991040706635, -0.48407721519470215, 0.21844269335269928, -0.1122383326292038, 0.27197545766830444, -0.18958772718906403, 0.18776826560497284, 0.0987580344080925, 0.2837068736553192, -0.4486690163612366, 0.4816776514053345, 0.02947971224784851, -0.1337375044822693, -0.39750921726226807, -0.08513020724058151, -0.054023586213588715, -0.3943594992160797, 0.23918119072914124, -0.12466679513454437, 0.09935147315263748, 0.31858691573143005, 0.48585832118988037, -0.6416525840759277, -0.15164820849895477, -0.4693508744239807, -0.13071806728839874], 'latents_std': [1.5792087316513062, 1.5769503116607666, 1.5864241123199463, 1.6454921960830688, 1.5336694717407227, 1.5587652921676636, 1.5838669538497925, 1.5659377574920654, 1.6860467195510864, 1.5192310810089111, 1.573639988899231, 1.5953549146652222, 1.5271092653274536, 1.6246271133422852, 1.7054023742675781, 1.607722282409668, 1.558642864227295, 1.5824549198150635, 1.6202995777130127, 1.6206320524215698, 1.6379750967025757, 1.6527063846588135, 1.498811960220337, 1.5706247091293335, 1.5854856967926025, 1.4828169345855713, 1.5693111419677734, 1.692481517791748, 1.6409776210784912, 1.6216280460357666, 1.6087706089019775, 1.5776633024215698]} were passed to AsymmetricAutoencoderKL, but are not expected and will be ignored. Please verify your config.json configuration file.\n",
+      "Перенос весов: 100%|██████████| 248/248 [00:00<00:00, 30616.26it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Статистика переноса: {'перенесено': 248, 'несовпадение_размеров': 0, 'пропущено': 0}\n",
+      "Неперенесенные ключи в новой модели:\n",
+      "decoder.condition_encoder.layers.0.bias\n",
+      "decoder.condition_encoder.layers.0.weight\n",
+      "decoder.condition_encoder.layers.1.bias\n",
+      "decoder.condition_encoder.layers.1.weight\n",
+      "decoder.condition_encoder.layers.2.bias\n",
+      "decoder.condition_encoder.layers.2.weight\n",
+      "decoder.condition_encoder.layers.3.bias\n",
+      "decoder.condition_encoder.layers.3.weight\n",
+      "decoder.condition_encoder.layers.4.bias\n",
+      "decoder.condition_encoder.layers.4.weight\n",
+      "decoder.up_blocks.3.upsamplers.0.conv.bias\n",
+      "decoder.up_blocks.3.upsamplers.0.conv.weight\n",
+      "decoder.up_blocks.4.resnets.0.conv1.bias\n",
+      "decoder.up_blocks.4.resnets.0.conv1.weight\n",
+      "decoder.up_blocks.4.resnets.0.conv2.bias\n",
+      "decoder.up_blocks.4.resnets.0.conv2.weight\n",
+      "decoder.up_blocks.4.resnets.0.norm1.bias\n",
+      "decoder.up_blocks.4.resnets.0.norm1.weight\n",
+      "decoder.up_blocks.4.resnets.0.norm2.bias\n",
+      "decoder.up_blocks.4.resnets.0.norm2.weight\n",
+      "decoder.up_blocks.4.resnets.1.conv1.bias\n",
+      "decoder.up_blocks.4.resnets.1.conv1.weight\n",
+      "decoder.up_blocks.4.resnets.1.conv2.bias\n",
+      "decoder.up_blocks.4.resnets.1.conv2.weight\n",
+      "decoder.up_blocks.4.resnets.1.norm1.bias\n",
+      "decoder.up_blocks.4.resnets.1.norm1.weight\n",
+      "decoder.up_blocks.4.resnets.1.norm2.bias\n",
+      "decoder.up_blocks.4.resnets.1.norm2.weight\n",
+      "decoder.up_blocks.4.resnets.2.conv1.bias\n",
+      "decoder.up_blocks.4.resnets.2.conv1.weight\n",
+      "decoder.up_blocks.4.resnets.2.conv2.bias\n",
+      "decoder.up_blocks.4.resnets.2.conv2.weight\n",
+      "decoder.up_blocks.4.resnets.2.norm1.bias\n",
+      "decoder.up_blocks.4.resnets.2.norm1.weight\n",
+      "decoder.up_blocks.4.resnets.2.norm2.bias\n",
+      "decoder.up_blocks.4.resnets.2.norm2.weight\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "from diffusers import AsymmetricAutoencoderKL,AutoencoderKL\n",
+    "from tqdm import tqdm\n",
+    "import torch.nn.init as init\n",
+    "\n",
+    "def log(message):\n",
+    "    print(message)\n",
+    "\n",
+    "def main():\n",
+    "    checkpoint_path_old = \"sdxs-1b/vae\"\n",
+    "    checkpoint_path_new = \"asymmetric_vae_empty\"\n",
+    "    device = \"cuda\"\n",
+    "    dtype = torch.float32\n",
+    "\n",
+    "    # Загрузка моделей\n",
+    "    old_unet = AutoencoderKL.from_pretrained(checkpoint_path_old).to(device, dtype=dtype)\n",
+    "    new_unet = AsymmetricAutoencoderKL.from_pretrained(checkpoint_path_new).to(device, dtype=dtype)\n",
+    "\n",
+    "    old_state_dict = old_unet.state_dict()\n",
+    "    new_state_dict = new_unet.state_dict()\n",
+    "\n",
+    "    transferred_state_dict = {}\n",
+    "    transfer_stats = {\n",
+    "        \"перенесено\": 0,\n",
+    "        \"несовпадение_размеров\": 0,\n",
+    "        \"пропущено\": 0\n",
+    "    }\n",
+    "\n",
+    "    transferred_keys = set()\n",
+    "\n",
+    "    # Обрабатываем каждый ключ старой модели\n",
+    "    for old_key in tqdm(old_state_dict.keys(), desc=\"Перенос весов\"):\n",
+    "        new_key = old_key\n",
+    "\n",
+    "        if new_key in new_state_dict:\n",
+    "            if old_state_dict[old_key].shape == new_state_dict[new_key].shape:\n",
+    "                transferred_state_dict[new_key] = old_state_dict[old_key].clone()\n",
+    "                transferred_keys.add(new_key)\n",
+    "                transfer_stats[\"перенесено\"] += 1\n",
+    "            else:\n",
+    "                log(f\"✗ Несовпадение размеров: {old_key} ({old_state_dict[old_key].shape}) -> {new_key} ({new_state_dict[new_key].shape})\")\n",
+    "                transfer_stats[\"несовпадение_размеров\"] += 1\n",
+    "        else:\n",
+    "            log(f\"? Ключ не найден в новой модели: {old_key} -> {old_state_dict[old_key].shape}\")\n",
+    "            transfer_stats[\"пропущено\"] += 1\n",
+    "\n",
+    "    # Обновляем состояние новой модели перенесенными весами\n",
+    "    new_state_dict.update(transferred_state_dict)\n",
+    "    \n",
+    "    # Инициализируем веса для нового mid блока\n",
+    "    #new_state_dict = initialize_mid_block_weights(new_state_dict, device, dtype)\n",
+    "    \n",
+    "    new_unet.load_state_dict(new_state_dict)\n",
+    "    new_unet.save_pretrained(\"asymmetric_vae\")\n",
+    "\n",
+    "    # Получаем список неперенесенных ключей\n",
+    "    non_transferred_keys = sorted(set(new_state_dict.keys()) - transferred_keys)\n",
+    "\n",
+    "    print(\"Статистика переноса:\", transfer_stats)\n",
+    "    print(\"Неперенесенные ключи в новой модели:\")\n",
+    "    for key in non_transferred_keys:\n",
+    "        print(key)\n",
+    "\n",
+    "if __name__ == \"__main__\":\n",
+    "    main()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "b316ee6c-d295-4396-9177-78e39a53055b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "The config attributes {'block_out_channels': [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": [
+      "ok\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "\n",
+    "from torchvision import transforms, utils\n",
+    "\n",
+    "import diffusers\n",
+    "from diffusers import AsymmetricAutoencoderKL\n",
+    "\n",
+    "from diffusers.utils import load_image\n",
+    "\n",
+    "def crop_image_to_nearest_divisible_by_8(img):\n",
+    "    # Check if the image height and width are divisible by 8\n",
+    "    if img.shape[1] % 8 == 0 and img.shape[2] % 8 == 0:\n",
+    "        return img\n",
+    "    else:\n",
+    "        # Calculate the closest lower resolution divisible by 8\n",
+    "        new_height = img.shape[1] - (img.shape[1] % 8)\n",
+    "        new_width = img.shape[2] - (img.shape[2] % 8)\n",
+    "        \n",
+    "        # Use CenterCrop to crop the image\n",
+    "        transform = transforms.CenterCrop((new_height, new_width), interpolation=transforms.InterpolationMode.BILINEAR)\n",
+    "        img = transform(img).to(torch.float32).clamp(-1, 1)\n",
+    "        \n",
+    "        return img\n",
+    "        \n",
+    "to_tensor = transforms.ToTensor()\n",
+    "\n",
+    "device = \"cuda\"\n",
+    "dtype=torch.float16\n",
+    "vae = AsymmetricAutoencoderKL.from_pretrained(\"asymmetric_vae\",torch_dtype=dtype).to(device).eval()\n",
+    "\n",
+    "image = load_image(\"123456789.jpg\")\n",
+    "\n",
+    "image = crop_image_to_nearest_divisible_by_8(to_tensor(image)).unsqueeze(0).to(device,dtype=dtype)\n",
+    "\n",
+    "upscaled_image = vae(image).sample\n",
+    "#vae.config.scaled_factor\n",
+    "# Save the reconstructed image\n",
+    "utils.save_image(upscaled_image, \"test.png\")\n",
+    "print('ok')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5a01b8e9-73c9-4da7-a097-e334019bd8e9",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python3 (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.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

create_asymmetric.ipynb

@@ -0,0 +1,620 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "407171be-ab46-442b-a0bd-83ca75173eba",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "AsymmetricAutoencoderKL(\n",
+      "  (encoder): Encoder(\n",
+      "    (conv_in): Conv2d(3, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "    (down_blocks): ModuleList(\n",
+      "      (0): DownEncoderBlock2D(\n",
+      "        (resnets): ModuleList(\n",
+      "          (0-1): 2 x ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "          )\n",
+      "        )\n",
+      "        (downsamplers): ModuleList(\n",
+      "          (0): Downsample2D(\n",
+      "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2))\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "      (1): DownEncoderBlock2D(\n",
+      "        (resnets): ModuleList(\n",
+      "          (0): ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "            (conv_shortcut): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1))\n",
+      "          )\n",
+      "          (1): ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "          )\n",
+      "        )\n",
+      "        (downsamplers): ModuleList(\n",
+      "          (0): Downsample2D(\n",
+      "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2))\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "      (2): DownEncoderBlock2D(\n",
+      "        (resnets): ModuleList(\n",
+      "          (0): ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "            (conv_shortcut): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1))\n",
+      "          )\n",
+      "          (1): ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "          )\n",
+      "        )\n",
+      "        (downsamplers): ModuleList(\n",
+      "          (0): Downsample2D(\n",
+      "            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2))\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "      (3): DownEncoderBlock2D(\n",
+      "        (resnets): ModuleList(\n",
+      "          (0-1): 2 x ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "    )\n",
+      "    (mid_block): UNetMidBlock2D(\n",
+      "      (attentions): ModuleList(\n",
+      "        (0): Attention(\n",
+      "          (group_norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "          (to_q): Linear(in_features=512, out_features=512, bias=True)\n",
+      "          (to_k): Linear(in_features=512, out_features=512, bias=True)\n",
+      "          (to_v): Linear(in_features=512, out_features=512, bias=True)\n",
+      "          (to_out): ModuleList(\n",
+      "            (0): Linear(in_features=512, out_features=512, bias=True)\n",
+      "            (1): Dropout(p=0.0, inplace=False)\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "      (resnets): ModuleList(\n",
+      "        (0-1): 2 x ResnetBlock2D(\n",
+      "          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "          (dropout): Dropout(p=0.0, inplace=False)\n",
+      "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "          (nonlinearity): SiLU()\n",
+      "        )\n",
+      "      )\n",
+      "    )\n",
+      "    (conv_norm_out): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "    (conv_act): SiLU()\n",
+      "    (conv_out): Conv2d(512, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "  )\n",
+      "  (decoder): MaskConditionDecoder(\n",
+      "    (conv_in): Conv2d(32, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "    (up_blocks): ModuleList(\n",
+      "      (0-1): 2 x UpDecoderBlock2D(\n",
+      "        (resnets): ModuleList(\n",
+      "          (0-2): 3 x ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "          )\n",
+      "        )\n",
+      "        (upsamplers): ModuleList(\n",
+      "          (0): Upsample2D(\n",
+      "            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "      (2): UpDecoderBlock2D(\n",
+      "        (resnets): ModuleList(\n",
+      "          (0): ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "            (conv_shortcut): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n",
+      "          )\n",
+      "          (1-2): 2 x ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "          )\n",
+      "        )\n",
+      "        (upsamplers): ModuleList(\n",
+      "          (0): Upsample2D(\n",
+      "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "      (3): UpDecoderBlock2D(\n",
+      "        (resnets): ModuleList(\n",
+      "          (0): ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "            (conv_shortcut): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1))\n",
+      "          )\n",
+      "          (1-2): 2 x ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "          )\n",
+      "        )\n",
+      "        (upsamplers): ModuleList(\n",
+      "          (0): Upsample2D(\n",
+      "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "      (4): UpDecoderBlock2D(\n",
+      "        (resnets): ModuleList(\n",
+      "          (0-2): 3 x ResnetBlock2D(\n",
+      "            (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "            (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "            (dropout): Dropout(p=0.0, inplace=False)\n",
+      "            (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "            (nonlinearity): SiLU()\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "    )\n",
+      "    (mid_block): UNetMidBlock2D(\n",
+      "      (attentions): ModuleList(\n",
+      "        (0): Attention(\n",
+      "          (group_norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "          (to_q): Linear(in_features=512, out_features=512, bias=True)\n",
+      "          (to_k): Linear(in_features=512, out_features=512, bias=True)\n",
+      "          (to_v): Linear(in_features=512, out_features=512, bias=True)\n",
+      "          (to_out): ModuleList(\n",
+      "            (0): Linear(in_features=512, out_features=512, bias=True)\n",
+      "            (1): Dropout(p=0.0, inplace=False)\n",
+      "          )\n",
+      "        )\n",
+      "      )\n",
+      "      (resnets): ModuleList(\n",
+      "        (0-1): 2 x ResnetBlock2D(\n",
+      "          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
+      "          (dropout): Dropout(p=0.0, inplace=False)\n",
+      "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "          (nonlinearity): SiLU()\n",
+      "        )\n",
+      "      )\n",
+      "    )\n",
+      "    (condition_encoder): MaskConditionEncoder(\n",
+      "      (layers): Sequential(\n",
+      "        (0): Conv2d(3, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "        (1): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "        (2): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
+      "        (3): Conv2d(512, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
+      "        (4): Conv2d(512, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
+      "      )\n",
+      "    )\n",
+      "    (conv_norm_out): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
+      "    (conv_act): SiLU()\n",
+      "    (conv_out): Conv2d(128, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+      "  )\n",
+      "  (quant_conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n",
+      "  (post_quant_conv): Conv2d(32, 32, kernel_size=(1, 1), stride=(1, 1))\n",
+      ")\n"
+     ]
+    }
+   ],
+   "source": [
+    "from diffusers.models import AsymmetricAutoencoderKL\n",
+    "import torch\n",
+    "\n",
+    "config = {\n",
+    "  \"_class_name\": \"AutoencoderKL\",\n",
+    "  \"_diffusers_version\": \"0.36.0\",\n",
+    "  \"act_fn\": \"silu\",\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",
+    "  \"down_block_out_channels\": [128, 256, 512, 512],\n",
+    "  \"down_block_types\": [\n",
+    "    \"DownEncoderBlock2D\",\n",
+    "    \"DownEncoderBlock2D\",\n",
+    "    \"DownEncoderBlock2D\",\n",
+    "    \"DownEncoderBlock2D\"\n",
+    "  ],\n",
+    "  \"force_upcast\": False,\n",
+    "  \"in_channels\": 3,\n",
+    "  \"latent_channels\": 32,\n",
+    "  \"latents_mean\": [\n",
+    "    -0.03542253375053406,\n",
+    "    0.20086465775966644,\n",
+    "    -0.016413161531090736,\n",
+    "    -0.0956302210688591,\n",
+    "    -0.2672063112258911,\n",
+    "    0.2609933018684387,\n",
+    "    -0.07806991040706635,\n",
+    "    -0.48407721519470215,\n",
+    "    0.21844269335269928,\n",
+    "    -0.1122383326292038,\n",
+    "    0.27197545766830444,\n",
+    "    -0.18958772718906403,\n",
+    "    0.18776826560497284,\n",
+    "    0.0987580344080925,\n",
+    "    0.2837068736553192,\n",
+    "    -0.4486690163612366,\n",
+    "    0.4816776514053345,\n",
+    "    0.02947971224784851,\n",
+    "    -0.1337375044822693,\n",
+    "    -0.39750921726226807,\n",
+    "    -0.08513020724058151,\n",
+    "    -0.054023586213588715,\n",
+    "    -0.3943594992160797,\n",
+    "    0.23918119072914124,\n",
+    "    -0.12466679513454437,\n",
+    "    0.09935147315263748,\n",
+    "    0.31858691573143005,\n",
+    "    0.48585832118988037,\n",
+    "    -0.6416525840759277,\n",
+    "    -0.15164820849895477,\n",
+    "    -0.4693508744239807,\n",
+    "    -0.13071806728839874\n",
+    "  ],\n",
+    "  \"latents_std\": [\n",
+    "    1.5792087316513062,\n",
+    "    1.5769503116607666,\n",
+    "    1.5864241123199463,\n",
+    "    1.6454921960830688,\n",
+    "    1.5336694717407227,\n",
+    "    1.5587652921676636,\n",
+    "    1.5838669538497925,\n",
+    "    1.5659377574920654,\n",
+    "    1.6860467195510864,\n",
+    "    1.5192310810089111,\n",
+    "    1.573639988899231,\n",
+    "    1.5953549146652222,\n",
+    "    1.5271092653274536,\n",
+    "    1.6246271133422852,\n",
+    "    1.7054023742675781,\n",
+    "    1.607722282409668,\n",
+    "    1.558642864227295,\n",
+    "    1.5824549198150635,\n",
+    "    1.6202995777130127,\n",
+    "    1.6206320524215698,\n",
+    "    1.6379750967025757,\n",
+    "    1.6527063846588135,\n",
+    "    1.498811960220337,\n",
+    "    1.5706247091293335,\n",
+    "    1.5854856967926025,\n",
+    "    1.4828169345855713,\n",
+    "    1.5693111419677734,\n",
+    "    1.692481517791748,\n",
+    "    1.6409776210784912,\n",
+    "    1.6216280460357666,\n",
+    "    1.6087706089019775,\n",
+    "    1.5776633024215698\n",
+    "  ],\n",
+    "  \"layers_per_block\": 2,\n",
+    "  \"mid_block_add_attention\": True,\n",
+    "  \"norm_num_groups\": 32,\n",
+    "  \"out_channels\": 3,\n",
+    "  \"sample_size\": 32,\n",
+    "  \"scaling_factor\": 1.0,\n",
+    "  \"shift_factor\": 0.0,\n",
+    "  \"use_post_quant_conv\": True,\n",
+    "  \"use_quant_conv\": True\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",
+    "    up_block_out_channels=config[\"up_block_out_channels\"],\n",
+    "    up_block_types=config[\"up_block_types\"],\n",
+    "    in_channels=config[\"in_channels\"],\n",
+    "    latent_channels=config[\"latent_channels\"],\n",
+    "    norm_num_groups=config[\"norm_num_groups\"],\n",
+    "    out_channels=config[\"out_channels\"],\n",
+    "    sample_size=config[\"sample_size\"],\n",
+    "    scaling_factor=config[\"latents_mean\"],\n",
+    "    #shift_factor=config[\"shift_factor\"],\n",
+    "    layers_per_down_block=config[\"layers_per_block\"],\n",
+    "    layers_per_up_block=config[\"layers_per_block\"],\n",
+    "    #latents_mean = config[\"latents_mean\"],\n",
+    "    #force_upcast=config[\"force_upcast\"],\n",
+    "    #use_post_quant_conv=config[\"use_post_quant_conv\"],\n",
+    "    #use_quant_conv=config[\"use_quant_conv\"]\n",
+    ")\n",
+    "#if hasattr(vae, 'latents_mean'):\n",
+    "#vae.latents_mean = torch.tensor(config[\"latents_mean\"])\n",
+    "#if hasattr(vae, 'latents_std'):\n",
+    "#vae.latents_std = torch.tensor(config[\"latents_std\"])\n",
+    "\n",
+    "# 👇 Добавляем новые поля в config diffusers\n",
+    "vae.register_to_config(\n",
+    "    latents_mean=config[\"latents_mean\"],\n",
+    "    latents_std=config[\"latents_std\"],\n",
+    ")\n",
+    "\n",
+    "# 👇 создаём tensor версии (удобно для инференса)\n",
+    "vae.latents_mean = torch.tensor(vae.config.latents_mean)\n",
+    "vae.latents_std = torch.tensor(vae.config.latents_std)\n",
+    "\n",
+    "vae.save_pretrained(\"asymmetric_vae_empty\")\n",
+    "print(vae)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "a2950158-5203-42b9-8791-e231ddbf1063",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "The config attributes {'block_out_channels': [128, 128, 256, 512, 512], 'force_upcast': False, 'latents_mean': [-0.03542253375053406, 0.20086465775966644, -0.016413161531090736, -0.0956302210688591, -0.2672063112258911, 0.2609933018684387, -0.07806991040706635, -0.48407721519470215, 0.21844269335269928, -0.1122383326292038, 0.27197545766830444, -0.18958772718906403, 0.18776826560497284, 0.0987580344080925, 0.2837068736553192, -0.4486690163612366, 0.4816776514053345, 0.02947971224784851, -0.1337375044822693, -0.39750921726226807, -0.08513020724058151, -0.054023586213588715, -0.3943594992160797, 0.23918119072914124, -0.12466679513454437, 0.09935147315263748, 0.31858691573143005, 0.48585832118988037, -0.6416525840759277, -0.15164820849895477, -0.4693508744239807, -0.13071806728839874], 'latents_std': [1.5792087316513062, 1.5769503116607666, 1.5864241123199463, 1.6454921960830688, 1.5336694717407227, 1.5587652921676636, 1.5838669538497925, 1.5659377574920654, 1.6860467195510864, 1.5192310810089111, 1.573639988899231, 1.5953549146652222, 1.5271092653274536, 1.6246271133422852, 1.7054023742675781, 1.607722282409668, 1.558642864227295, 1.5824549198150635, 1.6202995777130127, 1.6206320524215698, 1.6379750967025757, 1.6527063846588135, 1.498811960220337, 1.5706247091293335, 1.5854856967926025, 1.4828169345855713, 1.5693111419677734, 1.692481517791748, 1.6409776210784912, 1.6216280460357666, 1.6087706089019775, 1.5776633024215698]} were passed to AsymmetricAutoencoderKL, but are not expected and will be ignored. Please verify your config.json configuration file.\n",
+      "Перенос весов: 100%|██████████| 248/248 [00:00<00:00, 30616.26it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Статистика переноса: {'перенесено': 248, 'несовпадение_размеров': 0, 'пропущено': 0}\n",
+      "Неперенесенные ключи в новой модели:\n",
+      "decoder.condition_encoder.layers.0.bias\n",
+      "decoder.condition_encoder.layers.0.weight\n",
+      "decoder.condition_encoder.layers.1.bias\n",
+      "decoder.condition_encoder.layers.1.weight\n",
+      "decoder.condition_encoder.layers.2.bias\n",
+      "decoder.condition_encoder.layers.2.weight\n",
+      "decoder.condition_encoder.layers.3.bias\n",
+      "decoder.condition_encoder.layers.3.weight\n",
+      "decoder.condition_encoder.layers.4.bias\n",
+      "decoder.condition_encoder.layers.4.weight\n",
+      "decoder.up_blocks.3.upsamplers.0.conv.bias\n",
+      "decoder.up_blocks.3.upsamplers.0.conv.weight\n",
+      "decoder.up_blocks.4.resnets.0.conv1.bias\n",
+      "decoder.up_blocks.4.resnets.0.conv1.weight\n",
+      "decoder.up_blocks.4.resnets.0.conv2.bias\n",
+      "decoder.up_blocks.4.resnets.0.conv2.weight\n",
+      "decoder.up_blocks.4.resnets.0.norm1.bias\n",
+      "decoder.up_blocks.4.resnets.0.norm1.weight\n",
+      "decoder.up_blocks.4.resnets.0.norm2.bias\n",
+      "decoder.up_blocks.4.resnets.0.norm2.weight\n",
+      "decoder.up_blocks.4.resnets.1.conv1.bias\n",
+      "decoder.up_blocks.4.resnets.1.conv1.weight\n",
+      "decoder.up_blocks.4.resnets.1.conv2.bias\n",
+      "decoder.up_blocks.4.resnets.1.conv2.weight\n",
+      "decoder.up_blocks.4.resnets.1.norm1.bias\n",
+      "decoder.up_blocks.4.resnets.1.norm1.weight\n",
+      "decoder.up_blocks.4.resnets.1.norm2.bias\n",
+      "decoder.up_blocks.4.resnets.1.norm2.weight\n",
+      "decoder.up_blocks.4.resnets.2.conv1.bias\n",
+      "decoder.up_blocks.4.resnets.2.conv1.weight\n",
+      "decoder.up_blocks.4.resnets.2.conv2.bias\n",
+      "decoder.up_blocks.4.resnets.2.conv2.weight\n",
+      "decoder.up_blocks.4.resnets.2.norm1.bias\n",
+      "decoder.up_blocks.4.resnets.2.norm1.weight\n",
+      "decoder.up_blocks.4.resnets.2.norm2.bias\n",
+      "decoder.up_blocks.4.resnets.2.norm2.weight\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "from diffusers import AsymmetricAutoencoderKL,AutoencoderKL\n",
+    "from tqdm import tqdm\n",
+    "import torch.nn.init as init\n",
+    "\n",
+    "def log(message):\n",
+    "    print(message)\n",
+    "\n",
+    "def main():\n",
+    "    checkpoint_path_old = \"sdxs-1b/vae\"\n",
+    "    checkpoint_path_new = \"asymmetric_vae_empty\"\n",
+    "    device = \"cuda\"\n",
+    "    dtype = torch.float32\n",
+    "\n",
+    "    # Загрузка моделей\n",
+    "    old_unet = AutoencoderKL.from_pretrained(checkpoint_path_old).to(device, dtype=dtype)\n",
+    "    new_unet = AsymmetricAutoencoderKL.from_pretrained(checkpoint_path_new).to(device, dtype=dtype)\n",
+    "\n",
+    "    old_state_dict = old_unet.state_dict()\n",
+    "    new_state_dict = new_unet.state_dict()\n",
+    "\n",
+    "    transferred_state_dict = {}\n",
+    "    transfer_stats = {\n",
+    "        \"перенесено\": 0,\n",
+    "        \"несовпадение_размеров\": 0,\n",
+    "        \"пропущено\": 0\n",
+    "    }\n",
+    "\n",
+    "    transferred_keys = set()\n",
+    "\n",
+    "    # Обрабатываем каждый ключ старой модели\n",
+    "    for old_key in tqdm(old_state_dict.keys(), desc=\"Перенос весов\"):\n",
+    "        new_key = old_key\n",
+    "\n",
+    "        if new_key in new_state_dict:\n",
+    "            if old_state_dict[old_key].shape == new_state_dict[new_key].shape:\n",
+    "                transferred_state_dict[new_key] = old_state_dict[old_key].clone()\n",
+    "                transferred_keys.add(new_key)\n",
+    "                transfer_stats[\"перенесено\"] += 1\n",
+    "            else:\n",
+    "                log(f\"✗ Несовпадение размеров: {old_key} ({old_state_dict[old_key].shape}) -> {new_key} ({new_state_dict[new_key].shape})\")\n",
+    "                transfer_stats[\"несовпадение_размеров\"] += 1\n",
+    "        else:\n",
+    "            log(f\"? Ключ не найден в новой модели: {old_key} -> {old_state_dict[old_key].shape}\")\n",
+    "            transfer_stats[\"пропущено\"] += 1\n",
+    "\n",
+    "    # Обновляем состояние новой модели перенесенными весами\n",
+    "    new_state_dict.update(transferred_state_dict)\n",
+    "    \n",
+    "    # Инициализируем веса для нового mid блока\n",
+    "    #new_state_dict = initialize_mid_block_weights(new_state_dict, device, dtype)\n",
+    "    \n",
+    "    new_unet.load_state_dict(new_state_dict)\n",
+    "    new_unet.save_pretrained(\"asymmetric_vae\")\n",
+    "\n",
+    "    # Получаем список неперенесенных ключей\n",
+    "    non_transferred_keys = sorted(set(new_state_dict.keys()) - transferred_keys)\n",
+    "\n",
+    "    print(\"Статистика переноса:\", transfer_stats)\n",
+    "    print(\"Неперенесенные ключи в новой модели:\")\n",
+    "    for key in non_transferred_keys:\n",
+    "        print(key)\n",
+    "\n",
+    "if __name__ == \"__main__\":\n",
+    "    main()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "b316ee6c-d295-4396-9177-78e39a53055b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "The config attributes {'block_out_channels': [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": [
+      "ok\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "\n",
+    "from torchvision import transforms, utils\n",
+    "\n",
+    "import diffusers\n",
+    "from diffusers import AsymmetricAutoencoderKL\n",
+    "\n",
+    "from diffusers.utils import load_image\n",
+    "\n",
+    "def crop_image_to_nearest_divisible_by_8(img):\n",
+    "    # Check if the image height and width are divisible by 8\n",
+    "    if img.shape[1] % 8 == 0 and img.shape[2] % 8 == 0:\n",
+    "        return img\n",
+    "    else:\n",
+    "        # Calculate the closest lower resolution divisible by 8\n",
+    "        new_height = img.shape[1] - (img.shape[1] % 8)\n",
+    "        new_width = img.shape[2] - (img.shape[2] % 8)\n",
+    "        \n",
+    "        # Use CenterCrop to crop the image\n",
+    "        transform = transforms.CenterCrop((new_height, new_width), interpolation=transforms.InterpolationMode.BILINEAR)\n",
+    "        img = transform(img).to(torch.float32).clamp(-1, 1)\n",
+    "        \n",
+    "        return img\n",
+    "        \n",
+    "to_tensor = transforms.ToTensor()\n",
+    "\n",
+    "device = \"cuda\"\n",
+    "dtype=torch.float16\n",
+    "vae = AsymmetricAutoencoderKL.from_pretrained(\"asymmetric_vae\",torch_dtype=dtype).to(device).eval()\n",
+    "\n",
+    "image = load_image(\"123456789.jpg\")\n",
+    "\n",
+    "image = crop_image_to_nearest_divisible_by_8(to_tensor(image)).unsqueeze(0).to(device,dtype=dtype)\n",
+    "\n",
+    "upscaled_image = vae(image).sample\n",
+    "#vae.config.scaled_factor\n",
+    "# Save the reconstructed image\n",
+    "utils.save_image(upscaled_image, \"test.png\")\n",
+    "print('ok')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5a01b8e9-73c9-4da7-a097-e334019bd8e9",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python3 (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.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5af0ea59571d8825d073fc31cac736a05e0a534c6ba7bcdd7e099f8d8119fac0
+oid sha256:d349936a1faab1555bf81e68ba1e6fd2b84f6a6a46ffd11470079581d48bdfea
 size 343311604

down.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+
+TARGET_DIR="/workspace/d23"
+mkdir -p "$TARGET_DIR"
+
+BASE_URL="https://huggingface.co/datasets/AI-Art-Collab/dtasettar23/resolve/main/d23.tar."
+
+(
+  # Устанавливаем `set -e` внутри subshell, чтобы он завершился при первой ошибке curl
+  set -e
+  # Попробуем от 'a' до 'z' для первого символа суффикса
+  for c1 in {a..z}; do
+    # Попробуем от 'a' до 'z' для второго символа суффикса
+    for c2 in {a..z}; do
+      suffix="${c1}${c2}"
+      url="${BASE_URL}${suffix}"
+      echo "Fetching: $url" >&2
+      # Качаем часть архива. --fail заставит curl завершиться с ошибкой, если файла нет.
+      curl -LsS --fail "$url"
+    done
+  done
+) 2>/dev/null | tar -xv -C "$TARGET_DIR" --wildcards '*.png'
+#    └─ 1 ─┘   └────────── 2 ──────────┘ └─────────── 3 ───────────┘
+
+echo "Extraction of PNG files finished. Check $TARGET_DIR"

train_vae_16x.py

@@ -29,7 +29,7 @@ from collections import deque
 
 # --------------------------- Параметры ---------------------------
 ds_path            = "/workspace/d23"
-project            = "vae16x32ch"
+project            = "vae16x32ch_new"
 batch_size         = 1
 base_learning_rate = 6e-6
 min_learning_rate  = 7e-7
@@ -68,7 +68,7 @@ torch.backends.cuda.enable_math_sdp(False)
 train_decoder_only = False
 train_up_only      = False
 full_training      = True  # если True — учим весь VAE и добавляем KL (ниже)
-kl_ratio           = 0.00
+kl_ratio           = 0.001
 
 # Доли лоссов
 loss_ratios = {
@@ -76,8 +76,8 @@ loss_ratios = {
     "fdl" :  0.10,#0.25,   
     "edge":  0.05,
     "mse":   0.10,
-    "mae":   0.05,
-    "kl":    0.00,  # активируем при full_training=True
+    "mae":   0.049,
+    "kl":    0.001,
 }
 median_coeff_steps = 250
 

train_vae_fdl.py

@@ -0,0 +1,624 @@
+# -*- coding: utf-8 -*-
+import os
+import math
+import re
+import torch
+import numpy as np
+import random
+import gc
+from datetime import datetime
+from pathlib import Path
+
+import torchvision.transforms as transforms
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+from torch.optim.lr_scheduler import LambdaLR
+from diffusers import AutoencoderKL, AsymmetricAutoencoderKL
+# QWEN: импорт класса
+from diffusers import AutoencoderKLQwenImage
+from diffusers import AutoencoderKLWan
+
+from accelerate import Accelerator
+from PIL import Image, UnidentifiedImageError
+from tqdm import tqdm
+import bitsandbytes as bnb
+import wandb
+import lpips   # pip install lpips
+from FDL_pytorch import FDL_loss # pip install fdl-pytorch
+from collections import deque
+
+# --------------------------- Параметры ---------------------------
+ds_path            = "/workspace/d23"
+project            = "asymmetric_vae"
+batch_size         = 1
+base_learning_rate = 6e-6
+min_learning_rate  = 7e-7
+num_epochs         = 2
+sample_interval_share = 30
+use_wandb          = True
+save_model         = True
+use_decay          = True
+optimizer_type     = "adam8bit"
+dtype              = torch.float32
+
+model_resolution   = 384 #448 #288
+high_resolution    = 768 #896 #576 
+limit              = 0
+save_barrier       = 1.3
+warmup_percent     = 0.005
+percentile_clipping = 99
+beta2              = 0.997
+eps                = 1e-8
+clip_grad_norm     = 1.0
+mixed_precision    = "no"
+gradient_accumulation_steps = 1
+generated_folder   = "samples"
+save_as            = "asymmetric_vae_new"
+num_workers        = 0
+device = None
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+# Включение Flash Attention 2/SDPA #MAX_JOBS=4 pip install flash-attn --no-build-isolation
+torch.backends.cuda.enable_flash_sdp(True)
+torch.backends.cuda.enable_mem_efficient_sdp(True)
+torch.backends.cuda.enable_math_sdp(False) 
+
+# --- Режимы обучения ---
+# QWEN: учим только декодер
+train_decoder_only = True
+train_up_only      = False
+full_training      = False  # если True — учим весь VAE и добавляем KL (ниже)
+kl_ratio           = 0.00
+
+# Доли лоссов
+loss_ratios = {
+    "lpips": 0.70,#0.50,
+    "fdl" :  0.10,#0.25,   
+    "edge":  0.05,
+    "mse":   0.10,
+    "mae":   0.05,
+    "kl":    0.00,  # активируем при full_training=True
+}
+median_coeff_steps = 250
+
+resize_long_side = 1280  # ресайз длинной стороны исходных картинок
+
+# QWEN: конфиг загрузки модели
+vae_kind      = "kl"  # "qwen" или "kl" (обычный)
+
+Path(generated_folder).mkdir(parents=True, exist_ok=True)
+
+accelerator = Accelerator(
+    mixed_precision=mixed_precision,
+    gradient_accumulation_steps=gradient_accumulation_steps
+)
+device = accelerator.device
+
+# reproducibility
+seed = int(datetime.now().strftime("%Y%m%d")) + 13
+torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
+torch.backends.cudnn.benchmark = False
+
+# --------------------------- WandB ---------------------------
+if use_wandb and accelerator.is_main_process:
+    wandb.init(project=project, config={
+        "batch_size": batch_size,
+        "base_learning_rate": base_learning_rate,
+        "num_epochs": num_epochs,
+        "optimizer_type": optimizer_type,
+        "model_resolution": model_resolution,
+        "high_resolution": high_resolution,
+        "gradient_accumulation_steps": gradient_accumulation_steps,
+        "train_decoder_only": train_decoder_only,
+        "full_training": full_training,
+        "kl_ratio": kl_ratio,
+        "vae_kind": vae_kind,
+    })
+
+# --------------------------- VAE ---------------------------
+def get_core_model(model):
+    m = model
+    # если модель уже обёрнута torch.compile
+    if hasattr(m, "_orig_mod"):
+        m = m._orig_mod
+    return m
+
+def is_video_vae(model) -> bool:
+    # WAN/Qwen — это видео-VAEs
+    if vae_kind in ("wan", "qwen"):
+        return True
+    # fallback по структуре (если понадобится)
+    try:
+        core = get_core_model(model)
+        enc = getattr(core, "encoder", None)
+        conv_in = getattr(enc, "conv_in", None)
+        w = getattr(conv_in, "weight", None)
+        if isinstance(w, torch.nn.Parameter):
+            return w.ndim == 5
+    except Exception:
+        pass
+    return False
+
+# загрузка
+if vae_kind == "qwen":
+    vae = AutoencoderKLQwenImage.from_pretrained("Qwen/Qwen-Image", subfolder="vae")
+else:
+    if vae_kind == "wan":
+        vae = AutoencoderKLWan.from_pretrained(project)
+    else:
+        # старое поведение (пример)
+        if model_resolution==high_resolution:
+            vae = AutoencoderKL.from_pretrained(project)
+        else:
+            vae = AsymmetricAutoencoderKL.from_pretrained(project)
+
+vae = vae.to(dtype)
+
+# torch.compile (опционально)
+if hasattr(torch, "compile"):
+    try:
+        vae = torch.compile(vae)
+    except Exception as e:
+        print(f"[WARN] torch.compile failed: {e}")
+
+# --------------------------- Freeze/Unfreeze ---------------------------
+core = get_core_model(vae)
+
+for p in core.parameters():
+    p.requires_grad = False
+
+unfrozen_param_names = []
+
+if full_training and not train_decoder_only:
+    for name, p in core.named_parameters():
+        p.requires_grad = True
+        unfrozen_param_names.append(name)
+        loss_ratios["kl"] = float(kl_ratio)
+        trainable_module = core
+else:
+    # учим только 0-й блок декодера + post_quant_conv
+    if hasattr(core, "decoder"):
+        if train_up_only:#hasattr(core.decoder, "up_blocks") and len(core.decoder.up_blocks) > 0:
+            # --- только 0-й up_block ---
+            for name, p in core.decoder.up_blocks[0].named_parameters():
+                p.requires_grad = True
+                unfrozen_param_names.append(f"{name}")
+        else:
+            print("Decoder — fallback to full decoder")
+            for name, p in core.decoder.named_parameters():
+                p.requires_grad = True
+                unfrozen_param_names.append(f"decoder.{name}")
+    if hasattr(core, "post_quant_conv"):
+        for name, p in core.post_quant_conv.named_parameters():
+            p.requires_grad = True
+            unfrozen_param_names.append(f"post_quant_conv.{name}")
+    trainable_module = core.decoder if hasattr(core, "decoder") else core
+
+
+print(f"[INFO] Разморожено параметров: {len(unfrozen_param_names)}. Первые 200 имён:")
+for nm in unfrozen_param_names[:200]:
+    print(" ", nm)
+
+# --------------------------- Датасет ---------------------------
+class PngFolderDataset(Dataset):
+    def __init__(self, root_dir, min_exts=('.png',), resolution=1024, limit=0):
+        self.root_dir = root_dir
+        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))
+        if limit:
+            self.paths = self.paths[:limit]
+        valid = []
+        for p in self.paths:
+            try:
+                with Image.open(p) as im:
+                    im.verify()
+                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}")
+        random.shuffle(self.paths)
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, idx):
+        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)
+
+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))
+
+tfm = transforms.Compose([
+    transforms.ToTensor(),
+    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:
+        img = random_crop(img, high_resolution)
+        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
+)
+
+# --------------------------- Оптимизатор ---------------------------
+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 = [], []
+    for n, p in module.named_parameters():
+        if not p.requires_grad:
+            continue
+        n_l = n.lower()
+        if any(t in n_l for t in no_decay_tokens):
+            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 create_optimizer(name, param_groups):
+    if name == "adam8bit":
+        return bnb.optim.AdamW8bit(param_groups, lr=base_learning_rate, betas=(0.9, beta2), eps=eps)
+    raise ValueError(name)
+
+param_groups = get_param_groups(get_core_model(vae), weight_decay=0.001)
+optimizer = create_optimizer(optimizer_type, param_groups)
+
+# --------------------------- LR schedule ---------------------------
+batches_per_epoch = len(dataloader)
+steps_per_epoch = int(math.ceil(batches_per_epoch / float(gradient_accumulation_steps)))
+total_steps = steps_per_epoch * num_epochs
+
+def lr_lambda(step):
+    if not use_decay:
+        return 1.0
+    x = float(step) / float(max(1, total_steps))
+    warmup = float(warmup_percent)
+    min_ratio = float(min_learning_rate) / float(base_learning_rate)
+    if x < warmup:
+        return min_ratio + (1.0 - min_ratio) * (x / warmup)
+    decay_ratio = (x - warmup) / (1.0 - warmup)
+    return min_ratio + 0.5 * (1.0 - min_ratio) * (1.0 + math.cos(math.pi * decay_ratio))
+
+scheduler = LambdaLR(optimizer, lr_lambda)
+
+# Подготовка
+dataloader, vae, optimizer, scheduler = accelerator.prepare(dataloader, vae, optimizer, scheduler)
+trainable_params = [p for p in vae.parameters() if p.requires_grad]
+
+# fdl
+fdl_loss = FDL_loss()
+fdl_loss = fdl_loss.to(accelerator.device)
+
+# --------------------------- LPIPS и вспомогательные ---------------------------
+_lpips_net = None
+def _get_lpips():
+    global _lpips_net
+    if _lpips_net is None:
+        _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):
+        s = sum(desired_ratios.values())
+        self.ratios = {k: (v / s) if s > 0 else 0.0 for k, v in desired_ratios.items()}
+        self.buffers = {k: deque(maxlen=window_steps) for k in self.ratios.keys()}
+        self.window = window_steps
+
+    def update_and_total(self, abs_losses: dict):
+        for k, v in abs_losses.items():
+            if k in self.buffers:
+                self.buffers[k].append(float(v.detach().abs().cpu()))
+        meds = {k: (np.median(self.buffers[k]) if len(self.buffers[k]) > 0 else 1.0) for k in self.buffers}
+        coeffs = {k: (self.ratios[k] / max(meds[k], 1e-12)) for k in self.ratios}
+        total = sum(coeffs[k] * abs_losses[k] for k in abs_losses if k in coeffs)
+        return total, coeffs, meds
+
+if full_training and not train_decoder_only:
+    loss_ratios["kl"] = float(kl_ratio)
+normalizer = MedianLossNormalizer(loss_ratios, median_coeff_steps)
+
+# --------------------------- Сэмплы ---------------------------
+@torch.no_grad()
+def get_fixed_samples(n=3):
+    idx = random.sample(range(len(dataset)), min(n, len(dataset)))
+    pil_imgs = [dataset[i] for i in idx]
+    tensors = []
+    for img in pil_imgs:
+        img = random_crop(img, high_resolution)
+        tensors.append(tfm(img))
+    return torch.stack(tensors).to(accelerator.device, dtype)
+
+fixed_samples = get_fixed_samples()
+
+@torch.no_grad()
+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()
+        if hasattr(vae, "module"):
+            # Если это DDP или DistributedDataParallel
+            unwrapped_vae = vae.module
+        else:
+            unwrapped_vae = vae
+        
+        # Если использовался torch.compile, достаем оригинал
+        if hasattr(unwrapped_vae, "_orig_mod"):
+            temp_vae = unwrapped_vae._orig_mod
+        else:
+            temp_vae = unwrapped_vae
+        
+        temp_vae = temp_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}.png")
+            dec_img.save(f"{generated_folder}/sample_decoded_{i}.png")
+
+        # 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}.png", caption=f"real_{i}")
+                log_data[f"sample/decoded_{i}"] = wandb.Image(f"{generated_folder}/sample_decoded_{i}.png", caption=f"decoded_{i}")
+            wandb.log(log_data, step=step)
+
+    finally:
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+if accelerator.is_main_process and save_model:
+    print("Генерация сэмплов до старта обучения...")
+    generate_and_save_samples(0)
+
+accelerator.wait_for_everyone()
+
+# --------------------------- Тренировка ---------------------------
+progress = tqdm(total=total_steps, disable=not accelerator.is_local_main_process)
+global_step = 0
+min_loss = float("inf")
+sample_interval = max(1, total_steps // max(1, sample_interval_share * num_epochs))
+
+for epoch in range(num_epochs):
+    vae.train()
+    batch_losses, batch_grads = [], []
+    track_losses = {k: [] for k in loss_ratios.keys()}
+
+    for imgs in dataloader:
+        with accelerator.accumulate(vae):
+            imgs = imgs.to(accelerator.device)
+
+            if high_resolution != model_resolution:
+                imgs_low = F.interpolate(imgs, size=(model_resolution, model_resolution),mode="area") # mode="bilinear", align_corners=False)
+            else:
+                imgs_low = imgs
+
+            model_dtype = next(vae.parameters()).dtype
+            imgs_low_model = imgs_low.to(dtype=model_dtype) if imgs_low.dtype != model_dtype else imgs_low
+
+            # Вместо: current_vae = accelerator.unwrap_model(vae)
+            unwrapped = vae.module if hasattr(vae, "module") else vae
+            current_vae = getattr(unwrapped, "_orig_mod", unwrapped)
+
+
+            # QWEN: encode/decode с T=1
+            if is_video_vae(current_vae):
+                x_in = imgs_low_model.unsqueeze(2)             # [B,3,1,H,W]
+                enc = current_vae.encode(x_in)
+                latents = enc.latent_dist.mean if train_decoder_only else enc.latent_dist.sample()
+                dec = current_vae.decode(latents).sample               # [B,3,1,H,W]
+                rec = dec.squeeze(2)                           # [B,3,H,W]
+            else:
+                enc = current_vae.encode(imgs_low_model)
+                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),
+                "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:
+                mean   = enc.latent_dist.mean
+                logvar = enc.latent_dist.logvar
+                kl = -0.5 * torch.mean(1 + logvar - mean.pow(2) - logvar.exp())
+                abs_losses["kl"] = kl
+            else:
+                abs_losses["kl"] = torch.tensor(0.0, device=accelerator.device, dtype=torch.float32)
+
+            total_loss, coeffs, meds = normalizer.update_and_total(abs_losses)
+
+            if torch.isnan(total_loss) or torch.isinf(total_loss):
+                raise RuntimeError("NaN/Inf loss")
+
+            accelerator.backward(total_loss)
+
+            grad_norm = torch.tensor(0.0, device=accelerator.device)
+            if accelerator.sync_gradients:
+                grad_norm = accelerator.clip_grad_norm_(trainable_params, clip_grad_norm)
+                optimizer.step()
+                scheduler.step()
+                optimizer.zero_grad(set_to_none=True)
+                global_step += 1
+                progress.update(1)
+
+            if accelerator.is_main_process:
+                try:
+                    current_lr = optimizer.param_groups[0]["lr"]
+                except Exception:
+                    current_lr = scheduler.get_last_lr()[0]
+
+                batch_losses.append(total_loss.detach().item())
+                batch_grads.append(float(grad_norm.detach().cpu().item()) if isinstance(grad_norm, torch.Tensor) else float(grad_norm))
+                for k, v in abs_losses.items():
+                    track_losses[k].append(float(v.detach().item()))
+
+                if use_wandb and accelerator.sync_gradients:
+                    log_dict = {
+                        "total_loss": float(total_loss.detach().item()),
+                        "learning_rate": current_lr,
+                        "epoch": epoch,
+                        "grad_norm": batch_grads[-1],
+                    }
+                    for k, v in abs_losses.items():
+                        log_dict[f"loss_{k}"] = float(v.detach().item())
+                    for k in coeffs:
+                        log_dict[f"coeff_{k}"] = float(coeffs[k])
+                        log_dict[f"median_{k}"] = float(meds[k])
+                    wandb.log(log_dict, step=global_step)
+
+            if global_step > 0 and global_step % sample_interval == 0:
+                if accelerator.is_main_process:
+                    generate_and_save_samples(global_step)
+                accelerator.wait_for_everyone()
+
+                n_micro = sample_interval * gradient_accumulation_steps
+                avg_loss = float(np.mean(batch_losses[-n_micro:])) if len(batch_losses) >= n_micro else float(np.mean(batch_losses)) if batch_losses else float("nan")
+                avg_grad = float(np.mean(batch_grads[-n_micro:])) if len(batch_grads) >= 1 else float(np.mean(batch_grads)) if batch_grads else 0.0
+
+                if accelerator.is_main_process:
+                    print(f"Epoch {epoch} step {global_step} loss: {avg_loss:.6f}, grad_norm: {avg_grad:.6f}, lr: {current_lr:.9f}")
+                    if save_model and avg_loss < min_loss * save_barrier:
+                        min_loss = avg_loss
+                        unwrapped = vae.module if hasattr(vae, "module") else vae
+                        current_vae = getattr(unwrapped, "_orig_mod", unwrapped)
+                        current_vae.save_pretrained(save_as)
+                    if use_wandb:
+                        wandb.log({"interm_loss": avg_loss, "interm_grad": avg_grad}, step=global_step)
+
+    if accelerator.is_main_process:
+        epoch_avg = float(np.mean(batch_losses)) if batch_losses else float("nan")
+        print(f"Epoch {epoch} done, avg loss {epoch_avg:.6f}")
+        if use_wandb:
+            wandb.log({"epoch_loss": epoch_avg, "epoch": epoch + 1}, step=global_step)
+
+# --------------------------- Финальное сохранение ---------------------------
+if accelerator.is_main_process:
+    print("Training finished – saving final model")
+    if save_model:
+        unwrapped = vae.module if hasattr(vae, "module") else vae
+        current_vae = getattr(unwrapped, "_orig_mod", unwrapped)
+        current_vae.save_pretrained(save_as)
+
+accelerator.free_memory()
+if torch.distributed.is_initialized():
+    torch.distributed.destroy_process_group()
+print("Готово!")