cagliostro-forge-colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "[![visitor][visitor-badge]][visitor-stats]\n",
        "\n",
        "# **Cagliostro Forge Colab**\n",
        "Rise from the ashes, reborn and empowered by [lllyasviel/stable-diffusion-webui-forge](https://github.com/lllyasviel/stable-diffusion-webui-forge)\n",
        "\n",
        "**Version 1.0.0** | [Github][link-to-github] | [License](https://github.com/cagliostrolab/forge-colab/blob/main/LICENSE)\n",
        "\n",
        "[visitor-badge]: https://api.visitorbadge.io/api/visitors?path=cagliostro-forge-colab&label=Visitors&labelColor=%2334495E&countColor=%231ABC9C&style=flat&labelStyle=none\n",
        "[visitor-stats]: https://visitorbadge.io/status?path=cagliostro-forge-colab\n",
        "[link-to-github]: https://github.com/cagliostrolab/forge-colab/blob/main/cagliostro-forge-colab.ipynb"
      ],
      "metadata": {
        "id": "GnmYPpV3o719"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title ## **Install Environment**\n",
        "import sys\n",
        "import subprocess\n",
        "import os\n",
        "import time\n",
        "import json\n",
        "import shutil\n",
        "import random\n",
        "import string\n",
        "from pathlib import Path\n",
        "from tqdm import tqdm\n",
        "from pydantic import BaseModel\n",
        "\n",
        "python_version  = \".\".join(sys.version.split(\".\")[:2])\n",
        "python_path     = Path(f\"/usr/local/lib/python{python_version}/dist-packages/\")\n",
        "colablib_path   = python_path / \"colablib\"\n",
        "if not colablib_path.exists():\n",
        "    subprocess.run(['pip', 'install', '--upgrade', 'git+https://github.com/Linaqruf/colablib'], check=True)\n",
        "\n",
        "from colablib.colored_print import cprint, print_line\n",
        "from colablib.utils import py_utils, package_utils, config_utils\n",
        "from colablib.sd_models.downloader import aria2_download, download\n",
        "from colablib.utils.git_utils import update_repo, reset_repo, validate_repo, batch_update\n",
        "from colablib.utils.py_utils import get_filename\n",
        "\n",
        "################################\n",
        "# COLAB ARGUMENTS GOES HERE\n",
        "################################\n",
        "\n",
        "# It ain't much, but it's honest work.\n",
        "class CustomDirs(BaseModel):\n",
        "    url: str\n",
        "    dst: str\n",
        "\n",
        "# @markdown ### **Drive Config**\n",
        "mount_drive          = False  # @param {type: 'boolean'}\n",
        "output_drive_folder  = \"cagliostro-colab-forge\"  # @param {type: 'string'}\n",
        "\n",
        "# @markdown ### **Repo Config**\n",
        "update_webui         = True  # @param {type: 'boolean'}\n",
        "update_extensions    = True  # @param {type: 'boolean'}\n",
        "commit_hash          = \"\"  # @param {type: 'string'}\n",
        "\n",
        "# @markdown ### **Download Config**\n",
        "# @markdown > Check only the options you need\n",
        "animagine_xl_3_1     = True  # @param {type: 'boolean'}\n",
        "rae_diffusion_xl_v2  = False  # @param {type: 'boolean'}\n",
        "kivotos_xl_v2_0      = False  # @param {type: 'boolean'}\n",
        "urangdiffusion_2_0   = False  # @param {type: 'boolean'}\n",
        "\n",
        "# @markdown > **Note:**\n",
        "# @markdown - For multiple URLs, use comma separation (e.g. `url1, url2, url3`)\n",
        "# @markdown - Forge supports FLUX, SD, and SDXL, but this notebook focuses only on SDXL\n",
        "# @markdown - **Highly Recommended:** Use Hugging Face links whenever possible\n",
        "custom_model_url     = \"\"  # @param {'type': 'string'}\n",
        "custom_vae_url       = \"https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/resolve/main/sdxl.vae.safetensors\"  # @param {'type': 'string'}\n",
        "custom_lora_url      = \"\"  # @param {'type': 'string'}\n",
        "\n",
        "# @markdown ### **Tunnel Config**\n",
        "# @markdown > Default to `--share` until `ngrok_token` is not `None`\n",
        "ngrok_token          = \"\"  # @param {type: 'string'}\n",
        "ngrok_region         = \"ap\"  # @param [\"us\", \"eu\", \"au\", \"ap\", \"sa\", \"jp\", \"in\"]\n",
        "\n",
        "# @markdown ### **UI/UX Config**\n",
        "gradio_theme         = \"remilia/Ghostly\"  # @param [\"Default\", \"gradio/base\", \"gradio/glass\", \"gradio/monochrome\", \"gradio/seafoam\", \"gradio/soft\", \"gradio/dracula_test\", \"abidlabs/dracula_test\", \"abidlabs/Lime\", \"abidlabs/pakistan\", \"Ama434/neutral-barlow\", \"dawood/microsoft_windows\", \"finlaymacklon/smooth_slate\", \"Franklisi/darkmode\", \"freddyaboulton/dracula_revamped\", \"freddyaboulton/test-blue\", \"gstaff/xkcd\", \"Insuz/Mocha\", \"Insuz/SimpleIndigo\", \"JohnSmith9982/small_and_pretty\", \"nota-ai/theme\", \"nuttea/Softblue\", \"ParityError/Anime\", \"reilnuud/polite\", \"remilia/Ghostly\", \"rottenlittlecreature/Moon_Goblin\", \"step-3-profit/Midnight-Deep\", \"Taithrah/Minimal\", \"ysharma/huggingface\", \"ysharma/steampunk\", \"NoCrypt/miku\"]\n",
        "# @markdown Set `use_preset` for using default prompt, resolution, sampler, and other settings\n",
        "use_presets          = True  # @param {type: 'boolean'}\n",
        "\n",
        "# @markdown ### **Launch Arguments**\n",
        "use_gradio_auth      = False  # @param {type: 'boolean'}\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 --opt-sdp-attention\"  # @param {type: 'string'}\n",
        "\n",
        "################################\n",
        "# GLOBAL VARIABLES GOES HERE\n",
        "################################\n",
        "\n",
        "# GRADIO AUTH\n",
        "user      = \"cagliostro\"\n",
        "password  = \"\".join(random.choices(string.ascii_letters + string.digits, k=6))\n",
        "\n",
        "# ROOT DIR\n",
        "root_dir        = Path(\"/content\")\n",
        "drive_dir       = root_dir / \"drive\" / \"MyDrive\"\n",
        "repo_dir        = root_dir / \"stable-diffusion-webui-forge\"\n",
        "tmp_dir         = root_dir / \"tmp\"\n",
        "\n",
        "models_dir      = repo_dir / \"models\"\n",
        "extensions_dir  = repo_dir / \"extensions\"\n",
        "ckpt_dir        = models_dir / \"Stable-diffusion\"\n",
        "vae_dir         = models_dir / \"VAE\"\n",
        "lora_dir        = models_dir / \"Lora\"\n",
        "output_subdir   = [\"txt2img-samples\", \"img2img-samples\", \"extras-samples\", \"txt2img-grids\", \"img2img-grids\"]\n",
        "\n",
        "config_file_path    = repo_dir / \"config.json\"\n",
        "ui_config_file_path = repo_dir / \"ui-config.json\"\n",
        "\n",
        "package_url = [\n",
        "    \"https://huggingface.co/Linaqruf/fast-repo/resolve/main/webui-forge.tar.lz4\",\n",
        "    \"https://huggingface.co/Linaqruf/fast-repo/resolve/main/webui-forge-deps.tar.lz4\",\n",
        "]\n",
        "\n",
        "custom_dirs = {\n",
        "    \"model\" : CustomDirs(url=custom_model_url, dst=str(ckpt_dir)),\n",
        "    \"vae\"   : CustomDirs(url=custom_vae_url, dst=str(vae_dir)),\n",
        "    \"lora\"  : CustomDirs(url=custom_lora_url, dst=str(lora_dir)),\n",
        "}\n",
        "\n",
        "default_model_urls = {\n",
        "    \"animagine_xl_3_1\"      : \"https://huggingface.co/cagliostrolab/animagine-xl-3.1/resolve/main/animagine-xl-3.1.safetensors\",\n",
        "    \"rae_diffusion_xl_v2\"   : \"https://huggingface.co/Raelina/Rae-Diffusion-XL-V2/resolve/main/RaeDiffusion-XL-v2.safetensors\",\n",
        "    \"kivotos_xl_v2_0\"       : \"https://huggingface.co/yodayo-ai/kivotos-xl-2.0/resolve/main/kivotos-xl-2.0.safetensors\",\n",
        "    \"urangdiffusion_2_0\"    : \"https://huggingface.co/kayfahaarukku/UrangDiffusion-2.0/resolve/main/UrangDiffusion-2.0.safetensors\",\n",
        "}\n",
        "\n",
        "################################\n",
        "# HELPER FUNCTIONS STARTS HERE\n",
        "################################\n",
        "\n",
        "def mount_drive_function(directory):\n",
        "    output_dir = repo_dir / \"outputs\"\n",
        "\n",
        "    if mount_drive:\n",
        "        print_line(80, color=\"green\")\n",
        "        if not directory.exists():\n",
        "            from google.colab import drive\n",
        "            cprint(\"Mounting google drive...\", color=\"green\", reset=False)\n",
        "            drive.mount(str(directory.parent))\n",
        "        output_dir = 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 [ckpt_dir, vae_dir, lora_dir]:\n",
        "        dir.mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "def pre_download(dir, urls, desc, overwrite=False):\n",
        "    ffmpy_path = python_path / \"ffmpy-0.3.0.dist-info\"\n",
        "\n",
        "    for url in tqdm(urls, desc=desc):\n",
        "        filename = Path(url).name\n",
        "        aria2_download(dir, filename, url, quiet=True)\n",
        "        if filename == \"webui-forge-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(dir / filename)\n",
        "\n",
        "    subprocess.run([\"rm\", \"-rf\", str(ffmpy_path)])\n",
        "    subprocess.run([\"pip\", \"install\", \"--force-reinstall\", \"ffmpy\"], check=True)\n",
        "\n",
        "def install_dependencies():\n",
        "    ubuntu_deps = [\"aria2\", \"lz4\"]\n",
        "    cprint(\"Installing ubuntu dependencies\", color=\"green\")\n",
        "    subprocess.run([\"apt\", \"install\", \"-y\"] + ubuntu_deps, check=True)\n",
        "\n",
        "def install_webui(repo_dir, desc):\n",
        "    if not repo_dir.exists():\n",
        "        pre_download(root_dir, package_url, desc, overwrite=False)\n",
        "    else:\n",
        "        cprint(\"Stable Diffusion Web UI Forge already installed, skipping...\", color=\"green\")\n",
        "\n",
        "def configure_output_path(config_path, output_dir, output_subdir):\n",
        "    try:\n",
        "        config = config_utils.read_config(str(config_path))\n",
        "    except (FileNotFoundError, json.JSONDecodeError):\n",
        "        config = {}\n",
        "\n",
        "    config_updates = {\n",
        "        f\"outdir_{subdir.split('-')[0]}_{'_'.join(subdir.split('-')[1:])}\": str(output_dir / subdir)\n",
        "        for subdir in output_subdir\n",
        "    }\n",
        "    config.update(config_updates)\n",
        "\n",
        "    config_path.parent.mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "    config_utils.write_config(str(config_path), config)\n",
        "\n",
        "    for dir in output_subdir:\n",
        "        (output_dir / dir).mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "def prepare_environment():\n",
        "    cprint(\"Preparing environment...\", color=\"green\")\n",
        "    os.environ['PYTORCH_CUDA_ALLOC_CONF']   = \"garbage_collection_threshold:0.9,max_split_size_mb:512\"\n",
        "    os.environ[\"TF_CPP_MIN_LOG_LEVEL\"]      = \"3\"\n",
        "    os.environ[\"PYTHONWARNINGS\"]            = \"ignore\"\n",
        "\n",
        "def custom_download(custom_dirs):\n",
        "    filtered_urls = filter_dict_items(default_model_urls)\n",
        "\n",
        "    for key, value in custom_dirs.items():\n",
        "        urls = value.url.split(\",\")\n",
        "        dst = value.dst\n",
        "\n",
        "        if key == \"model\":\n",
        "            urls.extend(filtered_urls)\n",
        "\n",
        "        if urls[0]:\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()\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 + Path(get_filename(url)).suffix\n",
        "                else:\n",
        "                    filename = get_filename(url)\n",
        "\n",
        "                download(url=url, filename=filename, dst=dst, quiet=False)\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(url)\n",
        "    return result_list\n",
        "\n",
        "def auto_select_file(target_dir, config_key, file_types):\n",
        "    valid_files = [f for f in os.listdir(target_dir) if f.endswith(file_types)]\n",
        "    if valid_files:\n",
        "        file_path = random.choice(valid_files)\n",
        "\n",
        "        if Path(target_dir).joinpath(file_path).exists():\n",
        "            config = config_utils.read_config(str(config_file_path))\n",
        "            config[config_key] = file_path\n",
        "            config_utils.write_config(str(config_file_path), config)\n",
        "        return file_path\n",
        "    else:\n",
        "        return None\n",
        "\n",
        "def ui_config_presets():\n",
        "    preset_prompt = \"masterpiece, best quality, very aesthetic, absurdres\"\n",
        "    preset_negative_prompt = \"nsfw, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]\"\n",
        "\n",
        "    return {\n",
        "        \"txt2img/Prompt/value\"              : preset_prompt,\n",
        "        \"txt2img/Negative prompt/value\"     : preset_negative_prompt,\n",
        "        \"img2img/Prompt/value\"              : preset_prompt,\n",
        "        \"img2img/Negative prompt/value\"     : preset_negative_prompt,\n",
        "        \"customscript/sampler.py/txt2img/Sampling method/value\" : \"Euler a\",\n",
        "        \"customscript/sampler.py/txt2img/Sampling steps/value\"  : 28,\n",
        "        \"customscript/sampler.py/txt2img/Scheduler/value\"       : \"Automatic\",\n",
        "    }\n",
        "\n",
        "def ui_config_settings(ui_config_file: str):\n",
        "    config = config_utils.read_config(str(ui_config_file))\n",
        "    preset_config = ui_config_presets()\n",
        "\n",
        "    for key, value in preset_config.items():\n",
        "        config[key] = value\n",
        "\n",
        "    config_utils.write_config(str(ui_config_file), config)\n",
        "\n",
        "def general_config_presets(config_file: str, lora_dir: str, use_presets: bool, ui_config_file: str):\n",
        "    config = config_utils.read_config(str(config_file))\n",
        "\n",
        "    config.update({\n",
        "        \"CLIP_stop_at_last_layers\"      : 2,\n",
        "        \"show_progress_every_n_steps\"   : 10,\n",
        "        \"show_progressbar\"              : True,\n",
        "        \"samples_filename_pattern\"      : \"[model_name]_[seed]\",\n",
        "        \"show_progress_type\"            : \"Approx NN\",\n",
        "        \"live_preview_content\"          : \"Prompt\",\n",
        "        \"forge_preset\"                  : \"xl\",\n",
        "        \"xl_t2i_width\"                  : 832,\n",
        "        \"xl_t2i_height\"                 : 1216,\n",
        "        \"xl_t2i_cfg\"                    : 7,\n",
        "        \"xl_t2i_hr_cfg\"                 : 7,\n",
        "        \"xl_t2i_sampler\"                : \"Euler a\",\n",
        "        \"xl_t2i_scheduler\"              : \"Automatic\",\n",
        "        \"gradio_theme\"                  : gradio_theme,\n",
        "    })\n",
        "\n",
        "    config_utils.write_config(str(config_file), config)\n",
        "\n",
        "    if use_presets:\n",
        "        ui_config_settings(ui_config_file)\n",
        "\n",
        "def is_valid(target_dir, file_types):\n",
        "    return any(f.endswith(file_types) for f in os.listdir(target_dir))\n",
        "\n",
        "def parse_args(config):\n",
        "    args = []\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args.append(f'\"{v}\"')\n",
        "        elif isinstance(v, str):\n",
        "            args.append(f'--{k}=\"{v}\"')\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args.append(f\"--{k}\")\n",
        "        elif isinstance(v, (float, int)) and not isinstance(v, bool):\n",
        "            args.append(f\"--{k}={v}\")\n",
        "    return \" \".join(args)\n",
        "\n",
        "def main():\n",
        "    global output_dir, auto_select_model, auto_select_vae\n",
        "\n",
        "    ################################\n",
        "    # MAIN EXECUTION\n",
        "    ################################\n",
        "\n",
        "    os.chdir(root_dir)\n",
        "    start_time = time.time()\n",
        "    output_dir = mount_drive_function(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",
        "    try:\n",
        "        install_dependencies()\n",
        "\n",
        "        print_line(80, color=\"green\")\n",
        "        install_webui(repo_dir, cprint(\"Unpacking Web UI Forge\", color=\"green\", tqdm_desc=True))\n",
        "        prepare_environment()\n",
        "\n",
        "        configure_output_path(config_file_path, 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",
        "        elif commit_hash:\n",
        "            reset_repo(repo_dir, commit_hash)\n",
        "\n",
        "        setup_directories()\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",
        "        if update_extensions:\n",
        "            print_line(80, color=\"green\")\n",
        "            batch_update(fetch=True, directory=extensions_dir, desc=cprint(\"Updating extensions\", color=\"green\", tqdm_desc=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",
        "    except Exception as e:\n",
        "        cprint(f\"An error occurred: {str(e)}\", color=\"red\")\n",
        "        print_line(80, color=\"red\")\n",
        "        cprint(\"Setup failed. Please check the error message above and try again.\", color=\"red\")\n",
        "        print_line(80, color=\"red\")\n",
        "        return\n",
        "\n",
        "    start_time = time.time()\n",
        "\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",
        "    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(ckpt_dir, ('.ckpt', '.safetensors')):\n",
        "        cprint(f\"No checkpoints were found in the directory '{ckpt_dir}'.\", color=\"yellow\")\n",
        "        url = \"https://huggingface.co/cagliostrolab/animagine-xl-3.1/blob/main/animagine-xl-3.1.safetensors\"\n",
        "        filename = get_filename(url)\n",
        "        aria2_download(url=url, download_dir=ckpt_dir, filename=filename)\n",
        "        print_line(80, color=\"green\")\n",
        "        auto_select_model = True\n",
        "\n",
        "    if not is_valid(vae_dir, ('.vae.pt', '.vae.safetensors', '.pt', '.ckpt')):\n",
        "        cprint(f\"No VAEs were found in the directory '{vae_dir}'.\", color=\"yellow\")\n",
        "        url = \"https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/blob/main/sdxl.vae.safetensors\"\n",
        "        filename = get_filename(url)\n",
        "        aria2_download(url=url, download_dir=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(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(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",
        "    general_config_presets(config_file_path, lora_dir, use_presets, ui_config_file_path)\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",
        "        \"share\"                           : True if not ngrok_token else False,\n",
        "        \"ngrok\"                           : ngrok_token if ngrok_token else None,\n",
        "        \"ngrok-region\"                    : ngrok_region if ngrok_token else None,\n",
        "        \"gradio-auth\"                     : f\"{user}:{password}\" if use_gradio_auth else None,\n",
        "        \"no-hashing\"                      : True,\n",
        "        \"disable-console-progressbars\"    : True,\n",
        "        \"lowram\"                          : True,\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\"                        : ckpt_dir,\n",
        "        \"vae-dir\"                         : vae_dir,\n",
        "        \"lora-dir\"                        : 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",
        "    ! {final_args}\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()"
      ],
      "metadata": {
        "id": "YYzHDlgEkkrY",
        "cellView": "form"
      },
      "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 zipfile\n",
        "from pathlib import Path\n",
        "from google.colab import auth\n",
        "from pydrive.auth import GoogleAuth\n",
        "from pydrive.drive import GoogleDrive\n",
        "from oauth2client.client import GoogleCredentials\n",
        "from colablib.colored_print import cprint\n",
        "\n",
        "os.chdir(output_dir)\n",
        "\n",
        "use_drive = False  # @param {type:\"boolean\"}\n",
        "folder_name = \"cagliostro-forge-colab\"  # @param {type: \"string\"}\n",
        "filename = \"waifu.zip\"  # @param {type: \"string\"}\n",
        "save_as = filename\n",
        "\n",
        "def get_unique_filename(base_filename):\n",
        "    path = Path(base_filename)\n",
        "    if not path.exists():\n",
        "        return path\n",
        "    i = 1\n",
        "    while True:\n",
        "        new_path = path.with_name(f\"{path.stem}({i}){path.suffix}\")\n",
        "        if not new_path.exists():\n",
        "            return new_path\n",
        "        i += 1\n",
        "\n",
        "filename = get_unique_filename(filename)\n",
        "\n",
        "def zip_directory(directory, zipname):\n",
        "    with zipfile.ZipFile(zipname, 'w', zipfile.ZIP_DEFLATED) as zipf:\n",
        "        for file_path in directory.rglob('*'):\n",
        "            if file_path.is_file():\n",
        "                zipf.write(file_path, file_path.relative_to(directory.parent))\n",
        "\n",
        "zip_directory(output_dir, Path('/content/outputs.zip'))\n",
        "\n",
        "if use_drive:\n",
        "    auth.authenticate_user()\n",
        "    gauth = GoogleAuth()\n",
        "    gauth.credentials = GoogleCredentials.get_application_default()\n",
        "    drive_service = GoogleDrive(gauth)\n",
        "\n",
        "    def create_folder(folder_name):\n",
        "        query = f\"title='{folder_name}' and mimeType='application/vnd.google-apps.folder' and trashed=false\"\n",
        "        file_list = drive_service.ListFile({\"q\": query}).GetList()\n",
        "        if file_list:\n",
        "            cprint(\"Debug: Folder exists\", color=\"green\")\n",
        "            return file_list[0][\"id\"]\n",
        "        else:\n",
        "            cprint(\"Debug: Creating folder\", color=\"green\")\n",
        "            folder = drive_service.CreateFile({\n",
        "                \"title\": folder_name,\n",
        "                \"mimeType\": \"application/vnd.google-apps.folder\"\n",
        "            })\n",
        "            folder.Upload()\n",
        "            return folder[\"id\"]\n",
        "\n",
        "    def upload_file(file_path, folder_id, save_as):\n",
        "        save_as = get_unique_filename(save_as)\n",
        "        file = drive_service.CreateFile({\"title\": save_as.name, \"parents\": [{\"id\": folder_id}]})\n",
        "        file.SetContentFile(str(file_path))\n",
        "        file.Upload()\n",
        "        file.InsertPermission({\"type\": \"anyone\", \"value\": \"anyone\", \"role\": \"reader\"})\n",
        "        return file[\"id\"]\n",
        "\n",
        "    folder_id = create_folder(folder_name)\n",
        "    file_id = upload_file(Path('/content/outputs.zip'), folder_id, Path(save_as))\n",
        "    sharing_link = f\"https://drive.google.com/file/d/{file_id}/view?usp=sharing\"\n",
        "    cprint(f\"Your sharing link: {sharing_link}\", color=\"green\")\n",
        "else:\n",
        "    cprint(\"Files zipped locally. Download manually from the files tab.\", color=\"yellow\")\n"
      ],
      "metadata": {
        "id": "UyTKsCa1qUL4",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}