test_warhol_ui.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3e0c28d6",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7c933d8c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Imports done\n"
     ]
    }
   ],
   "source": [
    "import argparse\n",
    "from pathlib import Path\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "from ipywidgets import widgets\n",
    "import torch\n",
    "import torchvision\n",
    "from torch.utils.data import DataLoader\n",
    "from einops import repeat\n",
    "import clip\n",
    "import numpy as np\n",
    "from tqdm import tqdm \n",
    "import torch.nn.functional as F\n",
    "from PIL import Image\n",
    "from torchvision.utils import make_grid, save_image\n",
    "from dalle_pytorch import DiscreteVAE, OpenAIDiscreteVAE, VQGanVAE, WARHOL\n",
    "from dalle_pytorch.tokenizer import tokenizer, HugTokenizer, YttmTokenizer, ChineseTokenizer\n",
    "from dalle_pytorch.loader import TextImageDataset\n",
    "import os\n",
    "import locale\n",
    "os.environ[\"PYTHONIOENCODING\"] = \"utf-8\"\n",
    "print(\"Imports done\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35054773",
   "metadata": {},
   "source": [
    "## Load models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c3fa8a80",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eb35f4ec760d40958b28d3e4617b86b6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Text(value='f.name', description='Username:', placeholder='Type something')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import os\n",
    "username = widgets.Text(\n",
    "    value='f.name',\n",
    "    placeholder='Type something',\n",
    "    description='Username:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "display(username)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8179da3b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "get: `mytheresa_clip.tar': File exists\n",
      "data downloaded\n",
      "get: `02_warhol_shoes150k.pt': File exists\n",
      "get: `01_warhol_fashion200k.pt': File exists\n",
      "get: `01_warhol_mytheresa.pt': File exists\n",
      "get: `02_warhol_mytheresa_timeline.pt': File exists\n",
      "get: `vqgan_imagenet_f16_16384_model.ckpt': File exists\n",
      "get: `vqgan_imagenet_f16_16384_configs.yaml': File exists\n",
      "fatal: destination path 'Real-ESRGAN' already exists and is not an empty directory.\n",
      "--2022-07-25 16:38:05--  https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth\n",
      "Resolving github.com (github.com)... 140.82.121.3\n",
      "Connecting to github.com (github.com)|140.82.121.3|:443... connected.\n",
      "HTTP request sent, awaiting response... 302 Found\n",
      "Location: https://objects.githubusercontent.com/github-production-release-asset-2e65be/387326890/08f0e941-ebb7-48f0-9d6a-73e87b710e7e?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20220725%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20220725T163805Z&X-Amz-Expires=300&X-Amz-Signature=f66d6608378598f9bef9557590550c845154e5dc83d0cd1358fed660d2b4ea4f&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=387326890&response-content-disposition=attachment%3B%20filename%3DRealESRGAN_x4plus.pth&response-content-type=application%2Foctet-stream [following]\n",
      "--2022-07-25 16:38:05--  https://objects.githubusercontent.com/github-production-release-asset-2e65be/387326890/08f0e941-ebb7-48f0-9d6a-73e87b710e7e?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20220725%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20220725T163805Z&X-Amz-Expires=300&X-Amz-Signature=f66d6608378598f9bef9557590550c845154e5dc83d0cd1358fed660d2b4ea4f&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=387326890&response-content-disposition=attachment%3B%20filename%3DRealESRGAN_x4plus.pth&response-content-type=application%2Foctet-stream\n",
      "Resolving objects.githubusercontent.com (objects.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.111.133, ...\n",
      "Connecting to objects.githubusercontent.com (objects.githubusercontent.com)|185.199.108.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 67040989 (64M) [application/octet-stream]\n",
      "Saving to: 'experiments/pretrained_models/RealESRGAN_x4plus.pth.4'\n",
      "\n",
      "100%[======================================>] 67,040,989  --.-K/s   in 0.1s    \n",
      "\n",
      "2022-07-25 16:38:05 (445 MB/s) - 'experiments/pretrained_models/RealESRGAN_x4plus.pth.4' saved [67040989/67040989]\n",
      "\n",
      "model downloaded\n",
      "/home/u.tanielian/temp_data/Real-ESRGAN\n",
      "done\n"
     ]
    }
   ],
   "source": [
    "path = \"/home/\"+username.value+\"/temp_data/\"\n",
    "folder_for_gifs = \"/home/\"+username.value+\"/temp_data/\"\n",
    "\n",
    "if not os.path.exists(path):\n",
    "    os.mkdir(path)\n",
    "os.chdir(path)\n",
    "!hdfs dfs -get \"/user/u.tanielian/mytheresa_clip.tar\"\n",
    "!tar -xf mytheresa_clip.tar\n",
    "dataset_path = path + \"mytheresa_dalle_format/\"\n",
    "print(\"data downloaded\")\n",
    "\n",
    "os.chdir(path)\n",
    "!hdfs dfs -get \"/user/u.tanielian/02_warhol_shoes150k.pt\"\n",
    "!hdfs dfs -get \"/user/u.tanielian/01_warhol_fashion200k.pt\"\n",
    "!hdfs dfs -get \"/user/u.tanielian/01_warhol_mytheresa.pt\"\n",
    "!hdfs dfs -get \"/user/u.tanielian/02_warhol_mytheresa_timeline.pt\"\n",
    "!hdfs dfs -get \"/user/u.tanielian/vqgan_imagenet_f16_16384_model.ckpt\"\n",
    "!hdfs dfs -get \"/user/u.tanielian/vqgan_imagenet_f16_16384_configs.yaml\"\n",
    "\n",
    "os.chdir(path)\n",
    "!git clone https://github.com/tanouch/Real-ESRGAN.git\n",
    "folder = os.path.join(path, \"Real-ESRGAN\")\n",
    "os.chdir(folder)\n",
    "!wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P experiments/pretrained_models\n",
    "print(\"model downloaded\")\n",
    "    \n",
    "folder = os.path.join(path, \"Real-ESRGAN\")\n",
    "os.chdir(folder)\n",
    "print(folder)\n",
    "from basicsr.archs.rrdbnet_arch import RRDBNet\n",
    "from realesrgan import RealESRGANer\n",
    "    \n",
    "os.chdir(path)\n",
    "warhol_path = path + \"01_warhol_fashion200k.pt\"\n",
    "warhol_timeline_path = path + \"02_warhol_mytheresa_timeline.pt\"\n",
    "vqgan_model_path  = path + \"vqgan_imagenet_f16_16384_model.ckpt\"\n",
    "vqgan_config_path = path + \"vqgan_imagenet_f16_16384_configs.yaml\"\n",
    "srgan_model_path = path + \"Real-ESRGAN/experiments/pretrained_models/RealESRGAN_x4plus.pth\"\n",
    "print(\"done\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a9ee04d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['RN50', 'RN101', 'RN50x4', 'RN50x16', 'RN50x64', 'ViT-B/32', 'ViT-B/16', 'ViT-L/14', 'ViT-L/14@336px']\n",
      "Model parameters: 151,277,313\n",
      "Context length: 77\n",
      "Vocab size: 49408\n",
      "CLIP loaded\n",
      "Working with z of shape (1, 256, 16, 16) = 65536 dimensions.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading: \"https://download.pytorch.org/models/vgg16-397923af.pth\" to /home/u.tanielian/.cache/torch/hub/checkpoints/vgg16-397923af.pth\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cede471953b3409a984c21259a2054d4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0.00/528M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading vgg_lpips model from https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1 to taming/modules/autoencoder/lpips/vgg.pth\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "8.19kB [00:00, 912kB/s]                    \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loaded pretrained LPIPS loss from taming/modules/autoencoder/lpips/vgg.pth\n",
      "VQLPIPSWithDiscriminator running with hinge loss.\n",
      "Loaded VQGAN from /home/u.tanielian/temp_data/vqgan_imagenet_f16_16384_model.ckpt and /home/u.tanielian/temp_data/vqgan_imagenet_f16_16384_configs.yaml\n",
      "True\n",
      "True\n",
      "Working with z of shape (1, 256, 16, 16) = 65536 dimensions.\n",
      "loaded pretrained LPIPS loss from taming/modules/autoencoder/lpips/vgg.pth\n",
      "VQLPIPSWithDiscriminator running with hinge loss.\n",
      "Loaded VQGAN from /home/u.tanielian/temp_data/vqgan_imagenet_f16_16384_model.ckpt and /home/u.tanielian/temp_data/vqgan_imagenet_f16_16384_configs.yaml\n",
      "WARHOL loaded\n",
      "files loaded\n"
     ]
    }
   ],
   "source": [
    "def load_warhol_model():\n",
    "    loaded_obj = torch.load(warhol_path, map_location='cuda')\n",
    "    warhol_params, vae_params, weights = loaded_obj['hparams'], loaded_obj['vae_params'], loaded_obj['weights']\n",
    "    opt_state = loaded_obj.get('opt_state')\n",
    "    scheduler_state = loaded_obj.get('scheduler_state')\n",
    "    warhol_params = dict(**warhol_params)\n",
    "    resume_epoch = loaded_obj.get('epoch', 0)\n",
    "    USE_NEG_SAMPLES = False\n",
    "    USE_NEXT_PROD_MODULE = False\n",
    "    warhol_params[\"use_neg_samples\"] = USE_NEG_SAMPLES\n",
    "    warhol_params[\"use_next_prod_module\"] = USE_NEXT_PROD_MODULE\n",
    "\n",
    "    vae = VQGanVAE(vqgan_model_path, vqgan_config_path)\n",
    "    IMAGE_SIZE = vae.image_size\n",
    "\n",
    "    warhol = WARHOL(vae=vae, **warhol_params)\n",
    "    warhol.load_state_dict(weights)\n",
    "    warhol = warhol.cuda()\n",
    "    return warhol\n",
    "\n",
    "def load_warhol_timeline():\n",
    "    loaded_obj = torch.load(warhol_timeline_path, map_location='cuda')\n",
    "    warhol_params, vae_params, weights = loaded_obj['hparams'], loaded_obj['vae_params'], loaded_obj['weights']\n",
    "    opt_state = loaded_obj.get('opt_state')\n",
    "    scheduler_state = loaded_obj.get('scheduler_state')\n",
    "    warhol_params = dict(**warhol_params)\n",
    "    resume_epoch = loaded_obj.get('epoch', 0)\n",
    "    print(warhol_params[\"use_neg_samples\"])\n",
    "    print(warhol_params[\"use_next_prod_module\"])\n",
    "\n",
    "    vae = VQGanVAE(vqgan_model_path, vqgan_config_path)\n",
    "    IMAGE_SIZE = vae.image_size\n",
    "\n",
    "    warholTimeline = WARHOL(vae=vae, **warhol_params)\n",
    "    warholTimeline.load_state_dict(weights)\n",
    "    warholTimeline = warholTimeline.cuda()\n",
    "    return warholTimeline\n",
    "\n",
    "def load_clip_model():\n",
    "    MODELS = {\n",
    "        \"ViT-B/32\": \"https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt\",\n",
    "    }\n",
    "    print(clip.available_models())\n",
    "    model, preprocess = clip.load(\"ViT-B/32\")\n",
    "    model.to('cuda')\n",
    "    input_resolution = model.visual.input_resolution    \n",
    "    print(\"Model parameters:\", f\"{np.sum([int(np.prod(p.shape)) for p in model.parameters()]):,}\")\n",
    "    context_length = model.context_length\n",
    "    print(\"Context length:\", context_length)\n",
    "    vocab_size = model.vocab_size\n",
    "    print(\"Vocab size:\", str(vocab_size))\n",
    "    return model , input_resolution, context_length, vocab_size\n",
    "\n",
    "def infer_with_real_srgan(inputt, model_path=srgan_model_path, netscale=4, outscale=4, tile=0, tile_pad=10, pre_pad=0, \\\n",
    "                          face_enhance=False, half=True, block=23):\n",
    "        \n",
    "    if 'RealESRGAN_x4plus_anime_6B.pth' in model_path:\n",
    "        block = 6\n",
    "    elif 'RealESRGAN_x2plus.pth' in model_path:\n",
    "        netscale = 2\n",
    "\n",
    "    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=block, num_grow_ch=32, scale=netscale)\n",
    "    \n",
    "    upsampler = RealESRGANer(\n",
    "        scale=netscale,\n",
    "        model_path=model_path,\n",
    "        model=model,\n",
    "        tile=tile,\n",
    "        tile_pad=tile_pad,\n",
    "        pre_pad=pre_pad,\n",
    "        half=half)\n",
    "\n",
    "    if face_enhance:\n",
    "        from gfpgan import GFPGANer\n",
    "        face_enhancer = GFPGANer(\n",
    "            model_path='https://github.com/TencentARC/GFPGAN/releases/download/v0.2.0/GFPGANCleanv1-NoCE-C2.pth',\n",
    "            upscale=outscale,\n",
    "            arch='clean',\n",
    "            channel_multiplier=2,\n",
    "            bg_upsampler=upsampler)\n",
    "    \n",
    "    improved_images = list()\n",
    "    for idx, img in enumerate(inputt):\n",
    "        if len(img.shape) == 3 and img.shape[2] == 4:\n",
    "            img_mode = 'RGBA'\n",
    "        else:\n",
    "            img_mode = None\n",
    "\n",
    "        h, w = img.shape[0:2]\n",
    "        if max(h, w) > 1000 and netscale == 4:\n",
    "            import warnings\n",
    "            warnings.warn('The input image is large, try X2 model for better performance.')\n",
    "        if max(h, w) < 500 and netscale == 2:\n",
    "            import warnings\n",
    "            warnings.warn('The input image is small, try X4 model for better performance.')\n",
    "        \n",
    "        img = img*256\n",
    "        img = img.cpu().numpy()\n",
    "        img = np.transpose(img, (1, 2, 0))\n",
    "        if face_enhance:\n",
    "            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)\n",
    "        else:\n",
    "            output, _ = upsampler.enhance(img, outscale=outscale)\n",
    "        \n",
    "        output = np.transpose(output, (2, 0, 1))/256\n",
    "        output = torch.unsqueeze(torch.from_numpy(output), 0)\n",
    "        improved_images.append(output)\n",
    "        \n",
    "    improved_images = torch.cat(improved_images, dim=0)\n",
    "    return improved_images\n",
    "\n",
    "clip_model, input_resolution, context_length, vocab_size = load_clip_model()\n",
    "print(\"CLIP loaded\")\n",
    "warhol = load_warhol_model()\n",
    "warholTimeline = load_warhol_timeline()\n",
    "print(\"WARHOL loaded\")\n",
    "\n",
    "batch_size = 1\n",
    "from pathlib import Path\n",
    "path_folder = Path(dataset_path)\n",
    "clip_files = [*path_folder.glob('**/*.npy')]\n",
    "clip_files = sorted(list({clip_file.stem: clip_file for clip_file in clip_files}))\n",
    "im_clip_embs = torch.cat([torch.from_numpy(np.load(dataset_path+file+\".npy\"))[:512].unsqueeze(1) for file in clip_files], dim=1).cuda()\n",
    "im_files = [*path_folder.glob('**/*.jpg')]\n",
    "im_files = sorted(list({im_file.stem: im_file for im_file in im_files}))\n",
    "im_files = [dataset_path+file+\".jpg\" for file in im_files]\n",
    "print(\"files loaded\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c83de46-23d7-4e3d-879c-217e7128cebb",
   "metadata": {},
   "source": [
    "## Demo !!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "69af608f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6b321c7970cc4ff6bed8d9cebdc0faf0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(Text(value='Red high heels shoes', description='query'), IntSlider(value=4, description=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8d71c7cffd884a54ad07eaeb1d97378e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Checkbox(value=False, description='Content based', indent=False), Checkbox(value=False, descrip…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "home_folder = os.getcwd()\n",
    "from PIL import Image\n",
    "from ipywidgets import interact, interactive, fixed, interact_manual\n",
    "from IPython.display import display, Markdown, clear_output\n",
    "import ipywidgets as widgets\n",
    "from PIL import Image\n",
    "\n",
    "def showbig2(img, topk):\n",
    "    npimg = img.cpu().numpy()\n",
    "    npimg = np.clip(npimg, 0, 1)\n",
    "    fig = plt.figure(figsize=(topk*10, 3.))\n",
    "    ax = fig.add_subplot(111)\n",
    "    data = np.transpose(npimg, (1, 2, 0))\n",
    "    ax.imshow(data)\n",
    "    ax.set_axis_off()\n",
    "    ax.set_aspect(1)\n",
    "\n",
    "    \n",
    "def get_text_embedding_from_prompt(model, text_prompt):\n",
    "    texts = []\n",
    "    text_tokens = []\n",
    "    \n",
    "    text_tokens.append(tokenizer.encode(text_prompt))\n",
    "    texts.append(str(text_prompt))\n",
    "\n",
    "    text_input = torch.zeros(1, model.context_length, dtype=torch.long)\n",
    "    sot_token = tokenizer.encoder['<|startoftext|>']\n",
    "    eot_token = tokenizer.encoder['<|endoftext|>']\n",
    "\n",
    "    for i, tokens in enumerate(text_tokens):\n",
    "        tokens = [sot_token] + tokens + [eot_token]\n",
    "        text_input[i, :len(tokens)] = torch.tensor(tokens)\n",
    "\n",
    "    text_input = text_input.cuda()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        text_features = model.encode_text(text_input).float()\n",
    "\n",
    "    text_features /= text_features.norm(dim=-1, keepdim=True)\n",
    "    text_embedding = text_features\n",
    "    return text_embedding\n",
    "    \n",
    "    \n",
    "def output_txt_and_im_emb(im_clip_embs, prompt_text, topk):\n",
    "    ###GET QUERY EMB\n",
    "    txt_emb = get_text_embedding_from_prompt(clip_model, prompt_text).cuda()\n",
    "    im_clip_embs = im_clip_embs.float()\n",
    "    dot_product = torch.reshape(torch.matmul(txt_emb, im_clip_embs), (-1,))\n",
    "    (values, indices) = torch.topk(dot_product, topk)\n",
    "    target_img = [transforms.ToTensor()(Image.open(im_files[index])).unsqueeze_(0).cuda() for index in indices]\n",
    "    target_img = torch.cat(target_img, dim=0)\n",
    "    target_im_emb = torch.t(torch.index_select(im_clip_embs, dim=1, index=indices)).unsqueeze(1)\n",
    "    txt_emb = torch.cat([txt_emb.unsqueeze(0)]*topk, dim=0)\n",
    "    return target_img, target_im_emb, txt_emb\n",
    "\n",
    "    \n",
    "def get_closest_prods(query,topk):\n",
    "    target_img, target_im_emb, txt_emb = output_txt_and_im_emb(im_clip_embs, prompt_text=query, topk=topk)\n",
    "    qgrid = make_grid(target_img, nrow=topk)\n",
    "    showbig2(qgrid, topk)\n",
    "    return target_img, target_im_emb, txt_emb, topk, query\n",
    "    \n",
    "    \n",
    "def generate_prods(model_list, target_img, target_im_emb, txt_emb, temp, topk, \\\n",
    "                   num_images_sampled, nrow, prefix):\n",
    "    gen_images = list()\n",
    "    gen_images_np = list()\n",
    "    gen_texts = list()\n",
    "    clip_sims = list()\n",
    "    \n",
    "    for i in tqdm(range(num_images_sampled)): \n",
    "        for m in range(len(model_list)):\n",
    "            model = model_list[m]\n",
    "            text_tokens, gen_text = model.generate_texts(target_im_emb, txt_emb, temperature = temp)\n",
    "            gen_image = model.generate_images(target_im_emb, txt_emb, text_tokens, temperature = temp)\n",
    "            gen_image = F.interpolate(gen_image, size=224)\n",
    "            gen_images.append(gen_image)\n",
    "            gen_images_np.append(gen_image.cpu().numpy())\n",
    "\n",
    "            #Get CLIP embedding of gen_image and compute sim with query img and query txt\n",
    "            im_embed = clip_model.encode_image(gen_image).float()\n",
    "            im_embed = im_embed/im_embed.norm(dim=-1, keepdim=True)\n",
    "            img_clip_sim = (torch.diag(torch.matmul(im_embed, torch.t(target_im_emb.squeeze(1)))).cpu().detach().numpy())\n",
    "            txt_clip_sim = (torch.diag(torch.matmul(im_embed, torch.t(txt_emb.squeeze(1)))).cpu().detach().numpy())\n",
    "            clip_sims.append(txt_clip_sim)\n",
    "                \n",
    "            #Get generated text\n",
    "            new_texts = list()\n",
    "            every = 5\n",
    "            for i in range(len(gen_image)):\n",
    "                text = gen_text[i]\n",
    "                splitted_str = text.split(\" \")[:20]\n",
    "                splitted_str = [\" \".join(splitted_str[i:i+every]) for i in range(0, len(splitted_str), every)]\n",
    "                splitted_str = \"Cosine.Sim: \" + str(img_clip_sim[i]) +\"\\n\" + \"\\n\".join(splitted_str)+\"[...]\"\n",
    "                new_texts.append(splitted_str)\n",
    "            gen_texts.append(new_texts)\n",
    "            \n",
    "            #Display current generation\n",
    "            width = nrow*5.5\n",
    "            height = nrow*2.5\n",
    "            fig = plt.figure(figsize=(width, 6))\n",
    "            if m==0:\n",
    "                print(\"\\nContent-optimized generation: Products are optimized to be visually similar with the query product\\n\\n\")\n",
    "            else:\n",
    "                print(\"\\nUser-optimized generation: Products are optimized to be relevant to the users of the query product\\n\\n\")\n",
    "            \n",
    "            gen_image = infer_with_real_srgan(gen_image, netscale=4, outscale=3.5)\n",
    "            for i in range(len(gen_image)):\n",
    "                ax = fig.add_subplot(1, nrow, i+1)\n",
    "                ax.set_title(new_texts[i], fontsize=nrow*4.5)\n",
    "                ax.set_axis_off()\n",
    "\n",
    "                npimg = gen_image.cpu().numpy()[i]\n",
    "                npimg = np.clip(npimg, 0, 1)\n",
    "                img_data = np.transpose(npimg, (1, 2, 0))\n",
    "                ax.imshow(img_data)\n",
    "            plt.tight_layout()\n",
    "            plt.show()\n",
    "            \n",
    "    #Final image\n",
    "    if len(model_list)==1 and num_images_sampled>1:\n",
    "        fig = plt.figure(figsize=(width, 6))\n",
    "        print(\"\\n Showing the best generation !!\\n\")\n",
    "        for i in range(len(gen_image)):\n",
    "            j = np.argmax(np.array(clip_sims)[:,i])\n",
    "            ax = fig.add_subplot(1, nrow, i+1)\n",
    "            ax.set_title(np.array(gen_texts)[j,i], fontsize=nrow*4.5)\n",
    "            ax.set_axis_off()\n",
    "\n",
    "            npimg = np.array(gen_images_np)[j, i]\n",
    "            npimg = np.clip(npimg, 0, 1)\n",
    "            img_data = np.transpose(npimg, (1, 2, 0))\n",
    "            ax.imshow(img_data)\n",
    "        plt.tight_layout()\n",
    "        plt.show()\n",
    "    \n",
    "def generate_from_ui(query_tp, content_based, user_based, num_results, temp):\n",
    "    models = list()\n",
    "    if content_based:\n",
    "        models.append(warhol)\n",
    "    if user_based:\n",
    "        models.append(warholTimeline)\n",
    "    target_img, target_img_emb, txt_emb, topk, query = my_result.result\n",
    "    generate_prods(models, target_img, target_img_emb, txt_emb, \\\n",
    "                   temp, topk=topk, num_images_sampled=num_results, nrow=topk, prefix=\"WARHOL_\")\n",
    "    \n",
    "topk_slider = widgets.IntSlider(\n",
    "         value=4,\n",
    "         description='Num items:',\n",
    "         min=3,\n",
    "         max=6,\n",
    "         step=1)\n",
    "\n",
    "my_result = interactive(get_closest_prods, query=(\"Red high heels shoes\"), \\\n",
    "                        topk=topk_slider)\n",
    "display(my_result)\n",
    "\n",
    "num_results_slider = widgets.FloatSlider(\n",
    "         value=1,\n",
    "         description='Samples:',\n",
    "         min=1,\n",
    "         max=10,\n",
    "         step=1)\n",
    "\n",
    "temp_slider = widgets.FloatSlider(\n",
    "         value=0.4,\n",
    "         description='Temperature:',\n",
    "         min=0.1,\n",
    "         max=1.0,\n",
    "         step=0.1,)\n",
    "\n",
    "query_with_prompt = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description='Query with text prompt',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "content_based = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description='Content based',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "user_based = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description='User based',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "button = widgets.Button(description='Generate!')\n",
    "button100 = widgets.Button(description='Generate gifs !')\n",
    "out = widgets.Output()\n",
    "\n",
    "def on_button_clicked(_):\n",
    "    with out:\n",
    "        clear_output()\n",
    "        generate_from_ui(query_with_prompt.value, content_based.value, user_based.value, \\\n",
    "                                  int(num_results_slider.value), temp_slider.value)\n",
    "\n",
    "# linking button and function together using a button's method\n",
    "button.on_click(on_button_clicked)\n",
    "\n",
    "# displaying button and its output together\n",
    "widgets.VBox([content_based, user_based, num_results_slider, temp_slider, button, out])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b670dd8-4adc-45ae-b9c7-298117e477c3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "explore",
   "language": "python",
   "name": "venv"
  },
  "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}