notebooks/stable_diffusion/text_to_image_sd2.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Text-to-Image Generation on IPUs using Stable Diffusion 2.0 - Inference\n",
    "\n",
    "This notebook demonstrates how a Stable Diffusion 2.0 inference pipeline can be run on Graphcore IPUs. Stable Diffusion is a latent text-to-image diffusion model capable of generating photo-realistic images defined by text input. The improvements in this version of the model are described in the [release blog](https://stability.ai/blog/stable-diffusion-v2-release).\n",
    "\n",
    "![Text to image with SD2](sample_images/text_to_image_sd2.png)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "|  Domain | Tasks | Model | Datasets | Workflow |   Number of IPUs   | Execution time |\n",
    "|---------|-------|-------|----------|----------|--------------|--------------|\n",
    "| Image processing | Text to image generation | Stable Diffusion 2.0 | N/A | Inference | recommended: 4 | 24min |\n",
    "\n",
    "[![Join our Slack Community](https://img.shields.io/badge/Slack-Join%20Graphcore's%20Community-blue?style=flat-square&logo=slack)](https://www.graphcore.ai/join-community)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment setup\n",
    "\n",
    "The best way to run this demo is on Paperspace Gradient's cloud IPUs because everything is already set up for you.\n",
    "\n",
    "[![Run on Gradient](https://assets.paperspace.io/img/gradient-badge.svg)](https://ipu.dev/3X3IMDh)\n",
    "\n",
    "To run the demo using other IPU hardware, you need to have the Poplar SDK enabled. Refer to the [Getting Started guide](https://docs.graphcore.ai/en/latest/getting-started.html#getting-started) for your system for details on how to enable the Poplar SDK. Also refer to the [Jupyter Quick Start guide](https://docs.graphcore.ai/projects/jupyter-notebook-quick-start/en/latest/index.html) for how to set up Jupyter to be able to run this notebook on a remote IPU machine."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to improve usability and support for future users, Graphcore would like to collect information about the\n",
    "applications and code being run in this notebook. The following information will be anonymised before being sent to Graphcore:\n",
    "\n",
    "- User progression through the notebook\n",
    "- Notebook details: number of cells, code being run and the output of the cells\n",
    "- Environment details\n",
    "\n",
    "You can disable logging at any time by running `%unload_ext graphcore_cloud_tools.notebook_logging.gc_logger` from any cell."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dependencies and configuration"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Install the dependencies for this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -r requirements.txt\n",
    "!pip install \"ipywidgets>=7,<8\"\n",
    "%load_ext graphcore_cloud_tools.notebook_logging.gc_logger"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Values for the virtual IPU Pod size and the cache directories can be configured through environment variables or directly in the notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "n_ipu = int(os.getenv(\"NUM_AVAILABLE_IPU\", 4))\n",
    "executable_cache_dir = os.getenv(\"POPLAR_EXECUTABLE_CACHE_DIR\", \"/tmp/exe_cache/\") + \"/stablediffusion2_text2img\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To download the pre-trained Stable-Diffusion-v2 checkpoint, we must first authenticate to the 🤗 Hub:\n",
    "\n",
    "1. Create a read access token on the [🤗 website]((https://huggingface.co/settings/tokens)). [Sign up to 🤗](https://huggingface.co/join) if you haven't already.\n",
    "2. Execute the following cell and input your username and token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you have not done so already, you will need to accept the User License on the [model page](https://huggingface.co/stabilityai/stable-diffusion-2)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pipeline creation"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now ready to import and run the pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from diffusers import DPMSolverMultistepScheduler\n",
    "\n",
    "from optimum.graphcore.diffusers import IPUStableDiffusionPipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe = IPUStableDiffusionPipeline.from_pretrained(\n",
    "    \"stabilityai/stable-diffusion-2\",\n",
    "    revision=\"fp16\", \n",
    "    torch_dtype=torch.float16,\n",
    "    requires_safety_checker=False,\n",
    "    n_ipu=n_ipu,\n",
    "    num_prompts=1,\n",
    "    num_images_per_prompt=1,\n",
    "    common_ipu_config_kwargs={\"executable_cache_dir\": executable_cache_dir}\n",
    ")\n",
    "pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We run a dummy generation step to trigger the one-time compilation process. This should take on the order of 20 minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_width = os.getenv(\"STABLE_DIFFUSION_2_TXT2IMG_DEFAULT_WIDTH\", default=768)\n",
    "image_height = os.getenv(\"STABLE_DIFFUSION_2_TXT2IMG_DEFAULT_HEIGHT\", default=768)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "pipe(\"apple\", height=image_height, width=image_width, num_inference_steps=25, guidance_scale=9);"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Image generation"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We show some example prompts in the cells below. We encourage you to try your own!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"a shiba inu in a zen garden, acrylic painting\"\n",
    "pipe(prompt, height=image_height, width=image_width, num_inference_steps=25, guidance_scale=9).images[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"a photograph of an astronaut riding a horse\"\n",
    "pipe(prompt, height=image_height, width=image_width, num_inference_steps=25, guidance_scale=9).images[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"the living room of a cozy wooden house with a fireplace\"\n",
    "out = pipe(prompt, height=image_height, width=image_width, num_inference_steps=25, guidance_scale=9)\n",
    "out.images[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "fig, ax = plt.subplots(1,1)\n",
    "fig.set_size_inches(9, 9)\n",
    "ax.imshow(out.images[0])\n",
    "\n",
    "ax.set_title(f\"Prompt: {prompt}\")\n",
    "ax.axis(\"off\")\n",
    "fig.savefig(\"sample_images/text_to_image_sd2.png\", dpi=150)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Release IPUs in use\n",
    "\n",
    "The IPython kernel has a lock on the IPUs used in running the model, preventing other users from using them. For example, if you wish to use other notebooks after running this notebook, it may be necessary to manually run the cell below to release the IPUs you have been using. This will happen by default if you use the `Run All` notebook option. More information can be found in the notebook on \"managing IPU resources\" - `useful-tips/managing_ipu_resources.ipynb`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe.detach_from_device()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "You can try out the other notebooks using the Stable Diffusion version 1.5 model:\n",
    "\n",
    "* Text to image generation - `text_to_image.ipynb`\n",
    "* Image to image generation - `image_to_image.ipynb`\n",
    "* Inpainting - `inpainting.ipynb`\n",
    "\n",
    "Also, Try out the other [IPU-powered Jupyter Notebooks](https://www.graphcore.ai/ipu-jupyter-notebooks) to see how how IPUs perform on other tasks."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "bb566221b6fb66eb17be1e54d5bad447b448f6085074d60633cde4e94a153eec"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}