{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "68d83bd8-f0ae-4118-8005-ada7d8b0b3cf",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-19T08:59:29.479790Z",
     "iopub.status.busy": "2024-02-19T08:59:29.479500Z",
     "iopub.status.idle": "2024-02-19T08:59:38.923903Z",
     "shell.execute_reply": "2024-02-19T08:59:38.923356Z",
     "shell.execute_reply.started": "2024-02-19T08:59:29.479771Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "正克隆到 'U-ViT'...\n",
      "remote: Enumerating objects: 135, done.\u001b[K\n",
      "remote: Counting objects: 100% (79/79), done.\u001b[K\n",
      "remote: Compressing objects: 100% (26/26), done.\u001b[K\n",
      "remote: Total 135 (delta 68), reused 53 (delta 53), pack-reused 56\u001b[K\n",
      "接收对象中: 100% (135/135), 7.82 MiB | 2.75 MiB/s, 完成.\n",
      "处理 delta 中: 100% (82/82), 完成.\n",
      "Looking in indexes: https://mirrors.aliyun.com/pypi/simple\n",
      "Requirement already satisfied: einops in /opt/conda/lib/python3.10/site-packages (0.7.0)\n",
      "\u001b[33mDEPRECATION: pytorch-lightning 1.7.7 has a non-standard dependency specifier torch>=1.9.*. pip 24.0 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.3.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!git clone https://github.com/baofff/U-ViT\n",
    "!pip install einops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5a57ae81-d9fa-4ddd-a8f3-4d3e88e40d06",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-19T11:33:34.876466Z",
     "iopub.status.busy": "2024-02-19T11:33:34.876128Z",
     "iopub.status.idle": "2024-02-19T11:33:34.996801Z",
     "shell.execute_reply": "2024-02-19T11:33:34.996215Z",
     "shell.execute_reply.started": "2024-02-19T11:33:34.876447Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "attention mode is flash\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.chdir('/mnt/workspace/U-ViT')\n",
    "os.environ['PYTHONPATH'] = '/env/python:/content/U-ViT'\n",
    "\n",
    "import torch\n",
    "from dpm_solver_pp import NoiseScheduleVP, DPM_Solver\n",
    "import libs.autoencoder\n",
    "from libs.uvit import UViT\n",
    "import einops\n",
    "from torchvision.utils import save_image\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b457d379-0e44-4127-ae70-75b1c0866985",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-19T11:33:36.464889Z",
     "iopub.status.busy": "2024-02-19T11:33:36.464451Z",
     "iopub.status.idle": "2024-02-19T11:33:36.467697Z",
     "shell.execute_reply": "2024-02-19T11:33:36.467121Z",
     "shell.execute_reply.started": "2024-02-19T11:33:36.464870Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from modelscope.hub.file_download import model_file_download"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3c518405-82c0-44b4-b0ea-1720b2838874",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-02-19T11:33:37.878912Z",
     "iopub.status.busy": "2024-02-19T11:33:37.878608Z",
     "iopub.status.idle": "2024-02-19T11:33:56.952396Z",
     "shell.execute_reply": "2024-02-19T11:33:56.951707Z",
     "shell.execute_reply.started": "2024-02-19T11:33:37.878895Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading: 100%|█████████▉| 1.87G/1.87G [00:05<00:00, 354MB/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "UViT(\n",
       "  (patch_embed): PatchEmbed(\n",
       "    (proj): Conv2d(4, 1152, kernel_size=(2, 2), stride=(2, 2))\n",
       "  )\n",
       "  (time_embed): Identity()\n",
       "  (label_emb): Embedding(1001, 1152)\n",
       "  (in_blocks): ModuleList(\n",
       "    (0-13): 14 x Block(\n",
       "      (norm1): LayerNorm((1152,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): Attention(\n",
       "        (qkv): Linear(in_features=1152, out_features=3456, bias=False)\n",
       "        (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "        (proj): Linear(in_features=1152, out_features=1152, bias=True)\n",
       "        (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "      )\n",
       "      (norm2): LayerNorm((1152,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): Mlp(\n",
       "        (fc1): Linear(in_features=1152, out_features=4608, bias=True)\n",
       "        (act): GELU(approximate='none')\n",
       "        (fc2): Linear(in_features=4608, out_features=1152, bias=True)\n",
       "        (drop): Dropout(p=0.0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (mid_block): Block(\n",
       "    (norm1): LayerNorm((1152,), eps=1e-05, elementwise_affine=True)\n",
       "    (attn): Attention(\n",
       "      (qkv): Linear(in_features=1152, out_features=3456, bias=False)\n",
       "      (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "      (proj): Linear(in_features=1152, out_features=1152, bias=True)\n",
       "      (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "    )\n",
       "    (norm2): LayerNorm((1152,), eps=1e-05, elementwise_affine=True)\n",
       "    (mlp): Mlp(\n",
       "      (fc1): Linear(in_features=1152, out_features=4608, bias=True)\n",
       "      (act): GELU(approximate='none')\n",
       "      (fc2): Linear(in_features=4608, out_features=1152, bias=True)\n",
       "      (drop): Dropout(p=0.0, inplace=False)\n",
       "    )\n",
       "  )\n",
       "  (out_blocks): ModuleList(\n",
       "    (0-13): 14 x Block(\n",
       "      (norm1): LayerNorm((1152,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): Attention(\n",
       "        (qkv): Linear(in_features=1152, out_features=3456, bias=False)\n",
       "        (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "        (proj): Linear(in_features=1152, out_features=1152, bias=True)\n",
       "        (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "      )\n",
       "      (norm2): LayerNorm((1152,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): Mlp(\n",
       "        (fc1): Linear(in_features=1152, out_features=4608, bias=True)\n",
       "        (act): GELU(approximate='none')\n",
       "        (fc2): Linear(in_features=4608, out_features=1152, bias=True)\n",
       "        (drop): Dropout(p=0.0, inplace=False)\n",
       "      )\n",
       "      (skip_linear): Linear(in_features=2304, out_features=1152, bias=True)\n",
       "    )\n",
       "  )\n",
       "  (norm): LayerNorm((1152,), eps=1e-05, elementwise_affine=True)\n",
       "  (decoder_pred): Linear(in_features=1152, out_features=16, bias=True)\n",
       "  (final_layer): Identity()\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "image_size = \"256\" #@param [256, 512]\n",
    "image_size = int(image_size)\n",
    "\n",
    "if image_size == 256:\n",
    "    model_file_download(model_id='thu-ml/imagenet256_uvit_huge',file_path='imagenet256_uvit_huge.pth', cache_dir='/mnt/workspace')\n",
    "    !mv /mnt/workspace/thu-ml/imagenet256_uvit_huge/imagenet256_uvit_huge.pth /mnt/workspace/U-ViT\n",
    "else:\n",
    "    model_file_download(model_id='thu-ml/imagenet512_uvit_huge',file_path='imagenet512_uvit_huge.pth', cache_dir='/mnt/workspace')\n",
    "    !mv /mnt/workspace/thu-ml/imagenet512_uvit_huge/imagenet512_uvit_huge.pth /mnt/workspace/U-ViT\n",
    " \n",
    "z_size = image_size // 8\n",
    "patch_size = 2 if image_size == 256 else 4\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "nnet = UViT(img_size=z_size,\n",
    "       patch_size=patch_size,\n",
    "       in_chans=4,\n",
    "       embed_dim=1152,\n",
    "       depth=28,\n",
    "       num_heads=16,\n",
    "       num_classes=1001,\n",
    "       conv=False)\n",
    "\n",
    "nnet.to(device)\n",
    "nnet.load_state_dict(torch.load(f'imagenet{image_size}_uvit_huge.pth', map_location='cpu'))\n",
    "nnet.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "47b3cf27-4593-4abc-9b27-6fd9e3507204",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-02-19T11:34:01.179601Z",
     "iopub.status.busy": "2024-02-19T11:34:01.179298Z",
     "iopub.status.idle": "2024-02-19T11:34:05.051089Z",
     "shell.execute_reply": "2024-02-19T11:34:05.050547Z",
     "shell.execute_reply.started": "2024-02-19T11:34:01.179581Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading: 100%|██████████| 319M/319M [00:01<00:00, 207MB/s] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Create autoencoder with scale_factor=0.18215\n",
      "making attention of type 'vanilla' with 512 in_channels\n",
      "Working with z of shape (1, 4, 32, 32) = 4096 dimensions.\n",
      "making attention of type 'vanilla' with 512 in_channels\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "FrozenAutoencoderKL(\n",
       "  (encoder): Encoder(\n",
       "    (conv_in): Conv2d(3, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (down): ModuleList(\n",
       "      (0): Module(\n",
       "        (block): ModuleList(\n",
       "          (0-1): 2 x ResnetBlock(\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",
       "          )\n",
       "        )\n",
       "        (attn): ModuleList()\n",
       "        (downsample): Downsample(\n",
       "          (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2))\n",
       "        )\n",
       "      )\n",
       "      (1): Module(\n",
       "        (block): ModuleList(\n",
       "          (0): ResnetBlock(\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",
       "            (nin_shortcut): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "          )\n",
       "          (1): ResnetBlock(\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",
       "          )\n",
       "        )\n",
       "        (attn): ModuleList()\n",
       "        (downsample): Downsample(\n",
       "          (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2))\n",
       "        )\n",
       "      )\n",
       "      (2): Module(\n",
       "        (block): ModuleList(\n",
       "          (0): ResnetBlock(\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",
       "            (nin_shortcut): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "          )\n",
       "          (1): ResnetBlock(\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",
       "          )\n",
       "        )\n",
       "        (attn): ModuleList()\n",
       "        (downsample): Downsample(\n",
       "          (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2))\n",
       "        )\n",
       "      )\n",
       "      (3): Module(\n",
       "        (block): ModuleList(\n",
       "          (0-1): 2 x ResnetBlock(\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",
       "          )\n",
       "        )\n",
       "        (attn): ModuleList()\n",
       "      )\n",
       "    )\n",
       "    (mid): Module(\n",
       "      (block_1): ResnetBlock(\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",
       "      )\n",
       "      (attn_1): AttnBlock(\n",
       "        (norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (q): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (k): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (v): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (proj_out): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (block_2): ResnetBlock(\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",
       "      )\n",
       "    )\n",
       "    (norm_out): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "    (conv_out): Conv2d(512, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  )\n",
       "  (decoder): Decoder(\n",
       "    (conv_in): Conv2d(4, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (mid): Module(\n",
       "      (block_1): ResnetBlock(\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",
       "      )\n",
       "      (attn_1): AttnBlock(\n",
       "        (norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "        (q): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (k): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (v): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (proj_out): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (block_2): ResnetBlock(\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",
       "      )\n",
       "    )\n",
       "    (up): ModuleList(\n",
       "      (0): Module(\n",
       "        (block): ModuleList(\n",
       "          (0): ResnetBlock(\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",
       "            (nin_shortcut): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "          )\n",
       "          (1-2): 2 x ResnetBlock(\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",
       "          )\n",
       "        )\n",
       "        (attn): ModuleList()\n",
       "      )\n",
       "      (1): Module(\n",
       "        (block): ModuleList(\n",
       "          (0): ResnetBlock(\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",
       "            (nin_shortcut): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "          )\n",
       "          (1-2): 2 x ResnetBlock(\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",
       "          )\n",
       "        )\n",
       "        (attn): ModuleList()\n",
       "        (upsample): Upsample(\n",
       "          (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        )\n",
       "      )\n",
       "      (2-3): 2 x Module(\n",
       "        (block): ModuleList(\n",
       "          (0-2): 3 x ResnetBlock(\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",
       "          )\n",
       "        )\n",
       "        (attn): ModuleList()\n",
       "        (upsample): Upsample(\n",
       "          (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (norm_out): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "    (conv_out): Conv2d(128, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  )\n",
       "  (quant_conv): Conv2d(8, 8, kernel_size=(1, 1), stride=(1, 1))\n",
       "  (post_quant_conv): Conv2d(4, 4, kernel_size=(1, 1), stride=(1, 1))\n",
       ")"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_file_download(model_id='AI-ModelScope/autoencoder_kl_ema',file_path='autoencoder_kl_ema.pth', cache_dir='/mnt/workspace')\n",
    "!mv /mnt/workspace/AI-ModelScope/autoencoder_kl_ema/autoencoder_kl_ema.pth /mnt/workspace/U-ViT\n",
    "autoencoder = libs.autoencoder.get_model('autoencoder_kl_ema.pth')\n",
    "autoencoder.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "038b90cc-3884-44e3-87e3-ab3a0f0cd87d",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-02-19T11:34:10.013253Z",
     "iopub.status.busy": "2024-02-19T11:34:10.012921Z",
     "iopub.status.idle": "2024-02-19T11:34:24.747234Z",
     "shell.execute_reply": "2024-02-19T11:34:24.746758Z",
     "shell.execute_reply.started": "2024-02-19T11:34:10.013221Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 464588 characters]",
      "image/png": "[encoded data, 5239812 characters]",
      "text/plain": [
       "<PIL.PngImagePlugin.PngImageFile image mode=RGB size=2048x1024>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "seed = 4321 #@param {type:\"number\"}\n",
    "steps = 25 #@param {type:\"slider\", min:0, max:1000, step:1}\n",
    "cfg_scale = 3 #@param {type:\"slider\", min:0, max:10, step:0.1}\n",
    "class_labels = 207, 360, 387, 974, 88, 979, 417, 279 #@param {type:\"raw\"}\n",
    "samples_per_row = 4 #@param {type:\"number\"}\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "def stable_diffusion_beta_schedule(linear_start=0.00085, linear_end=0.0120, n_timestep=1000):\n",
    "    _betas = (\n",
    "        torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2\n",
    "    )\n",
    "    return _betas.numpy()\n",
    "\n",
    "\n",
    "_betas = stable_diffusion_beta_schedule()  # set the noise schedule\n",
    "noise_schedule = NoiseScheduleVP(schedule='discrete', betas=torch.tensor(_betas, device=device).float())\n",
    "\n",
    "\n",
    "y = torch.tensor(class_labels, device=device)\n",
    "y = einops.repeat(y, 'B -> (B N)', N=samples_per_row)\n",
    "\n",
    "def model_fn(x, t_continuous):\n",
    "    t = t_continuous * len(_betas)\n",
    "    _cond = nnet(x, t, y=y)\n",
    "    _uncond = nnet(x, t, y=torch.tensor([1000] * x.size(0), device=device))\n",
    "    return _cond + cfg_scale * (_cond - _uncond)  # classifier free guidance\n",
    "\n",
    "\n",
    "z_init = torch.randn(len(y), 4, z_size, z_size, device=device)\n",
    "dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=True, thresholding=False)\n",
    "\n",
    "with torch.no_grad():\n",
    "  with torch.cuda.amp.autocast():  # inference with mixed precision\n",
    "    z = dpm_solver.sample(z_init, steps=steps, eps=1. / len(_betas), T=1.)\n",
    "    samples = autoencoder.decode(z)\n",
    "samples = 0.5 * (samples + 1.)\n",
    "samples.clamp_(0., 1.)\n",
    "save_image(samples, \"sample.png\", nrow=samples_per_row * 2, padding=0)\n",
    "samples = Image.open(\"sample.png\")\n",
    "display(samples)"
   ]
  }
 ],
 "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
