cagliostro-colab-ui.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Linaqruf/sd-notebook-collection/blob/main/cagliostro-colab-ui.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "[![visitor][visitor-badge]][visitor-stats]\n",
        "[![ko-fi][ko-fi-badge]][ko-fi-link]\n",
        "\n",
        "# **Cagliostro Colab UI**\n",
        "All-in-One, Customizable and Flexible Stable Diffusion for Google Colab.\n",
        "\n",
        "**Version 3.0.0** | [Github][link-to-github] | [What's New?][README]\n",
        "\n",
        "<!-- [visitor-badge]: https://visitor-badge.glitch.me/badge?page_id=linaqruf.cag-webui -->\n",
        "[visitor-badge]: https://api.visitorbadge.io/api/visitors?path=Cagliostro%20Colab%20UI&label=Visitors&labelColor=%2334495E&countColor=%231ABC9C&style=flat&labelStyle=none\n",
        "[visitor-stats]: https://visitorbadge.io/status?path=Cagliostro%20Colab%20UI\n",
        "[ko-fi-badge]: https://img.shields.io/badge/Support%20me%20on%20Ko--fi-F16061?logo=ko-fi&logoColor=white&style=flat\n",
        "[ko-fi-link]: https://ko-fi.com/linaqruf\n",
        "[link-to-github]: https://github.com/Linaqruf/sd-notebook-collection/blob/main/cagliostro-colab-ui.ipynb\n",
        "[README]: https://github.com/Linaqruf/sd-notebook-collection/blob/main/README.md#whats-new\n",
        "[MANUAL]: https://github.com/Linaqruf/sd-notebook-collection/blob/main/MANUAL.md#cagliostro-colab-ui-user-manual\n",
        "\n"
      ],
      "metadata": {
        "id": "WgQr3s96015a"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title ## **Install Cagliostro Colab UI**\n",
        "import os\n",
        "import time\n",
        "import json\n",
        "import base64\n",
        "import shutil\n",
        "import subprocess\n",
        "import threading\n",
        "import sys\n",
        "from IPython.display import display, HTML\n",
        "from google.colab.output import eval_js\n",
        "from IPython.utils import capture\n",
        "from tqdm import tqdm\n",
        "\n",
        "python_version      = \".\".join(sys.version.split(\".\")[:2])\n",
        "colablib_path       = f\"/usr/local/lib/python{python_version}/dist-packages/colablib\"\n",
        "if not os.path.exists(colablib_path):\n",
        "    subprocess.run(['pip', 'install', 'git+https://github.com/Linaqruf/colablib'])\n",
        "\n",
        "from colablib.colored_print import cprint, print_line\n",
        "from colablib.utils import py_utils, config_utils, package_utils\n",
        "from colablib.utils.config_utils import pastebin_reader as read\n",
        "from colablib.utils.ubuntu_utils import ubuntu_deps\n",
        "from colablib.sd_models.downloader import aria2_download\n",
        "from colablib.utils.git_utils import update_repo, batch_update, validate_repo, reset_repo, patch_repo\n",
        "\n",
        "%store -r\n",
        "\n",
        "################################\n",
        "# COLAB ARGUMENTS GOES HERE\n",
        "################################\n",
        "\n",
        "# @markdown ### **Drive Config**\n",
        "mount_drive         = False  # @param {type:'boolean'}\n",
        "output_drive_folder = \"cagliostro-colab-ui/outputs\" #@param {type:'string'}\n",
        "# @markdown ### **Repo Config**\n",
        "repo_type           = \"Anapnoe\" #@param [\"AUTOMATIC1111\", \"AUTOMATIC1111-Dev\", \"Anapnoe\"]\n",
        "update_webui        = False  # @param {type:'boolean'}\n",
        "update_extensions   = False  # @param {type:'boolean'}\n",
        "commit_hash         = \"\"  # @param {type:'string'}\n",
        "dpmpp_2m_v2_patch   = True  # @param {type:'boolean'}\n",
        "# @markdown ### **Optimization Config**\n",
        "# @markdown > It's not recommended to set params below to `True` if you have **Colab Pro** subscription.\n",
        "colab_optimizations = True  # @param {type:'boolean'}\n",
        "# @markdown > Specify `mobile_optimizations` to keep colab tab alive for mobile users\n",
        "mobile_optimizations = False  # @param {type:'boolean'}\n",
        "\n",
        "################################\n",
        "# DIRECTORY CONFIG\n",
        "################################\n",
        "\n",
        "# VAR\n",
        "voldemort, voldy = read(\"kq6ZmHFU\")[:2]\n",
        "\n",
        "# ROOT DIR\n",
        "root_dir            = \"/content\"\n",
        "drive_dir           = os.path.join(root_dir, \"drive\", \"MyDrive\")\n",
        "repo_dir            = os.path.join(root_dir, \"cagliostro-colab-ui\")\n",
        "tmp_dir             = os.path.join(root_dir, \"tmp\")\n",
        "patches_dir         = os.path.join(root_dir, \"patches\")\n",
        "deps_dir            = os.path.join(root_dir, \"deps\")\n",
        "fused_dir           = os.path.join(root_dir, \"fused\")\n",
        "\n",
        "# REPO DIR\n",
        "models_dir          = os.path.join(repo_dir, \"models\", \"Stable-diffusion\")\n",
        "vaes_dir            = os.path.join(repo_dir, \"models\", \"VAE\")\n",
        "hypernetworks_dir   = os.path.join(repo_dir, \"models\", \"hypernetworks\")\n",
        "lora_dir            = os.path.join(repo_dir, \"models\", \"Lora\")\n",
        "control_dir         = os.path.join(repo_dir, \"models\", \"ControlNet\")\n",
        "esrgan_dir          = os.path.join(repo_dir, \"models\", \"ESRGAN\")\n",
        "embeddings_dir      = os.path.join(repo_dir, \"embeddings\")\n",
        "extensions_dir      = os.path.join(repo_dir, \"extensions\")\n",
        "annotator_dir       = os.path.join(extensions_dir, f\"{voldy}-controlnet\", \"annotator\")\n",
        "output_subdir       = [\"txt2img-images\", \"img2img-images\", \"extras-images\", \"txt2img-grids\", \"img2img-grids\"]\n",
        "\n",
        "# CONFIG\n",
        "config_file         = os.path.join(repo_dir, \"config.json\")\n",
        "ui_config_file      = os.path.join(repo_dir, \"ui-config.json\")\n",
        "style_path          = os.path.join(repo_dir, \"style.css\")\n",
        "download_list       = os.path.join(root_dir, \"download_list.txt\")\n",
        "\n",
        "\n",
        "################################\n",
        "# REPO TYPE CONFIG\n",
        "################################\n",
        "\n",
        "repo_type_lower = repo_type.lower()\n",
        "\n",
        "package_url = [\n",
        "    f\"https://huggingface.co/Linaqruf/fast-repo/resolve/main/{repo_type_lower}-webui.tar.lz4\",\n",
        "    f\"https://huggingface.co/Linaqruf/fast-repo/resolve/main/{repo_type_lower}-webui-deps.tar.lz4\",\n",
        "    f\"https://huggingface.co/Linaqruf/fast-repo/resolve/main/{repo_type_lower}-webui-cache.tar.lz4\",\n",
        "]\n",
        "\n",
        "repo_type_to_repo_name = {\n",
        "    \"anapnoe\"           : f\"anapnoe/{voldemort}-ux\",\n",
        "    \"automatic1111\"     : f\"AUTOMATIC1111/{voldemort}\",\n",
        "    \"automatic1111-dev\" : f\"AUTOMATIC1111/{voldemort}\",\n",
        "}\n",
        "\n",
        "branch_type_to_branch = {\n",
        "    \"automatic1111\"     : \"master\",\n",
        "    \"automatic1111-dev\" : \"dev\"\n",
        "}\n",
        "\n",
        "with capture.capture_output() as cap:\n",
        "    for dir in  [\"root_dir\", \"fused_dir\", \"repo_dir\", \"tmp_dir\", \"models_dir\", \"vaes_dir\", \"hypernetworks_dir\", \"embeddings_dir\", \"extensions_dir\", \"lora_dir\", \"control_dir\", \"esrgan_dir\"]:\n",
        "        %store {dir}\n",
        "    for file in [\"config_file\", \"ui_config_file\", \"style_path\", \"download_list\"]:\n",
        "        %store {file}\n",
        "    for var in  [\"voldemort\", \"voldy\"]:\n",
        "        %store {var}\n",
        "    del cap\n",
        "\n",
        "def mount_func(directory):\n",
        "    output_dir = os.path.join(repo_dir, \"outputs\")\n",
        "\n",
        "    if mount_drive:\n",
        "        print_line(80, color=\"green\")\n",
        "        if not os.path.exists(directory):\n",
        "            from google.colab import drive\n",
        "            cprint(\"Mounting google drive...\", color=\"green\", reset=False)\n",
        "            drive.mount(os.path.dirname(directory))\n",
        "        output_dir  = os.path.join(directory, output_drive_folder)\n",
        "        cprint(\"Set default output path to:\", output_dir, color=\"green\")\n",
        "\n",
        "    return output_dir\n",
        "\n",
        "def setup_directories():\n",
        "    for dir in [fused_dir, models_dir, vaes_dir,\n",
        "                hypernetworks_dir, embeddings_dir, extensions_dir,\n",
        "                lora_dir, control_dir, esrgan_dir]:\n",
        "        os.makedirs(dir, exist_ok=True)\n",
        "\n",
        "def pre_download(dir, urls, desc, overwrite=False):\n",
        "    gpu_info          = py_utils.get_gpu_info()\n",
        "    version           = py_utils.get_python_version().split()[0]\n",
        "    major_minor       = \".\".join(version.split(\".\")[:2])\n",
        "    python_path       = f\"/usr/local/lib/python{major_minor}/dist-packages/\"\n",
        "    ffmpy_path        = os.path.join(python_path, \"ffmpy-0.3.0.dist-info\")\n",
        "\n",
        "    for url in tqdm(urls, desc=desc):\n",
        "        filename  = py_utils.get_filename(url)\n",
        "        aria2_download(dir, filename, url, quiet=True)\n",
        "        if filename == f\"{repo_type.lower()}-webui-deps.tar.lz4\":\n",
        "            package_utils.extract_package(filename, python_path, overwrite=True)\n",
        "        else:\n",
        "            package_utils.extract_package(filename, \"/\", overwrite=overwrite)\n",
        "        os.remove(filename)\n",
        "\n",
        "    if os.path.exists(ffmpy_path):\n",
        "        shutil.rmtree(ffmpy_path)\n",
        "\n",
        "    subprocess.run(['pip', 'uninstall', '-y', 'xformers'], check=True)\n",
        "\n",
        "def install_dependencies():\n",
        "    ubuntu_deps = [\"aria2\", \"lz4\", \"unionfs-fuse\"]\n",
        "    cprint(\"Installing ubuntu dependencies\", color=\"green\")\n",
        "    subprocess.run([\"apt\", \"install\"] + ubuntu_deps)\n",
        "\n",
        "def install_webui(repo_dir, desc):\n",
        "    try:\n",
        "        if not os.path.exists(repo_dir):\n",
        "            pre_download(root_dir, package_url, desc, overwrite=False)\n",
        "            return\n",
        "\n",
        "        repo_name, _, current_branch = validate_repo(repo_dir)\n",
        "        repo_type_lower = repo_type.lower()\n",
        "        expected_repo_name = repo_type_to_repo_name.get(repo_type_lower)\n",
        "\n",
        "        if expected_repo_name == repo_name:\n",
        "            expected_branch = branch_type_to_branch.get(repo_type_lower)\n",
        "            if expected_branch is None or expected_branch == current_branch:\n",
        "                cprint(f\"'{repo_name}' {current_branch if expected_branch else ''} already installed, skipping...\", color=\"green\")\n",
        "                return\n",
        "\n",
        "        cprint(f\"Another repository exist. Uninstall '{repo_name}'...\", color=\"green\")\n",
        "        shutil.rmtree(repo_dir)\n",
        "        pre_download(root_dir, package_url, desc)\n",
        "    except Exception as e:\n",
        "        cprint(f\"An error occurred: {e}\", color=\"green\")\n",
        "\n",
        "def configure_output_path(config_path, output_dir, output_subdir):\n",
        "    config = config_utils.read_config(config_path)\n",
        "    config_updates = {\n",
        "        \"outdir_txt2img_samples\"  : os.path.join(output_dir, output_subdir[0]),\n",
        "        \"outdir_img2img_samples\"  : os.path.join(output_dir, output_subdir[1]),\n",
        "        \"outdir_extras_samples\"   : os.path.join(output_dir, output_subdir[2]),\n",
        "        \"outdir_txt2img_grids\"    : os.path.join(output_dir, output_subdir[3]),\n",
        "        \"outdir_img2img_grids\"    : os.path.join(output_dir, output_subdir[4])\n",
        "    }\n",
        "\n",
        "    config.update(config_updates)\n",
        "    config_utils.write_config(config_path, config)\n",
        "\n",
        "    for dir in output_subdir:\n",
        "        os.makedirs(os.path.join(output_dir, dir), exist_ok=True)\n",
        "\n",
        "def prepare_environment():\n",
        "    cprint(f\"Preparing environment...\", color=\"green\")\n",
        "\n",
        "    os.environ[\"colab_url\"]               = eval_js(\"google.colab.kernel.proxyPort(7860, {'cache': false})\")\n",
        "    os.environ[\"TF_CPP_MIN_LOG_LEVEL\"]    = \"3\"\n",
        "    os.environ[\"SAFETENSORS_FAST_GPU\"]    = \"1\"\n",
        "    os.environ['PYTORCH_CUDA_ALLOC_CONF'] = \"garbage_collection_threshold:0.9,max_split_size_mb:512\"\n",
        "    os.environ[\"PYTHONWARNINGS\"]          = \"ignore\"\n",
        "\n",
        "def play_audio(url):\n",
        "    display(HTML(f'<audio src=\"{url}\" controls autoplay style=\"display:none\"></audio>'))\n",
        "\n",
        "def main():\n",
        "    global output_dir\n",
        "\n",
        "    os.chdir(root_dir)\n",
        "    start_time = time.time()\n",
        "\n",
        "    output_dir = mount_func(drive_dir)\n",
        "\n",
        "    gpu_info    = py_utils.get_gpu_info(get_gpu_name=True)\n",
        "    python_info = py_utils.get_python_version()\n",
        "    torch_info  = py_utils.get_torch_version()\n",
        "\n",
        "    print_line(80, color=\"green\")\n",
        "    cprint(f\" [-] Current GPU:\", gpu_info, color=\"flat_yellow\")\n",
        "    cprint(f\" [-] Python\", python_info, color=\"flat_yellow\")\n",
        "    cprint(f\" [-] Torch\", torch_info, color=\"flat_yellow\")\n",
        "    print_line(80, color=\"green\")\n",
        "\n",
        "    install_dependencies()\n",
        "\n",
        "    print_line(80, color=\"green\")\n",
        "    install_webui(repo_dir, cprint(f\"Unpacking {repo_type} Webui\", color=\"green\", tqdm_desc=True))\n",
        "    prepare_environment()\n",
        "\n",
        "    configure_output_path(config_file, output_dir, output_subdir)\n",
        "\n",
        "    print_line(80, color=\"green\")\n",
        "    if update_webui and not commit_hash:\n",
        "        update_repo(cwd=repo_dir, args=\"-X theirs --rebase --autostash\")\n",
        "\n",
        "    setup_directories ()\n",
        "\n",
        "    if commit_hash:\n",
        "        reset_repo(repo_dir, commit_hash)\n",
        "\n",
        "    repo_name, current_commit_hash, current_branch = validate_repo(repo_dir)\n",
        "    cprint(f\"Using '{repo_name}' repository...\", color=\"green\")\n",
        "    cprint(f\"Branch: {current_branch}, Commit hash: {current_commit_hash}\", color=\"green\")\n",
        "\n",
        "    print_line(80, color=\"green\")\n",
        "    cprint(\"Hotfixes and Optimization:\", color=\"green\")\n",
        "\n",
        "    if dpmpp_2m_v2_patch:\n",
        "        dpmpp_2m_v2_url  = \"https://gist.githubusercontent.com/Linaqruf/514d40676e97a70ffc3a2451bbf51555/raw/3fa447ebfac6b98a25485374b70447f848267589/01-add-DPMPP-2M-V2.patch\"\n",
        "        patch_repo(url=dpmpp_2m_v2_url, dir=patches_dir, cwd=repo_dir, whitespace_fix=True, quiet=True)\n",
        "        shutil.rmtree(patches_dir)\n",
        "        cprint(\" [-] DPM++ 2m V2 and DPM++ 2m Karras V2 patch done!\", color=\"green\")\n",
        "\n",
        "    if colab_optimizations:\n",
        "        lowram_patch_url = \"https://raw.githubusercontent.com/ddPn08/automatic1111-colab/main/patches/stablediffusion-lowram.patch\"\n",
        "        stable_diffusion_repo_dir = os.path.join(repo_dir, \"repositories/stable-diffusion-stability-ai\")\n",
        "        patch_repo(url=lowram_patch_url, dir=patches_dir, cwd=stable_diffusion_repo_dir, quiet=True)\n",
        "        shutil.rmtree(patches_dir)\n",
        "        cprint(\" [-] Stable Diffusion V2.x lowram patch done!\", color=\"green\")\n",
        "\n",
        "        subprocess.run([\"sed\", \"-i\", f\"s@os.path.splitext(checkpoint_file)@os.path.splitext(checkpoint_file); map_location='cuda'@\", os.path.join(repo_dir, \"modules\", \"sd_models.py\")])\n",
        "        subprocess.run([\"sed\", \"-i\", f\"s@ui.create_ui().*@ui.create_ui();shared.demo.queue(concurrency_count=999999,status_update_rate=0.1)@\", os.path.join(repo_dir, \"webui.py\")])\n",
        "        subprocess.run([\"sed\", \"-i\", f\"s@map_location='cpu'@map_location='cuda'@\", os.path.join(repo_dir, \"modules\", \"extras.py\")])\n",
        "        cprint(\" [-] TheLastben's colab optimization done!\", color=\"green\")\n",
        "\n",
        "    if mobile_optimizations:\n",
        "        audio_url    = \"https://raw.githubusercontent.com/KoboldAI/KoboldAI-Client/main/colab/silence.m4a\"\n",
        "        audio_thread = threading.Thread(target=play_audio, args=(audio_url,))\n",
        "        audio_thread.start()\n",
        "        cprint(\" [-] Mobile Optimization done!\", color=\"green\")\n",
        "\n",
        "    if \"anapnoe\" in repo_name and \"9931e861dfb128735c4a928a7beb5b5c0af30593\" in current_commit_hash:\n",
        "        hires_prompt_fix = \"https://gist.githubusercontent.com/Linaqruf/8fef456d53604f8c3bcd16722ea7d2f6/raw/a3382087c6e32f9a171f4b5e8aeb572a61682801/0001-Add-New-Label-for-Hires-Prompt.patch\"\n",
        "        patch_repo(url=hires_prompt_fix, dir=patches_dir, cwd=repo_dir, whitespace_fix=True, quiet=True)\n",
        "        shutil.rmtree(patches_dir)\n",
        "        cprint(\" [-] Hires Prompt patch done!\", color=\"green\")\n",
        "\n",
        "    print_line(80, color=\"green\")\n",
        "\n",
        "    if update_extensions:\n",
        "        batch_update(fetch=True, directory=extensions_dir, desc=cprint(f\"Updating extensions\", color=\"green\", tqdm_desc=True))\n",
        "\n",
        "    if not os.path.exists(download_list):\n",
        "        download_list_url = \"https://raw.githubusercontent.com/Linaqruf/sd-notebook-collection/main/config/download_list.txt\"\n",
        "        aria2_download(os.path.dirname(download_list), os.path.basename(download_list), download_list_url, quiet=True)\n",
        "\n",
        "    elapsed_time = py_utils.calculate_elapsed_time(start_time)\n",
        "    print_line(80, color=\"green\")\n",
        "    cprint(f\"Finished installation. Took {elapsed_time}.\", color=\"flat_yellow\")\n",
        "    cprint(\"All is done! Go to the next step.\", color=\"flat_yellow\")\n",
        "    print_line(80, color=\"green\")\n",
        "\n",
        "main()"
      ],
      "metadata": {
        "id": "A6c7-qjDdb0X",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title ## **Download Model and VAE**\n",
        "import os\n",
        "import time\n",
        "from colablib.utils import py_utils\n",
        "from colablib.colored_print import cprint, print_line\n",
        "from colablib.sd_models.downloader import aria2_download, get_modelname\n",
        "\n",
        "%store -r\n",
        "\n",
        "################################\n",
        "# COLAB ARGUMENTS GOES HERE\n",
        "################################\n",
        "\n",
        "# @markdown ### **Stable Diffusion v1.x Model**\n",
        "Anything_V3_0         = False  # @param {type: 'boolean'}\n",
        "AnyLoRA_Default       = False  # @param {type: 'boolean'}\n",
        "AnyLoRA_Anime_Mix     = True  # @param {type: 'boolean'}\n",
        "Ghost_Note_Delta      = False  # @param {type: 'boolean'}\n",
        "SDHK_V3               = False  # @param {type: 'boolean'}\n",
        "Majic_Mix_V5          = False  # @param {type: 'boolean'}\n",
        "# @markdown ### **Stable Diffusion v2.x Model**\n",
        "Replicant_V3          = False  # @param {type: 'boolean'}\n",
        "Illuminati_Diffusion  = False  # @param {type: 'boolean'}\n",
        "# @markdown ### **VAE Model**\n",
        "Anime                 = True  # @param {type: 'boolean'}\n",
        "Blessed               = False  # @param {type: 'boolean'}\n",
        "Waifu_Diffusion       = False  # @param {type: 'boolean'}\n",
        "Stable_Diffusion      = False  # @param {type: 'boolean'}\n",
        "\n",
        "# VAR\n",
        "read_token  = \"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\"\n",
        "user_header = f\"Authorization: Bearer {read_token}\"\n",
        "\n",
        "################################\n",
        "# URL DICT GOES HERE\n",
        "################################\n",
        "\n",
        "model_dict = {\n",
        "    \"Anything_V3_0\"         : \"https://huggingface.co/Linaqruf/stolen/resolve/main/pruned-models/anything-v3-0-pruned.ckpt\",\n",
        "    \"AnyLoRA_Default\"       : \"https://huggingface.co/Linaqruf/stolen/resolve/main/pruned-models/AnyLoRA_noVae_fp16-pruned.safetensors\",\n",
        "    \"AnyLoRA_Anime_Mix\"     : \"https://huggingface.co/Linaqruf/stolen/resolve/main/fp16/aamAnyloraAnimeMixAnime_v10-fp16-pruned.safetensors\",\n",
        "    \"Ghost_Note_Delta\"      : \"https://huggingface.co/Linaqruf/stolen/resolve/main/fp16/GhostNoteDelta_m0528_fp16.safetensors\",\n",
        "    \"SDHK_V3\"               : \"https://huggingface.co/Linaqruf/stolen/resolve/main/fp16/sdhk_v30.safetensors\",\n",
        "    \"Majic_Mix_V5\"          : \"https://huggingface.co/Linaqruf/stolen/resolve/main/fp16/majicmixRealistic_v5.safetensors\",\n",
        "    \"Replicant_V3\"          : \"https://huggingface.co/gsdf/Replicant-V3.0/resolve/main/Replicant-V3.0_fp16.safetensors\",\n",
        "    \"Illuminati_Diffusion\"  : \"https://huggingface.co/Linaqruf/stolen/resolve/main/pruned-models/illuminatiDiffusionV1_v11.safetensors\",\n",
        "}\n",
        "\n",
        "vae_dict = {\n",
        "    \"Anime\"                 : \"https://huggingface.co/NoCrypt/resources/resolve/main/VAE/any.vae.safetensors\",\n",
        "    \"Blessed\"               : \"https://huggingface.co/NoCrypt/resources/resolve/main/VAE/blessed2.vae.safetensors\",\n",
        "    \"Waifu_Diffusion\"       : \"https://huggingface.co/NoCrypt/resources/resolve/main/VAE/wd.vae.safetensors\",\n",
        "    \"Stable_Diffusion\"      : \"https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors\",\n",
        "}\n",
        "\n",
        "def filter_dict_items(dict_items):\n",
        "    result_list = []\n",
        "    for key, url in dict_items.items():\n",
        "        if globals().get(key):\n",
        "            result_list.append((key, url))\n",
        "    return result_list\n",
        "\n",
        "def main():\n",
        "    start_time = time.time()\n",
        "\n",
        "    download_list = [\n",
        "        (filter_dict_items(model_dict), models_dir),\n",
        "        (filter_dict_items(vae_dict), vaes_dir)\n",
        "    ]\n",
        "\n",
        "    print_line(80, color=\"green\")\n",
        "    cprint(\" [-] Downloading Stable Diffusion Models and VAEs...\", color=\"flat_yellow\")\n",
        "    for lst, dst in download_list:\n",
        "        for key, url in lst:\n",
        "            print_line(80, color=\"green\")\n",
        "            extensions = os.path.splitext(get_modelname(url))[1]\n",
        "            if dst == vaes_dir:\n",
        "                extensions = \".vae\" + extensions\n",
        "            aria2_download(url=url, download_dir=dst, filename=key + extensions, user_header=user_header)\n",
        "\n",
        "    elapsed_time = py_utils.calculate_elapsed_time(start_time)\n",
        "    print_line(80, color=\"green\")\n",
        "    cprint(f\"Download finished. Took {elapsed_time}.\", color=\"flat_yellow\")\n",
        "    cprint(\"All is done! Go to the next step.\", color=\"flat_yellow\")\n",
        "    print_line(80, color=\"green\")\n",
        "\n",
        "main()"
      ],
      "metadata": {
        "id": "7nUM-wRFhBa3",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title ## **ControlNet v1.1**\n",
        "import time\n",
        "import shutil\n",
        "import os\n",
        "import subprocess\n",
        "from tqdm import tqdm\n",
        "from colablib.sd_models.downloader import aria2_download, download\n",
        "from colablib.utils import config_utils, py_utils, git_utils\n",
        "from colablib.colored_print import cprint, print_line\n",
        "from colablib.utils.py_utils import get_filename\n",
        "\n",
        "%store -r\n",
        "\n",
        "# @markdown ### **ControlNet Annotator**\n",
        "pre_download_annotator      = True   # @param {type: 'boolean'}\n",
        "# @markdown ### **SDv1.x ControlNet Model**\n",
        "control_v11_sd15_model      = True   # @param {type: 'boolean'}\n",
        "t2i_adapter_model           = False  # @param {type: 'boolean'}\n",
        "# @markdown ### **SDv2.x ControlNet Model**\n",
        "control_v11_sd21_model      = False  # @param {type: 'boolean'}\n",
        "# @markdown ### **Custom ControlNet Model**\n",
        "# @markdown - Make sure your custom controlnet model has `sd15`/`sd21` in the filename.\n",
        "# @markdown - Use comma separation for multiple URLs, e.g. `url1, url2, url3`.\n",
        "custom_controlnet_url = \"\" #@param [\"\", \"https://huggingface.co/ioclab/ioc-controlnet/resolve/main/models/control_v1p_sd15_illumination.safetensors\", \"https://huggingface.co/ioclab/ioc-controlnet/resolve/main/models/control_v1p_sd15_brightness.safetensors\"] {allow-input: true}\n",
        "# @markdown ### **ControlNet Config**\n",
        "control_net_max_models_num  = 2      # @param {type:\"slider\", min:1, max:10, step:1}\n",
        "\n",
        "custom_controlnet_dict = {\n",
        "    \"control_v1p_sd15_illumination\" : \"https://huggingface.co/ioclab/ioc-controlnet/resolve/main/models/control_v1p_sd15_illumination.safetensors\",\n",
        "    \"control_v1p_sd15_brightness\"   : \"https://huggingface.co/ioclab/ioc-controlnet/resolve/main/models/control_v1p_sd15_brightness.safetensors\"\n",
        "}\n",
        "\n",
        "annotator_dict = {\n",
        "    \"midas\"         : \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt\",\n",
        "    \"leres\"         : [\n",
        "        \"https://huggingface.co/lllyasviel/Annotators/resolve/main/res101.pth\",\n",
        "        \"https://huggingface.co/lllyasviel/Annotators/resolve/main/latest_net_G.pth\"\n",
        "    ],\n",
        "    \"hed\"           : \"https://huggingface.co/lllyasviel/Annotators/resolve/main/ControlNetHED.pth\",\n",
        "    \"mlsd\"          : \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/mlsd_large_512_fp32.pth\",\n",
        "    \"openpose\"      : [\n",
        "        \"https://huggingface.co/lllyasviel/Annotators/resolve/main/body_pose_model.pth\",\n",
        "        \"https://huggingface.co/lllyasviel/Annotators/resolve/main/hand_pose_model.pth\",\n",
        "        \"https://huggingface.co/lllyasviel/Annotators/resolve/main/facenet.pth\"\n",
        "    ],\n",
        "    \"clip_vision\"   : \"https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/pytorch_model.bin\",\n",
        "    \"pidinet\"       : \"https://huggingface.co/lllyasviel/Annotators/resolve/main/table5_pidinet.pth\",\n",
        "    \"uniformer\"     : \"https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/upernet_global_small.pth\",\n",
        "    \"zoedepth\"      : \"https://huggingface.co/lllyasviel/Annotators/resolve/main/ZoeD_M12_N.pt\",\n",
        "    \"normal_bae\"    : \"https://huggingface.co/lllyasviel/Annotators/resolve/main/scannet.pt\",\n",
        "    \"oneformer\"     : [\n",
        "        \"https://huggingface.co/lllyasviel/Annotators/resolve/main/150_16_swin_l_oneformer_coco_100ep.pth\",\n",
        "        \"https://huggingface.co/lllyasviel/Annotators/resolve/main/250_16_swin_l_oneformer_ade20k_160k.pth\"\n",
        "    ],\n",
        "    \"lineart\"       : [\n",
        "        \"https://huggingface.co/lllyasviel/Annotators/resolve/main/sk_model.pth\",\n",
        "        \"https://huggingface.co/lllyasviel/Annotators/resolve/main/sk_model2.pth\"\n",
        "    ],\n",
        "    \"lineart_anime\" : \"https://huggingface.co/lllyasviel/Annotators/resolve/main/netG.pth\",\n",
        "    \"manga_line\"    : \"https://huggingface.co/lllyasviel/Annotators/resolve/main/erika.pth\"\n",
        "}\n",
        "\n",
        "control_v11_sd15_url = [\n",
        "    \"https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11e_sd15_ip2p_fp16.safetensors\",\n",
        "    \"https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11e_sd15_shuffle_fp16.safetensors\",\n",
        "    \"https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15_canny_fp16.safetensors\",\n",
        "    \"https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11f1p_sd15_depth_fp16.safetensors\",\n",
        "    \"https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15_inpaint_fp16.safetensors\",\n",
        "    \"https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15_lineart_fp16.safetensors\",\n",
        "    \"https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15_mlsd_fp16.safetensors\",\n",
        "    \"https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15_normalbae_fp16.safetensors\",\n",
        "    \"https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15_openpose_fp16.safetensors\",\n",
        "    \"https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15_scribble_fp16.safetensors\",\n",
        "    \"https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15_seg_fp16.safetensors\",\n",
        "    \"https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15_softedge_fp16.safetensors\",\n",
        "    \"https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15s2_lineart_anime_fp16.safetensors\",\n",
        "    \"https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11f1e_sd15_tile_fp16.safetensors\",\n",
        "]\n",
        "\n",
        "control_v11_sd21_url = [\n",
        "    \"https://huggingface.co/thibaud/controlnet-sd21/resolve/main/control_v11p_sd21_ade20k.safetensors\",\n",
        "    \"https://huggingface.co/thibaud/controlnet-sd21/resolve/main/control_v11p_sd21_canny.safetensors\",\n",
        "    \"https://huggingface.co/thibaud/controlnet-sd21/resolve/main/control_v11p_sd21_color.safetensors\",\n",
        "    \"https://huggingface.co/thibaud/controlnet-sd21/resolve/main/control_v11p_sd21_depth.safetensors\",\n",
        "    \"https://huggingface.co/thibaud/controlnet-sd21/resolve/main/control_v11p_sd21_hed.safetensors\",\n",
        "    \"https://huggingface.co/thibaud/controlnet-sd21/resolve/main/control_v11p_sd21_lineart.safetensors\",\n",
        "    \"https://huggingface.co/thibaud/controlnet-sd21/resolve/main/control_v11p_sd21_normalbae.safetensors\",\n",
        "    \"https://huggingface.co/thibaud/controlnet-sd21/resolve/main/control_v11p_sd21_openpose.safetensors\",\n",
        "    \"https://huggingface.co/thibaud/controlnet-sd21/resolve/main/control_v11p_sd21_openposev2.safetensors\",\n",
        "    \"https://huggingface.co/thibaud/controlnet-sd21/resolve/main/control_v11p_sd21_scribble.safetensors\",\n",
        "    \"https://huggingface.co/thibaud/controlnet-sd21/resolve/main/control_v11p_sd21_zoedepth.safetensors\"\n",
        "]\n",
        "\n",
        "t2i_adapter_url = [\n",
        "    \"https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_canny_sd14v1.pth\",\n",
        "    \"https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_canny_sd15v2.pth\",\n",
        "    \"https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_color_sd14v1.pth\",\n",
        "    \"https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_depth_sd14v1.pth\",\n",
        "    \"https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_depth_sd15v2.pth\",\n",
        "    \"https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_keypose_sd14v1.pth\",\n",
        "    \"https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_openpose_sd14v1.pth\",\n",
        "    \"https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_seg_sd14v1.pth\",\n",
        "    \"https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_sketch_sd14v1.pth\",\n",
        "    \"https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_sketch_sd15v2.pth\",\n",
        "    \"https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_style_sd14v1.pth\",\n",
        "    \"https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_zoedepth_sd15v1.pth\"\n",
        "]\n",
        "\n",
        "def cldm_config_path(destination_path):\n",
        "    if \"control\" in destination_path:\n",
        "        if \"sd15\" in destination_path:\n",
        "            return \"control_v11e_sd15_shuffle.yaml\" if \"_shuffle_\" in destination_path else \"cldm_v15.yaml\"\n",
        "        if \"sd21\" in destination_path:\n",
        "            return \"cldm_v21.yaml\"\n",
        "    elif \"t2i\" in destination_path:\n",
        "        adapter_name = os.path.splitext(os.path.basename(destination_path))[0]\n",
        "        return adapter_name + \".yaml\"\n",
        "    return None\n",
        "\n",
        "def cldm_config(destination_path):\n",
        "    repo_name, _, _ = git_utils.validate_repo(repo_dir)\n",
        "\n",
        "    control_net_model_config = cldm_config_path(destination_path)\n",
        "    if control_net_model_config is not None:\n",
        "        cldm_config_src = os.path.join(extensions_dir, os.path.join(f\"{voldy}-controlnet\", \"models\", control_net_model_config))\n",
        "        cldm_config_dst = os.path.splitext(destination_path)[0] + \".yaml\"\n",
        "        if not os.path.exists(cldm_config_dst):\n",
        "            shutil.copy(cldm_config_src, cldm_config_dst)\n",
        "\n",
        "def batch_download(urls, dst, desc=None, quiet=False, cldm_model=False):\n",
        "    for url in tqdm(urls, disable=quiet, desc=cprint(desc, color=\"green\", tqdm_desc=True)):\n",
        "        filename = get_filename(url)\n",
        "        aria2_download(url=url, download_dir=dst, filename=filename, quiet=True)\n",
        "        if cldm_model:\n",
        "            cldm_config(os.path.join(dst, filename))\n",
        "\n",
        "\n",
        "def custom_controlnet_download(urls, dst):\n",
        "    for url in urls.split(\",\"):\n",
        "        url = url.strip()\n",
        "        if url != \"\":\n",
        "            print_line(80, color=\"green\")\n",
        "            filename = get_filename(url)\n",
        "            download(url=url, filename=filename, dst=control_dir)\n",
        "            cldm_config(os.path.join(dst, filename))\n",
        "\n",
        "def download_annotator(directory, desc):\n",
        "    for category, urls in tqdm(annotator_dict.items(), desc=cprint(desc, color=\"green\", tqdm_desc=True)):\n",
        "        if category == \"clip_vision\":\n",
        "            dst = os.path.join(directory, \"clip_vision\")\n",
        "        else:\n",
        "            dst = os.path.join(directory, \"downloads\", category)\n",
        "        os.makedirs(dst, exist_ok=True)\n",
        "        urls = [urls] if isinstance(urls, str) else urls\n",
        "        batch_download(urls, dst, quiet=True)\n",
        "\n",
        "def main():\n",
        "    start_time = time.time()\n",
        "\n",
        "    repo_name, _, _ = git_utils.validate_repo(repo_dir)\n",
        "\n",
        "    config = config_utils.read_config(config_file)\n",
        "    config[\"control_net_max_models_num\"]        = control_net_max_models_num\n",
        "    config[\"control_net_models_path\"]           = control_dir\n",
        "    config[\"control_net_allow_script_control\"]  = True\n",
        "    config_utils.write_config(config_file, config)\n",
        "\n",
        "    print_line(80, color=\"green\")\n",
        "    cprint(\" [-] Downloading ControlNet Models...\", color=\"flat_yellow\")\n",
        "    print_line(80, color=\"green\")\n",
        "\n",
        "    if pre_download_annotator:\n",
        "        download_annotator(annotator_dir, \"Downloading ControlNet Annotator/Preprocessor\")\n",
        "    if control_v11_sd15_model:\n",
        "        batch_download(control_v11_sd15_url, control_dir, \"Downloading SDv1.x ControlNet Model\", cldm_model=True)\n",
        "    if control_v11_sd21_model:\n",
        "        batch_download(control_v11_sd21_url, control_dir, \"Downloading SDv2.x ControlNet Model\", cldm_model=True)\n",
        "    if t2i_adapter_model:\n",
        "        batch_download(t2i_adapter_url, control_dir, \"Downloading SDv1.x Text2Image Adapter Model\", cldm_model=True)\n",
        "    if custom_controlnet_url:\n",
        "        print_line(80, color=\"green\")\n",
        "        cprint(\" [-] Downloading Custom ControlNet Models...\", color=\"flat_yellow\")\n",
        "        custom_controlnet_download(custom_controlnet_url, control_dir)\n",
        "\n",
        "    print_line(80, color=\"green\")\n",
        "    elapsed_time = py_utils.calculate_elapsed_time(start_time)\n",
        "    cprint(f\"Download finished. Took {elapsed_time}.\", color=\"flat_yellow\")\n",
        "    cprint(\"All is done! Go to the next step.\", color=\"flat_yellow\")\n",
        "    print_line(80, color=\"green\")\n",
        "\n",
        "main()"
      ],
      "metadata": {
        "id": "K28QYTFhf7Cu",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title ## **Custom Download Corner**\n",
        "import os\n",
        "import time\n",
        "from pydantic import BaseModel\n",
        "from colablib.utils.py_utils import get_filename\n",
        "from colablib.sd_models.downloader import aria2_download, download\n",
        "from colablib.utils.ubuntu_utils import unionfuse\n",
        "from colablib.utils.git_utils import clone_repo\n",
        "from colablib.colored_print import cprint, print_line\n",
        "from colablib.utils.config_utils import read_config\n",
        "\n",
        "%store -r\n",
        "\n",
        "# @markdown\n",
        "# @markdown ### **Download from Custom URLs**\n",
        "# @markdown - Use comma separation for multiple URLs, e.g. `url1, url2, url3`.\n",
        "# @markdown - To load Google Drive, use `fuse:` followed by path, e.g. `fuse:/content/MyDrive/LoRA`.\n",
        "# @markdown - Copy your model path from Google Drive to URL fields to copy your model to the web UI models directory.\n",
        "custom_model_url        = \"\"  # @param {'type': 'string'}\n",
        "custom_vae_url          = \"\"  # @param {'type': 'string'}\n",
        "custom_embedding_url    = \"\"  # @param {'type': 'string'}\n",
        "custom_LoRA_url         = \"\"  # @param {'type': 'string'}\n",
        "custom_hypernetwork_url = \"\"  # @param {'type': 'string'}\n",
        "custom_extensions_url   = \"\"  # @param {'type': 'string'}\n",
        "custom_upscaler_url     = \"\"  # @param {'type': 'string'}\n",
        "# @markdown ### <br>`NEW` **Download from Textfile**\n",
        "# @markdown - Provide a custom download URL for a `.txt` file instead of using the URL field. Edit the file: `/content/download_list.txt`.\n",
        "# @markdown - Available hashtags: `#model`, `#vae`, `#embedding`, `#lora`, `#hypernetwork`, `#extensions`, `#upscaler`.\n",
        "# @markdown - Or you can input your `.txt` file in `custom_download_list_url` below. Works for `pastebin`.\n",
        "custom_download_list_url = \"\"  # @param {'type': 'string'}\n",
        "\n",
        "class CustomDirs(BaseModel):\n",
        "    url: str\n",
        "    dst: str\n",
        "\n",
        "custom_dirs = {\n",
        "    \"model\"       : CustomDirs(url=custom_model_url, dst=models_dir),\n",
        "    \"vae\"         : CustomDirs(url=custom_vae_url, dst=vaes_dir),\n",
        "    \"embedding\"   : CustomDirs(url=custom_embedding_url, dst=embeddings_dir),\n",
        "    \"lora\"        : CustomDirs(url=custom_LoRA_url, dst=lora_dir),\n",
        "    \"hypernetwork\": CustomDirs(url=custom_hypernetwork_url, dst=hypernetworks_dir),\n",
        "    \"extensions\"  : CustomDirs(url=custom_extensions_url, dst=extensions_dir),\n",
        "    \"upscaler\"    : CustomDirs(url=custom_upscaler_url, dst=esrgan_dir)\n",
        "}\n",
        "\n",
        "def fuse(url, key, dst):\n",
        "    if \"extensions\" in key:\n",
        "        cprint(f\"Folder can't be fused, skipping...\")\n",
        "        return\n",
        "\n",
        "    path = url.split(\"fuse:\")[1].strip()\n",
        "    category_dir = os.path.join(fused_dir, key)\n",
        "    if os.path.exists(category_dir):\n",
        "        cprint(f\"Folder '{category_dir}' is already fused, skipping...\", color=\"yellow\")\n",
        "        return\n",
        "\n",
        "    cprint(f\"Fusing process started for PATH: '{path}'\", color=\"green\")\n",
        "    unionfuse(category_dir, path, dst)\n",
        "    cprint(f\"Fusing process completed. Valid '{key}' folder located at: '{category_dir}' \", color=\"green\")\n",
        "\n",
        "def parse_urls(filename):\n",
        "    content = read_config(filename)\n",
        "    lines   = content.strip().split('\\n')\n",
        "    result  = {}\n",
        "    key     = ''\n",
        "    for line in lines:\n",
        "        if not line.strip():\n",
        "            continue\n",
        "        if line.startswith('//'):\n",
        "            continue\n",
        "        if line.startswith('#'):\n",
        "            key = line[1:].lower()\n",
        "            result[key] = []\n",
        "        else:\n",
        "            urls = [url.strip() for url in line.split(',') if url.strip() != '']\n",
        "            result[key].extend(urls)\n",
        "    return result\n",
        "\n",
        "def custom_download(custom_dirs):\n",
        "    for key, value in custom_dirs.items():\n",
        "        urls     = value.url.split(\",\")  # Split the comma-separated URLs\n",
        "        dst      = value.dst\n",
        "\n",
        "        if value.url:\n",
        "            print_line(80, color=\"green\")\n",
        "            cprint(f\" [-] Downloading Custom {key}...\", color=\"flat_yellow\")\n",
        "\n",
        "        for url in urls:\n",
        "            url = url.strip()  # Remove leading/trailing whitespaces from each URL\n",
        "            if url != \"\":\n",
        "                print_line(80, color=\"green\")\n",
        "                if \"|\" in url:\n",
        "                    url, filename = map(str.strip, url.split(\"|\"))\n",
        "                    if not filename.endswith((\".safetensors\", \".ckpt\", \".pt\", \"pth\")):\n",
        "                        filename = filename + os.path.splitext(get_filename(url))[1]\n",
        "                else:\n",
        "                    if not url.startswith(\"fuse:\"):\n",
        "                        filename = get_filename(url)\n",
        "\n",
        "                if url.startswith(\"fuse:\"):\n",
        "                    fuse(url, key, dst)\n",
        "                elif key == \"extensions\":\n",
        "                    clone_repo(url, cwd=dst)\n",
        "                else:\n",
        "                    download(url=url, filename=filename, dst=dst, quiet=False)\n",
        "\n",
        "def download_from_textfile(filename):\n",
        "    for key, urls in parse_urls(filename).items():\n",
        "        key_lower = key.lower()\n",
        "        if key_lower in custom_dirs:\n",
        "            if custom_dirs[key_lower].url:\n",
        "                custom_dirs[key_lower].url += ',' + ','.join(urls)\n",
        "            else:\n",
        "                custom_dirs[key_lower].url = ','.join(urls)\n",
        "        else:\n",
        "            cprint(f\"Warning: Category '{key}' from the file is not found in custom_dirs.\", color=\"yellow\")\n",
        "\n",
        "def custom_download_list(url):\n",
        "    filename = \"custom_download_list.txt\"\n",
        "    filepath = os.path.join(root_dir, filename)\n",
        "    if os.path.exists(filepath):\n",
        "        os.remove(filepath)\n",
        "    if 'pastebin.com' in url:\n",
        "        if 'raw' not in url:\n",
        "            url = url.replace('pastebin.com', 'pastebin.com/raw')\n",
        "    download(url=url, filename=filename, dst=root_dir, quiet=True)\n",
        "    return filepath\n",
        "\n",
        "def main():\n",
        "    start_time    = time.time()\n",
        "    textfile_path = download_list\n",
        "    if custom_download_list_url:\n",
        "        textfile_path = custom_download_list(custom_download_list_url)\n",
        "    download_from_textfile(textfile_path)\n",
        "    custom_download(custom_dirs)\n",
        "\n",
        "    elapsed_time  = py_utils.calculate_elapsed_time(start_time)\n",
        "    print_line(80, color=\"green\")\n",
        "    cprint(f\"Download finished. Took {elapsed_time}.\", color=\"flat_yellow\")\n",
        "    cprint(\"All is done! Go to the next step.\", color=\"flat_yellow\")\n",
        "    print_line(80, color=\"green\")\n",
        "\n",
        "main()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "SoZcrfQk6Y2z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ## **Start Cagliostro Colab UI**\n",
        "import random\n",
        "import string\n",
        "from pydantic import BaseModel\n",
        "from typing import List, Optional\n",
        "from colablib.utils import config_utils\n",
        "from colablib.colored_print import cprint, print_line\n",
        "from colablib.utils.git_utils import validate_repo\n",
        "\n",
        "%store -r\n",
        "\n",
        "################################\n",
        "# COLAB ARGUMENTS GOES HERE\n",
        "################################\n",
        "\n",
        "# @markdown ### **Alternative Tunnel**\n",
        "# @markdown > Recommended Tunnels: `ngrok` > `gradio` > `cloudflared` > `remotemoe` > `localhostrun` > `googleusercontent`\n",
        "select_tunnel         = \"multiple\" # @param ['gradio', 'multiple','cloudflared', 'localhostrun', 'remotemoe', \"googleusercontent\"]\n",
        "# @markdown > Get your `ngrok_token` [here](https://dashboard.ngrok.com/get-started/your-authtoken)\n",
        "ngrok_token           = \"\" # @param {type: 'string'}\n",
        "ngrok_region          = \"ap\" # @param [\"us\", \"eu\", \"au\", \"ap\", \"sa\", \"jp\", \"in\"]\n",
        "# @markdown ### **UI/UX Config**\n",
        "select_theme          = \"minimal_orange\" # @param ['moonlight', 'ogxRed', 'fun', 'ogxCyan', 'ogxCyanInvert', 'ogxBGreen', 'default_orange', 'tron2', 'd-230-52-94', 'minimal', 'ogxRedYellow', 'retrog', 'ogxRedPurple', 'ogxGreen', 'tron', 'default_cyan', 'default', 'backup', 'minimal_orange', 'Golde']\n",
        "# @markdown Set `use_preset` for using default prompt, resolution, sampler, and other settings\n",
        "use_presets           = True # @param {type: 'boolean'}\n",
        "# @markdown ### **Arguments**\n",
        "use_gradio_auth       = False # @param {type: 'boolean'}\n",
        "accelerator           = \"opt-sdp-attention\" # @param ['xformers', 'opt-sdp-attention', 'opt-sdp-no-mem-attention', 'opt-split-attention']\n",
        "auto_select_model     = False # @param {type: 'boolean'}\n",
        "auto_select_vae       = True # @param {type: 'boolean'}\n",
        "additional_arguments  = \"--lowram --theme dark --no-half-vae\" #@param {type: 'string'}\n",
        "\n",
        "# GRADIO AUTH\n",
        "user                  = \"cagliostro\"\n",
        "password              = \"\".join(random.choices(string.ascii_letters + string.digits, k=6))\n",
        "\n",
        "def change_theme(filename):\n",
        "    themes_folder   = os.path.join(repo_dir, \"extensions-builtin\", \"sd_theme_editor\", \"themes\")\n",
        "    themes_file     = os.path.join(themes_folder, f\"{filename}.css\")\n",
        "\n",
        "    style_config    = config_utils.read_config(style_path)\n",
        "    style_contents  = style_config.split(\"/*BREAKPOINT_CSS_CONTENT*/\")[1]\n",
        "\n",
        "    theme_config    = config_utils.read_config(themes_file)\n",
        "    style_data      = \":host{\" + theme_config + \"}\" + \"/*BREAKPOINT_CSS_CONTENT*/\" + style_contents\n",
        "    config_utils.write_config(style_path, style_data)\n",
        "\n",
        "def is_valid(valid_dir, file_types):\n",
        "    return [f for f in os.listdir(valid_dir) if f.endswith(file_types)]\n",
        "\n",
        "def auto_select_file(valid_dir, config_key, file_types):\n",
        "    valid_files = is_valid(valid_dir, file_types)\n",
        "    if valid_files:\n",
        "        file_path = random.choice(valid_files)\n",
        "        if os.path.exists(os.path.join(valid_dir, file_path)):\n",
        "            config = config_utils.read_config(config_file)\n",
        "            config[config_key] = file_path\n",
        "            config_utils.write_config(config_file, config)\n",
        "        return file_path\n",
        "    else:\n",
        "        return None\n",
        "\n",
        "def ui_preset_config():\n",
        "    global default_upscaler, default_sampler_v2\n",
        "\n",
        "    default_prompt        = \"masterpiece, best quality,\"\n",
        "    default_neg_prompt    = \"(worst quality, low quality:1.4)\"\n",
        "    default_sampler       = \"DPM++ 2M Karras\"\n",
        "    default_steps         = 20\n",
        "    default_width         = 512\n",
        "    default_height        = 768\n",
        "    default_strength      = 0.55\n",
        "    default_cfg_scale     = 7\n",
        "    default_upscaler       = \"Latent (nearest-exact)\"\n",
        "\n",
        "    config = {\n",
        "        \"Prompt/value\"              : default_prompt,\n",
        "        \"Negative prompt/value\"     : default_neg_prompt,\n",
        "        \"Sampling method/value\"     : default_sampler,\n",
        "        \"Sampling steps/value\"      : default_steps,\n",
        "        \"Width/value\"               : default_width,\n",
        "        \"Height/value\"              : default_height,\n",
        "        \"Denoising strength/value\"  : default_strength,\n",
        "        \"CFG Scale/value\"           : default_cfg_scale\n",
        "    }\n",
        "\n",
        "    return config\n",
        "\n",
        "def configure_main_settings(config_file: str, lora_dir: str, use_presets: bool, ui_config_file: str):\n",
        "    config = config_utils.read_config(config_file)\n",
        "\n",
        "    config[\"additional_networks_extra_lora_path\"] = lora_dir\n",
        "    config[\"CLIP_stop_at_last_layers\"] = 2\n",
        "    config[\"eta_noise_seed_delta\"] = 0\n",
        "    config[\"show_progress_every_n_steps\"] = 10\n",
        "    config[\"show_progressbar\"] = True\n",
        "    config[\"samples_filename_pattern\"] = \"[model_name]_[seed]\"\n",
        "    config[\"show_progress_type\"] = \"Approx NN\" # Full, Approx NN, TAESD, Approx cheap\n",
        "    config[\"live_preview_content\"] = \"Prompt\" # Combined, Prompt, Negative Prompt\n",
        "    config[\"hires_fix_show_sampler\"] = True\n",
        "    config[\"hires_fix_show_prompts\"] = True\n",
        "    config[\"state\"] = [\"tabs\"]\n",
        "    config[\"state_txt2img\"] = [\"prompt\", \"negative_prompt\", \"styles\", \"sampling\", \"sampling_steps\", \"width\", \"height\", \"batch_count\", \"batch_size\", \"hires_resize_y\", \"hires_resize_x\", \"hires_scale\", \"hires_steps\", \"hires_upscaler\", \"hires_fix\", \"tiling\", \"restore_faces\", \"cfg_scale\", \"hires_denoising_strength\"]\n",
        "    config[\"state_img2img\"] = [\"prompt\", \"negative_prompt\", \"styles\", \"sampling\", \"resize_mode\", \"sampling_steps\", \"tiling\", \"restore_faces\", \"width\", \"height\", \"batch_count\", \"batch_size\", \"cfg_scale\", \"denoising_strength\"]\n",
        "    config[\"state_extensions\"] = [\"control-net\"]\n",
        "\n",
        "    quicksettings_values = [\"sd_model_checkpoint\", \"sd_vae\", \"CLIP_stop_at_last_layers\",\n",
        "                            \"use_old_karras_scheduler_sigmas\", \"always_discard_next_to_last_sigma\",\n",
        "                            \"token_merging_ratio\", \"s_min_uncond\"]\n",
        "\n",
        "    if \"quicksettings\" in config:\n",
        "        config[\"quicksettings\"] = \", \".join(quicksettings_values)\n",
        "    elif \"quicksettings_list\" in config:\n",
        "        config[\"quicksettings_list\"] = quicksettings_values\n",
        "\n",
        "    config_utils.write_config(config_file, config)\n",
        "\n",
        "    if use_presets:\n",
        "        configure_ui_settings(ui_config_file)\n",
        "\n",
        "def configure_ui_settings(ui_config_file: str):\n",
        "    config = config_utils.read_config(ui_config_file)\n",
        "    preset_config = ui_preset_config()\n",
        "    for key in [\"txt2img\", \"img2img\"]:\n",
        "        for subkey, value in preset_config.items():\n",
        "            config[f\"{key}/{subkey}\"] = value\n",
        "\n",
        "    config[\"txt2img/Upscaler/value\"] = default_upscaler\n",
        "    config_utils.write_config(ui_config_file, config)\n",
        "\n",
        "def is_dir_exist(cloned_dir, original_dir):\n",
        "    if os.path.exists(cloned_dir):\n",
        "        return cloned_dir\n",
        "    else:\n",
        "        return original_dir\n",
        "\n",
        "def parse_args(config):\n",
        "    args = \"\"\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args += f'\"{v}\" '\n",
        "        elif isinstance(v, str):\n",
        "            args += f'--{k}=\"{v}\" '\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args += f\"--{k} \"\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "\n",
        "    return args\n",
        "\n",
        "def main():\n",
        "    global auto_select_model, auto_select_vae\n",
        "\n",
        "    repo_name, _, _ = validate_repo(repo_dir)\n",
        "    if \"anapnoe\" in repo_name:\n",
        "        change_theme(select_theme)\n",
        "\n",
        "    valid_ckpt_dir          = is_dir_exist(os.path.join(fused_dir, \"model\"), models_dir)\n",
        "    valid_vae_dir           = is_dir_exist(os.path.join(fused_dir, \"vae\"), vaes_dir)\n",
        "    valid_embedding_dir     = is_dir_exist(os.path.join(fused_dir, \"embedding\"), embeddings_dir)\n",
        "    valid_lora_dir          = is_dir_exist(os.path.join(fused_dir, \"lora\"), lora_dir)\n",
        "    valid_hypernetwork_dir  = is_dir_exist(os.path.join(fused_dir, \"hypernetwork\"), hypernetworks_dir)\n",
        "\n",
        "    print_line(80, color=\"green\")\n",
        "    cprint(f\"Launching '{repo_name}'\", color=\"flat_yellow\")\n",
        "    print_line(80, color=\"green\")\n",
        "\n",
        "    if not is_valid(valid_ckpt_dir, ('.ckpt', '.safetensors')):\n",
        "        cprint(f\"No checkpoints were found in the directory '{valid_ckpt_dir}'.\", color=\"yellow\")\n",
        "        url = \"https://huggingface.co/Linaqruf/stolen/resolve/main/fp16/aamAnyloraAnimeMixAnime_v10-fp16-pruned.safetensors\"\n",
        "        filename = \"AnyLoRA_Anime_mix.safetensors\"\n",
        "        aria2_download(url=url, download_dir=valid_ckpt_dir, filename=filename)\n",
        "        print_line(80, color=\"green\")\n",
        "        auto_select_model = True\n",
        "\n",
        "    if not is_valid(valid_vae_dir, ('.vae.pt', '.vae.safetensors', '.pt', '.ckpt')):\n",
        "        cprint(f\"No VAEs were found in the directory '{valid_vae_dir}'.\", color=\"yellow\")\n",
        "        url = \"https://huggingface.co/NoCrypt/resources/resolve/main/any.vae.safetensors\"\n",
        "        filename = \"Anime.vae.safetensors\"\n",
        "        aria2_download(url=url, download_dir=valid_vae_dir, filename=filename)\n",
        "        print_line(80, color=\"green\")\n",
        "        auto_select_vae = True\n",
        "\n",
        "    if auto_select_model:\n",
        "        selected_model  = auto_select_file(valid_ckpt_dir, \"sd_model_checkpoint\", ('.ckpt', '.safetensors'))\n",
        "        cprint(f\"Selected Model: {selected_model}\", color=\"green\")\n",
        "\n",
        "    if auto_select_vae:\n",
        "        selected_vae    = auto_select_file(valid_vae_dir, \"sd_vae\", ('.vae.pt', '.vae.safetensors', '.pt', '.ckpt'))\n",
        "        cprint(f\"Selected VAE: {selected_vae}\", color=\"green\")\n",
        "\n",
        "    print_line(80, color=\"green\")\n",
        "\n",
        "    configure_main_settings(config_file, valid_lora_dir, use_presets, ui_config_file)\n",
        "\n",
        "    if use_gradio_auth:\n",
        "      cprint(\"Gradio Auth (use this account to login):\", color=\"green\")\n",
        "      cprint(\"[-] Username: cagliostro\", color=\"green\")\n",
        "      cprint(\"[-] Password:\", password, color=\"green\")\n",
        "      print_line(80, color=\"green\")\n",
        "\n",
        "    config = {\n",
        "        \"enable-insecure-extension-access\": True,\n",
        "        \"disable-safe-unpickle\"           : True,\n",
        "        f\"{accelerator}\"                  : True,\n",
        "        f\"{select_tunnel}\"                : True if not select_tunnel == \"gradio\" and not ngrok_token else False,\n",
        "        \"share\"                           : True if not ngrok_token else False,\n",
        "        \"gradio-auth\"                     : f\"{user}:{password}\" if use_gradio_auth else None,\n",
        "        \"no-hashing\"                      : True,\n",
        "        \"disable-console-progressbars\"    : True,\n",
        "        \"ngrok\"                           : ngrok_token if ngrok_token else None,\n",
        "        \"ngrok-region\"                    : ngrok_region if ngrok_token else None,\n",
        "        \"opt-sub-quad-attention\"          : True,\n",
        "        \"opt-channelslast\"                : True,\n",
        "        \"no-download-sd-model\"            : True,\n",
        "        \"gradio-queue\"                    : True,\n",
        "        \"listen\"                          : True,\n",
        "        \"ckpt-dir\"                        : valid_ckpt_dir,\n",
        "        \"vae-dir\"                         : valid_vae_dir,\n",
        "        \"hypernetwork-dir\"                : valid_hypernetwork_dir,\n",
        "        \"embeddings-dir\"                  : valid_embedding_dir,\n",
        "        \"lora-dir\"                        : valid_lora_dir,\n",
        "        \"lyco-dir\"                        : valid_lora_dir,\n",
        "    }\n",
        "\n",
        "    args = parse_args(config)\n",
        "    final_args = f\"python launch.py {args} {additional_arguments}\"\n",
        "\n",
        "    cprint()\n",
        "    os.chdir(repo_dir)\n",
        "    os.environ[\"TCMALLOC_AGGRESSIVE_DECOMMIT\"] = \"t\"\n",
        "    !{final_args}\n",
        "\n",
        "main()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "Oyrwg8cMyDXj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title ## **Download Generated Images**\n",
        "# @markdown Download file manually from files tab or save to Google Drive\n",
        "import os\n",
        "\n",
        "from pydrive.auth import GoogleAuth\n",
        "from pydrive.drive import GoogleDrive\n",
        "from google.colab import auth, drive\n",
        "from oauth2client.client import GoogleCredentials\n",
        "from colablib.colored_print import cprint, print_line\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(output_dir)\n",
        "\n",
        "use_drive = False  # @param {type:\"boolean\"}\n",
        "folder_name = \"cagliostro-colab-ui\"  # @param {type: \"string\"}\n",
        "filename = \"waifu.zip\"  # @param {type: \"string\"}\n",
        "save_as = filename\n",
        "\n",
        "if os.path.exists(filename):\n",
        "    i = 1\n",
        "    while os.path.exists(f\"waifu({i}).zip\"):\n",
        "        i += 1\n",
        "    filename = f\"waifu({i}).zip\"\n",
        "\n",
        "os.system('zip -r /content/outputs.zip .')\n",
        "\n",
        "if use_drive:\n",
        "    auth.authenticate_user()\n",
        "    gauth = GoogleAuth()\n",
        "    gauth.credentials = GoogleCredentials.get_application_default()\n",
        "    drive = GoogleDrive(gauth)\n",
        "\n",
        "    def create_folder(folder_name):\n",
        "        file_list = drive.ListFile({\n",
        "            \"q\": f\"title='{folder_name}' and mimeType='application/vnd.google-apps.folder' and trashed=false\"\n",
        "        }).GetList()\n",
        "        if file_list:\n",
        "            cprint(\"Debug: Folder exists\", color=\"green\")\n",
        "            folder_id = file_list[0][\"id\"]\n",
        "        else:\n",
        "            cprint(\"Debug: Creating folder\", color=\"green\")\n",
        "            file = drive.CreateFile({\n",
        "                \"title\": folder_name,\n",
        "                \"mimeType\": \"application/vnd.google-apps.folder\"\n",
        "            })\n",
        "            file.Upload()\n",
        "            folder_id = file.attr[\"metadata\"][\"id\"]\n",
        "        return folder_id\n",
        "\n",
        "    def upload_file(file_name, folder_id, save_as):\n",
        "        file_list = drive.ListFile({\"q\": f\"title='{save_as}' and trashed=false\"}).GetList()\n",
        "        if file_list:\n",
        "            cprint(\"Debug: File already exists\", color=\"green\")\n",
        "            i = 1\n",
        "            while True:\n",
        "                new_name = f\"{os.path.splitext(save_as)[0]}({i}){os.path.splitext(save_as)[1]}\"\n",
        "                file_list = drive.ListFile({\"q\": f\"title='{new_name}' and trashed=false\"}).GetList()\n",
        "                if not file_list:\n",
        "                    save_as = new_name\n",
        "                    break\n",
        "                i += 1\n",
        "        file = drive.CreateFile({\"title\": save_as, \"parents\": [{\"id\": folder_id}]})\n",
        "        file.SetContentFile(file_name)\n",
        "        file.Upload()\n",
        "        file.InsertPermission({\"type\": \"anyone\", \"value\": \"anyone\", \"role\": \"reader\"})\n",
        "        return file.attr[\"metadata\"][\"id\"]\n",
        "\n",
        "    file_id = upload_file(\"/content/outputs.zip\", create_folder(folder_name), save_as)\n",
        "    cprint(f\"Your sharing link: https://drive.google.com/file/d/{file_id}/view?usp=sharing\", color=\"green\")\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "Ojb4nAieATxI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Extras"
      ],
      "metadata": {
        "id": "4SUHPtGLz2m4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title ## **Download Generated Images V2**\n",
        "import shutil\n",
        "import os\n",
        "from IPython.utils import capture\n",
        "from huggingface_hub import login, HfApi\n",
        "from huggingface_hub.utils import validate_repo_id, HfHubHTTPError\n",
        "from colablib.colored_print import cprint, print_line\n",
        "\n",
        "# @markdown Download your output by upload it to **Huggingface** instead of Google Drive.\n",
        "# @markdown > Get **your** huggingface `WRITE` token [here](https://huggingface.co/settings/tokens)\n",
        "write_token = \"\"  # @param {type:\"string\"}\n",
        "# @markdown Specify where is your repo located, it will automatically create your repo if you didn't have one.\n",
        "repo_name = \"cagliostro-colab-ui\"  # @param{type:\"string\"}\n",
        "private_repo = False  # @param{type:\"boolean\"}\n",
        "# @markdown This will be compressed to zip and uploaded to datasets repo\n",
        "project_name = \"waifu\"  # @param {type :\"string\"}\n",
        "\n",
        "repo_name = repo_name.replace(\" \", \"-\")\n",
        "project_name = project_name.replace(\" \", \"_\")\n",
        "\n",
        "if not project_name:\n",
        "    project_name = \"waifu\"\n",
        "\n",
        "dataset_zip = f\"{project_name}.zip\"\n",
        "output_path = os.path.join(root_dir, dataset_zip)\n",
        "commit_message = f\"Feat: Upload {dataset_zip} with Cagliostro Colab UI\"\n",
        "\n",
        "def create_or_validate_repo(api, datasets_repo):\n",
        "    try:\n",
        "        validate_repo_id(datasets_repo)\n",
        "        api.create_repo(\n",
        "            repo_id=datasets_repo, repo_type=\"dataset\", private=private_repo\n",
        "        )\n",
        "        cprint(f\"Repo created, located at \"\n",
        "              f\"https://huggingface.co/datasets/{datasets_repo}\", color=\"green\")\n",
        "\n",
        "    except HfHubHTTPError:\n",
        "        cprint(f\"Repo exists, skipping...\", color=\"green\")\n",
        "\n",
        "def compress_to_zip(output_path):\n",
        "    os.chdir(output_dir)\n",
        "    cprint(f\"Compressing to ZIP...\", color=\"green\")\n",
        "    with capture.capture_output() as cap:\n",
        "        !zip -rv {output_path} .\n",
        "\n",
        "def upload_and_cleanup(api, output_path, datasets_repo):\n",
        "    cprint(f\"Uploading generated images... Please wait...\", color=\"green\")\n",
        "\n",
        "    api.upload_file(\n",
        "        path_or_fileobj=output_path,\n",
        "        path_in_repo=dataset_zip,\n",
        "        repo_id=datasets_repo,\n",
        "        repo_type=\"dataset\",\n",
        "        commit_message=commit_message,\n",
        "    )\n",
        "\n",
        "    cprint(f\"Upload success, download directly at \"\n",
        "          f\"https://huggingface.co/datasets/{datasets_repo}/resolve/main/{dataset_zip}\", color=\"green\")\n",
        "\n",
        "    os.remove(output_path)\n",
        "\n",
        "def main():\n",
        "    with capture.capture_output() as cap:\n",
        "        login(write_token, add_to_git_credential=True)\n",
        "    output = cap.stdout.strip()\n",
        "\n",
        "    if \"Token is valid.\" in output:\n",
        "        cprint(f\"Login Successful.\", color=\"green\")\n",
        "\n",
        "    api = HfApi()\n",
        "    user = api.whoami(write_token)\n",
        "    datasets_repo = f\"{user['name']}/{repo_name.strip()}\"\n",
        "\n",
        "    if repo_name:\n",
        "        create_or_validate_repo(api, datasets_repo)\n",
        "        compress_to_zip(output_path)\n",
        "        upload_and_cleanup(api, output_path, datasets_repo)\n",
        "\n",
        "main()\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "EGXqJLXwnJQB"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}