Diffusion_WebUI.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "private_outputs": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/acheong08/Diffusion-ColabUI/blob/main/Diffusion_WebUI.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Diffusion WebUI\n",
        "Choose your models and just run it. No extra steps"
      ],
      "metadata": {
        "id": "QOmvhMFXLjbI"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### [GitHub](https://github.com/acheong08/NovelAI-Colab)\n",
        "# Quick start: \n",
        "Run the Unified code.\n",
        "\n",
        "(If there are any issues, report them on GitHub)"
      ],
      "metadata": {
        "id": "lOC4X1umLtAi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Check dependencies (Run this if you get errors)\n",
        "import os\n",
        "\n",
        "!nvidia-smi\n",
        "!pip install pytorch_lightning\n",
        "import pytorch_lightning as pl\n",
        "print(pl.__version__)\n"
      ],
      "metadata": {
        "id": "twBYi6Y4aj_a",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Unified code"
      ],
      "metadata": {
        "id": "TL9o4lz5XK2y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ## Add models\n",
        "#@markdown - You can run this multiple times to add more models to the list\n",
        "#@markdown - Leave input blank to ignore\n",
        "#@markdown - No spaces allowed in the name\n",
        "\n",
        "try:\n",
        "  customModels\n",
        "except NameError:\n",
        "  customModels = []\n",
        "else:\n",
        "  pass\n",
        "\n",
        "#@markdown ### Custom\n",
        "#@markdown The model URL should be a direct download link, Google Drive URL or Civitai download API URL\n",
        "modelName = \"\" #@param {'type': 'string'}\n",
        "modelURL = \"\" #@param [\"\", \"https://huggingface.co/ShadoWxShinigamI/MidJourney-PaperCut/resolve/main/Mdjrny-pprct_step_7000.ckpt\", \"https://huggingface.co/hesw23168/SD-Elysium-Model/resolve/main/Elysium_V1.ckpt\", \"https://huggingface.co/acheong08/nutmegmix/resolve/main/nutmegmix.ckpt\", \"https://huggingface.co/hesw23168/SD-Elysium-Model/resolve/main/Elysium_Anime_V2.ckpt\", \"https://huggingface.co/Linaqruf/anything-v3.0/resolve/main/Anything-V3.0-pruned.ckpt\", \"https://huggingface.co/prompthero/linkedin-diffusion/resolve/main/lnkdn.ckpt\"] {allow-input: true}\n",
        "\n",
        "#@markdown ## Defaults\n",
        "#@markdown SDv2.0 and mdjrny-v4 style has been added\n",
        "defaultModelURLs = [\"\" , \"https://huggingface.co/XpucT/Deliberate/resolve/main/Deliberate_v2.safetensors\",\"https://huggingface.co/acheong08/secretAI/resolve/main/stableckpt/animefull-final-pruned/model.ckpt\", \"https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt\", \"https://huggingface.co/hakurei/waifu-diffusion-v1-3/resolve/main/wd-v1-3-float32.ckpt\", \"https://huggingface.co/Deltaadams/HentaiDiffusion/resolve/main/HD-16.ckpt\", \"https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.ckpt\", \"https://huggingface.co/runwayml/stable-diffusion-inpainting/resolve/main/sd-v1-5-inpainting.ckpt\", \"https://huggingface.co/acheong08/f222/resolve/main/f222.ckpt\", \"https://huggingface.co/andite/anything-v4.0/resolve/main/anything-v4.0-pruned-fp32.ckpt\", \"https://huggingface.co/andite/anything-v4.0/resolve/main/anything-v4.5-pruned-fp32.ckpt\", \"https://huggingface.co/acheong08/r34/resolve/main/r34_e4.ckpt\", \"https://huggingface.co/hassanblend/HassanBlend1.5.1.2/resolve/main/HassanBlend1.5.1.2-pruned.safetensors\", \"https://huggingface.co/stabilityai/stable-diffusion-2/resolve/main/768-v-ema.ckpt\", \"https://huggingface.co/prompthero/openjourney/resolve/main/mdjrny-v4.ckpt\", \"https://huggingface.co/WarriorMama777/OrangeMixs/resolve/main/Models/AbyssOrangeMix2/AbyssOrangeMix2_nsfw.safetensors\"]\n",
        "defaultModels = [\"\" , \"deliberate\",\"NovelAI\", \"SD1-4\", \"WaifuDiffusion\", \"HentaiDiffusion\", \"SD1-5\", \"SD-inpainting\", \"f222-NSFW\", \"Anything-V4\", \"Anything-V4.5\", \"Rule34\", \"HassansBlend\", \"SDv2-0\", \"OpenMidjourney\", \"AbyssOrangeMix2_nsfw\"]\n",
        "defaultModel = \"deliberate\" #@param [\"\" ,\"deliberate\", \"NovelAI\", \"SD1-4\", \"WaifuDiffusion\", \"HentaiDiffusion\", \"SD1-5\", \"SD-inpainting\", \"f222-NSFW\", \"Anything-V4\", \"Anything-V4.5\", \"Rule34\", \"HassansBlend\", \"SDv2-0\", \"OpenMidjourney\", \"AbyssOrangeMix2_nsfw\"]\n",
        "\n",
        "#@markdown # Parameters\n",
        "#@markdown Select this if the model is based on SD2.0\n",
        "SD2 = \"None\" #@param [\"None\", \"V2-512\", \"V2-768\", \"V2-inpainting\", \"V2-upscaling\"]\n",
        "\n",
        "if modelName == \"\" or modelURL == \"\":\n",
        "  pass\n",
        "else:\n",
        "  if modelURL.endswith(\"safetensors\"):\n",
        "    modelName = modelName + \".safetensors\"\n",
        "  elif modelURL.endswith(\"ckpt\"):\n",
        "    modelName = modelName + \".ckpt\"\n",
        "  else:\n",
        "    pass\n",
        "  customModels.append((modelName, modelURL, SD2))\n",
        "\n",
        "if defaultModel != \"\":\n",
        "  if defaultModelURLs[defaultModels.index(defaultModel)].endswith(\"safetensors\"):\n",
        "    defaultModelName = defaultModel + \".safetensors\"\n",
        "  elif defaultModel.endswith(\".safetensors\"):\n",
        "    defaultModelName = defaultModel\n",
        "  else:\n",
        "    defaultModelName = defaultModel + \".ckpt\"\n",
        "  # Map model to URL\n",
        "  customModels.append((defaultModelName, defaultModelURLs[defaultModels.index(defaultModel)], SD2))\n",
        "\n",
        "print(customModels)\n",
        "from IPython import display\n",
        "display.HTML(\"<p>Want faster speed? <a href=\\\"https://www.youtube.com/watch?v=dQw4w9WgXcQ\\\">Here</a>'s a better version of this Colab!</p>\")\n"
      ],
      "metadata": {
        "id": "dtrNFeQiYDz-",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ## Clear model list\n",
        "customModels = []"
      ],
      "metadata": {
        "cellView": "form",
        "id": "LEPOSoQ3_xNT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Start WebUI\n",
        "#@markdown # Instructions:\n",
        "#@markdown Select your favorite model (Or all of them)\n",
        "#@markdown 0. Set models in the cell above\n",
        "#@markdown 1. Run this cell.\n",
        "#@markdown  <br> - Ignore alerts about disk space. You got plenty\n",
        "#@markdown 2. Wait\n",
        "#@markdown 3. Open gradio link\n",
        "\n",
        "#@markdown #### The default username is `webui` and password is `diffusion`\n",
        "\n",
        "!pip3 install numpy --upgrade\n",
        "#@markdown # Options\n",
        "from google.colab import drive\n",
        "import os\n",
        "import time\n",
        "from subprocess import getoutput\n",
        "vae_args = \"\"\n",
        "save_to_gdrive = False\n",
        "\n",
        "root_dir = \"/content\"\n",
        "\n",
        "#@markdown Import LoRAs stored on your Google Drive. Path is relative to your drive's root.\n",
        "use_gdrive_loras = False #@param {'type':'boolean'}\n",
        "gdrive_loras_path = \"StableDiffusion/LoRAs\" #@param {'type': 'string'}\n",
        "\n",
        "if use_gdrive_loras:\n",
        "  drive.mount('/content/drive')\n",
        "\n",
        "if save_to_gdrive:\n",
        "  !mkdir -p {'/content/drive/MyDrive/AI/output'}\n",
        "  output_path = \"--img-output-dir /content/drive/MyDrive/AI/output\"\n",
        "else:\n",
        "  root_dir = \"/content\"\n",
        "  output_path = \"\"\n",
        "\n",
        "%cd {root_dir}\n",
        "\n",
        "def get_gdrive_loras():\n",
        "  path_loras = root_dir + '/stable-diffusion-webui/models/Lora'\n",
        "  if os.path.exists(path_loras) and not os.path.islink(path_loras):\n",
        "    !rmdir {root_dir}/stable-diffusion-webui/models/Lora\n",
        "  if not os.path.exists(path_loras):\n",
        "    os.symlink('/content/drive/MyDrive/' + gdrive_loras_path, root_dir + '/stable-diffusion-webui/models/Lora')\n",
        "\n",
        "def get_hypernetworks():\n",
        "  !mkdir -p {root_dir}/stable-diffusion-webui/models/hypernetworks\n",
        "  hypernetworks = ['anime_2.pt', 'anime.pt', 'anime_3.pt', 'furry_2.pt', 'furry_3.pt', 'furry_protogen.pt', 'furry_transformation.pt', 'furry_scalie.pt', 'pony.pt', 'aini.pt', 'furry.pt', 'furry_kemono.pt']\n",
        "  for network in hypernetworks:\n",
        "    !wget -c https://huggingface.co/acheong08/secretAI/resolve/main/stableckpt/modules/modules/{network} -O {root_dir}/stable-diffusion-webui/models/hypernetworks/{network}\n",
        "\n",
        "def install_aria():\n",
        "  if not os.path.exists('/usr/bin/aria2c'):\n",
        "    !apt install -y -qq aria2\n",
        "\n",
        "def custom_model(checkpoint_name, url, SD2):\n",
        "  inference_url = \"https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/\"\n",
        "  if SD2 == \"V2-768\":\n",
        "    inference_url += \"v2-inference-v.yaml\"\n",
        "  elif SD2 == \"V2-512\":\n",
        "    inference_url += \"v2-inference.yaml\"\n",
        "  elif SD2 == \"V2-inpainting\":\n",
        "    inference_url += \"v2-inpainting-inference.yaml\"\n",
        "  elif SD2 == \"V2-upscaling\":\n",
        "    inference_url += \"x4-upscaling.yaml\"\n",
        "    \n",
        "  if SD2.startswith(\"V2\"):\n",
        "    checkpoint_config = checkpoint_name.split('.')[0]\n",
        "    !wget {inference_url} -O {root_dir}/stable-diffusion-webui/models/Stable-diffusion/{checkpoint_config}.yaml\n",
        "\n",
        "  if url.startswith(\"https://drive.google.com\"):\n",
        "    !gdown --fuzzy -O \"{root_dir}/stable-diffusion-webui/models/Stable-diffusion/{checkpoint_name}\" \"{url}\"\n",
        "  elif url.startswith(\"magnet:?\"):\n",
        "    install_aria()\n",
        "    !aria2c --summary-interval=10 -c -x 10 -k 1M -s 10 -o {root_dir}/stable-diffusion-webui/models/Stable-diffusion/{checkpoint_name} \"{url}\"\n",
        "  elif url.startswith(\"https://civitai.com/api/download\"):\n",
        "    !wget -c --header=\"User-Agent: Diffusion_WebUI/Fork_Ganbat\" \"{url}\" -O {root_dir}/stable-diffusion-webui/models/Stable-diffusion/{checkpoint_name}\n",
        "  else:\n",
        "    #@markdown You can use the default token or enter your own for private models\n",
        "    user_token = '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' #@param {'type': 'string'}\n",
        "    user_header = f\"\\\"Authorization: Bearer {user_token}\\\"\"\n",
        "    \n",
        "    !wget -c --header={user_header} {url} -O {root_dir}/stable-diffusion-webui/models/Stable-diffusion/{checkpoint_name}\n",
        "\n",
        "def install_deps():\n",
        "  %cd {root_dir}\n",
        "  global new_ui\n",
        "  new_ui = False #@param {'type': 'boolean'}\n",
        "  !git clone https://github.com/acheong08/stable-diffusion-webui --recursive --branch={\"v2\" if new_ui else \"main\"}\n",
        "  %cd {root_dir}/stable-diffusion-webui/\n",
        "  !mkdir -p {root_dir}/stable-diffusion-webui/models/Stable-diffusion/\n",
        "  !git pull\n",
        "  %cd {root_dir}\n",
        "  use_hypernetworks = False #@param {'type':'boolean'}\n",
        "\n",
        "  # Get models\n",
        "  for model in customModels:\n",
        "    custom_model(model[0], model[1], model[2])\n",
        "  if use_hypernetworks:\n",
        "    get_hypernetworks()\n",
        "  if use_gdrive_loras:\n",
        "    get_gdrive_loras()\n",
        "  \n",
        "  use_controlnet = False #@param {'type':'boolean'}\n",
        "  if use_controlnet:\n",
        "    !git clone \"https://huggingface.co/lllyasviel/ControlNet-v1-1\" ControlNetModels --depth 1 -j 12\n",
        "    !mv {root_dir}/ControlNetModels/* {root_dir}/stable-diffusion-webui/extensions-builtin/sd-webui-controlnet/models/\n",
        "  \n",
        "def run_webui():\n",
        "  vae = \"Stable Diffusion\" #@param [\"\", \"Anime\", \"Stable Diffusion\", \"Anything-V4\"]\n",
        "  if vae == \"Anime\":\n",
        "    !wget -c https://huggingface.co/acheong08/secretAI/resolve/main/stableckpt/animevae.pt -O {root_dir}/stable-diffusion-webui/models/Stable-diffusion/novelAI.vae.pt\n",
        "    vae_args = str(\"--vae \" if not new_ui else \"--vae-path \") + root_dir + \"/stable-diffusion-webui/models/Stable-diffusion/novelAI.vae.pt\"\n",
        "  elif vae == \"Stable Diffusion\":\n",
        "    !wget -c https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt -O {root_dir}/stable-diffusion-webui/models/Stable-diffusion/sd-v1-5.vae.pt\n",
        "    vae_args = str(\"--vae \" if not new_ui else \"--vae-path \") + root_dir + \"/stable-diffusion-webui/models/Stable-diffusion/sd-v1-5.vae.pt\"\n",
        "  elif vae == \"Anything-V4\":\n",
        "    !wget -c https://huggingface.co/andite/anything-v4.0/resolve/main/anything-v4.0.vae.pt -O {root_dir}/stable-diffusion-webui/models/Stable-diffusion/anything-v4.vae.pt\n",
        "    vae_args = str(\"--vae \" if not new_ui else \"--vae-path \") + root_dir + \"/stable-diffusion-webui/models/Stable-diffusion/anything-v4.vae.pt\"\n",
        "\n",
        "  %cd {root_dir}/stable-diffusion-webui/\n",
        "  vram = \"\" #@param [\"--medvram\", \"--lowvram\", \"\"]\n",
        "  other_args = \"--share\" #@param [\"\"] {'allow-input': true}\n",
        "  #@markdown Tick this if you will be handling models greater than 7GB or merging multiple models\n",
        "  large_models = False #@param {'type': 'boolean'}\n",
        "  !IGNORE_CMD_ARGS_ERRORS=\"true\" COMMANDLINE_ARGS=\"{other_args} {vae_args} {vram}\" REQS_FILE=\"requirements.txt\" python launch.py\n",
        "install_deps()\n",
        "run_webui()"
      ],
      "metadata": {
        "id": "Y4qjIc1XXKWw",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Saving your images"
      ],
      "metadata": {
        "id": "_xPLSbzNTgF2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Download file manually from files tab or save to Google Drive\n",
        "%cd /content/stable-diffusion-webui/\n",
        "!zip -r /content/output.zip outputs\n",
        "\n",
        "from pydrive.auth import GoogleAuth\n",
        "from google.colab import drive\n",
        "from pydrive.drive import GoogleDrive\n",
        "from google.colab import auth\n",
        "from oauth2client.client import GoogleCredentials\n",
        "def create_folder(folder_name):\n",
        "    # Check if folder exists\n",
        "    file_list = drive.ListFile({'q': \"title='{}' and mimeType='application/vnd.google-apps.folder' and trashed=false\".format(folder_name)}).GetList()\n",
        "    if len(file_list) > 0:\n",
        "        # Folder exists\n",
        "        print('Debug: Folder exists')\n",
        "        folder_id = file_list[0]['id']\n",
        "    else:\n",
        "        print('Debug: Creating folder')\n",
        "        file = drive.CreateFile({'title': folder_name, 'mimeType': 'application/vnd.google-apps.folder'})\n",
        "        file.Upload()\n",
        "        folder_id = file.attr['metadata']['id']\n",
        "    # return folder id\n",
        "    return folder_id\n",
        "# Upload file to Google Drive\n",
        "def upload_file(file_name, folder_id, save_as):\n",
        "    # Check if file exists\n",
        "    file_list = drive.ListFile({'q': \"title='{}' and trashed=false\".format(save_as)}).GetList()\n",
        "    if len(file_list) > 0:\n",
        "        print('Debug: File already exists')\n",
        "        # Change file name to avoid overwriting\n",
        "        save_as = save_as + ' (1)'\n",
        "    file = drive.CreateFile({'title': save_as, 'parents': [{'id': folder_id}]})\n",
        "    file.SetContentFile(file_name)\n",
        "    # Upload and set permission to public\n",
        "    file.Upload()\n",
        "    file.InsertPermission({'type': 'anyone', 'value': 'anyone', 'role': 'reader'})\n",
        "    # return file id\n",
        "    return file.attr['metadata']['id']\n",
        "\n",
        "use_drive = True #@param {type:\"boolean\"}\n",
        "folder_name = \"AI_pic_archive\" #@param {type: \"string\"}\n",
        "save_as = \"oni.zip\" #@param {type: \"string\"}\n",
        "\n",
        "if use_drive:\n",
        "  auth.authenticate_user()\n",
        "  gauth = GoogleAuth()\n",
        "  gauth.credentials = GoogleCredentials.get_application_default()\n",
        "  drive = GoogleDrive(gauth)\n",
        "  file_id = upload_file('/content/output.zip', create_folder(folder_name), save_as)\n",
        "  print(\"Your sharing link: https://drive.google.com/file/d/\" + file_id + \"/view?usp=sharing\")  "
      ],
      "metadata": {
        "id": "lpiP4AkYTic0",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content\n",
        "!rm -rf stable-diffusion-webui"
      ],
      "metadata": {
        "id": "2wc3I3oj9i5l"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!git clone https://huggingface.co/lllyasviel/ControlNet-v1-1 ControlNetModels --depth 1"
      ],
      "metadata": {
        "id": "n0Asxkvy9n--"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!mv ControlNetModels/* {root_dir}/stable-diffusion-webui/extensions-builtin/sd-webui-controlnet/models/"
      ],
      "metadata": {
        "id": "7bFBeUV9IuZN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content"
      ],
      "metadata": {
        "id": "Z8e8U21MHVDb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!rm -rf ControlNetModels"
      ],
      "metadata": {
        "id": "OFmbjTIuHJcm"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}