{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Running Peacasso - Python API and UI\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Requirements and Installation\n",
    "\n",
    "- Step 1: **Access to Weights via HuggingFace**\n",
    "\n",
    "  Access to the diffusion model weights requires a HuggingFace model account and access token. Please create an account at [huggingface.co](https://huggingface.co/), get an [access token](https://huggingface.co/settings/tokens) and agree to the model terms [here](https://huggingface.co/CompVis/stable-diffusion-v1-4). Next, create a `HF_API_TOKEN` environment variable containing your token. `export HF_API_TOKEN=your_token`. Note that the first time you run peacasso, the weights for the SD model are [cached locally](https://huggingface.co/transformers/v4.0.1/installation.html#caching-models) on your machine. In theory, you can download the weights, and run peacasso by pointing to the folder with the weights.\n",
    "\n",
    "- Step 2: **Verify Environment - Pythong 3.7+ and CUDA**\n",
    "  Setup and verify that your python environment is `python 3.7` or higher (preferably, use Conda). Also verify that you have CUDA installed correctly (`torch.cuda.is_available()` is true) and your GPU has about [7GB of VRAM memory](https://stability.ai/blog/stable-diffusion-public-release).\n",
    "\n",
    "Once requirements are met, run the following command to install the library:\n",
    "\n",
    "```bash\n",
    "pip install peacasso\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install -q peacasso -U"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# os.environ['HF_API_TOKEN'] = \"your HF Token Here\" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from peacasso.generator import ImageGenerator\n",
    "from peacasso.datamodel import GeneratorConfig\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "\n",
    "# helper function to plot images\n",
    "def plot_images(images, title=\"\"):\n",
    "  plt.figure(figsize=(18,10))\n",
    "  plt.imshow(images);\n",
    "  plt.title(title)\n",
    "  plt.axis('off');\n",
    "  plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "{'trained_betas'} was not found in config. Values will be initialized to default values.\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# load token from .env variable\n",
    "token = os.environ.get(\"HF_API_TOKEN\")\n",
    "gen = ImageGenerator(token=token, model=\"CompVis/stable-diffusion-v1-4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prompt Based Generation\n",
    "\n",
    "Generating images conditioned on text representation. The following code generates 3 images conditioned on the text prompt `\"portrait photo of a asia old warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes, 50mm portrait photography, hard rim lighting photography--beta\"`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.014884233474731445,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": null,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": null,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a4fb811fa2d7400180dcaf3c4f68b5a4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 2835500 characters]",
      "text/plain": [
       "<Figure size 1800x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "prompt = \"portrait photo of a asia old warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes, 50mm portrait photography, hard rim lighting photography--beta\"\n",
    "prompt_config = GeneratorConfig(\n",
    "    prompt=prompt,\n",
    "    num_images=3,\n",
    "    width=512,\n",
    "    height=512,\n",
    "    guidance_scale=7.5,\n",
    "    num_inference_steps=20,\n",
    "    mode=\"prompt\",\n",
    "    return_intermediates=True,\n",
    ")\n",
    "prompt_result = gen.generate(prompt_config)\n",
    "plot_images(np.asarray(np.concatenate(prompt_result[\"images\"], axis=1)), title=\"Text to Image Generation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 630728 characters]",
      "text/plain": [
       "<Figure size 1800x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 627320 characters]",
      "text/plain": [
       "<Figure size 1800x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 629608 characters]",
      "text/plain": [
       "<Figure size 1800x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i in range(prompt_config.num_images):\n",
    "  plot_images(np.asarray(np.concatenate([x[i] for x in prompt_result[\"intermediates\"]], axis=1)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Image to Image (Variations)\n",
    "\n",
    "We can also generate images based on existing images. This is achieved by creating an initial latent representation based on the existing image and then conditioning the generation process on the provided text prompt. We do this by setting the `init_image` field in prompt_config and setting `mode='image'`. Ideally, this should give us variations of the same image.\n",
    "\n",
    "In the section below, we will each image generated in the previous section and generate new images based on the same prompt. Note that the generated images are variations of their respective source images image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.014462947845458984,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": null,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": null,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a24f89d385cb478fb3cf893a8d1ca4fb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.015017032623291016,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": null,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": null,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4a150a60828e43fd8cf4ac603187885a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.014709949493408203,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": null,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": null,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fc6778f4c7a6484d90ba4a9f6db11eea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 2938748 characters]",
      "text/plain": [
       "<Figure size 1800x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 2473268 characters]",
      "text/plain": [
       "<Figure size 1800x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 2730892 characters]",
      "text/plain": [
       "<Figure size 1800x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i in range(prompt_config.num_images):\n",
    "    image_prompt_config = GeneratorConfig(\n",
    "        prompt=prompt,\n",
    "        num_images=3,\n",
    "        width=512,\n",
    "        height=512,\n",
    "        guidance_scale=7.5,\n",
    "        num_inference_steps=20,\n",
    "        mode=\"image\",\n",
    "        return_intermediates=True,\n",
    "        init_image = prompt_result[\"images\"][i]\n",
    "    )\n",
    "    image_result = gen.generate(image_prompt_config)\n",
    "    plot_images(np.asarray(np.concatenate(image_result[\"images\"], axis=1)), title=\"Image to Image Generation\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Image to Making with Masking \n",
    "\n",
    "We can also generate images based on _parts of_ existing images. This is achieved by creating an initial latent representation based on the existing image and a mask and then conditioning the generation process on the provided text prompt. We do this by setting the `init_image` field in prompt_config and setting `mode='mask'`. Ideally, this should give us variations of the same image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 2066600 characters]",
      "text/plain": [
       "<Figure size 1800x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "init_mask = Image.open(\"../docs/images/mask.png\") \n",
    "init_image =  prompt_result[\"images\"][0]\n",
    "# concatenate init_mask and init_image and plot  \n",
    "plot_images(np.concatenate((init_mask.convert(\"RGB\"), init_image), axis=1), title=\"Mask + Image (Black Pixels are Inpainted)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.014477014541625977,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": null,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": null,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "14bc0fe760d941eab0cb32f68ae46441",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.014795303344726562,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": null,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": null,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "812582c1c39a4be6ade06f38fa59c1f9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.014929771423339844,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": null,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": null,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "32205429fc39495d8385a33a45d9e2ee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 2453788 characters]",
      "text/plain": [
       "<Figure size 1800x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 2072476 characters]",
      "text/plain": [
       "<Figure size 1800x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 2171476 characters]",
      "text/plain": [
       "<Figure size 1800x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i in range(image_prompt_config.num_images):\n",
    "    mask_prompt_config = GeneratorConfig(\n",
    "        prompt=prompt,\n",
    "        num_images=3,\n",
    "        width=512,\n",
    "        height=512,\n",
    "        guidance_scale=7.5,\n",
    "        num_inference_steps=50,\n",
    "        init_image=prompt_result[\"images\"][i],\n",
    "        mode=\"image\",\n",
    "        return_intermediates=True,\n",
    "        # seed=2147483647,\n",
    "        attention_slice=None,\n",
    "        mask_image=init_mask,\n",
    "    ) \n",
    "    mask_result = gen.generate(mask_prompt_config)\n",
    "    plot_images(np.asarray(np.concatenate([prompt_result[\"images\"][i]] + mask_result[\"images\"], axis=1)), title=\"In Painting (First Image is Original | Upper Half of Face is Inpainted)\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run UI in Background Thread"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import threading\n",
    "import uvicorn\n",
    "\n",
    "def ui(\n",
    "    host: str = \"127.0.0.1\", port: int = 8081, workers: int = 1, reload: bool = True\n",
    "):\n",
    "    \"\"\"\n",
    "    Launch the peacasso UI.Pass in parameters host, port, workers, and reload to override the default values.\n",
    "    \"\"\"\n",
    "    uvicorn.run(\n",
    "        \"peacasso.web.backend.app:app\",\n",
    "        host=host,\n",
    "        port=port,\n",
    "        workers=workers, \n",
    "    )\n",
    "\n",
    "def start_server():\n",
    "    ui()\n",
    "\n",
    "#uncomment the next line\n",
    "# threading.Thread(target=start_server).start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Note \n",
    "Wait a few minutes before you run the cell below. This is to allow the model to load into memory.\n",
    "\n",
    "![image.png](data:image/png;base64,[encoded data, 556228 characters])You should see the UI. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import IFrame\n",
    "IFrame(\"http://127.0.0.1:8081\", 900,1200)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.13 ('py37')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "ef6267dc51b05090b38b35ee4b4e0c15ffdecf0643e9d054740389ea5c66ad97"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
