stable_fast_colab.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/camenduru/stable-fast-colab/blob/main/stable_fast_colab.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# !pip install -q https://download.pytorch.org/whl/cu118/torchaudio-2.1.0%2Bcu118-cp310-cp310-linux_x86_64.whl\n",
        "# !pip install -q https://download.pytorch.org/whl/torchdata-0.7.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n",
        "# !pip install -q https://download.pytorch.org/whl/torchtext-0.16.0%2Bcpu-cp310-cp310-linux_x86_64.whl\n",
        "# !pip install -q https://download.pytorch.org/whl/cu118/torchvision-0.16.0%2Bcu118-cp310-cp310-linux_x86_64.whl\n",
        "# !pip install -q fastai==2.7.13\n",
        "!pip install -q diffusers transformers accelerate\n",
        "!pip install -q https://download.pytorch.org/whl/cu118/torch-2.1.0%2Bcu118-cp310-cp310-linux_x86_64.whl\n",
        "!pip install -q https://download.pytorch.org/whl/triton-2.1.0-0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl\n",
        "!pip install -q https://download.pytorch.org/whl/cu121/xformers-0.0.22.post7-cp310-cp310-manylinux2014_x86_64.whl\n",
        "!pip install -q https://github.com/camenduru/wheels/releases/download/colab/stable_fast-0.0.2-cp310-cp310-linux_x86_64.whl\n",
        "\n",
        "# !pip install -q ninja\n",
        "# !pip install -q git+https://github.com/chengzeyi/stable-fast.git@main#egg=stable-fast\n",
        "\n",
        "# !apt -y update -qq\n",
        "# !wget https://github.com/camenduru/gperftools/releases/download/v1.0/libtcmalloc_minimal.so.4 -O /content/libtcmalloc_minimal.so.4\n",
        "# %env LD_PRELOAD=/content/libtcmalloc_minimal.so.4\n",
        "\n",
        "!ldconfig /usr/lib64-nvidia"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch, xformers, triton\n",
        "from diffusers import StableDiffusionPipeline\n",
        "from sfast.compilers.stable_diffusion_pipeline_compiler import (compile, CompilationConfig)\n",
        "torch.backends.cuda.matmul.allow_tf32 = True\n",
        "model = StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5', torch_dtype=torch.float16, safety_checker=None).to('cuda')\n",
        "config = CompilationConfig.Default()\n",
        "config.enable_xformers = True\n",
        "config.enable_triton = True\n",
        "config.enable_cuda_graph = True # After capturing, the model only accepts one fixed image size. If you want the model to be dynamic, don't enable it.\n",
        "compiled_model = compile(model, config)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "kwarg_inputs = dict(\n",
        "    prompt=\n",
        "    '(masterpiece:1,2), best quality, masterpiece, best detail face, lineart, monochrome, a beautiful girl',\n",
        "    height=512,\n",
        "    width=512,\n",
        "    num_inference_steps=50,\n",
        "    num_images_per_prompt=1,\n",
        ")\n",
        "output_image = compiled_model(**kwarg_inputs).images[0] # The first call will trigger compilation and might be very slow. # After the first call, it should be very fast.\n",
        "output_image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "output_image = compiled_model(**kwarg_inputs).images[0]\n",
        "output_image"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}