{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "eade093b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install rudalle==0.0.1rc5 > /dev/null"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dc79e298",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-11-05 15:07:58.768069: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1\n"
     ]
    }
   ],
   "source": [
    "from rudalle.pipelines import generate_images, show, super_resolution, cherry_pick_by_clip\n",
    "from rudalle.image_prompts import ImagePrompts\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": "ae1ea3fd",
   "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"
     ]
    }
   ],
   "source": [
    "device = 'cuda'\n",
    "dalle = get_rudalle_model('Malevich', pretrained=True, fp16=True, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a2741b22",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x4 --> ready\n",
      "tokenizer --> 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",
    "tokenizer = get_tokenizer()\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": "12ba6a35",
   "metadata": {},
   "source": [
    "## image prompts to dress manequins by ruDALLE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ddea6c2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import requests\n",
    "from PIL import Image\n",
    "import torch\n",
    "\n",
    "mannequin_woman_url = 'https://manekens.ru/storage/app/uploads/public/34/5d25eabf21090114407732.jpg'\n",
    "mannequin_man_url = 'https://ncm-torg.ru/wp-content/uploads/2019/08/CGE-11.jpg'\n",
    "\n",
    "mannequin_woman = Image.open(requests.get(mannequin_woman_url, stream=True).raw).resize((256, 256))\n",
    "mannequin_man = Image.open(requests.get(mannequin_man_url, stream=True).raw).resize((256, 256))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "48ae46e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/user/conda/lib/python3.7/site-packages/taming/modules/diffusionmodules/model.py:137: UserWarning: Mixed memory format inputs detected while calling the operator. The operator will output channels_last tensor even if some of the inputs are not in channels_last format. (Triggered internally at  /pytorch/aten/src/ATen/native/TensorIterator.cpp:924.)\n",
      "  return x+h\n"
     ]
    }
   ],
   "source": [
    "pt_items = {\n",
    "    'woman': {\n",
    "        'prompt': ImagePrompts(\n",
    "            mannequin_woman, \n",
    "            {'up': 5, 'left': 0, 'right': 0, 'down': 0}, \n",
    "            vae, \n",
    "            torch.device('cuda'), \n",
    "            crop_first=True\n",
    "        ),\n",
    "        'texts': [\n",
    "            'красивое красное платье на манекене',\n",
    "            'зеленый женский наряд на манекене',\n",
    "        ]\n",
    "    },\n",
    "    'man': {\n",
    "        'prompt': ImagePrompts(\n",
    "            mannequin_man, \n",
    "            {'up': 7, 'left': 0, 'right': 0, 'down': 0}, \n",
    "            vae, \n",
    "            torch.device('cuda'), \n",
    "            crop_first=True\n",
    "        ),\n",
    "        'texts': [\n",
    "            'белый праздничный мужской костюм на манекене',\n",
    "            'черный смокинг на манекене',\n",
    "        ]\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "217f0266",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "974013b0b65141e59c0fff3e16fa48d9",
       "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": "fb6e615862a641e581b91c0cdbf572e2",
       "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": "804f4c648694405ba8cbafe19db0321c",
       "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": "6d182b93e40f4757b3b3fc8243fd444e",
       "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": "2db36b4b7dec4e0ba6e7dea15ee0eefb",
       "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": "18074fe8db864f98865c43c64e56548f",
       "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": "55b0846e56b04970bdf5ed6aec2de7a9",
       "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": "4186e497a74a4ec592137cff3e7ecbed",
       "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": "f3329b9fee784901b4eab13d9e287f20",
       "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": "6d7a8116efe845a3bfd0c7829c80be2f",
       "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": "88ec256fc2fe4cd5ba8772f901d60d50",
       "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": "75353665c0b0405c96ed89149925a9ef",
       "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": "561d269fcc43451b86f3f07d602b0eb5",
       "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": "fdb0f7ca365c4e159d8685352ce18526",
       "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": "e0abf04508e34625b1210a9bc6a548a9",
       "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": "ddfd020cc51a4b2d9536270a5fa271ea",
       "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": "7e111a7ecfcf49d6aebd3dad9e4b9a10",
       "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": "b2c09f9d3e674f309c5b9d4748ea7428",
       "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": "274b693b511e4a58aec62ef0214b380a",
       "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": "36abfc2b373e43a6a47d8a060a75e5e9",
       "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": "e12f01f534934b6ba45a5ace102c1d5a",
       "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": "9479a5bf0b844cd79df16ee6ffa0566f",
       "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": "931defb08dbe453584aed72c5d975ee6",
       "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": "c5113b8cbd7c4aa4b06280feec5dd2ae",
       "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": "47b99b16c83547de8e33840cf2e91d90",
       "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": "1fdeb0d0769b482c8fffd00c9aff33bd",
       "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": "0858b94badda408eb859a9a58c821535",
       "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": "66ee10d0347648c082779a45dbc1a10e",
       "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": "cd53437ee74d40098c3b8a33c420ba10",
       "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": "de4bc0b27b174b16ba5799ab5acf16f2",
       "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": "e77c623e328e4e9090488ab7bbafe605",
       "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": "a339cb89299b4e219d7f59e3a1da7406",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1024 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "all_images = defaultdict(list)\n",
    "\n",
    "for name, item in pt_items.items():\n",
    "    seed_everything(42)\n",
    "    for text in item['texts']:\n",
    "        pil_images = []\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, _ = generate_images(\n",
    "                text,\n",
    "                tokenizer,\n",
    "                dalle,\n",
    "                vae,\n",
    "                top_k=top_k,\n",
    "                images_num=images_num,\n",
    "                image_prompts=item['prompt'],\n",
    "                top_p=top_p,\n",
    "                use_cache=True\n",
    "            )\n",
    "            pil_images += _pil_images\n",
    "        top_images, _ = cherry_pick_by_clip(pil_images, text, ruclip, ruclip_processor, device=device, count=6)\n",
    "        all_images[name] += super_resolution(top_images, realesrgan)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1c98c37e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 271392 characters]\n",
      "text/plain": [
       "<Figure size 1008x1008 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "show(all_images['woman'], 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f5f57589",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 155608 characters]\n",
      "text/plain": [
       "<Figure size 1008x1008 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "show(all_images['man'], 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60f829d7",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
