{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "18009882-afb5-4ca1-ae6c-455404b3e2f6",
   "metadata": {},
   "source": [
    "# Stable Diffusion with diffusers\n",
    "\n",
    "[![Open In Studio Lab](https://studiolab.sagemaker.aws/studiolab.svg)](https://studiolab.sagemaker.aws/import/github/aws/studio-lab-examples/blob/main/custom-environments/diffusers.ipynb)\n",
    "\n",
    "This is sample notebook of running [Stable Diffusion v2](https://huggingface.co/stabilityai/stable-diffusion-2) with [diffusers](https://github.com/huggingface/diffusers) on SageMaker Studio Lab.\n",
    "\n",
    "This notebook is designed to work in both CPU and GPU environment.\n",
    "\n",
    "It is recommended to run this notebook on GPU runtime since it is very slow on CPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f3cc787-1f30-45bd-9094-d7f62ce07f9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from diffusers import StableDiffusionPipeline\n",
    "import IPython.display as display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fbdbddf-fd30-4a38-af67-246afc6773a1",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_name = \"stabilityai/stable-diffusion-2\"\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(device)\n",
    "\n",
    "if device == \"cuda\":\n",
    "    pipe = StableDiffusionPipeline.from_pretrained(\n",
    "        model_name, torch_dtype=torch.float16, revision=\"fp16\"\n",
    "    )\n",
    "else:\n",
    "    pipe = StableDiffusionPipeline.from_pretrained(model_name)\n",
    "pipe = pipe.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d68c0769-44a4-4bd2-9748-93285564923b",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"a photo of an astronaut riding a horse on mars\"\n",
    "image = pipe(prompt, height=768, width=768).images[0]\n",
    "image.save(\"test.png\")\n",
    "display.display(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab600abf-0314-4e93-bb95-b8008d2b6aa0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "diffusers:Python",
   "language": "python",
   "name": "conda-env-diffusers-py"
  },
  "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}