{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fdac3238",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install rudalle==0.0.1rc4 > /dev/null"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9264c4b3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-11-03 19:22:10.306367: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0\n"
     ]
    }
   ],
   "source": [
    "from rudalle.pipelines import generate_images, show, super_resolution, cherry_pick_by_clip\n",
    "from rudalle import get_rudalle_model, get_tokenizer, get_vae, get_realesrgan, get_ruclip\n",
    "from rudalle.utils import seed_everything"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c76c31da",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tokenizer --> ready\n",
      "◼️ Malevich is 1.3 billion params model from the family GPT3-like, that uses Russian language and text+image multi-modality.\n"
     ]
    }
   ],
   "source": [
    "device = 'cuda'\n",
    "tokenizer = get_tokenizer()\n",
    "dalle = get_rudalle_model('Malevich', pretrained=True, fp16=True, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f010cf3c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x4 --> ready\n",
      "Working with z of shape (1, 256, 32, 32) = 262144 dimensions.\n",
      "vae --> ready\n",
      "ruclip --> ready\n"
     ]
    }
   ],
   "source": [
    "realesrgan = get_realesrgan('x4', device=device)\n",
    "vae = get_vae().to(device)\n",
    "ruclip, ruclip_processor = get_ruclip('ruclip-vit-base-patch32-v5')\n",
    "ruclip = ruclip.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0dea7796",
   "metadata": {},
   "source": [
    "## generation by ruDALLE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3a4bf4b5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1024/1024 [00:30<00:00, 33.04it/s]\n",
      "100%|██████████| 1024/1024 [00:30<00:00, 33.47it/s]\n",
      "100%|██████████| 1024/1024 [00:30<00:00, 33.73it/s]\n",
      "100%|██████████| 1024/1024 [00:30<00:00, 33.64it/s]\n",
      "100%|██████████| 1024/1024 [00:30<00:00, 33.43it/s]\n",
      "100%|██████████| 1024/1024 [00:30<00:00, 33.45it/s]\n",
      "100%|██████████| 1024/1024 [00:30<00:00, 33.38it/s]\n",
      "100%|██████████| 1024/1024 [00:30<00:00, 33.52it/s]\n"
     ]
    }
   ],
   "source": [
    "text = 'изображение радуги на фоне ночного города'\n",
    "pil_images = []\n",
    "scores = []\n",
    "\n",
    "seed_everything(42)\n",
    "\n",
    "for top_k, top_p, images_num in [\n",
    "    (2048, 0.995, 3),\n",
    "    (1536, 0.99, 3),\n",
    "    (1024, 0.99, 3),\n",
    "    (1024, 0.98, 3),\n",
    "    (512, 0.97, 3),\n",
    "    (384, 0.96, 3),\n",
    "    (256, 0.95, 3),\n",
    "    (128, 0.95, 3), \n",
    "]:\n",
    "    _pil_images, _scores = generate_images(text, tokenizer, dalle, vae, top_k=top_k, images_num=images_num, top_p=top_p)\n",
    "    pil_images += _pil_images\n",
    "    scores += _scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b0780246",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 877476 characters]\n",
      "text/plain": [
       "<Figure size 1008x1008 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "show([pil_image for pil_image, score in sorted(zip(pil_images, scores), key=lambda x: -x[1])] , 6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0ec9628f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 779736 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=device, count=6)\n",
    "show(top_images, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "dbb1f38e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 706184 characters]\n",
      "text/plain": [
       "<Figure size 1008x1008 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sr_images = super_resolution(top_images, realesrgan)\n",
    "show(sr_images, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8c1bb5a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
