camenduru's_tensorrt_diffusers.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/notebooks/blob/main/camenduru's_tensorrt_diffusers.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "%cd /content\n",
        "!mkdir /content/external\n",
        "%cd /content/external\n",
        "!git clone https://github.com/pybind/pybind11\n",
        "!wget https://www.python.org/ftp/python/3.10.12/Python-3.10.12.tgz\n",
        "!tar -xvf Python-3.10.12.tgz\n",
        "!mkdir -p /content/external/python3.10/include\n",
        "!cp -r /content/external/Python-3.10.12/Include/* /content/external/python3.10/include\n",
        "!wget http://http.us.debian.org/debian/pool/main/p/python3.10/libpython3.10-dev_3.10.13-1_amd64.deb\n",
        "!ar x libpython3.10-dev*.deb\n",
        "!mkdir /content/external/debian\n",
        "!tar -xf data.tar.xz -C /content/external/debian\n",
        "!cp /content/external/debian/usr/include/x86_64-linux-gnu/python3.10/pyconfig.h /content/external/python3.10/include\n",
        "%cd /content/TensorRT/python\n",
        "!TENSORRT_MODULE=tensorrt PYTHON_MAJOR_VERSION=3 PYTHON_MINOR_VERSION=10 TARGET_ARCHITECTURE=x86_64 ./build.sh"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "%cd /content\n",
        "!git clone -b release/9.1 --recursive https://github.com/NVIDIA/TensorRT\n",
        "%cd /content/TensorRT\n",
        "\n",
        "%cd /content\n",
        "!wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/9.1.0/tars/tensorrt-9.1.0.4.linux.x86_64-gnu.cuda-11.8.tar.gz\n",
        "!mkdir trt\n",
        "%cd /content/trt\n",
        "!tar -xvzf /content/tensorrt-9.1.0.4.linux.x86_64-gnu.cuda-11.8.tar.gz\n",
        "\n",
        "%cd /content/TensorRT\n",
        "!mkdir -p /content/TensorRT/build\n",
        "%cd /content/TensorRT/build\n",
        "!cmake .. -DTRT_LIB_DIR=/content/trt/TensorRT-9.1.0.4/lib -DTRT_OUT_DIR=/content/TensorRT/build/out -DCUDNN_ROOT_DIR=/usr/local/cuda -DGPU_ARCHS=\"75\"\n",
        "!make -j$(nproc)\n",
        "# !/content/TensorRT/build/out/trtexec -h"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# !pip install -q diffusers transformers accelerate colored cuda-python ftfy nvtx controlnet_aux\n",
        "# !pip install -q --extra-index-url https://pypi.ngc.nvidia.com \"onnx\" \"onnxruntime\" \"onnx-graphsurgeon\" \"polygraphy\"\n",
        "!pip install -q diffusers==0.19.3 transformers==4.31.0 accelerate colored cuda-python ftfy nvtx controlnet_aux==0.0.6 onnx==1.14.0 onnxruntime==1.15.1 -U\n",
        "!pip install -q --extra-index-url https://pypi.ngc.nvidia.com \"onnx-graphsurgeon>=0.3.27\" \"polygraphy\" -U\n",
        "!pip install -q https://huggingface.co/tensorrt/build/resolve/main/tensorrt-9.1.0.post11.dev4-cp310-none-linux_x86_64.whl -U"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "%env LD_LIBRARY_PATH=/content/trt/TensorRT-9.1.0.4/lib:/usr/lib64-nvidia\n",
        "!ln -s /content/trt/TensorRT-9.1.0.4/lib/* /usr/lib\n",
        "# !echo $LD_LIBRARY_PATH\n",
        "%cd /content/TensorRT/demo/Diffusion\n",
        "# !LD_PRELOAD=/content/trt/TensorRT-9.1.0.4/lib/libnvinfer.so.9 python demo_txt2img_xl.py -h\n",
        "# !LD_PRELOAD=/content/trt/TensorRT-9.1.0.4/lib/libnvinfer.so.9 python demo_txt2img_xl.py \"a beautiful photograph of Mt. Fuji during cherry blossom\"\n",
        "!LD_PRELOAD=/content/trt/TensorRT-9.1.0.4/lib/libnvinfer.so.9 python demo_txt2img.py \"a beautiful photograph of Mt. Fuji during cherry blossom\" # sd1.5"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n_Kgy0tjscUC"
      },
      "outputs": [],
      "source": [
        "# https://github.com/NVIDIA/TensorRT/blob/release/8.6/demo/experimental/HuggingFace-Diffusers/TensorRT-diffusers-txt2img.ipynb\n",
        "# https://github.com/huggingface/diffusers/tree/main/examples/community#tensorrt-text2image-stable-diffusion-pipeline\n",
        "# https://github.com/huggingface/diffusers/tree/main/examples/community#tensorrt-image2image-stable-diffusion-pipeline\n",
        "# https://github.com/huggingface/diffusers/tree/main/examples/community#tensorrt-inpainting-stable-diffusion-pipeline\n",
        "\n",
        "# https://github.com/huggingface/diffusers/blob/main/examples/community/stable_diffusion_tensorrt_txt2img.py\n",
        "# https://github.com/huggingface/diffusers/blob/main/examples/community/stable_diffusion_tensorrt_img2img.py\n",
        "# https://github.com/huggingface/diffusers/blob/main/examples/community/stable_diffusion_tensorrt_inpaint.py\n",
        "# https://github.com/huggingface/diffusers/blob/main/examples/community/run_tensorrt_controlnet.py\n",
        "\n",
        "!pip install -q tensorrt diffusers transformers accelerate\n",
        "!pip install -q --extra-index-url https://pypi.ngc.nvidia.com \"onnx-graphsurgeon\" \"onnxruntime\" \"polygraphy\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# !pip install -q https://huggingface.co/tensorrt/build/resolve/main/tensorrt-9.1.0.post11.dev4-cp310-none-linux_x86_64.whl\n",
        "# %cd /content\n",
        "# !wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/9.1.0/tars/tensorrt-9.1.0.4.linux.x86_64-gnu.cuda-11.8.tar.gz\n",
        "# !mkdir trt\n",
        "# %cd /content/trt\n",
        "# !tar -xvzf /content/tensorrt-9.1.0.4.linux.x86_64-gnu.cuda-11.8.tar.gz\n",
        "# !ln -s /content/trt/TensorRT-9.1.0.4/lib/* /usr/lib\n",
        "# !pip install -q diffusers transformers accelerate colored cuda-python ftfy nvtx controlnet_aux onnx onnxruntime\n",
        "# !pip install -q --extra-index-url https://pypi.ngc.nvidia.com \"onnx-graphsurgeon\" \"polygraphy\"\n",
        "\n",
        "import torch\n",
        "from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline\n",
        "from diffusers import DDIMScheduler\n",
        "\n",
        "scheduler = DDIMScheduler.from_pretrained(\"stabilityai/stable-diffusion-2-1\", subfolder=\"scheduler\")\n",
        "pipe_trt = StableDiffusionPipeline.from_pretrained(\"stabilityai/stable-diffusion-2-1\", custom_pipeline=\"stable_diffusion_tensorrt_txt2img\", revision='fp16',\n",
        "    torch_dtype=torch.float16, scheduler=scheduler, image_height=512, image_width=512)\n",
        "\n",
        "pipe_trt.set_cached_folder(\"stabilityai/stable-diffusion-2-1\", revision='fp16')\n",
        "\n",
        "pipe_trt = pipe_trt.to(\"cuda\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "prompt = \"a beautiful photograph of Mt. Fuji during cherry blossom\"\n",
        "\n",
        "# warm up runs to stabilize performance benchmarking\n",
        "num_warm_up_steps=5\n",
        "for _ in range(num_warm_up_steps):\n",
        "    _ = pipe_trt(prompt)\n",
        "\n",
        "image = pipe_trt(prompt).images[0]\n",
        "display(image)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}