Stable Diffusion ControlNet Google Colab

StableDiffusion_ControlNet_Gradio.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "47kV9o1Ni8GH"
      },
      "source": [
        "# **From https://github.com/lllyasviel/ControlNet**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_u3q60di584x",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title ## 1.1. Clone ControlNet\n",
        "#@markdown Clone ControlNet from GitHub and check for updates. Use textbox below if you want to checkout other branch or old commit. Leave it empty to stay the HEAD on main.\n",
        "\n",
        "import os\n",
        "%store -r\n",
        "\n",
        "!nvidia-smi\n",
        "\n",
        "root_dir = \"/content\"\n",
        "%store root_dir\n",
        "repo_dir = str(root_dir)+\"/ControlNet\"\n",
        "%store repo_dir\n",
        "models_dir = str(root_dir)+\"/ControlNet/models\"\n",
        "%store models_dir\n",
        "\n",
        "\n",
        "branch = \"\" #@param {type: \"string\"}\n",
        "repo_url = \"https://github.com/lllyasviel/ControlNet\"\n",
        "\n",
        "def clone_repo():\n",
        "  if os.path.isdir(repo_dir):\n",
        "    print(\"The repository folder already exists, will do a !git pull instead\\n\")\n",
        "    %cd {repo_dir}\n",
        "    !git pull origin {branch} if branch else !git pull\n",
        "  else:\n",
        "    %cd {root_dir}\n",
        "    !git clone {repo_url} {repo_dir}\n",
        "\n",
        "if not os.path.isdir(repo_dir):\n",
        "  clone_repo()\n",
        "\n",
        "%cd {root_dir}\n",
        "os.makedirs(repo_dir, exist_ok=True)\n",
        "os.makedirs(models_dir, exist_ok=True)\n",
        "\n",
        "if branch:\n",
        "  %cd {repo_dir}\n",
        "  status = os.system(f\"git checkout {branch}\")\n",
        "  if status != 0:\n",
        "    raise Exception(\"Failed to checkout branch or commit\")\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ## 1.2. Installing Dependencies\n",
        "#@markdown This will install required Python packages\n",
        "import os\n",
        "%store -r\n",
        "\n",
        "install_xformers = True #@param {'type':'boolean'}\n",
        "\n",
        "%cd {repo_dir}\n",
        "\n",
        "def install_dependencies():\n",
        "    if install_xformers:\n",
        "        !pip install gradio -q #gradio isn't required to be mentioned on requirements for apps on HF space\n",
        "        !sudo apt-get install aria2\n",
        "        !pip install timm==0.6.12\n",
        "        !pip install triton==2.0.0\n",
        "        #!pip install opencv-contrib-python==4.4.0.46 #4.3.0.36\n",
        "        !pip install pytorch-lightning==1.9.1\n",
        "        !pip install einops==0.3.0\n",
        "        !pip install open_clip_torch==2.0.2\n",
        "        !pip install omegaconf==2.1.1\n",
        "        !pip install transformers==4.19.2\n",
        "        !pip install xformers --no-cache-dir\n",
        "        #!pip install --pre -U triton\n",
        "        !sudo apt-get install aria2\n",
        "        !sed -i \"s@block.launch(server_name='0.0.0.0')@block.launch(debug=True, share=True)@\" /content/ControlNet/gradio_canny2image.py\n",
        "        !sed -i \"s@block.launch(server_name='0.0.0.0')@block.launch(debug=True, share=True)@\" /content/ControlNet/gradio_annotator.py\n",
        "        !sed -i \"s@block.launch(server_name='0.0.0.0')@block.launch(debug=True, share=True)@\" /content/ControlNet/gradio_depth2image.py\n",
        "        !sed -i \"s@block.launch(server_name='0.0.0.0')@block.launch(debug=True, share=True)@\" /content/ControlNet/gradio_fake_scribble2image.py\n",
        "        !sed -i \"s@block.launch(server_name='0.0.0.0')@block.launch(debug=True, share=True)@\" /content/ControlNet/gradio_hed2image.py\n",
        "        !sed -i \"s@block.launch(server_name='0.0.0.0')@block.launch(debug=True, share=True)@\" /content/ControlNet/gradio_hough2image.py\n",
        "        !sed -i \"s@block.launch(server_name='0.0.0.0')@block.launch(debug=True, share=True)@\" /content/ControlNet/gradio_normal2image.py\n",
        "        !sed -i \"s@block.launch(server_name='0.0.0.0')@block.launch(debug=True, share=True)@\" /content/ControlNet/gradio_pose2image.py\n",
        "        !sed -i \"s@block.launch(server_name='0.0.0.0')@block.launch(debug=True, share=True)@\" /content/ControlNet/gradio_scribble2image.py\n",
        "        !sed -i \"s@block.launch(server_name='0.0.0.0')@block.launch(debug=True, share=True)@\" /content/ControlNet/gradio_seg2image.py\n",
        "        !sed -i \"s@block.launch(server_name='0.0.0.0')@block.launch(debug=True, share=True)@\" /content/ControlNet/gradio_scribble2image_interactive.py\n",
        "        !sed -i \"s@save_memory = False@save_memory = True@\" /content/ControlNet/config.py\n",
        "\n",
        "install_dependencies()"
      ],
      "metadata": {
        "id": "DExW5D960idK",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ## 2.1. Download ControlNet Model\n",
        "import os\n",
        "%store -r\n",
        "\n",
        "%cd {root_dir}\n",
        "\n",
        "installModels = []\n",
        "installv2Models = []\n",
        "\n",
        "#@markdown ### Available Model\n",
        "#@markdown You must download the model before launching the Gradio interface\n",
        "#@markdown Select one of available model to download:\n",
        "\n",
        "\n",
        "modelUrl = [\"\", \\\n",
        "            \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_canny.pth\", \\\n",
        "            \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_depth.pth\", \\\n",
        "            \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_hed.pth\", \\\n",
        "            \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_mlsd.pth\", \\\n",
        "            \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_normal.pth\", \\\n",
        "            \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_openpose.pth\", \\\n",
        "            \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_scribble.pth\", \\\n",
        "            \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_seg.pth\"]\n",
        "modelList = [\"\", \\\n",
        "             \"control_sd15_canny\", \\\n",
        "             \"control_sd15_depth\", \\\n",
        "             \"control_sd15_hed\", \\\n",
        "             \"control_sd15_mlsd\", \\\n",
        "             \"control_sd15_normal\", \\\n",
        "             \"control_sd15_openpose\", \\\n",
        "             \"control_sd15_scribble\", \\\n",
        "             \"control_sd15_seg\" ]\n",
        "modelName = \"control_sd15_canny\" #@param [\"control_sd15_canny\", \"control_sd15_depth\", \"control_sd15_hed\", \"control_sd15_mlsd\", \"control_sd15_normal\", \"control_sd15_openpose\", \"control_sd15_scribble\", \"control_sd15_seg\"]\n",
        "\n",
        "\n",
        "if modelName != \"\":\n",
        "  installModels.append((modelName, modelUrl[modelList.index(modelName)]))\n",
        "\n",
        "def install(checkpoint_name, url):\n",
        "  ext = \"pth\" if url.endswith(\".pth\") else \"pt\"\n",
        "\n",
        "  hf_token = ''\n",
        "  user_header = f\"\\\"Authorization: Bearer {hf_token}\\\"\"\n",
        "  !aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 -d {repo_dir}/models -o {checkpoint_name}.{ext} \"{url}\"\n",
        "\n",
        "def install_checkpoint():\n",
        "  for model in installModels:\n",
        "    install(model[0], model[1])\n",
        "\n",
        "install_checkpoint()"
      ],
      "metadata": {
        "id": "YnpIRGP6z0EC",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ## 2.1. Download Additional model\n",
        "#@markdown You must download the model before launching the Gradio interface\n",
        "\n",
        "!aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d /content/ControlNet/annotator/ckpts -o body_pose_model.pth \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/body_pose_model.pth\"\n",
        "!aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d /content/ControlNet/annotator/ckpts -o dpt_hybrid-midas-501f0c75.pt \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt\"\n",
        "!aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d /content/ControlNet/annotator/ckpts -o hand_pose_model.pth \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/hand_pose_model.pth\"\n",
        "!aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d /content/ControlNet/annotator/ckpts -o mlsd_large_512_fp32.pth \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/mlsd_large_512_fp32.pth\"\n",
        "!aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d /content/ControlNet/annotator/ckpts -o mlsd_tiny_512_fp32.ptht \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/mlsd_tiny_512_fp32.pth\"\n",
        "!aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d /content/ControlNet/annotator/ckpts -o network-bsds500.pth \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/network-bsds500.pth\"\n",
        "!aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d /content/ControlNet/annotator/ckpts -o upernet_global_small.pth \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/upernet_global_small.pth\""
      ],
      "metadata": {
        "cellView": "form",
        "id": "NTKm3RNOXrgg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Start **ControlNet** ##"
      ],
      "metadata": {
        "id": "yyhHbK0NEfzY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%cd {repo_dir}\n",
        "\n",
        "#@markdown Example : If you downloaded the **control_sd15_canny** model you must select **gradio_canny2image.py**\n",
        "\n",
        "gradioLink = \"gradio_canny2image.py\" #@param [\"gradio_canny2image.py\", \"gradio_annotator.py\", \"gradio_depth2image.py\", \"gradio_fake_scribble2image.py\", \"gradio_hed2image.py\", \"gradio_hough2image.py\", \"gradio_normal2image.py\", \"gradio_pose2image.py\", \"gradio_scribble2image.py\"]\n",
        "\n",
        "\n",
        "if gradioLink != \"\":\n",
        "   !python {gradioLink}"
      ],
      "metadata": {
        "cellView": "form",
        "id": "kJ8THSaY_rIO"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}