Stable_Diffusion_XL_Turbo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyMsPmM0OaZ/MgCmOve3SHDC",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sihwapark/SDXL_Turbo_Demo_Notebook/blob/main/Stable_Diffusion_XL_Turbo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Stable Diffusion XL Turbo Image-to-Image Demo Notebook\n",
        "By Sihwa Park (shpark@yorku.ca) <br/>\n",
        "December 2, 2023\n",
        "\n",
        "References:<br/>\n",
        "- [SDXL Turbo Online Demo](https://stablediffusionweb.com/SDXL-Turbo)\n",
        "- [SDXL Turbo Model Card](https://huggingface.co/stabilityai/sdxl-turbo)\n",
        "- [SDXL Turbo Documentation](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/sdxl_turbo)\n",
        "- [Jupyter widgets](https://github.com/jupyter-widgets/ipywidgets)\n",
        "- [redromnon's notebook](https://colab.research.google.com/github/redromnon/stable-diffusion-interactive-notebook/blob/main/stable_diffusion_interactive_notebook.ipynb) for the use of Jupyter widgets\n"
      ],
      "metadata": {
        "id": "DMf78SoOHjqc"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Install Prerequisites"
      ],
      "metadata": {
        "id": "KfnJRvITKXt0"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3kU7cSlcUwFH"
      },
      "outputs": [],
      "source": [
        "!pip -q install torch diffusers transformers accelerate scipy safetensors xformers mediapy ipywidgets==7.7.1"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Import Libraries"
      ],
      "metadata": {
        "id": "tNbuWynnMSw3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import ipywidgets as widgets, mediapy, random\n",
        "import importlib\n",
        "from pathlib import Path\n",
        "\n",
        "#Enable third party widget support\n",
        "from google.colab import output\n",
        "output.enable_custom_widget_manager()\n",
        "\n",
        "from diffusers import AutoPipelineForImage2Image\n",
        "from diffusers.utils import load_image\n",
        "\n",
        "# See https://github.com/tin2tin/Pallaidium/issues/72\n",
        "# WARNING:xformers:WARNING[XFORMERS]: xFormers can't load C++/CUDA extensions. xFormers was built for:\n",
        "  # PyTorch 2.1.0+cu121 with CUDA 1201 (you have 2.1.0+cu118)\n",
        "  # Python  3.10.13 (you have 3.10.12)"
      ],
      "metadata": {
        "id": "m779caUAVnMR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load Pipeline and Model"
      ],
      "metadata": {
        "id": "iKFm1SUQK00d"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pipe = AutoPipelineForImage2Image.from_pretrained(\"stabilityai/sdxl-turbo\", torch_dtype=torch.float16, variant=\"fp16\")\n",
        "\n",
        "# To resolve the error: \"RuntimeError: \"LayerNormKernelImpl\" not implemented for 'Half'\"\n",
        "# It is caused when trying to load the diffusion model in float16(Half) format on CPU which is not supported.\n",
        "# For float16 format, GPU needs to be used. For CPU run the model in float32 format.\n",
        "# https://stackoverflow.com/questions/75641074/i-run-stable-diffusion-its-wrong-runtimeerror-layernormkernelimpl-not-implem\n",
        "pipe = pipe.to(\"cuda\")\n",
        "\n",
        "# or call .enable_model_cpu_offload() that offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance.\n",
        "# https://huggingface.co/docs/diffusers/v0.24.0/en/api/pipelines/stable_diffusion/gligen#diffusers.StableDiffusionGLIGENTextImagePipeline.enable_model_cpu_offload\n",
        "# pipe.enable_model_cpu_offload()"
      ],
      "metadata": {
        "id": "m9EdsDsFjAu7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Image-to-Image Generation with GUI"
      ],
      "metadata": {
        "id": "fgWW4_tIL6hO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title\n",
        "\n",
        "init_image = None;\n",
        "\n",
        "image_file_path = widgets.Text(\n",
        "    value=\"/content/1.jpeg\",\n",
        "    description=\"Image:\",\n",
        "    layout=widgets.Layout(width=\"80%\")\n",
        ")\n",
        "\n",
        "button_select = widgets.Button(\n",
        "    description=\"Select\",\n",
        "    button_style=\"success\",\n",
        "    layout=widgets.Layout(width=\"20%\")\n",
        ")\n",
        "\n",
        "reference_image = widgets.Output()\n",
        "\n",
        "def preload_image(e):\n",
        "  reference_image.clear_output()\n",
        "\n",
        "  global init_image\n",
        "  with reference_image:\n",
        "    init_image = load_image(image_file_path.value).resize((512, 512))\n",
        "    mediapy.show_image(init_image)\n",
        "\n",
        "button_select.on_click(preload_image)\n",
        "\n",
        "# When using SDXL-Turbo for image-to-image generation, make sure that num_inference_steps * strength is larger or equal to 1.\n",
        "# The image-to-image pipeline will run for int(num_inference_steps * strength) steps, e.g. 0.5 * 2.0 = 1 step in our example below.\n",
        "steps = widgets.IntSlider(\n",
        "    value=5,\n",
        "    min=1,\n",
        "    max=10,\n",
        "    step=1,\n",
        "    orientation='horizontal',\n",
        "    readout_format='d',\n",
        "    description=\"Steps:\",\n",
        "    layout=widgets.Layout(width=\"99%\")\n",
        ")\n",
        "\n",
        "# strength (float, optional, defaults to 0.3) — Conceptually, indicates how much to transform the reference image.\n",
        "# Must be between 0 and 1. the reference image will be used as a starting point, adding more noise to it the larger the strength.\n",
        "# The number of denoising steps depends on the amount of noise initially added.\n",
        "# When strength is 1, added noise will be maximum and the denoising process will run for the full number of iterations specified in num_inference_steps.\n",
        "strength = widgets.FloatSlider(\n",
        "    value=0.2,\n",
        "    min=0,\n",
        "    max=1.0,\n",
        "    step=0.001,\n",
        "    orientation='horizontal',\n",
        "    readout_format='.3f',\n",
        "    description=\"Strength:\",\n",
        "    layout=widgets.Layout(width=\"99%\")\n",
        ")\n",
        "\n",
        "\n",
        "# guidance_scale (float, optional, defaults to 7.5) — Guidance scale as defined in Classifier-Free Diffusion Guidance. guidance_scale is defined as w of equation 2. of Imagen Paper.\n",
        "# Guidance scale is enabled by setting guidance_scale > 1.\n",
        "# Higher guidance scale encourages to generate images that are closely linked to the text prompt, usually at the expense of lower image quality.\n",
        "# NOTE: SDXL Turbo should disable guidance scale by setting guidance_scale=0.0\n",
        "# CFG = widgets.widgets.FloatSlider(\n",
        "#     value=7.5,\n",
        "#     min=0,\n",
        "#     max=1.0,\n",
        "#     step=0.001,\n",
        "#     orientation='horizontal',\n",
        "#     readout_format='.3f',\n",
        "#     description=\"CFG:\",\n",
        "#     layout=widgets.Layout(width=width)\n",
        "# )\n",
        "\n",
        "random_seed = widgets.IntSlider(\n",
        "    value=random.randint(0, 12013012031030),\n",
        "    min=0,\n",
        "    max=12013012031030,\n",
        "    step=1,\n",
        "    orientation='horizontal',\n",
        "    readout_format='d',\n",
        "    description=\"Seed:\",\n",
        "    layout=widgets.Layout(width=\"99%\")\n",
        ")\n",
        "\n",
        "prompt = widgets.Textarea(\n",
        "    value=\"\",\n",
        "    placeholder=\"Enter prompt\",\n",
        "    rows=1,\n",
        "    layout=widgets.Layout(width=\"80%\")\n",
        ")\n",
        "\n",
        "# neg_prompt = widgets.Textarea(\n",
        "#     value=\"\",\n",
        "#     placeholder=\"Enter negative prompt\",\n",
        "#     rows=5,\n",
        "#     layout=widgets.Layout(width=\"600px\")\n",
        "# )\n",
        "\n",
        "generate = widgets.Button(\n",
        "    description=\"Generate\",\n",
        "    button_style=\"primary\",\n",
        "    layout=widgets.Layout(width=\"20%\")\n",
        ")\n",
        "\n",
        "display_imgs = widgets.Output()\n",
        "\n",
        "process_info = widgets.Output() #widgets.HTML(value=\"\")\n",
        "\n",
        "generated_image = None;\n",
        "\n",
        "def generate_img(i):\n",
        "  global generated_image\n",
        "  #Clear output\n",
        "  process_info.clear_output()\n",
        "  display_imgs.clear_output()\n",
        "  generate.disabled = True\n",
        "  generated_image = None;\n",
        "\n",
        "  #Calculate seed\n",
        "  seed = random.randint(0, 12013012031030) if random_seed.value == -1 else random_seed.value\n",
        "  # print(init_image)\n",
        "  if init_image == None:\n",
        "    preload_image(None)\n",
        "\n",
        "  with process_info:\n",
        "    print(\"Running...\")\n",
        "    # process_info.value = \"Running...\"\n",
        "    # https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/sdxl_turbo\n",
        "    images = pipe(\n",
        "        prompt.value,\n",
        "        image=init_image,\n",
        "        num_inference_steps=steps.value,\n",
        "        strength=strength.value,\n",
        "        # num_images_per_prompt = 1,\n",
        "        # guidance_scale=CFG.value,\n",
        "        guidance_scale=0.0,\n",
        "        generator=torch.Generator(\"cuda\").manual_seed(seed)\n",
        "      ).images;\n",
        "\n",
        "    generated_image = images[0]\n",
        "    print(\"Done!\")\n",
        "\n",
        "  with display_imgs:\n",
        "    mediapy.show_images(images)\n",
        "\n",
        "  generate.disabled = False\n",
        "\n",
        "generate.on_click(generate_img)\n",
        "\n",
        "save = widgets.Button(\n",
        "    description=\"Save\",\n",
        "    button_style=\"primary\",\n",
        "    layout=widgets.Layout(width=\"20%\")\n",
        ")\n",
        "\n",
        "def save_image(b):\n",
        "  global generated_image\n",
        "  generated_image.save(\"result.png\")\n",
        "\n",
        "save.on_click(save_image);\n",
        "\n",
        "\n",
        "widgets.VBox(\n",
        "  [\n",
        "    widgets.AppLayout(\n",
        "      header=None,\n",
        "      # widgets.HTML(\n",
        "      #     value=\"<h2>Stable Diffusion XL Turbo</h2>\",\n",
        "      # ),\n",
        "      left_sidebar=widgets.VBox(\n",
        "          [\n",
        "            widgets.HBox([image_file_path, button_select]),\n",
        "            steps, strength, random_seed\n",
        "          ], layout=widgets.Layout(width=\"95%\")\n",
        "      ),\n",
        "      center=None,\n",
        "      right_sidebar=widgets.VBox(\n",
        "          [\n",
        "            widgets.HBox([prompt, generate]),\n",
        "            process_info\n",
        "          ],\n",
        "          layout=widgets.Layout(width=\"80%\")\n",
        "      ),\n",
        "      footer=None,\n",
        "      pane_widths=[1, 0, 2],\n",
        "    ),\n",
        "    widgets.AppLayout(\n",
        "      header=None,\n",
        "      left_sidebar=reference_image,\n",
        "      center=None,\n",
        "      right_sidebar=display_imgs,\n",
        "      footer=None,\n",
        "      pane_widths=[1, 0, 2],\n",
        "    ),\n",
        "    # save\n",
        "  ]\n",
        ")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "fXRWncVJZ0Zc"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}