stable_diffusion_interactive_notebook.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyNVqD2247RRHVMWeh9Znaqy"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard",
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Stable Diffusion Interactive Notebook 📓 🤖\n",
        "\n",
        "A widgets-based interactive notebook for Google Colab that lets users generate AI images from prompts (Text2Image) using [Stable Diffusion (by Stability AI, Runway & CompVis)](https://en.wikipedia.org/wiki/Stable_Diffusion). \n",
        "\n",
        "This notebook aims to be an alternative to WebUIs while offering a simple and lightweight GUI for anyone to get started with Stable Diffusion.\n",
        "\n",
        "Uses Stable Diffusion, [HuggingFace](https://huggingface.co/) Diffusers and [Jupyter widgets](https://github.com/jupyter-widgets/ipywidgets).\n",
        "\n",
        "<br/>\n",
        "\n",
        "Made with ❤️ by redromnon\n",
        "\n",
        "[GitHub](https://github.com/redromnon/stable-diffusion-interactive-notebook)"
      ],
      "metadata": {
        "id": "wILzWiWRfTX8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 👇 Installing dependencies { display-mode: \"form\" }\n",
        "#@markdown ---\n",
        "#@markdown Make sure to select **GPU** as the runtime type:<br/>\n",
        "#@markdown *Runtime->Change Runtime Type->Under Hardware accelerator, select GPU*\n",
        "#@markdown \n",
        "#@markdown ---\n",
        "\n",
        "!pip -q install torch diffusers transformers accelerate scipy safetensors xformers mediapy ipywidgets==7.7.1"
      ],
      "metadata": {
        "id": "vCR176NNfn0o"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 👇 Selecting Model { form-width: \"20%\", display-mode: \"form\" }\n",
        "#@markdown ---\n",
        "#@markdown - **Select Model** - A list of Stable Diffusion models to choose from.\n",
        "#@markdown - **Select Sampler** - A list of schedulers to choose from. Default is EulerAncestralScheduler.\n",
        "#@markdown - **Safety Checker** - Enable/Disable uncensored content\n",
        "#@markdown \n",
        "#@markdown ---\n",
        "\n",
        "from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler, DDIMScheduler, EulerDiscreteScheduler, UniPCMultistepScheduler\n",
        "from diffusers.models import AutoencoderKL\n",
        "import torch\n",
        "import ipywidgets as widgets\n",
        "import importlib\n",
        "\n",
        "#Enable third party widget support\n",
        "from google.colab import output\n",
        "output.enable_custom_widget_manager()\n",
        "\n",
        "#Pipe\n",
        "pipe = None\n",
        "\n",
        "#Models\n",
        "select_model = widgets.Dropdown(\n",
        "    options=[\n",
        "        (\"Stable Diffusion 2.1 Base\" , \"stabilityai/stable-diffusion-2-1-base\"),\n",
        "        (\"Stable Diffusion 2.1\" , \"stabilityai/stable-diffusion-2-1\"),\n",
        "        (\"Stable Diffusion 1.5\", \"runwayml/stable-diffusion-v1-5\"),\n",
        "        (\"Dreamlike Photoreal 2.0\" , \"dreamlike-art/dreamlike-photoreal-2.0\"),\n",
        "        (\"OpenJourney v4\" , \"prompthero/openjourney-v4\")\n",
        "    ],\n",
        "    description=\"Select Model:\"\n",
        ")\n",
        "\n",
        "#Schedulers\n",
        "select_sampler = widgets.Dropdown(\n",
        "    options=[\n",
        "        \"EulerAncestralDiscreteScheduler\",\n",
        "        \"EulerDiscreteScheduler\",\n",
        "        \"UniPCMultistepScheduler\",\n",
        "        \"DDIMScheduler\"\n",
        "    ],\n",
        "    description=\"Select Schedular:\"\n",
        ")\n",
        "select_sampler.style.description_width = \"auto\"\n",
        "\n",
        "#Safety Checker\n",
        "safety_check = widgets.Checkbox(\n",
        "    value=True,\n",
        "    description=\"Enable Safety Check\",\n",
        "    layout=widgets.Layout(margin=\"0px 0px 0px -85px\")\n",
        ")\n",
        "\n",
        "#Output\n",
        "out = widgets.Output()\n",
        "\n",
        "#Apply Settings\n",
        "apply_btn = widgets.Button(\n",
        "    description=\"Apply\",\n",
        "    button_style=\"info\"\n",
        ")\n",
        "\n",
        "\n",
        "#Get scheduler\n",
        "def get_scheduler(name):\n",
        "\n",
        "  match name:\n",
        "\n",
        "    case \"EulerAncestralDiscreteScheduler\":\n",
        "      return EulerAncestralDiscreteScheduler.from_pretrained(select_model.value, subfolder=\"scheduler\")\n",
        "\n",
        "    case \"DDIMScheduler\":\n",
        "      return DDIMScheduler.from_pretrained(select_model.value, subfolder=\"scheduler\")\n",
        "\n",
        "    case \"EulerDiscreteScheduler\":\n",
        "      return EulerDiscreteScheduler.from_pretrained(select_model.value, subfolder=\"scheduler\")\n",
        "\n",
        "    case \"UniPCMultistepScheduler\":\n",
        "      return UniPCMultistepScheduler.from_pretrained(select_model.value, subfolder=\"scheduler\")\n",
        "\n",
        "#Run pipeline\n",
        "def pipeline(p):\n",
        "\n",
        "  global pipe\n",
        "  \n",
        "  out.clear_output()\n",
        "  apply_btn.disabled = True\n",
        "\n",
        "  with out:\n",
        "\n",
        "    print(\"Running, please wait...\")\n",
        "    \n",
        "    pipe = StableDiffusionPipeline.from_pretrained(\n",
        "      select_model.value, \n",
        "      scheduler=get_scheduler(select_sampler.value),\n",
        "      torch_dtype=torch.float16,  \n",
        "      vae=AutoencoderKL.from_pretrained(\"stabilityai/sd-vae-ft-mse\", torch_dtype=torch.float16).to(\"cuda\")\n",
        "    ).to(\"cuda\")\n",
        "\n",
        "    if not safety_check.value:\n",
        "      pipe.safety_checker = None\n",
        "\n",
        "    pipe.enable_xformers_memory_efficient_attention()\n",
        "\n",
        "    print(\"Finished!\")\n",
        "\n",
        "  apply_btn.disabled = False\n",
        "\n",
        "\n",
        "#Display\n",
        "apply_btn.on_click(pipeline)\n",
        "\n",
        "widgets.VBox(\n",
        "    [\n",
        "      widgets.HTML(value=\"<h2>Configure Pipeline</h2>\"),\n",
        "      select_model, select_sampler, safety_check, apply_btn, out\n",
        "    ]\n",
        ")\n"
      ],
      "metadata": {
        "id": "CV_UTS40oD1k"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 👇 Generating Images { form-width: \"20%\", display-mode: \"form\" }\n",
        "#@markdown ---\n",
        "#@markdown - **Prompt** - Description of the image\n",
        "#@markdown - **Negative Prompt** - Things you don't want to see or ignore in the image\n",
        "#@markdown - **Steps** - Number of denoising steps. Higher steps may lead to better results but takes longer time to generate the image. Default is `30`.\n",
        "#@markdown - **CFG** - Guidance scale ranging from `0` to `20`. Lower values allow the AI to be more creative and less strict at following the prompt. Default is `7.5`.\n",
        "#@markdown - **Seed** - A random value that controls image generation. The same seed and prompt produce the same images. Set `-1` for using random seed values.\n",
        "#@markdown ---\n",
        "import ipywidgets as widgets, mediapy, random\n",
        "import IPython.display\n",
        "\n",
        "\n",
        "#PARAMETER WIDGETS \n",
        "width = \"300px\"\n",
        "\n",
        "prompt = widgets.Textarea(\n",
        "    value=\"\",\n",
        "    placeholder=\"Enter prompt\",\n",
        "    #description=\"Prompt:\",\n",
        "    rows=5,\n",
        "    layout=widgets.Layout(width=\"600px\")\n",
        ")\n",
        "\n",
        "neg_prompt = widgets.Textarea(\n",
        "    value=\"\",\n",
        "    placeholder=\"Enter negative prompt\",\n",
        "    #description=\"Negative Prompt:\",\n",
        "    rows=5,\n",
        "    layout=widgets.Layout(width=\"600px\")\n",
        ")\n",
        "\n",
        "num_images = widgets.IntText(\n",
        "    value=1,\n",
        "    description=\"Images:\",\n",
        "    layout=widgets.Layout(width=width),\n",
        ")\n",
        "\n",
        "steps = widgets.IntText(\n",
        "    value=30,\n",
        "    description=\"Steps:\",\n",
        "    layout=widgets.Layout(width=width)\n",
        ")\n",
        "\n",
        "CFG = widgets.FloatText(\n",
        "    value=7.5,\n",
        "    description=\"CFG:\",\n",
        "    layout=widgets.Layout(width=width)\n",
        ")\n",
        "\n",
        "img_height = widgets.Dropdown(\n",
        "    options=[('512px', 512), ('768px', 768)],\n",
        "    value=512,\n",
        "    description=\"Height:\",\n",
        "    layout=widgets.Layout(width=width)\n",
        ")\n",
        "\n",
        "img_width = widgets.Dropdown(\n",
        "    options=[('512px', 512), ('768px', 768)],\n",
        "    value=512,\n",
        "    description=\"Width:\",\n",
        "    layout=widgets.Layout(width=width)\n",
        ")\n",
        "\n",
        "random_seed = widgets.IntText(\n",
        "    value=-1,\n",
        "    description=\"Seed:\",\n",
        "    layout=widgets.Layout(width=width),\n",
        "    disabled=False\n",
        ")\n",
        "\n",
        "generate = widgets.Button(\n",
        "    description=\"Generate\",\n",
        "    disabled=False,\n",
        "    button_style=\"primary\"\n",
        ")\n",
        "\n",
        "display_imgs = widgets.Output()\n",
        "\n",
        "\n",
        "#RUN\n",
        "def generate_img(i):\n",
        "\n",
        "  #Clear output\n",
        "  display_imgs.clear_output()\n",
        "  generate.disabled = True\n",
        "\n",
        "  #Calculate seed\n",
        "  seed = random.randint(0, 2147483647) if random_seed.value == -1 else random_seed.value\n",
        "\n",
        "  with display_imgs: \n",
        "\n",
        "    print(\"Running...\")\n",
        "    \n",
        "    images = pipe(\n",
        "      prompt.value,\n",
        "      height = img_height.value,\n",
        "      width = img_width.value,\n",
        "      num_inference_steps = steps.value,\n",
        "      guidance_scale = CFG.value,\n",
        "      num_images_per_prompt = num_images.value,\n",
        "      negative_prompt = neg_prompt.value,\n",
        "      generator = torch.Generator(\"cuda\").manual_seed(seed),\n",
        "    ).images \n",
        "    mediapy.show_images(images)\n",
        "\n",
        "    print(f\"Seed:\\n{seed}\")\n",
        "\n",
        "  generate.disabled = False\n",
        "\n",
        "#Display\n",
        "generate.on_click(generate_img)\n",
        "\n",
        "widgets.VBox(\n",
        "    [\n",
        "      widgets.AppLayout(\n",
        "        header=widgets.HTML(\n",
        "            value=\"<h2>Stable Diffusion</h2>\",\n",
        "        ),\n",
        "        left_sidebar=widgets.VBox(\n",
        "            [num_images, steps, CFG, img_height, img_width, random_seed]\n",
        "        ),\n",
        "        center=widgets.VBox(\n",
        "            [prompt, neg_prompt, generate]\n",
        "        ),\n",
        "        right_sidebar=None,\n",
        "        footer=None\n",
        "      ),\n",
        "      display_imgs\n",
        "    ]\n",
        ")"
      ],
      "metadata": {
        "id": "atmx0PNQ78Wa"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}