{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "romantic-shopper",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1AoolDYePUpPkRCKIu0cP9zV7lX5QGD3Z?usp=sharing)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "atmospheric-benchmark",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mWARNING: Running pip as root will break packages and permissions. You should install packages reliably by using venv: https://pip.pypa.io/warnings/venv\u001b[0m\n",
      "\u001b[33mWARNING: You are using pip version 21.1.1; however, version 21.3.1 is available.\n",
      "You should consider upgrading via the '/opt/conda/bin/python -m pip install --upgrade pip' command.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install rudalle==0.0.1rc8 > /dev/null"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "authorized-temple",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ruDALL-E batch size: 1\n"
     ]
    }
   ],
   "source": [
    "ALLOWED_MEMORY = 3.5 # choose your GPU memory in GB  \n",
    "if ALLOWED_MEMORY < 4.5:\n",
    "    DALLE_BS = 1\n",
    "elif ALLOWED_MEMORY < 5.5:\n",
    "    DALLE_BS = 2\n",
    "elif ALLOWED_MEMORY < 6.5:\n",
    "    DALLE_BS = 3\n",
    "elif ALLOWED_MEMORY < 7.5:\n",
    "    DALLE_BS = 4\n",
    "elif ALLOWED_MEMORY < 8.5:\n",
    "    DALLE_BS = 5\n",
    "elif ALLOWED_MEMORY < 9.5:\n",
    "    DALLE_BS = 6\n",
    "elif ALLOWED_MEMORY < 10.5:\n",
    "    DALLE_BS = 7\n",
    "else:\n",
    "    DALLE_BS = 8\n",
    "    \n",
    "print('ruDALL-E batch size:', DALLE_BS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "molecular-butler",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total GPU RAM: 23.69 Gb\n",
      "CPU: 32\n",
      "RAM GB: 125.7\n",
      "PyTorch version: 1.9.0+cu111\n",
      "CUDA version: 11.1\n",
      "cuDNN version: 8005\n",
      "Allowed GPU RAM: 3.5 Gb\n",
      "GPU part 0.1477\n",
      "Wed Nov 10 16:50:38 2021       \n",
      "+-----------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 460.91.03    Driver Version: 460.91.03    CUDA Version: 11.2     |\n",
      "|-------------------------------+----------------------+----------------------+\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                               |                      |               MIG M. |\n",
      "|===============================+======================+======================|\n",
      "|   0  GeForce RTX 3090    Off  | 00000000:04:00.0 Off |                  N/A |\n",
      "| 53%   47C    P2   120W / 350W |   1807MiB / 24260MiB |     17%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "                                                                               \n",
      "+-----------------------------------------------------------------------------+\n",
      "| Processes:                                                                  |\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
      "|        ID   ID                                                   Usage      |\n",
      "|=============================================================================|\n",
      "+-----------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "import multiprocessing\n",
    "import torch\n",
    "from psutil import virtual_memory\n",
    "\n",
    "total_memory = torch.cuda.get_device_properties(0).total_memory / 2**30\n",
    "if total_memory < ALLOWED_MEMORY:\n",
    "    raise MemoryError\n",
    "print('Total GPU RAM:', round(total_memory, 2), 'Gb')\n",
    "\n",
    "ram_gb = round(virtual_memory().total / 1024**3, 1)\n",
    "print('CPU:', multiprocessing.cpu_count())\n",
    "print('RAM GB:', ram_gb)\n",
    "print(\"PyTorch version:\", torch.__version__)\n",
    "print(\"CUDA version:\", torch.version.cuda)\n",
    "print(\"cuDNN version:\", torch.backends.cudnn.version())\n",
    "\n",
    "if torch.__version__ >= '1.8.0':\n",
    "    k = ALLOWED_MEMORY/ total_memory\n",
    "    torch.cuda.set_per_process_memory_fraction(k, 0)\n",
    "    print('Allowed GPU RAM:', round(ALLOWED_MEMORY, 2), 'Gb')\n",
    "    print('GPU part', round(k, 4))\n",
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "useful-stake",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "◼️ Malevich is 1.3 billion params model from the family GPT3-like, that uses Russian language and text+image multi-modality.\n",
      "tokenizer --> ready\n",
      "Working with z of shape (1, 256, 32, 32) = 262144 dimensions.\n",
      "vae --> ready\n",
      "ruclip --> ready\n"
     ]
    }
   ],
   "source": [
    "import transformers\n",
    "import more_itertools\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from rudalle.pipelines import show, cherry_pick_by_clip\n",
    "from rudalle import get_rudalle_model, get_tokenizer, get_vae, get_ruclip\n",
    "from rudalle.utils import seed_everything, torch_tensors_to_pil_list\n",
    "\n",
    "# prepare models:\n",
    "device = 'cuda'\n",
    "dalle = get_rudalle_model('Malevich', pretrained=True, fp16=True, device=device)\n",
    "tokenizer = get_tokenizer()\n",
    "vae = get_vae(dwt=True)\n",
    "\n",
    "# prepare utils:\n",
    "ruclip, ruclip_processor = get_ruclip('ruclip-vit-base-patch32-v5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "further-forge",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_codebooks(text, tokenizer, dalle, top_k, top_p, images_num, image_prompts=None, temperature=1.0, bs=8,\n",
    "                    seed=None, use_cache=True):\n",
    "    vocab_size = dalle.get_param('vocab_size')\n",
    "    text_seq_length = dalle.get_param('text_seq_length')\n",
    "    image_seq_length = dalle.get_param('image_seq_length')\n",
    "    total_seq_length = dalle.get_param('total_seq_length')\n",
    "    device = dalle.get_param('device')\n",
    "    text = text.lower().strip()\n",
    "    input_ids = tokenizer.encode_text(text, text_seq_length=text_seq_length)\n",
    "    codebooks = []\n",
    "    for chunk in more_itertools.chunked(range(images_num), bs):\n",
    "        chunk_bs = len(chunk)\n",
    "        with torch.no_grad():\n",
    "            attention_mask = torch.tril(torch.ones((chunk_bs, 1, total_seq_length, total_seq_length), device=device))\n",
    "            out = input_ids.unsqueeze(0).repeat(chunk_bs, 1).to(device)\n",
    "            has_cache = False\n",
    "            if image_prompts is not None:\n",
    "                prompts_idx, prompts = image_prompts.image_prompts_idx, image_prompts.image_prompts\n",
    "                prompts = prompts.repeat(chunk_bs, 1)\n",
    "            for idx in tqdm(range(out.shape[1], total_seq_length)):\n",
    "                idx -= text_seq_length\n",
    "                if image_prompts is not None and idx in prompts_idx:\n",
    "                    out = torch.cat((out, prompts[:, idx].unsqueeze(1)), dim=-1)\n",
    "                else:\n",
    "                    logits, has_cache = dalle(out, attention_mask,\n",
    "                                              has_cache=has_cache, use_cache=use_cache, return_loss=False)\n",
    "                    logits = logits[:, -1, vocab_size:]\n",
    "                    logits /= temperature\n",
    "                    filtered_logits = transformers.top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)\n",
    "                    probs = torch.nn.functional.softmax(filtered_logits, dim=-1)\n",
    "                    sample = torch.multinomial(probs, 1)\n",
    "                    out = torch.cat((out, sample), dim=-1)\n",
    "            codebooks.append(out[:, -image_seq_length:].cpu())\n",
    "    return codebooks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "tight-priority",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "83d5b9ed02484cd7826ea133fc1f81f7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.8/site-packages/torch/_tensor.py:575: UserWarning: floor_divide is deprecated, and will be removed in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values.\n",
      "To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor'). (Triggered internally at  /pytorch/aten/src/ATen/native/BinaryOps.cpp:467.)\n",
      "  return torch.floor_divide(self, other)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2edf9c4cf03e4735934ac990d485155d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b633c20ce52e48e8839a3b6918ac5a08",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "624ef56e7cfb467fb2b8a59933e17abb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3d0f8d52fb724f3e9988103d556070bf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6ae304ae57684143b936d57dac700f75",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e18e7018f0874ab89cee3968ab5cfe57",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8ba79fe8d9b54cc094e5d33e641d9292",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c0dff1ba89d14c5cae07223dcb8079e2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4a0182f68c094d66bad0ba4665568021",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fd43ace8af394f918725dcff30a2e235",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3eff4f7ac4ba4ae38e956f6a9d78e176",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2fafab41113740f8a427077f5a352687",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d48dcb06c41f45bda0990abf7e303283",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1a098d3efb554c20bf02d57fc909d4f3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "da44a5e218924f52b543cdb9db77f983",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "text = 'Пингвины радуются - неизвестная картина Казимира Малевича'\n",
    "\n",
    "seed_everything(6955)\n",
    "codebooks = []\n",
    "for top_k, top_p, images_num in [\n",
    "    (2048, 0.995, 8),\n",
    "    (1536, 0.99, 8),\n",
    "#     (1024, 0.99, 8),\n",
    "#     (1024, 0.98, 8),\n",
    "#     (512, 0.97, 8),\n",
    "#     (384, 0.96, 8),\n",
    "#     (256, 0.95, 8),\n",
    "#     (128, 0.95, 8),\n",
    "]:\n",
    "    codebooks += generate_codebooks(text, tokenizer, dalle, top_k=top_k, images_num=images_num, top_p=top_p, bs=DALLE_BS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "advanced-costs",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e5ee4fdb8bb3408a90240134f426a8a4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/16 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pil_images = []\n",
    "for _codebooks in tqdm(torch.cat(codebooks).cpu()):\n",
    "    with torch.no_grad():\n",
    "        images = vae.decode(_codebooks.unsqueeze(0))\n",
    "        pil_images += torch_tensors_to_pil_list(images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "stone-smooth",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1169868 characters]\n",
      "text/plain": [
       "<Figure size 1008x1008 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "show(pil_images, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "classified-honduras",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 745356 characters]\n",
      "text/plain": [
       "<Figure size 1008x1008 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "top_images, clip_scores = cherry_pick_by_clip(pil_images, text, ruclip, ruclip_processor, device='cpu', count=6)\n",
    "show(top_images, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "chinese-violence",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
