{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JKTCrY9LU7Oq"
      },
      "outputs": [],
      "source": [
        "##~   LAUNCH CODE | BY: ANXETY   ~##\n",
        "\n",
        "import os\n",
        "import re\n",
        "import time\n",
        "import json\n",
        "import requests\n",
        "import cloudpickle as pickle\n",
        "from datetime import timedelta\n",
        "from IPython.display import clear_output\n",
        "\n",
        "\n",
        "#  ================= DETECT ENV =================\n",
        "def detect_environment():\n",
        "    free_plan = (os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES') / (1024 ** 3) <= 20)\n",
        "    environments = {\n",
        "        'COLAB_GPU': ('Google Colab', \"/root\" if free_plan else \"/content\"),\n",
        "        'KAGGLE_URL_BASE': ('Kaggle', \"/kaggle/working/content\")\n",
        "    }\n",
        "\n",
        "    for env_var, (environment, path) in environments.items():\n",
        "        if env_var in os.environ:\n",
        "            return environment, path, free_plan\n",
        "\n",
        "env, root_path, free_plan = detect_environment()\n",
        "webui_path = f\"{root_path}/sdw\"\n",
        "#  ----------------------------------------------\n",
        "\n",
        "def load_settings():\n",
        "    SETTINGS_FILE = f'{root_path}/settings.json'\n",
        "    if os.path.exists(SETTINGS_FILE):\n",
        "        with open(SETTINGS_FILE, 'r') as f:\n",
        "            settings = json.load(f)\n",
        "        return settings\n",
        "\n",
        "settings = load_settings()\n",
        "ngrok_token = settings['ngrok_token']\n",
        "zrok_token = settings['zrok_token']\n",
        "commandline_arguments = settings['commandline_arguments']\n",
        "change_webui = settings['change_webui']\n",
        "\n",
        "\n",
        "# ======================== TUNNEL V2 ========================\n",
        "print('Please Wait...')\n",
        "\n",
        "def get_public_ip(version='ipv4'):\n",
        "    try:\n",
        "        url = f'https://api64.ipify.org?format=json&{version}=true'\n",
        "        response = requests.get(url)\n",
        "        data = response.json()\n",
        "        public_ip = data['ip']\n",
        "        return public_ip\n",
        "    except Exception as e:\n",
        "        print(f\"Error getting public {version} address:\", e)\n",
        "\n",
        "# Check if public IP is already saved, if not then get it\n",
        "try:\n",
        "    with open(f\"{root_path}/public_ip.txt\", \"r\") as file:\n",
        "        public_ipv4 = file.read().strip()\n",
        "except FileNotFoundError:\n",
        "    public_ipv4 = get_public_ip(version='ipv4')\n",
        "    with open(f\"{root_path}/public_ip.txt\", \"w\") as file:\n",
        "        file.write(public_ipv4)\n",
        "\n",
        "tunnel_class = pickle.load(open(f\"{root_path}/new_tunnel\", \"rb\"), encoding=\"utf-8\")\n",
        "tunnel_port = 1734\n",
        "tunnel = tunnel_class(tunnel_port)\n",
        "tunnel.add_tunnel(command=\"cl tunnel --url localhost:{port}\", name=\"cl\", pattern=re.compile(r\"[\\w-]+\\.trycloudflare\\.com\"))\n",
        "tunnel.add_tunnel(command=\"lt --port {port}\", name=\"lt\", pattern=re.compile(r\"[\\w-]+\\.loca\\.lt\"), note=\"Password : \" + \"\\033[32m\" + public_ipv4 + \"\\033[0m\" + \" rerun cell if 404 error.\")\n",
        "\n",
        "''' add zrok tunnel '''\n",
        "if zrok_token:\n",
        "    get_ipython().system('zrok enable {zrok_token} &> /dev/null')\n",
        "    tunnel.add_tunnel(command=\"zrok share public http://localhost:{port}/ --headless\", name=\"zrok\", pattern=re.compile(r\"[\\w-]+\\.share\\.zrok\\.io\"))\n",
        "\n",
        "clear_output()\n",
        "\n",
        "\n",
        "# =============== Automatic Fixing Path V3 ===============\n",
        "paths_to_check = [\n",
        "    (\"tagger_hf_cache_dir\", f\"{webui_path}/models/interrogators/\"),\n",
        "    (\"additional_networks_extra_lora_path\", f\"{webui_path}/models/Lora/\"),\n",
        "    (\"ad_extra_models_dir\", f\"{webui_path}/models/adetailer/\"),\n",
        "    (\"sd_checkpoint_hash\", \"\"),\n",
        "    (\"sd_model_checkpoint\", \"\"),\n",
        "    (\"sd_vae\", \"None\")\n",
        "]\n",
        "\n",
        "config_path = f'{webui_path}/ui-config.json'\n",
        "\n",
        "with open(config_path, 'r') as file:\n",
        "    config_data = json.load(file)\n",
        "\n",
        "for key, value in paths_to_check:\n",
        "    if key in config_data and config_data[key] != value:\n",
        "        sed_command = f\"sed -i 's|\\\"{key}\\\": \\\".*\\\"|\\\"{key}\\\": \\\"{value}\\\"|' {config_path}\"\n",
        "        os.system(sed_command)\n",
        "\n",
        "# Additional check for Kaggle\n",
        "if env == 'Kaggle':\n",
        "   get_ipython().system('sed -i \\'s/\"civitai_interface\\\\/NSFW content\\\\/value\":.*/\"civitai_interface\\\\/NSFW content\\\\/value\": false/g\\' {webui_path}/ui-config.json')\n",
        "#  -------------------------------------------------------\n",
        "\n",
        "\n",
        "with tunnel:\n",
        "    %cd {webui_path}\n",
        "\n",
        "    commandline_arguments += f' --port={tunnel_port}'\n",
        "    if ngrok_token:\n",
        "        commandline_arguments += f' --ngrok {ngrok_token}'\n",
        "    if env != \"Google Colab\":\n",
        "        commandline_arguments += f' --encrypt-pass={tunnel_port} --api'\n",
        "\n",
        "    # -- FORGE --\n",
        "    if change_webui == 'Forge':\n",
        "        commandline_arguments += ' --cuda-stream --pin-shared-memory'\n",
        "\n",
        "    !COMMANDLINE_ARGS=\"{commandline_arguments}\" python launch.py\n",
        "\n",
        "\n",
        "# after runnig\n",
        "start_colab = float(open(f'{webui_path}/static/colabTimer.txt', 'r').read())\n",
        "time_since_start = str(timedelta(seconds=time.time()-start_colab)).split('.')[0]\n",
        "print(f\"\\n⌚️ \\033[0mВы проводите эту сессию в течение - \\033[33m{time_since_start}\\033[0m\\n\\n\")\n",
        "\n",
        "''' del zrok tunnel '''\n",
        "if zrok_token:\n",
        "    !zrok disable &> /dev/null"
      ]
    }
  ]
}