{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "eac29097-35bd-4739-b61a-c6cad7737a6c",
   "metadata": {},
   "source": [
    "# Stable Diffusion 2.1 on SageMaker Studio Notebooks\n",
    "Hey folks! I built this out on the Python 3 (Data Science) kernel, using an ml.g4dn.xlarge Studio notebook. Make sure you're using updated Python and PyTorch versions, otherwise you'll struggle to import the most recent objects from `diffusers`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4ee72acf-dae1-4a04-bda6-87e500a2702a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade diffusers transformers accelerate torch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d21b0e7e-a275-4eb7-a5be-f31bd691ed0e",
   "metadata": {},
   "source": [
    "### From the `diffusers` SDK from our friends at Hugging Face!\n",
    "https://github.com/huggingface/diffusers "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbe84989-bc4a-4796-9d23-5a5ea6a1fd4d",
   "metadata": {},
   "source": [
    "# Text to image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "0ee13807-11f8-4224-8283-7b71a0e96b6b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler\n",
    "\n",
    "model_id = \"stabilityai/stable-diffusion-2-1\"\n",
    "\n",
    "# Use the DPMSolverMultistepScheduler (DPM-Solver++) scheduler here instead\n",
    "pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)\n",
    "pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)\n",
    "pipe = pipe.to(\"cuda\")\n",
    "\n",
    "prompt = \"a photo of an astronaut riding a horse on mars\"\n",
    "image = pipe(prompt).images[0]\n",
    "    \n",
    "image.save(\"astronaut_rides_horse.png\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cce23f22-5baa-4b7f-b267-2ae3137aa159",
   "metadata": {
    "tags": []
   },
   "source": [
    "![](astronaut_rides_horse.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbb7b881-217a-45bd-a9d1-179a5ed7eb62",
   "metadata": {},
   "source": [
    "# Image-to-image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9cc5f409-b094-40ed-baa7-87ac63924e61",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ff4256d-3508-45da-b623-53208c32e726",
   "metadata": {},
   "source": [
    "![](sketch-mountains-input.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "bd4b7645-1b4e-480f-8b26-8fa56783ff1e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import requests\n",
    "import torch\n",
    "from PIL import Image\n",
    "from io import BytesIO\n",
    "\n",
    "from diffusers import StableDiffusionImg2ImgPipeline\n",
    "\n",
    "# load the pipeline\n",
    "device = \"cuda\"\n",
    "model_id_or_path = \"runwayml/stable-diffusion-v1-5\"\n",
    "pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)\n",
    "\n",
    "# or download via git clone https://huggingface.co/runwayml/stable-diffusion-v1-5\n",
    "# and pass `model_id_or_path=\"./stable-diffusion-v1-5\"`.\n",
    "pipe = pipe.to(device)\n",
    "\n",
    "# let's download an initial image\n",
    "url = \"https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg\"\n",
    "\n",
    "response = requests.get(url)\n",
    "init_image = Image.open(BytesIO(response.content)).convert(\"RGB\")\n",
    "init_image = init_image.resize((768, 512))\n",
    "\n",
    "prompt = \"A fantasy landscape with a castle, trending on artstation\"\n",
    "\n",
    "images = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images\n",
    "\n",
    "images[0].save(\"fantasy_landscape_v2.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b833e5ff-dffb-4e89-9d45-6168e7f62051",
   "metadata": {
    "tags": []
   },
   "source": [
    "![](fantasy_landscape_v2.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2a97f29-8bab-4491-bb7a-d18109ad3909",
   "metadata": {},
   "source": [
    "# In-painting with Stable Diffusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "65762a2b-76f6-46ad-ab66-87a4e70502c0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import PIL\n",
    "import requests\n",
    "import torch\n",
    "from io import BytesIO\n",
    "\n",
    "from diffusers import StableDiffusionInpaintPipeline\n",
    "\n",
    "def download_image(url):\n",
    "    response = requests.get(url)\n",
    "    return PIL.Image.open(BytesIO(response.content)).convert(\"RGB\")\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",
    "init_image = download_image(img_url).resize((512, 512))\n",
    "mask_image = download_image(mask_url).resize((512, 512))\n",
    "\n",
    "pipe = StableDiffusionInpaintPipeline.from_pretrained(\"runwayml/stable-diffusion-inpainting\", torch_dtype=torch.float16)\n",
    "pipe = pipe.to(\"cuda\")\n",
    "\n",
    "prompt = \"Face of a yellow cat, high resolution, sitting on a park bench\"\n",
    "image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]\n",
    "image.save('yellow_cat.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1136e0d-0b11-4b51-8455-5030bfbf8011",
   "metadata": {},
   "source": [
    "![](yellow_cat.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d129d6b6-6205-437d-9d02-4d1c7d516211",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "instance_type": "ml.g4dn.xlarge",
  "kernelspec": {
   "display_name": "Python 3 (Data Science)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/datascience-1.0"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}