optimized_colab.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "This colab requires you to have model.ckpt on your google drive (or you can download it on the next step)"
   ],
   "metadata": {
    "collapsed": false,
    "id": "f20_rzsJc1zU"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "id": "JtgQHPFgc1zZ",
    "cellView": "form"
   },
   "outputs": [],
   "source": [
    "#@title Install requirements\n",
    "!git clone https://github.com/neonsecret/stable-diffusion.git\n",
    "%cd /content/stable-diffusion\n",
    "!pip install gradio albumentations diffusers opencv-python pudb invisible-watermark imageio imageio-ffmpeg pytorch-lightning omegaconf test-tube streamlit einops torch-fidelity transformers torchmetrics kornia lpips yapf tb-nightly\n",
    "# !pip install taming\n",
    "!pip install -e git+https://github.com/openai/CLIP.git@main#egg=clip\n",
    "!pip install git+https://github.com/crowsonkb/k-diffusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "#@title Install requirements step 2\n",
    "%cd /content/stable-diffusion\n",
    "!pip install -e .\n",
    "!pip install Pillow==8.4.0 taming-transformers-rom1504\n",
    "!pip install --upgrade pytorch-lightning"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "id": "ksvqjvync1zb",
    "cellView": "form"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "#@title Mount Google Drive\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ],
   "metadata": {
    "id": "RJeVH0KtdVBh",
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "34e7ed19-71b6-4a0b-e80b-c8db1e3221e8"
   },
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Mounted at /content/drive\n"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "#@markdown # Installing xformers\n",
    "%%capture\n",
    "from IPython.display import HTML\n",
    "from subprocess import getoutput\n",
    "import os\n",
    "\n",
    "!pip install -e git+https://github.com/facebookresearch/xformers.git#egg=xformers\n",
    "\n",
    "if not os.path.exists('/content/gdrive/MyDrive/sd/stable-diffusion-webui/xformers'):\n",
    "  %cd /content/gdrive/MyDrive/sd/stable-diffusion/src\n",
    "  !git clone https://github.com/facebookresearch/xformers\n",
    "  !cp -R '/content/gdrive/MyDrive/sd/stable-diffusion/src/k-diffusion/k_diffusion' '/content/gdrive/MyDrive/sd/stable-diffusion-webui/'\n",
    "  !cp -R '/content/gdrive/MyDrive/sd/stable-diffusion/src/xformers/xformers' '/content/gdrive/MyDrive/sd/stable-diffusion-webui/'\n",
    "  !cp -R '/content/gdrive/MyDrive/sd/stable-diffusion/ldm' '/content/gdrive/MyDrive/sd/stable-diffusion-webui/'\n",
    "\n",
    "\n",
    "s = getoutput('nvidia-smi')\n",
    "if 'T4' in s:\n",
    "  gpu = 'T4'\n",
    "elif 'P100' in s:\n",
    "  gpu = 'P100'\n",
    "\n",
    "if gpu== 'T4':\n",
    "  %cd /content/\n",
    "  !git clone https://github.com/TheLastBen/fast-stable-diffusion\n",
    "  %cd /content/fast-stable-diffusion/precompiled\n",
    "  !mv /content/fast-stable-diffusion/precompiled/_C_flashattention.1 /content/fast-stable-diffusion/precompiled/_C_flashattention.7z.001\n",
    "  !mv /content/fast-stable-diffusion/precompiled/_C_flashattention.2 /content/fast-stable-diffusion/precompiled/_C_flashattention.7z.002\n",
    "  !7z x /content/fast-stable-diffusion/precompiled/_C_flashattention.7z.001\n",
    "  !mv -f /content/fast-stable-diffusion/precompiled/_C_flashattention.so /content/gdrive/MyDrive/sd/stable-diffusion-webui/xformers\n",
    "  !mv -f /content/fast-stable-diffusion/precompiled/_C.so /content/gdrive/MyDrive/sd/stable-diffusion-webui/xformers\n",
    "\n",
    "elif gpu== 'P100':\n",
    "  %cd /content/\n",
    "  !git clone https://github.com/TheLastBen/fast-stable-diffusion\n",
    "  %cd /content/fast-stable-diffusion/precompiled\n",
    "  !mv /content/fast-stable-diffusion/precompiled/_C_flashattention-p100.1 /content/fast-stable-diffusion/precompiled/_C_flashattention.7z.001\n",
    "  !mv /content/fast-stable-diffusion/precompiled/_C_flashattention-p100.2 /content/fast-stable-diffusion/precompiled/_C_flashattention.7z.002\n",
    "  !7z x /content/fast-stable-diffusion/precompiled/_C_flashattention.7z.001\n",
    "  !mv -f /content/fast-stable-diffusion/precompiled/_C.flashattention.so /content/gdrive/MyDrive/sd/stable-diffusion-webui/xformers/_C_flashattention.so\n",
    "  !mv -f /content/fast-stable-diffusion/precompiled/_C-p100.so /content/gdrive/MyDrive/sd/stable-diffusion-webui/xformers/_C.so\n",
    "\n",
    "\n",
    "%cd /content/gdrive/MyDrive/sd/stable-diffusion-webui/ldm/modules\n"
   ],
   "metadata": {
    "id": "a---cT2rwUQj",
    "cellView": "form"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "#@title Restarting the runtime here! Once it's done, just proceed to the next step.\n",
    "import os\n",
    "os.kill(os.getpid(), 9)"
   ],
   "metadata": {
    "cellView": "form",
    "id": "bFiD-_qJuNNu"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vkgBVo5OEpqn",
    "cellView": "form"
   },
   "outputs": [],
   "source": [
    "#@markdown # Load the stable-diffusion model\n",
    "\n",
    "#@markdown **Download the model if it isn't already in the 'models_path' folder**\n",
    "\n",
    "#@markdown To download the model, you need to have accepted the terms [HERE](https://huggingface.co/CompVis/stable-diffusion-v1-4)\n",
    "#@markdown and have copied a token from [HERE](https://huggingface.co/settings/tokens)\n",
    "download_if_missing = True #@param {type:\"boolean\"}\n",
    "token = \"\" #@param {type:\"string\"}\n",
    "\n",
    "#@markdown **Google Drive Path Variables**\n",
    "mount_google_drive = True #@param {type:\"boolean\"}\n",
    "force_remount = False\n",
    "\n",
    "%cd /content/\n",
    "import os\n",
    "mount_success = True\n",
    "if mount_google_drive:\n",
    "    from google.colab import drive\n",
    "    try:\n",
    "        drive_path = \"/content/drive\"\n",
    "        drive.mount(drive_path,force_remount=force_remount)\n",
    "        models_path_gdrive = \"/content/drive/MyDrive/\" #@param {type:\"string\"}\n",
    "        output_path_gdrive = \"/content/drive/MyDrive/outputs\" #@param {type:\"string\"}\n",
    "        models_path = models_path_gdrive\n",
    "        output_path = output_path_gdrive\n",
    "    except:\n",
    "        print(\"...error mounting drive or with drive path variables\")\n",
    "        print(\"...reverting to default path variables\")\n",
    "        mount_success = False\n",
    "        output_path = \"/content/outputs\"\n",
    "\n",
    "os.makedirs(models_path, exist_ok=True)\n",
    "os.makedirs(output_path, exist_ok=True)\n",
    "\n",
    "if download_if_missing:\n",
    "    if not mount_success:\n",
    "        print(\"Downloading model to \" + models_path + \" due to gdrive mount error\")\n",
    "    if token == \"\":\n",
    "        print(\"No token provided. Assuming model is already in \" + models_path)\n",
    "    elif not os.path.exists(models_path + '/sd-v1-4.ckpt'):\n",
    "        !git lfs install --system --skip-repo\n",
    "        !mkdir sd-model\n",
    "        %cd /content/sd-model/\n",
    "        !git init\n",
    "        !git remote add -f origin \"https://USER:{token}@huggingface.co/CompVis/stable-diffusion-v-1-4-original\"\n",
    "        !git config core.sparsecheckout true\n",
    "        !echo \"sd-v1-4.ckpt\" > .git/info/sparse-checkout\n",
    "        !git pull origin main\n",
    "        !mv '/content/sd-model/sd-v1-4.ckpt' '{models_path}/'\n",
    "    else:\n",
    "        print(\"Model already downloaded, moving to next step\")\n",
    "\n",
    "print(f\"models_path: {models_path}\")\n",
    "print(f\"output_path: {output_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "#@title Load the models\n",
    "%cd /content/stable-diffusion\n",
    "import gc\n",
    "gc.collect()\n",
    "import os\n",
    "import sys\n",
    "import git\n",
    "if not os.path.exists(\"CodeFormer/\"):\n",
    "    print(\"Installing CodeFormer..\")\n",
    "    git.Repo.clone_from(\"https://github.com/sczhou/CodeFormer/\", \"CodeFormer\")\n",
    "    os.chdir(\"CodeFormer\")\n",
    "    os.system(\"python basicsr/setup.py develop\")\n",
    "    os.chdir(\"..\")\n",
    "    print(\"Installation successful\")\n",
    "%cd /content/stable-diffusion\n",
    "sys.path.append('CodeFormer/')\n",
    "sys.path.append('../CodeFormer/')\n",
    "sys.path.append('optimizedSD/')\n",
    "\n",
    "import argparse\n",
    "import asyncio\n",
    "import logging\n",
    "import mimetypes\n",
    "import re\n",
    "import time\n",
    "from contextlib import nullcontext\n",
    "from itertools import islice\n",
    "from random import randint\n",
    "\n",
    "import gradio as gr\n",
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image\n",
    "from einops import rearrange, repeat\n",
    "from omegaconf import OmegaConf\n",
    "from pytorch_lightning import seed_everything\n",
    "from torch import autocast\n",
    "from torchvision.utils import make_grid\n",
    "from tqdm import tqdm, trange\n",
    "from transformers import logging as transformers_logging\n",
    "\n",
    "from ldm.util import instantiate_from_config\n",
    "from optimUtils import split_weighted_subprompts\n",
    "\n",
    "from basicsr.utils import img2tensor, tensor2img\n",
    "from basicsr.utils.download_util import load_file_from_url\n",
    "from facelib.utils.face_restoration_helper import FaceRestoreHelper\n",
    "from basicsr.archs.rrdbnet_arch import RRDBNet\n",
    "from basicsr.utils.realesrgan_utils import RealESRGANer\n",
    "\n",
    "from basicsr.utils.registry import ARCH_REGISTRY\n",
    "from torchvision.transforms.functional import normalize\n",
    "import cv2\n",
    "\n",
    "transformers_logging.set_verbosity_error()\n",
    "\n",
    "mimetypes.init()\n",
    "mimetypes.add_type(\"application/javascript\", \".js\")\n",
    "\n",
    "def load_model_from_config(ckpt, verbose=False):\n",
    "    print(f\"Loading model from {ckpt}\")\n",
    "    pl_sd = torch.load(ckpt, map_location=\"cpu\")\n",
    "    if \"global_step\" in pl_sd:\n",
    "        print(f\"Global Step: {pl_sd['global_step']}\")\n",
    "    sd = pl_sd[\"state_dict\"]\n",
    "    return sd\n",
    "\n",
    "def chunk(it, size):\n",
    "    it = iter(it)\n",
    "    return iter(lambda: tuple(islice(it, size)), ())\n",
    "\n",
    "\n",
    "def load_img(image, h0, w0):\n",
    "    image = image.convert(\"RGB\")\n",
    "    w, h = image.size\n",
    "    print(f\"loaded input image of size ({w}, {h})\")\n",
    "    if h0 is not None and w0 is not None:\n",
    "        h, w = h0, w0\n",
    "\n",
    "    w, h = map(lambda x: x - x % 64, (w, h))  # resize to integer multiple of 32\n",
    "\n",
    "    print(f\"New image size ({w}, {h})\")\n",
    "    image = image.resize((w, h), resample=Image.LANCZOS)\n",
    "    image = np.array(image).astype(np.float32) / 255.0\n",
    "    image = image[None].transpose(0, 3, 1, 2)\n",
    "    image = torch.from_numpy(image)\n",
    "    return 2.0 * image - 1.0\n",
    "\n",
    "\n",
    "def load_mask(mask, h0, w0, newH, newW, invert=False):\n",
    "    image = mask.convert(\"RGB\")\n",
    "    w, h = image.size\n",
    "    print(f\"loaded input mask of size ({w}, {h})\")\n",
    "    if h0 is not None and w0 is not None:\n",
    "        h, w = h0, w0\n",
    "\n",
    "    w, h = map(lambda x: x - x % 64, (w, h))  # resize to integer multiple of 32\n",
    "\n",
    "    print(f\"New mask size ({w}, {h})\")\n",
    "    image = image.resize((newW, newH), resample=Image.LANCZOS)\n",
    "    # image = image.resize((64, 64), resample=Image.LANCZOS)\n",
    "    image = np.array(image)\n",
    "\n",
    "    if invert:\n",
    "        print(\"inverted\")\n",
    "        where_0, where_1 = np.where(image == 0), np.where(image == 255)\n",
    "        image[where_0], image[where_1] = 255, 0\n",
    "    image = image.astype(np.float32) / 255.0\n",
    "    image = image[None].transpose(0, 3, 1, 2)\n",
    "    image = torch.from_numpy(image)\n",
    "    return image\n",
    "\n",
    "\n",
    "async def get_logs():\n",
    "    return \"\\n\".join([x for x in open(\"log.txt\", \"r\", encoding=\"utf8\").readlines()] +\n",
    "                     [y for y in open(\"tqdm.txt\", \"r\", encoding=\"utf8\").readlines()])\n",
    "\n",
    "\n",
    "async def get_nvidia_smi():\n",
    "    proc = await asyncio.create_subprocess_shell('nvidia-smi', stdout=asyncio.subprocess.PIPE)\n",
    "    stdout, stderr = await proc.communicate()\n",
    "    return str(stdout)\n",
    "\n",
    "def toImgOpenCV(imgPIL):  # Conver imgPIL to imgOpenCV\n",
    "    i = np.array(imgPIL)  # After mapping from PIL to numpy : [R,G,B,A]\n",
    "    # numpy Image Channel system: [B,G,R,A]\n",
    "    red = i[:, :, 0].copy()\n",
    "    i[:, :, 0] = i[:, :, 2].copy()\n",
    "    i[:, :, 2] = red\n",
    "    return i\n",
    "\n",
    "\n",
    "def generate_img2img(\n",
    "        image,\n",
    "        negative_prompt,\n",
    "        prompt,\n",
    "        strength,\n",
    "        ddim_steps,\n",
    "        n_iter,\n",
    "        batch_size,\n",
    "        Width,\n",
    "        Height,\n",
    "        scale,\n",
    "        ddim_eta,\n",
    "        unet_bs,\n",
    "        device,\n",
    "        seed,\n",
    "        outdir,\n",
    "        img_format,\n",
    "        turbo,\n",
    "        full_precision,\n",
    "        sampler,\n",
    "        speed_mp,\n",
    "):\n",
    "    logging.info(f\"prompt: {prompt}, W: {Width}, H: {Height}\")\n",
    "    try:\n",
    "        init_image = load_img(image['image'], Height, Width).to(device)\n",
    "    except:\n",
    "        init_image = load_img(image, Height, Width).to(device)\n",
    "    model.unet_bs = unet_bs\n",
    "    model.turbo = turbo\n",
    "    model.cdevice = device\n",
    "    modelCS.cond_stage_model.device = device\n",
    "\n",
    "    try:\n",
    "        seed = int(seed)\n",
    "    except:\n",
    "        seed = randint(0, 1000000)\n",
    "\n",
    "    if device != \"cpu\" and not full_precision:\n",
    "        model.half()\n",
    "        modelCS.half()\n",
    "        modelFS.half()\n",
    "        init_image = init_image.half()\n",
    "\n",
    "    tic = time.time()\n",
    "    os.makedirs(outdir, exist_ok=True)\n",
    "    outpath = outdir\n",
    "    sample_path = os.path.join(outpath, \"_\".join(re.split(\":| \", prompt)))[:150]\n",
    "    os.makedirs(sample_path, exist_ok=True)\n",
    "    base_count = len(os.listdir(sample_path))\n",
    "\n",
    "    # n_rows = opt.n_rows if opt.n_rows > 0 else batch_size\n",
    "    assert prompt is not None\n",
    "    data = [batch_size * [prompt]]\n",
    "\n",
    "    modelFS.to(device)\n",
    "\n",
    "    init_image = repeat(init_image, \"1 ... -> b ...\", b=batch_size)\n",
    "    init_latent = modelFS.get_first_stage_encoding(modelFS.encode_first_stage(init_image))  # move to latent space\n",
    "    try:\n",
    "        image['mask']\n",
    "        use_mask = True\n",
    "    except:\n",
    "        use_mask = False\n",
    "    if use_mask:\n",
    "        mask = load_mask(image['mask'], Height, Width, init_latent.shape[2], init_latent.shape[3], True).to(device)\n",
    "        mask = mask[0][0].unsqueeze(0).repeat(4, 1, 1).unsqueeze(0)\n",
    "        mask = repeat(mask, '1 ... -> b ...', b=batch_size)\n",
    "        if device != \"cpu\" and not full_precision:\n",
    "            mask = mask.half().to(device)\n",
    "\n",
    "    if device != \"cpu\":\n",
    "        mem = torch.cuda.memory_allocated() / 1e6\n",
    "        modelFS.to(\"cpu\")\n",
    "        while torch.cuda.memory_allocated() / 1e6 >= mem:\n",
    "            time.sleep(1)\n",
    "\n",
    "    assert 0.0 <= strength <= 1.0, \"can only work with strength in [0.0, 1.0]\"\n",
    "    t_enc = int(strength * ddim_steps)\n",
    "    print(f\"target t_enc is {t_enc} steps\")\n",
    "\n",
    "    if not full_precision and device != \"cpu\":\n",
    "        precision_scope = autocast\n",
    "    else:\n",
    "        precision_scope = nullcontext\n",
    "\n",
    "    all_samples = []\n",
    "    seeds = \"\"\n",
    "    with torch.no_grad():\n",
    "        for _ in trange(n_iter, desc=\"Sampling\"):\n",
    "            for prompts in tqdm(data, desc=\"data\"):\n",
    "                with precision_scope(\"cuda\"):\n",
    "                    modelCS.to(device)\n",
    "                    uc = None\n",
    "                    if scale != 1.0:\n",
    "                        uc = modelCS.get_learned_conditioning(\n",
    "                            batch_size * [negative_prompt if negative_prompt is not None else \"\"])\n",
    "                    if isinstance(prompts, tuple):\n",
    "                        prompts = list(prompts)\n",
    "\n",
    "                    subprompts, weights = split_weighted_subprompts(prompts[0])\n",
    "                    if len(subprompts) > 1:\n",
    "                        c = torch.zeros_like(uc)\n",
    "                        totalWeight = sum(weights)\n",
    "                        # normalize each \"sub prompt\" and add it\n",
    "                        for i in range(len(subprompts)):\n",
    "                            weight = weights[i]\n",
    "                            # if not skip_normalize:\n",
    "                            weight = weight / totalWeight\n",
    "                            c = torch.add(c, modelCS.get_learned_conditioning(subprompts[i]), alpha=weight)\n",
    "                    else:\n",
    "                        c = modelCS.get_learned_conditioning(prompts)\n",
    "\n",
    "                    if device != \"cpu\":\n",
    "                        mem = torch.cuda.memory_allocated() / 1e6\n",
    "                        modelCS.to(\"cpu\")\n",
    "                        while torch.cuda.memory_allocated() / 1e6 >= mem:\n",
    "                            time.sleep(1)\n",
    "\n",
    "                    # encode (scaled latent)\n",
    "                    z_enc = model.stochastic_encode(\n",
    "                        init_latent, torch.tensor([t_enc] * batch_size).to(device), seed, ddim_eta, ddim_steps\n",
    "                    )\n",
    "                    # decode it\n",
    "                    samples_ddim = model.sample(\n",
    "                        t_enc,\n",
    "                        c,\n",
    "                        z_enc,\n",
    "                        unconditional_guidance_scale=scale,\n",
    "                        unconditional_conditioning=uc,\n",
    "                        sampler=sampler,\n",
    "                        speed_mp=speed_mp,\n",
    "                        batch_size=batch_size,\n",
    "                        x_T=init_latent if use_mask else None,\n",
    "                        mask=mask if use_mask else None\n",
    "                    )\n",
    "\n",
    "                    modelFS.to(device)\n",
    "                    print(\"saving images\")\n",
    "                    for i in range(batch_size):\n",
    "                        x_samples_ddim = modelFS.decode_first_stage(samples_ddim[i].unsqueeze(0))\n",
    "                        x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)\n",
    "                        all_samples.append(x_sample.to(\"cpu\"))\n",
    "                        x_sample = 255.0 * rearrange(x_sample[0].cpu().numpy(), \"c h w -> h w c\")\n",
    "                        Image.fromarray(x_sample.astype(np.uint8)).save(\n",
    "                            os.path.join(sample_path, \"seed_\" + str(seed) + \"_\" + f\"{base_count:05}.{img_format}\")\n",
    "                        )\n",
    "                        seeds += str(seed) + \",\"\n",
    "                        seed += 1\n",
    "                        base_count += 1\n",
    "\n",
    "                    if device != \"cpu\":\n",
    "                        mem = torch.cuda.memory_allocated() / 1e6\n",
    "                        modelFS.to(\"cpu\")\n",
    "                        while torch.cuda.memory_allocated() / 1e6 >= mem:\n",
    "                            time.sleep(1)\n",
    "\n",
    "                    del samples_ddim\n",
    "                    del x_sample\n",
    "                    del x_samples_ddim\n",
    "                    print(\"memory_final = \", torch.cuda.memory_allocated() / 1e6)\n",
    "\n",
    "    toc = time.time()\n",
    "\n",
    "    time_taken = (toc - tic) / 60.0\n",
    "    grid = torch.cat(all_samples, 0)\n",
    "    grid = make_grid(grid, nrow=n_iter)\n",
    "    grid = 255.0 * rearrange(grid, \"c h w -> h w c\").cpu().numpy()\n",
    "\n",
    "    txt = (\n",
    "            \"Samples finished in \"\n",
    "            + str(round(time_taken, 3))\n",
    "            + \" minutes and exported to \\n\"\n",
    "            + sample_path\n",
    "            + \"\\nSeeds used = \"\n",
    "            + seeds[:-1]\n",
    "    )\n",
    "    return Image.fromarray(grid.astype(np.uint8)), txt\n",
    "\n",
    "\n",
    "def upscale2x(img):\n",
    "    img = Image.fromarray(upsampler.enhance(img, outscale=2)[0])\n",
    "    return img, f\"Upscaled to resolution: {img.size}\"\n",
    "\n",
    "\n",
    "def face_restore(img):\n",
    "    only_center_face = False\n",
    "    draw_box = False\n",
    "    codeformer_fidelity = 0.5\n",
    "    upscale = 2\n",
    "    face_upsample = True\n",
    "    detection_model = \"retinaface_resnet50\"\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    face_helper = FaceRestoreHelper(\n",
    "        True,\n",
    "        face_size=512,\n",
    "        crop_ratio=(1, 1),\n",
    "        det_model=detection_model,\n",
    "        save_ext=\"png\",\n",
    "        use_parse=True,\n",
    "        device=device,\n",
    "    )\n",
    "    codeformer_net.to(device)\n",
    "    bg_upsampler = upsampler\n",
    "    face_upsampler = upsampler\n",
    "    face_helper.read_image(img)\n",
    "    num_det_faces = face_helper.get_face_landmarks_5(\n",
    "        only_center_face=only_center_face, resize=640, eye_dist_threshold=5\n",
    "    )\n",
    "    print(f\"\\tdetect {num_det_faces} faces\")\n",
    "    # align and warp each face\n",
    "    face_helper.align_warp_face()\n",
    "\n",
    "    for idx, cropped_face in enumerate(face_helper.cropped_faces):\n",
    "        # prepare data\n",
    "        cropped_face_t = img2tensor(\n",
    "            cropped_face / 255.0, bgr2rgb=True, float32=True\n",
    "        )\n",
    "        normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)\n",
    "        cropped_face_t = cropped_face_t.unsqueeze(0).to(device)\n",
    "\n",
    "        try:\n",
    "            with torch.no_grad():\n",
    "                output = codeformer_net(\n",
    "                    cropped_face_t, w=codeformer_fidelity, adain=True\n",
    "                )[0]\n",
    "                restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))\n",
    "            del output\n",
    "            torch.cuda.empty_cache()\n",
    "        except Exception as error:\n",
    "            print(f\"\\tFailed inference for CodeFormer: {error}\")\n",
    "            restored_face = tensor2img(\n",
    "                cropped_face_t, rgb2bgr=True, min_max=(-1, 1)\n",
    "            )\n",
    "\n",
    "        restored_face = restored_face.astype(\"uint8\")\n",
    "        face_helper.add_restored_face(restored_face)\n",
    "\n",
    "    # paste_back\n",
    "    # upsample the background\n",
    "    if bg_upsampler is not None:\n",
    "        # Now only support RealESRGAN for upsampling background\n",
    "        bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]\n",
    "    else:\n",
    "        bg_img = None\n",
    "    face_helper.get_inverse_affine(None)\n",
    "    # paste each restored face to the input image\n",
    "    if face_upsample and face_upsampler is not None:\n",
    "        restored_img = face_helper.paste_faces_to_input_image(\n",
    "            upsample_img=bg_img,\n",
    "            draw_box=draw_box,\n",
    "            face_upsampler=face_upsampler,\n",
    "        )\n",
    "    else:\n",
    "        restored_img = face_helper.paste_faces_to_input_image(\n",
    "            upsample_img=bg_img, draw_box=draw_box\n",
    "        )\n",
    "    img = Image.fromarray(restored_img)\n",
    "    return img, f\"Fixed a face, new img size: {img.size}\"\n",
    "\n",
    "\n",
    "def generate_txt2img(\n",
    "        prompt,\n",
    "        negative_prompt,\n",
    "        ddim_steps,\n",
    "        n_iter,\n",
    "        batch_size,\n",
    "        Width,\n",
    "        Height,\n",
    "        scale,\n",
    "        ddim_eta,\n",
    "        unet_bs,\n",
    "        device,\n",
    "        seed,\n",
    "        outdir,\n",
    "        img_format,\n",
    "        turbo,\n",
    "        full_precision,\n",
    "        sampler,\n",
    "        speed_mp\n",
    "):\n",
    "    logging.info(f\"prompt: {prompt}, W: {Width}, H: {Height}\")\n",
    "    C = 4\n",
    "    f = 8\n",
    "    start_code = None\n",
    "    model.to(device)\n",
    "    model.unet_bs = unet_bs\n",
    "    model.turbo = turbo\n",
    "    model.cdevice = device\n",
    "    modelCS.cond_stage_model.device = device\n",
    "\n",
    "    if seed == \"\":\n",
    "        seed = randint(0, 1000000)\n",
    "    seed = int(seed)\n",
    "    seed_everything(seed)\n",
    "\n",
    "    if device != \"cpu\" and not full_precision:\n",
    "        model.half()\n",
    "        modelFS.half()\n",
    "        modelCS.half()\n",
    "\n",
    "    tic = time.time()\n",
    "    os.makedirs(outdir, exist_ok=True)\n",
    "    outpath = outdir\n",
    "    sample_path = os.path.join(outpath, \"_\".join(re.split(\":| \", prompt.replace(\"/\", \"\"))))[:150]\n",
    "    os.makedirs(sample_path, exist_ok=True)\n",
    "    base_count = len(os.listdir(sample_path))\n",
    "\n",
    "    # n_rows = opt.n_rows if opt.n_rows > 0 else batch_size\n",
    "    assert prompt is not None\n",
    "    data = [batch_size * [prompt]]\n",
    "\n",
    "    if device != \"cpu\" and not full_precision:\n",
    "        precision_scope = autocast\n",
    "    else:\n",
    "        precision_scope = nullcontext\n",
    "\n",
    "    seeds = \"\"\n",
    "    with torch.no_grad():\n",
    "        all_samples = list()\n",
    "        for _ in trange(n_iter, desc=\"Sampling\"):\n",
    "            for prompts in tqdm(data, desc=\"data\"):\n",
    "                with precision_scope(\"cuda\"):\n",
    "                    modelCS.to(device)\n",
    "                    uc = None\n",
    "                    if scale != 1.0:\n",
    "                        uc = modelCS.get_learned_conditioning(\n",
    "                            batch_size * [negative_prompt if negative_prompt is not None else \"\"])\n",
    "                    if isinstance(prompts, tuple):\n",
    "                        prompts = list(prompts)\n",
    "\n",
    "                    subprompts, weights = split_weighted_subprompts(prompts[0])\n",
    "                    if len(subprompts) > 1:\n",
    "                        c = torch.zeros_like(uc)\n",
    "                        totalWeight = sum(weights)\n",
    "                        # normalize each \"sub prompt\" and add it\n",
    "                        for i in range(len(subprompts)):\n",
    "                            weight = weights[i]\n",
    "                            # if not skip_normalize:\n",
    "                            weight = weight / totalWeight\n",
    "                            c = torch.add(c, modelCS.get_learned_conditioning(subprompts[i]), alpha=weight)\n",
    "                    else:\n",
    "                        c = modelCS.get_learned_conditioning(prompts)\n",
    "\n",
    "                    shape = [batch_size, C, Height // f, Width // f]\n",
    "\n",
    "                    if device != \"cpu\":\n",
    "                        mem = torch.cuda.memory_allocated() / 1e6\n",
    "                        modelCS.to(\"cpu\")\n",
    "                        while torch.cuda.memory_allocated() / 1e6 >= mem:\n",
    "                            time.sleep(1)\n",
    "                    samples_ddim = model.sample(\n",
    "                        S=ddim_steps,\n",
    "                        conditioning=c,\n",
    "                        seed=seed,\n",
    "                        shape=shape,\n",
    "                        verbose=False,\n",
    "                        unconditional_guidance_scale=scale,\n",
    "                        unconditional_conditioning=uc,\n",
    "                        eta=ddim_eta,\n",
    "                        x_T=start_code,\n",
    "                        sampler=sampler,\n",
    "                        speed_mp=speed_mp\n",
    "                    )\n",
    "\n",
    "                    modelFS.to(device)\n",
    "                    model.cpu()\n",
    "                    logging.info(\"saving images\")\n",
    "                    for i in range(batch_size):\n",
    "                        x_samples_ddim = modelFS.decode_first_stage(samples_ddim[i].unsqueeze(0))\n",
    "                        x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)\n",
    "                        all_samples.append(x_sample.to(\"cpu\"))\n",
    "                        x_sample = 255.0 * rearrange(x_sample[0].cpu().numpy(), \"c h w -> h w c\")\n",
    "                        Image.fromarray(x_sample.astype(np.uint8)).save(\n",
    "                            os.path.join(sample_path, \"seed_\" + str(seed) + \"_\" + f\"{base_count:05}.{img_format}\")\n",
    "                        )\n",
    "                        seeds += str(seed) + \",\"\n",
    "                        seed += 1\n",
    "                        base_count += 1\n",
    "\n",
    "                    if device != \"cpu\":\n",
    "                        mem = torch.cuda.memory_allocated() / 1e6\n",
    "                        modelFS.to(\"cpu\")\n",
    "                        while torch.cuda.memory_allocated() / 1e6 >= mem:\n",
    "                            time.sleep(1)\n",
    "\n",
    "                    del samples_ddim\n",
    "                    del x_sample\n",
    "                    del x_samples_ddim\n",
    "                    logging.info(str(\"memory_final = \" + str(torch.cuda.memory_allocated() / 1e6)))\n",
    "\n",
    "    toc = time.time()\n",
    "\n",
    "    time_taken = (toc - tic) / 60.0\n",
    "    grid = torch.cat(all_samples, 0)\n",
    "    grid = make_grid(grid, nrow=n_iter)\n",
    "    grid = 255.0 * rearrange(grid, \"c h w -> h w c\").cpu().numpy()\n",
    "    txt = (\n",
    "            \"Samples finished in \"\n",
    "            + str(round(time_taken, 3))\n",
    "            + \" minutes and exported to \"\n",
    "            + sample_path\n",
    "            + \"\\nSeeds used = \"\n",
    "            + seeds[:-1]\n",
    "    )\n",
    "    return Image.fromarray(grid.astype(np.uint8)), txt\n",
    "\n",
    "\n",
    "def generate_img2img_interp(\n",
    "        image,\n",
    "        prompt,\n",
    "        strength,\n",
    "        ddim_steps,\n",
    "        n_iter,\n",
    "        batch_size,\n",
    "        Width,\n",
    "        Height,\n",
    "        scale,\n",
    "        ddim_eta,\n",
    "        unet_bs,\n",
    "        device,\n",
    "        seed,\n",
    "        outdir,\n",
    "        img_format,\n",
    "        turbo,\n",
    "        full_precision,\n",
    "        sampler,\n",
    "        speed_mp,\n",
    "        n_interpolate_samples\n",
    "):\n",
    "    logging.info(f\"prompt: {prompt}, W: {Width}, H: {Height}\")\n",
    "    try:\n",
    "        init_image = load_img(image['image'], Height, Width).to(device)\n",
    "    except:\n",
    "        init_image = load_img(image, Height, Width).to(device)\n",
    "    model.unet_bs = unet_bs\n",
    "    model.turbo = turbo\n",
    "    model.cdevice = device\n",
    "    modelCS.cond_stage_model.device = device\n",
    "\n",
    "    try:\n",
    "        seed = int(seed)\n",
    "    except:\n",
    "        seed = randint(0, 1000000)\n",
    "\n",
    "    if device != \"cpu\" and not full_precision:\n",
    "        model.half()\n",
    "        modelCS.half()\n",
    "        modelFS.half()\n",
    "        init_image = init_image.half()\n",
    "\n",
    "    os.makedirs(outdir, exist_ok=True)\n",
    "    outpath = outdir\n",
    "    sample_path = os.path.join(outpath, \"_\".join(re.split(\":| \", prompt)))[:150]\n",
    "    os.makedirs(sample_path, exist_ok=True)\n",
    "    base_count = len(os.listdir(sample_path))\n",
    "\n",
    "    # n_rows = opt.n_rows if opt.n_rows > 0 else batch_size\n",
    "    assert prompt is not None\n",
    "    data = [batch_size * [prompt]]\n",
    "\n",
    "    modelFS.to(device)\n",
    "\n",
    "    init_image = repeat(init_image, \"1 ... -> b ...\", b=batch_size)\n",
    "    init_latent = modelFS.get_first_stage_encoding(modelFS.encode_first_stage(init_image))  # move to latent space\n",
    "    if device != \"cpu\":\n",
    "        mem = torch.cuda.memory_allocated() / 1e6\n",
    "        modelFS.to(\"cpu\")\n",
    "        while torch.cuda.memory_allocated() / 1e6 >= mem:\n",
    "            time.sleep(1)\n",
    "\n",
    "    assert 0.0 <= strength <= 1.0, \"can only work with strength in [0.0, 1.0]\"\n",
    "    t_enc = int(strength * ddim_steps)\n",
    "    print(f\"target t_enc is {t_enc} steps\")\n",
    "\n",
    "    if not full_precision and device != \"cpu\":\n",
    "        precision_scope = autocast\n",
    "    else:\n",
    "        precision_scope = nullcontext\n",
    "\n",
    "    seeds = \"\"\n",
    "    with torch.no_grad():\n",
    "        for _ in trange(n_iter, desc=\"Sampling\"):\n",
    "            for prompts in tqdm(data, desc=\"data\"):\n",
    "                with precision_scope(\"cuda\"):\n",
    "                    modelCS.to(device)\n",
    "                    uc = None\n",
    "                    if scale != 1.0:\n",
    "                        uc = modelCS.get_learned_conditioning(batch_size * [\"\"])\n",
    "                    if isinstance(prompts, tuple):\n",
    "                        prompts = list(prompts)\n",
    "\n",
    "                    subprompts, weights = split_weighted_subprompts(prompts[0])\n",
    "                    if len(subprompts) > 1:\n",
    "                        c = torch.zeros_like(uc)\n",
    "                        totalWeight = sum(weights)\n",
    "                        # normalize each \"sub prompt\" and add it\n",
    "                        for i in range(len(subprompts)):\n",
    "                            weight = weights[i]\n",
    "                            # if not skip_normalize:\n",
    "                            weight = weight / totalWeight\n",
    "                            c = torch.add(c, modelCS.get_learned_conditioning(subprompts[i]), alpha=weight)\n",
    "                    else:\n",
    "                        c = modelCS.get_learned_conditioning(prompts)\n",
    "\n",
    "                    if device != \"cpu\":\n",
    "                        mem = torch.cuda.memory_allocated() / 1e6\n",
    "                        modelCS.to(\"cpu\")\n",
    "                        while torch.cuda.memory_allocated() / 1e6 >= mem:\n",
    "                            time.sleep(1)\n",
    "\n",
    "                    # encode (scaled latent)\n",
    "                    true_z_enc = model.stochastic_encode(\n",
    "                        init_latent, torch.tensor([t_enc] * batch_size).to(device), seed, ddim_eta, ddim_steps\n",
    "                    )\n",
    "                    # decode it\n",
    "                    samples_ddim = model.sample(\n",
    "                        t_enc,\n",
    "                        c,\n",
    "                        true_z_enc,\n",
    "                        unconditional_guidance_scale=scale,\n",
    "                        unconditional_conditioning=uc,\n",
    "                        sampler=sampler,\n",
    "                        speed_mp=speed_mp\n",
    "                    )\n",
    "                    modelFS.to(device)\n",
    "                    print(\"saving frames\")\n",
    "                    all_time_samples = []\n",
    "                    for ij in range(n_interpolate_samples):\n",
    "                        temp_all_samples = []\n",
    "                        for i in range(batch_size):\n",
    "                            start0_sample = samples_ddim[i].unsqueeze(0)\n",
    "                            interp_sample = torch.lerp(init_latent, start0_sample, (ij / n_interpolate_samples))\n",
    "                            x_samples_ddim = modelFS.decode_first_stage(interp_sample)\n",
    "                            x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)\n",
    "                            temp_all_samples.append(x_sample.to(\"cpu\"))\n",
    "                            x_sample = 255.0 * rearrange(x_sample[0].cpu().numpy(), \"c h w -> h w c\")\n",
    "                            Image.fromarray(x_sample.astype(np.uint8)).save(\n",
    "                                os.path.join(sample_path, \"seed_\" + str(seed) + \"_\" + f\"{base_count:05}.{img_format}\")\n",
    "                            )\n",
    "                            seeds += str(seed) + \",\"\n",
    "                            base_count += 1\n",
    "                        grid = torch.cat(temp_all_samples, 0)\n",
    "                        grid = make_grid(grid, nrow=n_iter)\n",
    "                        grid = 255.0 * rearrange(grid, \"c h w -> h w c\").cpu().numpy()\n",
    "                        all_time_samples.append(Image.fromarray(grid.astype(np.uint8)))\n",
    "                    if device != \"cpu\":\n",
    "                        mem = torch.cuda.memory_allocated() / 1e6\n",
    "                        modelFS.to(\"cpu\")\n",
    "                        while torch.cuda.memory_allocated() / 1e6 >= mem:\n",
    "                            time.sleep(1)\n",
    "\n",
    "                    del samples_ddim\n",
    "                    del x_sample\n",
    "                    del x_samples_ddim\n",
    "                    print(\"memory_final = \", torch.cuda.memory_allocated() / 1e6)\n",
    "        # all_samples.append(all_time_samples)\n",
    "    all_time_samples = [toImgOpenCV(img) for img in all_time_samples]\n",
    "    out = cv2.VideoWriter(\"tempfile.mp4\", cv2.VideoWriter_fourcc(*'h264'), 15, (all_time_samples[0].shape[1], all_time_samples[0].shape[0]))\n",
    "    for img in all_time_samples:\n",
    "        out.write(img)\n",
    "    out.release()\n",
    "\n",
    "    return \"tempfile.mp4\", f\"yeah here's your video {Width}x{Height}\"\n",
    "\n",
    "\n",
    "def generate_double_triple(\n",
    "        prompt,\n",
    "        ddim_steps,\n",
    "        img2img_strength,\n",
    "        Width,\n",
    "        Height,\n",
    "        scale,\n",
    "        ddim_eta,\n",
    "        unet_bs,\n",
    "        device,\n",
    "        seed,\n",
    "        outdir,\n",
    "        img_format,\n",
    "        turbo,\n",
    "        full_precision,\n",
    "        sampler,\n",
    "        speed_mp,\n",
    "        upscale_reso\n",
    "):\n",
    "    C = 4\n",
    "    f = 8\n",
    "    start_code = None\n",
    "    model.unet_bs = unet_bs\n",
    "    model.turbo = turbo\n",
    "    model.cdevice = device\n",
    "    modelCS.cond_stage_model.device = device\n",
    "\n",
    "    if seed == \"\":\n",
    "        seed = randint(0, 1000000)\n",
    "    seed = int(seed)\n",
    "    seed_everything(seed)\n",
    "\n",
    "    if device != \"cpu\" and not full_precision:\n",
    "        model.half()\n",
    "        modelFS.half()\n",
    "        modelCS.half()\n",
    "\n",
    "    tic = time.time()\n",
    "    os.makedirs(outdir, exist_ok=True)\n",
    "    outpath = outdir\n",
    "    sample_path = os.path.join(outpath, \"_\".join(re.split(\":| \", prompt.replace(\"/\", \"\"))))[:150]\n",
    "    os.makedirs(sample_path, exist_ok=True)\n",
    "    base_count = len(os.listdir(sample_path))\n",
    "\n",
    "    # n_rows = opt.n_rows if opt.n_rows > 0 else batch_size\n",
    "    assert prompt is not None\n",
    "    data = [1 * [prompt]]\n",
    "\n",
    "    if device != \"cpu\" and not full_precision:\n",
    "        precision_scope = autocast\n",
    "    else:\n",
    "        precision_scope = nullcontext\n",
    "\n",
    "    seeds = \"\"\n",
    "    with torch.no_grad():\n",
    "        all_samples = list()\n",
    "        for prompts in tqdm(data, desc=\"data\"):\n",
    "            with precision_scope(\"cuda\"):\n",
    "                modelCS.to(device)\n",
    "                uc = None\n",
    "                if scale != 1.0:\n",
    "                    uc = modelCS.get_learned_conditioning(1 * [\"\"])\n",
    "                if isinstance(prompts, tuple):\n",
    "                    prompts = list(prompts)\n",
    "\n",
    "                subprompts, weights = split_weighted_subprompts(prompts[0])\n",
    "                if len(subprompts) > 1:\n",
    "                    c = torch.zeros_like(uc)\n",
    "                    totalWeight = sum(weights)\n",
    "                    # normalize each \"sub prompt\" and add it\n",
    "                    for i in range(len(subprompts)):\n",
    "                        weight = weights[i]\n",
    "                        # if not skip_normalize:\n",
    "                        weight = weight / totalWeight\n",
    "                        c = torch.add(c, modelCS.get_learned_conditioning(subprompts[i]), alpha=weight)\n",
    "                else:\n",
    "                    c = modelCS.get_learned_conditioning(prompts)\n",
    "\n",
    "                shape = [1, C, Height // f, Width // f]\n",
    "\n",
    "                if device != \"cpu\":\n",
    "                    mem = torch.cuda.memory_allocated() / 1e6\n",
    "                    modelCS.to(\"cpu\")\n",
    "                    while torch.cuda.memory_allocated() / 1e6 >= mem:\n",
    "                        time.sleep(1)\n",
    "\n",
    "                samples_ddim = model.sample(\n",
    "                    S=ddim_steps,\n",
    "                    conditioning=c,\n",
    "                    seed=seed,\n",
    "                    shape=shape,\n",
    "                    verbose=False,\n",
    "                    unconditional_guidance_scale=scale,\n",
    "                    unconditional_conditioning=uc,\n",
    "                    eta=ddim_eta,\n",
    "                    x_T=start_code,\n",
    "                    sampler=sampler,\n",
    "                    speed_mp=speed_mp\n",
    "                )\n",
    "\n",
    "                modelFS.to(device)\n",
    "\n",
    "                x_samples_ddim = modelFS.decode_first_stage(samples_ddim[0].unsqueeze(0))\n",
    "                x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)\n",
    "                x_sample = 255.0 * rearrange(x_sample[0].cpu().numpy(), \"c h w -> h w c\")\n",
    "                Image.fromarray(x_sample.astype(np.uint8)).save(\n",
    "                    os.path.join(sample_path, \"seed_\" + str(seed) + \"_step1_\" + f\"{base_count:05}.{img_format}\")\n",
    "                )\n",
    "\n",
    "                ### STEP 2\n",
    "\n",
    "                init_image = repeat(\n",
    "                    load_img(Image.fromarray(x_sample.astype(np.uint8)), Height * 2, Width * 2).to(device),\n",
    "                    \"1 ... -> b ...\", b=1)\n",
    "                init_latent = modelFS.get_first_stage_encoding(modelFS.encode_first_stage(init_image))\n",
    "\n",
    "                modelFS.cpu()\n",
    "                model.to(device)\n",
    "\n",
    "                z_enc = model.stochastic_encode(\n",
    "                    init_latent, torch.tensor([int(img2img_strength * ddim_steps)]).to(device), seed, ddim_eta,\n",
    "                    ddim_steps\n",
    "                ).to(device)\n",
    "\n",
    "                samples_ddim = model.sample(\n",
    "                    int(img2img_strength * ddim_steps // 2),\n",
    "                    c,\n",
    "                    z_enc,\n",
    "                    unconditional_guidance_scale=scale,\n",
    "                    unconditional_conditioning=uc,\n",
    "                    sampler=\"ddim\",\n",
    "                    speed_mp=speed_mp\n",
    "                )\n",
    "\n",
    "                modelFS.to(device)\n",
    "                model.cpu()\n",
    "                modelCS.to(\"cpu\")\n",
    "                torch.cuda.empty_cache()\n",
    "\n",
    "                x_samples_ddim = modelFS.decode_first_stage(samples_ddim[0].unsqueeze(0))\n",
    "                x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)\n",
    "                if upscale_reso < 3:\n",
    "                    all_samples.append(x_sample.cpu())\n",
    "                x_sample = 255.0 * rearrange(x_sample[0].cpu().numpy(), \"c h w -> h w c\")\n",
    "                Image.fromarray(x_sample.astype(np.uint8)).save(\n",
    "                    os.path.join(sample_path, \"seed_\" + str(seed) + \"_step2_\" + f\"{base_count:05}.{img_format}\")\n",
    "                )\n",
    "\n",
    "                ### STEP 3\n",
    "                if upscale_reso >= 3:\n",
    "                    init_image = repeat(\n",
    "                        load_img(Image.fromarray(x_sample.astype(np.uint8)), Height * 3, Width * 3).to(device),\n",
    "                        \"1 ... -> b ...\", b=1)\n",
    "                    init_latent = modelFS.get_first_stage_encoding(modelFS.encode_first_stage(init_image))\n",
    "\n",
    "                    modelFS.cpu()\n",
    "                    model.to(device)\n",
    "\n",
    "                    z_enc = model.stochastic_encode(\n",
    "                        init_latent, torch.tensor([int(img2img_strength * ddim_steps)]).to(device), seed, ddim_eta,\n",
    "                        ddim_steps\n",
    "                    ).to(device)\n",
    "\n",
    "                    samples_ddim = model.sample(\n",
    "                        int(img2img_strength * ddim_steps // 2),\n",
    "                        c,\n",
    "                        z_enc,\n",
    "                        unconditional_guidance_scale=scale,\n",
    "                        unconditional_conditioning=uc,\n",
    "                        sampler=\"ddim\",\n",
    "                        speed_mp=speed_mp\n",
    "                    )\n",
    "\n",
    "                    print(\"saving images\")\n",
    "                    model.cpu()\n",
    "                    modelFS.to(device)\n",
    "\n",
    "                    x_samples_ddim = modelFS.decode_first_stage(samples_ddim[0].unsqueeze(0))\n",
    "                    x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)\n",
    "                    all_samples.append(x_sample.to(\"cpu\"))\n",
    "                    x_sample = 255.0 * rearrange(x_sample[0].cpu().numpy(), \"c h w -> h w c\")\n",
    "                    Image.fromarray(x_sample.astype(np.uint8)).save(\n",
    "                        os.path.join(sample_path, \"seed_\" + str(seed) + \"_step3_\" + f\"{base_count:05}.{img_format}\")\n",
    "                    )\n",
    "\n",
    "                if device != \"cpu\":\n",
    "                    mem = torch.cuda.memory_allocated() / 1e6\n",
    "                    modelFS.to(\"cpu\")\n",
    "                    while torch.cuda.memory_allocated() / 1e6 >= mem:\n",
    "                        time.sleep(1)\n",
    "\n",
    "                del samples_ddim\n",
    "                del x_sample\n",
    "                del x_samples_ddim\n",
    "                print(\"memory_final = \", torch.cuda.memory_allocated() / 1e6)\n",
    "\n",
    "    toc = time.time()\n",
    "\n",
    "    time_taken = (toc - tic) / 60.0\n",
    "    grid = torch.cat(all_samples, 0)\n",
    "    grid = make_grid(grid, nrow=1)\n",
    "    grid = 255.0 * rearrange(grid, \"c h w -> h w c\").cpu().numpy()\n",
    "\n",
    "    txt = (\n",
    "            \"Samples finished in \"\n",
    "            + str(round(time_taken, 3))\n",
    "            + \" minutes and exported to \"\n",
    "            + sample_path\n",
    "            + \"\\nSeeds used = \"\n",
    "            + seeds[:-1]\n",
    "    )\n",
    "    return Image.fromarray(grid.astype(np.uint8)), txt\n",
    "\n",
    "\n",
    "def download_codeformer():\n",
    "    pretrain_model_url = {\n",
    "        'codeformer': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth',\n",
    "        'detection': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth',\n",
    "        'parsing': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth',\n",
    "        'realesrgan': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/RealESRGAN_x2plus.pth'\n",
    "    }\n",
    "    # download weights\n",
    "    if not os.path.exists(\"models/codeformer/\" + \"CodeFormer/codeformer.pth\"):\n",
    "        load_file_from_url(url=pretrain_model_url['codeformer'], model_dir=\"models/codeformer/\" + \"CodeFormer\",\n",
    "                           progress=True, file_name=None)\n",
    "    if not os.path.exists(\"models/codeformer/\" + \"facelib/detection_Resnet50_Final.pth\"):\n",
    "        load_file_from_url(url=pretrain_model_url['detection'], model_dir=\"models/codeformer/\" + \"facelib\",\n",
    "                           progress=True,\n",
    "                           file_name=None)\n",
    "    if not os.path.exists(\"models/codeformer/\" + \"facelib/parsing_parsenet.pth\"):\n",
    "        load_file_from_url(url=pretrain_model_url['parsing'], model_dir=\"models/codeformer/\" + \"facelib\", progress=True,\n",
    "                           file_name=None)\n",
    "    if not os.path.exists(\"models/codeformer/\" + \"realesrgan/RealESRGAN_x2plus.pth\"):\n",
    "        load_file_from_url(url=pretrain_model_url['realesrgan'], model_dir=\"models/codeformer/\" + \"realesrgan\",\n",
    "                           progress=True, file_name=None)\n",
    "\n",
    "\n",
    "# set enhancer with RealESRGAN\n",
    "def set_realesrgan():\n",
    "    half = True if torch.cuda.is_available() else False\n",
    "    model = RRDBNet(\n",
    "        num_in_ch=3,\n",
    "        num_out_ch=3,\n",
    "        num_feat=64,\n",
    "        num_block=23,\n",
    "        num_grow_ch=32,\n",
    "        scale=2,\n",
    "    )\n",
    "    upsampler = RealESRGANer(\n",
    "        scale=2,\n",
    "        model_path=\"models/codeformer/\" + \"realesrgan/RealESRGAN_x2plus.pth\",\n",
    "        model=model,\n",
    "        tile=400,\n",
    "        tile_pad=40,\n",
    "        pre_pad=0,\n",
    "        half=half,\n",
    "    )\n",
    "    return upsampler\n",
    "\n",
    "\n",
    "class TqdmLoggingHandler(logging.Handler):\n",
    "    def __init__(self, level=logging.NOTSET):\n",
    "        super().__init__(level)\n",
    "\n",
    "    def emit(self, record):\n",
    "        try:\n",
    "            msg = self.format(record)\n",
    "            tqdm.write(msg)\n",
    "            self.flush()\n",
    "        except Exception:\n",
    "            self.handleError(record)\n",
    "\n",
    "print(\"Downloading codeformer weights..\")\n",
    "download_codeformer()\n",
    "\n",
    "print(\"Loading realesr..\")\n",
    "upsampler = set_realesrgan()\n",
    "print(\"Loading codeformer..\")\n",
    "codeformer_net = ARCH_REGISTRY.get(\"CodeFormer\")(\n",
    "  dim_embd=512,\n",
    "  codebook_size=1024,\n",
    "  n_head=8,\n",
    "  n_layers=9,\n",
    "  connect_list=[\"32\", \"64\", \"128\", \"256\"],\n",
    ")\n",
    "checkpoint = torch.load(\"models/codeformer/\" + \"CodeFormer/codeformer.pth\")[\"params_ema\"]\n",
    "codeformer_net.load_state_dict(checkpoint)\n",
    "codeformer_net.eval().cuda()\n",
    "gc.collect()\n",
    "\n",
    "\n",
    "config = \"optimizedSD/v1-inference.yaml\"\n",
    "ckpt = '/content/drive/MyDrive/sd-v1-4.ckpt' #@param {type:\"string\"}\n",
    "sd = load_model_from_config(f\"{ckpt}\")\n",
    "li, lo = [], []\n",
    "for key, v_ in sd.items():\n",
    "    sp = key.split(\".\")\n",
    "    if (sp[0]) == \"model\":\n",
    "        if \"input_blocks\" in sp:\n",
    "            li.append(key)\n",
    "        elif \"middle_block\" in sp:\n",
    "            li.append(key)\n",
    "        elif \"time_embed\" in sp:\n",
    "            li.append(key)\n",
    "        else:\n",
    "            lo.append(key)\n",
    "for key in li:\n",
    "    sd[\"model1.\" + key[6:]] = sd.pop(key)\n",
    "for key in lo:\n",
    "    sd[\"model2.\" + key[6:]] = sd.pop(key)\n",
    "\n",
    "config = OmegaConf.load(f\"{config}\")\n",
    "\n",
    "model = instantiate_from_config(config.modelUNet)\n",
    "_, _ = model.load_state_dict(sd, strict=False)\n",
    "model.eval()\n",
    "\n",
    "modelCS = instantiate_from_config(config.modelCondStage)\n",
    "_, _ = modelCS.load_state_dict(sd, strict=False)\n",
    "modelCS.eval()\n",
    "\n",
    "modelFS = instantiate_from_config(config.modelFirstStage)\n",
    "_, _ = modelFS.load_state_dict(sd, strict=False)\n",
    "modelFS.eval()\n",
    "del sd\n",
    "gc.collect()"
   ],
   "metadata": {
    "id": "voKbhBMEe-yg",
    "cellView": "form"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Okay next run the ultimate interface"
   ],
   "metadata": {
    "id": "NyafrRNTF1Xe"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "#@title Ultimate Gradio Interface\n",
    "import gradio as gr\n",
    "demo = gr.Blocks()\n",
    "with demo:\n",
    "    with gr.Tab(\"txt2img\"):\n",
    "        with gr.Column():\n",
    "            gr.Markdown(\"# Generate images from text (neonsecret's adjustments)\")\n",
    "            with gr.Row():\n",
    "                with gr.Column():\n",
    "                    out_image = gr.Image(label=\"Output Image\")\n",
    "                    gen_res = gr.Text(label=\"Generation results\")\n",
    "                    b1 = gr.Button(\"Generate!\")\n",
    "                    b4 = gr.Button(\"Face correction\")\n",
    "                    b5 = gr.Button(\"Upscale 2x\")\n",
    "                with gr.Column():\n",
    "                    with gr.Box():\n",
    "                        b4.click(face_restore, inputs=[out_image], outputs=[out_image, gen_res])\n",
    "                        b5.click(upscale2x, inputs=[out_image], outputs=[out_image, gen_res])\n",
    "                        b1.click(generate_txt2img, inputs=[\n",
    "                            gr.Text(label=\"Your Prompt\"),\n",
    "                            gr.Text(label=\"Your Negative Prompt\"),\n",
    "                            gr.Slider(1, 200, value=50, label=\"Sampling Steps\"),\n",
    "                            gr.Slider(1, 100, step=1, label=\"Number of images\"),\n",
    "                            gr.Slider(1, 100, step=1, label=\"Batch size\"),\n",
    "                            gr.Slider(64, 4096, value=512, step=64, label=\"Width\"),\n",
    "                            gr.Slider(64, 4096, value=512, step=64, label=\"Height\"),\n",
    "                            gr.Slider(0, 50, value=7.5, step=0.1, label=\"Guidance scale\"),\n",
    "                            gr.Slider(0, 1, step=0.01, label=\"DDIM sampling ETA\"),\n",
    "                            gr.Slider(1, 2, value=1, step=1, label=\"U-Net batch size\"),\n",
    "                            gr.Radio([\"cuda\", \"cpu\"], value=\"cuda\", label=\"Device\"),\n",
    "                            gr.Text(label=\"Seed\"),\n",
    "                            gr.Text(value=output_path, label=\"Outputs path\"),\n",
    "                            gr.Radio([\"png\", \"jpg\"], value='png', label=\"Image format\"),\n",
    "                            gr.Checkbox(value=True, label=\"Turbo mode (better leave this on)\"),\n",
    "                            gr.Checkbox(label=\"Full precision mode (practically does nothing)\"),\n",
    "                            gr.Radio(\n",
    "                                [\"ddim\", \"plms\", \"k_dpm_2_a\", \"k_dpm_2\", \"k_euler_a\", \"k_euler\", \"k_heun\", \"k_lms\"],\n",
    "                                value=\"plms\", label=\"Sampler\"),\n",
    "                            gr.Checkbox(value=True, label=\"Lightning Attention (only on linux + xformers installed)\"),\n",
    "                        ], outputs=[out_image, gen_res])\n",
    "    with gr.Tab(\"img2img\"):\n",
    "        with gr.Column():\n",
    "            gr.Markdown(\"# Generate images from images (neonsecret's adjustments)\")\n",
    "            with gr.Row():\n",
    "                with gr.Column():\n",
    "                    out_image2 = gr.Image(label=\"Output Image\")\n",
    "                    gen_res2 = gr.Text(label=\"Generation results\")\n",
    "                    b1 = gr.Button(\"Generate!\")\n",
    "                    b4 = gr.Button(\"Face correction\")\n",
    "                    b5 = gr.Button(\"Upscale 2x\")\n",
    "                with gr.Column():\n",
    "                    with gr.Box():\n",
    "                        b4.click(face_restore, inputs=[out_image2], outputs=[out_image2, gen_res2])\n",
    "                        b5.click(upscale2x, inputs=[out_image2], outputs=[out_image2, gen_res2])\n",
    "                        b1.click(generate_img2img, inputs=[\n",
    "                            gr.Image(tool=\"editor\", type=\"pil\", label=\"Initial image\"),\n",
    "                            gr.Text(label=\"Your Prompt\"),\n",
    "                            gr.Text(label=\"Your Negative Prompt\"),\n",
    "                            gr.Slider(0, 1, value=0.75, label=\"Generated image strength\"),\n",
    "                            gr.Slider(1, 200, value=50, label=\"Sampling Steps\"),\n",
    "                            gr.Slider(1, 100, step=1, label=\"Number of images\"),\n",
    "                            gr.Slider(1, 100, step=1, label=\"Batch size\"),\n",
    "                            gr.Slider(64, 4096, value=512, step=64, label=\"Width\"),\n",
    "                            gr.Slider(64, 4096, value=512, step=64, label=\"Height\"),\n",
    "                            gr.Slider(0, 50, value=7.5, step=0.1, label=\"Guidance scale\"),\n",
    "                            gr.Slider(0, 1, step=0.01, label=\"DDIM sampling ETA\"),\n",
    "                            gr.Slider(1, 2, value=1, step=1, label=\"U-Net batch size\"),\n",
    "                            gr.Radio([\"cuda\", \"cpu\"], value=\"cuda\", label=\"Device\"),\n",
    "                            gr.Text(label=\"Seed\"),\n",
    "                            gr.Text(value=output_path, label=\"Outputs path\"),\n",
    "                            gr.Radio([\"png\", \"jpg\"], value='png', label=\"Image format\"),\n",
    "                            gr.Checkbox(value=True, label=\"Turbo mode (better leave this on)\"),\n",
    "                            gr.Checkbox(label=\"Full precision mode (practically does nothing)\"),\n",
    "                            gr.Radio(\n",
    "                                [\"ddim\", \"plms\", \"k_dpm_2_a\", \"k_dpm_2\", \"k_euler_a\", \"k_euler\", \"k_heun\", \"k_lms\"],\n",
    "                                value=\"ddim\", label=\"Sampler\"),\n",
    "                            gr.Checkbox(value=True, label=\"Lightning Attention (only on linux + xformers installed)\"),\n",
    "                        ], outputs=[out_image2, gen_res2])\n",
    "    with gr.Tab(\"img2img inpaint\"):\n",
    "        with gr.Column():\n",
    "            gr.Markdown(\"# Generate images from images (with a mask) (neonsecret's adjustments)\")\n",
    "            with gr.Row():\n",
    "                with gr.Column():\n",
    "                    out_image3 = gr.Image(label=\"Output Image\")\n",
    "                    gen_res3 = gr.Text(label=\"Generation results\")\n",
    "                    b1 = gr.Button(\"Generate!\")\n",
    "                    b4 = gr.Button(\"Face correction\")\n",
    "                    b5 = gr.Button(\"Upscale 2x\")\n",
    "                with gr.Column():\n",
    "                    with gr.Box():\n",
    "                        b4.click(face_restore, inputs=[out_image3], outputs=[out_image3, gen_res3])\n",
    "                        b5.click(upscale2x, inputs=[out_image3], outputs=[out_image3, gen_res3])\n",
    "                        b1.click(generate_img2img, inputs=[\n",
    "                            gr.Image(tool=\"sketch\", type=\"pil\", label=\"Initial image with a mask\"),\n",
    "                            gr.Text(label=\"Your Prompt\"),\n",
    "                            gr.Text(label=\"Your Negative Prompt\"),\n",
    "                            gr.Slider(0, 1, value=0.75, label=\"Generated image strength\"),\n",
    "                            gr.Slider(1, 200, value=50, label=\"Sampling Steps\"),\n",
    "                            gr.Slider(1, 100, step=1, label=\"Number of images\"),\n",
    "                            gr.Slider(1, 100, step=1, label=\"Batch size\"),\n",
    "                            gr.Slider(64, 4096, value=512, step=64, label=\"Width\"),\n",
    "                            gr.Slider(64, 4096, value=512, step=64, label=\"Height\"),\n",
    "                            gr.Slider(0, 50, value=7.5, step=0.1, label=\"Guidance scale\"),\n",
    "                            gr.Slider(0, 1, step=0.01, label=\"DDIM sampling ETA\"),\n",
    "                            gr.Slider(1, 2, value=1, step=1, label=\"U-Net batch size\"),\n",
    "                            gr.Radio([\"cuda\", \"cpu\"], value=\"cuda\", label=\"Device\"),\n",
    "                            gr.Text(label=\"Seed\"),\n",
    "                            gr.Text(value=output_path, label=\"Outputs path\"),\n",
    "                            gr.Radio([\"png\", \"jpg\"], value='png', label=\"Image format\"),\n",
    "                            gr.Checkbox(value=True, label=\"Turbo mode (better leave this on)\"),\n",
    "                            gr.Checkbox(label=\"Full precision mode (practically does nothing)\"),\n",
    "                            gr.Radio(\n",
    "                                [\"ddim\", \"plms\", \"k_dpm_2_a\", \"k_dpm_2\", \"k_euler_a\", \"k_euler\", \"k_heun\", \"k_lms\"],\n",
    "                                value=\"ddim\", label=\"Sampler\"),\n",
    "                            gr.Checkbox(value=True, label=\"Lightning Attention (only on linux + xformers installed)\"),\n",
    "                        ], outputs=[out_image3, gen_res3])\n",
    "    with gr.Tab(\"img2img interpolate\"):\n",
    "        with gr.Column():\n",
    "            gr.Markdown(\"# Generate a video interpolation from images\")\n",
    "            gr.Markdown(\"### Press 'generation status' button to get the model output logs\")\n",
    "            with gr.Row():\n",
    "                with gr.Column():\n",
    "                    out_video = gr.Video()\n",
    "                    gen_res4 = gr.Text(label=\"Generation results\")\n",
    "                    outs2 = [gr.Text(label=\"Logs\")]\n",
    "                    outs3 = [gr.Text(label=\"nvidia-smi\")]\n",
    "                    b1 = gr.Button(\"Generate!\")\n",
    "                    b2 = gr.Button(\"generation status\")\n",
    "                    b3 = gr.Button(\"nvidia-smi\")\n",
    "                with gr.Column():\n",
    "                    with gr.Box():\n",
    "                        b1.click(generate_img2img_interp, inputs=[\n",
    "                            gr.Image(tool=\"editor\", type=\"pil\", label=\"Initial image\"),\n",
    "                            gr.Text(label=\"Your Prompt\"),\n",
    "                            gr.Slider(0, 1, value=0.75, label=\"Generated image strength\"),\n",
    "                            gr.Slider(1, 200, value=50, label=\"Sampling Steps\"),\n",
    "                            gr.Slider(1, 100, step=1, label=\"Number of images\"),\n",
    "                            gr.Slider(1, 100, step=1, label=\"Batch size\"),\n",
    "                            gr.Slider(64, 4096, value=512, step=64, label=\"Width\"),\n",
    "                            gr.Slider(64, 4096, value=512, step=64, label=\"Height\"),\n",
    "                            gr.Slider(-25, 25, value=7.5, step=0.1, label=\"Guidance scale\"),\n",
    "                            gr.Slider(0, 1, step=0.01, label=\"DDIM sampling ETA\"),\n",
    "                            gr.Slider(1, 2, value=1, step=1, label=\"U-Net batch size\"),\n",
    "                            gr.Radio([\"cuda\", \"cpu\"], value=\"cuda\", label=\"Device\"),\n",
    "                            gr.Text(label=\"Seed\"),\n",
    "                            gr.Text(value=output_path, label=\"Outputs path\"),\n",
    "                            gr.Radio([\"png\", \"jpg\"], value='png', label=\"Image format\"),\n",
    "                            gr.Checkbox(value=True, label=\"Turbo mode (better leave this on)\"),\n",
    "                            gr.Checkbox(label=\"Full precision mode (practically does nothing)\"),\n",
    "                            gr.Radio(\n",
    "                                [\"ddim\", \"plms\", \"k_dpm_2_a\", \"k_dpm_2\", \"k_euler_a\", \"k_euler\", \"k_heun\", \"k_lms\"],\n",
    "                                value=\"ddim\", label=\"Sampler\"),\n",
    "                            gr.Checkbox(value=True, label=\"Lightning Attention (only on linux + xformers installed)\"),\n",
    "                            gr.Slider(1, 120, value=60, step=1, label=\"How smooth the video will be\"),\n",
    "                        ], outputs=[out_video, gen_res4])\n",
    "                        b2.click(get_logs, inputs=[], outputs=outs2)\n",
    "                        b3.click(get_nvidia_smi, inputs=[], outputs=outs3)\n",
    "    with gr.Tab(\"txt2img 2x-3x upscale\"):\n",
    "        with gr.Column():\n",
    "            gr.Markdown(\"# Generate images from text using SD upscaling\")\n",
    "            gr.Markdown(\"### Generate images in 2(3) steps - Wx -> 2Wx2H (-> 3Wx3H)\")\n",
    "            gr.Markdown(\"### Press 'generation status' button to get the model output logs\")\n",
    "            with gr.Row():\n",
    "                with gr.Column():\n",
    "                    out_image = gr.Image(label=\"Output Image\")\n",
    "                    gen_res = gr.Text(label=\"Generation results\")\n",
    "                    outs2 = [gr.Text(label=\"Logs\")]\n",
    "                    outs3 = gr.Text(label=\"nvidia-smi\")\n",
    "                    b1 = gr.Button(\"Generate!\")\n",
    "                    b4 = gr.Button(\"Face correction\")\n",
    "                    b5 = gr.Button(\"Upscale 2x\")\n",
    "                    b2 = gr.Button(\"generation status\")\n",
    "                    b3 = gr.Button(\"nvidia-smi\")\n",
    "                with gr.Column():\n",
    "                    with gr.Box():\n",
    "                        b4.click(face_restore, inputs=[out_image], outputs=[out_image, gen_res])\n",
    "                        b5.click(upscale2x, inputs=[out_image], outputs=[out_image, gen_res])\n",
    "                        b1.click(generate_double_triple, inputs=[\n",
    "                            gr.Text(label=\"Your Prompt\"),\n",
    "                            gr.Slider(1, 200, value=50, label=\"Sampling Steps\"),\n",
    "                            gr.Slider(0, 1, value=0.35, label=\"Upscaled image changes strength\"),\n",
    "                            gr.Slider(64, 4096, value=512, step=64, label=\"Width\"),\n",
    "                            gr.Slider(64, 4096, value=512, step=64, label=\"Height\"),\n",
    "                            gr.Slider(-25, 25, value=7.5, step=0.1, label=\"Guidance scale\"),\n",
    "                            gr.Slider(0, 1, step=0.01, label=\"DDIM sampling ETA\"),\n",
    "                            gr.Slider(1, 2, value=1, step=1, label=\"U-Net batch size\"),\n",
    "                            gr.Radio([\"cuda\", \"cpu\"], value=\"cuda\", label=\"Device\"),\n",
    "                            gr.Text(label=\"Seed\"),\n",
    "                            gr.Text(value=output_path, label=\"Outputs path\"),\n",
    "                            gr.Radio([\"png\", \"jpg\"], value='png', label=\"Image format\"),\n",
    "                            gr.Checkbox(value=True, label=\"Turbo mode (better leave this on)\"),\n",
    "                            gr.Checkbox(label=\"Full precision mode (practically does nothing)\"),\n",
    "                            gr.Radio(\n",
    "                                [\"ddim\", \"plms\", \"k_dpm_2_a\", \"k_dpm_2\", \"k_euler_a\", \"k_euler\", \"k_heun\", \"k_lms\"],\n",
    "                                value=\"plms\", label=\"Sampler\"),\n",
    "                            gr.Checkbox(value=True, label=\"Lightning Attention (only on linux + xformers installed)\"),\n",
    "                            gr.Slider(2, 3, value=2, step=1,\n",
    "                                      label=\"Neural scaling factor, 3 will take much longer\"),\n",
    "                        ], outputs=[out_image, gen_res])\n",
    "                        b2.click(get_logs, inputs=[], outputs=outs2)\n",
    "                        b3.click(get_nvidia_smi, inputs=[], outputs=[outs3])\n",
    "\n",
    "debug = False #@param {type:\"boolean\"}\n",
    "demo.launch(share=True, debug=debug)"
   ],
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "id": "gnYhjq2Gc1zc",
    "cellView": "form"
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  },
  "colab": {
   "provenance": [],
   "collapsed_sections": []
  },
  "accelerator": "GPU"
 },
 "nbformat": 4,
 "nbformat_minor": 0
}