notebooks/stable_diffusion/inpainting.ipynb (377 lines of code) (raw):
{
"cells": [
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"# Text-guided Inpainting on IPUs using Stable Diffusion - Inference\n",
"\n",
"This notebook demonstrates how a stable diffusion inference pipeline can be run on Graphcore IPUs to replace objects inside an image.\n",
"\n",
""
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"| Domain | Tasks | Model | Datasets | Workflow | Number of IPUs | Execution time |\n",
"|---------|-------|-------|----------|----------|--------------|--------------|\n",
"| Image processing | Inpainting | Stable Diffusion (stable-diffusion-v1-5) | N/A | Inference | recommended: 4 | 14 min |\n",
"\n",
"\n",
"[](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",
"[](https://ipu.dev/3W56b5R)\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": [
"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/\") + \"/stablediffusion_inpainting\""
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"To download the pre-trained Stable-Diffusion-v1-5 checkpoint, we must first authenticate to the 🤗 Hub. \n",
"\n",
"1. Create a read access token on the [Hugging Face 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 read 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/runwayml/stable-diffusion-inpainting)."
]
},
{
"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",
"\n",
"from optimum.graphcore.diffusers import IPUStableDiffusionInpaintPipeline"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"pipe = IPUStableDiffusionInpaintPipeline.from_pretrained(\n",
" \"runwayml/stable-diffusion-inpainting\", \n",
" revision=\"fp16\", \n",
" torch_dtype=torch.float16,\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",
")"
]
},
{
"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 15 minutes."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"image_width = os.getenv(\"STABLE_DIFFUSION_INPAINT_DEFAULT_WIDTH\", default=512)\n",
"image_height = os.getenv(\"STABLE_DIFFUSION_INPAINT_DEFAULT_HEIGHT\", default=512)\n",
"image_dimensions = (image_width, image_height)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"from PIL import Image\n",
"\n",
"pipe(\"apple\", \n",
" image=Image.new(\"RGB\", image_dimensions), \n",
" mask_image=Image.new(\"RGB\", image_dimensions, (255, 255, 255)), \n",
" guidance_scale=7.5\n",
");"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Image generation"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"We preprocess and visualize a context image which will be used to initialize the latents passed to the UNet."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import requests\n",
"from io import BytesIO\n",
"\n",
"img_url = \"https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png\"\n",
"mask_url = \"https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png\"\n",
"\n",
"def download_image(url):\n",
" response = requests.get(url)\n",
" return Image.open(BytesIO(response.content)).convert(\"RGB\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"image = download_image(img_url).resize(image_dimensions)\n",
"image"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"mask_image = download_image(mask_url).resize(image_dimensions)\n",
"mask_image"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"Below you will find an example prompt. We encourage you to try your own!"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"prompt = \"a mecha robot sitting on a bench\"\n",
"out = pipe(prompt, image=image, mask_image=mask_image, guidance_scale=7.5)\n",
"out.images[0]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from matplotlib import pyplot as plt"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"fig, axs = plt.subplots(1, 3)\n",
"fig.set_size_inches(12, 5)\n",
"axs[0].imshow(image)\n",
"axs[0].set_title(\"Starting image\")\n",
"axs[1].imshow(mask_image)\n",
"axs[1].set_title(\"Mask\")\n",
"axs[2].imshow(out.images[0])\n",
"axs[2].set_title(\"Generated image\")\n",
"for ax in axs:\n",
" ax.axis(\"off\")\n",
"fig.tight_layout()\n",
"fig.savefig(\"sample_images/inpainting.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",
"Try out our other Stable Diffusion notebooks:\n",
"\n",
"* Text to image generation (Stable Diffusion version 1.5) - `text_to_image.ipynb`\n",
"* Text to image generation (Stable Diffusion version 2.0) - `text_to_image2.ipynb`\n",
"* Image to image generation - `image_to_image.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.\n"
]
}
],
"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
}