{
  "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": [
        "### 1. ⚠️ подключить гуглодиск на панели слева"
      ],
      "metadata": {
        "id": "7FWvC7H5xcYm"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "AXREDL496gE-"
      },
      "outputs": [],
      "source": [
        "# @title 2. установить\n",
        "from urllib.parse import urlparse, unquote\n",
        "from tarfile import open as open_tar\n",
        "from pathlib import Path\n",
        "from zstandard import ZstdDecompressor\n",
        "from tqdm import tqdm\n",
        "from json import dumps, loads\n",
        "\n",
        "webui_type = \"a1111\" # @param [\"a1111\",\"sdnext\",\"forge\",\"reforge\",\"comfy\"]\n",
        "commnit_hash = \"\" # @param {\"type\":\"string\"}\n",
        "base_url = 'https://huggingface.co/2ch/repos/resolve/main'\n",
        "ui_url = f'{base_url}/{{ui}}.tar.zst'\n",
        "venv_url = f'{base_url}/{{ui}}_venv.tar.zst'\n",
        "output = Path('/content/drive/MyDrive/work/outputs')\n",
        "output.mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "%cd /content/\n",
        "\n",
        "get_ipython().system('mount -o remount,size=11G /dev/shm')\n",
        "\n",
        "\n",
        "def decompress(input_file, output_dir):\n",
        "    output_path = Path(output_dir).resolve()\n",
        "    output_path.mkdir(parents=True, exist_ok=True)\n",
        "    with open(input_file, 'rb') as source:\n",
        "        with ZstdDecompressor().stream_reader(source) as reader:\n",
        "            with open_tar(mode='r|', fileobj=reader) as tar:\n",
        "                pbar = tqdm(desc='распаковка', unit=' файлов')\n",
        "                member = tar.next()\n",
        "                while member is not None:\n",
        "                    if member.isfile():\n",
        "                        try:\n",
        "                            target_path = output_path / member.name\n",
        "                            target_path = target_path.resolve()\n",
        "                            if not target_path.is_relative_to(output_path):\n",
        "                                raise ValueError(f'попытка извлечения за пределы целевого каталога: {member.name}')\n",
        "                            target_path.parent.mkdir(parents=True, exist_ok=True)\n",
        "                            pbar.set_description(f'распаковано: ...{member.name[30:]}')\n",
        "                            tar.extract(member, output_path)\n",
        "                            pbar.update(1)\n",
        "                        except Exception as e:\n",
        "                            pbar.write(f'ошибка распаковки {member.name}: {e}')\n",
        "                    member = tar.next()\n",
        "                pbar.close()\n",
        "\n",
        "\n",
        "\n",
        "def dl_unpack(ui):\n",
        "    get_ipython().system('rm -rf /dev/shm/*')\n",
        "    for url in [venv_url.format(ui=ui), ui_url.format(ui=ui)]:\n",
        "        get_ipython().system(f'wget -nv -t 10 --show-progress -q --content-disposition {url} -P /dev/shm')\n",
        "        file = Path(unquote(urlparse(url).path))\n",
        "        archive, folder = Path('/dev/shm') / file.name, Path('/content')\n",
        "        webui_folder = folder / Path(file.stem).stem\n",
        "        decompress(archive, folder)\n",
        "        get_ipython().system('rm -rf /dev/shm/*')\n",
        "        if commnit_hash:\n",
        "            get_ipython().system(f'git -C {webui_folder.as_posix} restore .')\n",
        "            get_ipython().system(f'git -C {webui_folder.as_posix} checkout {commnit_hash}')\n",
        "        if webui_type != 'comfy':\n",
        "            get_ipython().system(f'git clone https://huggingface.co/prolapse/dl_models_addon {webui_folder / \"extensions\" / \"dl_models\"}')\n",
        "\n",
        "\n",
        "def output_config():\n",
        "    config_json_path = Path(f'/content/{webui_type}') / 'config.json'\n",
        "    if not config_json_path.exists():\n",
        "        config_json_path.touch()\n",
        "    default_outdirs = {\n",
        "        'outdir_samples': '',\n",
        "        'outdir_txt2img_samples': (output / 't2i').as_posix(),\n",
        "        'outdir_img2img_samples': (output / 'i2i').as_posix(),\n",
        "        'outdir_extras_samples': (output / 'upscaled_face_restored').as_posix(),\n",
        "        'outdir_grids': '',\n",
        "        'outdir_txt2img_grids': (output / 't2i_tiles').as_posix(),\n",
        "        'outdir_img2img_grids': (output / 'i2i_tiles').as_posix(),\n",
        "        'outdir_save': (output / 'selected').as_posix(),\n",
        "        'outdir_init_images': (output / 'init-images').as_posix(),\n",
        "    }\n",
        "    config_json_path.write_text(dumps(default_outdirs, indent=4), encoding='utf-8')\n",
        "\n",
        "\n",
        "dl_unpack(webui_type)\n",
        "\n",
        "if webui_type != 'comfy':\n",
        "    from json import dumps, loads\n",
        "    output_config()\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @title 3. запустить\n",
        "\n",
        "from pathlib import Path\n",
        "from random import randint\n",
        "from subprocess import PIPE, Popen, STDOUT, run\n",
        "from time import sleep\n",
        "from re import compile\n",
        "from requests import get\n",
        "\n",
        "webui_type = \"a1111\" # @param [\"a1111\",\"sdnext\",\"forge\",\"reforge\",\"comfy\"]\n",
        "\n",
        "output = Path('/content/drive/MyDrive/work/outputs')\n",
        "port = randint(7877, 8077)\n",
        "work_dir = f'/content/{webui_type}'\n",
        "python = f'/content/{webui_type}_venv/bin/python'\n",
        "\n",
        "\n",
        "def create_gradio_tunnel(port: int) -> str:\n",
        "    gradio_bin = Path('/content/.config/frpc_linux_amd64')\n",
        "    bin_url = 'https://cdn-media.huggingface.co/frpc-gradio-0.1/frpc_linux_amd64'\n",
        "    max_attempts = 3\n",
        "\n",
        "    for attempt in range(max_attempts):\n",
        "        try:\n",
        "            if not gradio_bin.exists():\n",
        "                gradio_bin.parent.mkdir(parents=True, exist_ok=True)\n",
        "                response = get(bin_url)\n",
        "                response.raise_for_status()\n",
        "                gradio_bin.write_bytes(response.content)\n",
        "                gradio_bin.chmod(0o777)\n",
        "\n",
        "            run(['pkill', '-f', gradio_bin.name], capture_output=True)\n",
        "\n",
        "            resp = get('https://api.gradio.app/v2/tunnel-request').json()\n",
        "            host, r_port = resp[0]['host'], resp[0]['port']\n",
        "\n",
        "            cmd = [\n",
        "                str(gradio_bin), 'http',\n",
        "                '-n', 'random',\n",
        "                '-l', str(port),\n",
        "                '-i', '127.0.0.1',\n",
        "                '--server_addr', f'{host}:{r_port}',\n",
        "                '--uc', '--sd', 'random', '--ue', '--disable_log_color'\n",
        "            ]\n",
        "            process = Popen(\n",
        "                cmd,\n",
        "                stdout=PIPE,\n",
        "                stderr=STDOUT,\n",
        "                text=True\n",
        "            )\n",
        "\n",
        "            url_pattern = compile(r'https://\\S+')\n",
        "            for _ in range(5):\n",
        "                line = process.stdout.readline()\n",
        "                if match := url_pattern.search(line):\n",
        "                    return match.group()\n",
        "\n",
        "                if process.poll() is not None:\n",
        "                    break\n",
        "\n",
        "            raise RuntimeError('url не найден в выводе')\n",
        "\n",
        "        except Exception as e:\n",
        "            if attempt == max_attempts - 1:\n",
        "                raise RuntimeError(f'не удалось создать туннель после {max_attempts} попыток: {str(e)}')\n",
        "\n",
        "            print(f'ошибка (попытка {attempt + 1}): {str(e)}')\n",
        "            sleep(2)\n",
        "\n",
        "    return ''\n",
        "\n",
        "\n",
        "\n",
        "%cd $work_dir\n",
        "\n",
        "gradio_args = [\n",
        "    '--opt-sdp-attention',\n",
        "    '--no-half-vae',\n",
        "    '--api',\n",
        "    '--no-download-sd-model',\n",
        "    '--disable-console-progressbars',\n",
        "    '--no-hashing',\n",
        "    '--skip-version-check',\n",
        "    '--skip-python-version-check',\n",
        "    '--skip-torch-cuda-test',\n",
        "    '--disable-safe-unpickle',\n",
        "    '--enable-insecure-extension-access',\n",
        "    f'--port {port}',\n",
        "    '--listen',\n",
        "    '--theme dark',\n",
        "    '--no-gradio-queue'\n",
        "]\n",
        "\n",
        "comfy_args = [\n",
        "    '--listen 0.0.0.0',\n",
        "    f'--port {port}',\n",
        "    f'--output-directory {output / \"cozy\"}',\n",
        "    '--disable-auto-launch',\n",
        "    '--dont-print-server',\n",
        "    '--fp8_e4m3fn-unet --fp8_e4m3fn-text-enc',  # для флюкс, сд3, т.п.\n",
        "    '--use-pytorch-cross-attention --disable-xformers',\n",
        "    '--lowvram'\n",
        "]\n",
        "\n",
        "\n",
        "url = create_gradio_tunnel(port)\n",
        "args = ' '.join(comfy_args if webui_type == 'comfy' else gradio_args)\n",
        "print(url)\n",
        "get_ipython().system(f'{python} launch.py {args}')\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "1kUYNYTMTXNb"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}