easy_stable_diffusion.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/nicarinai/nicarin/blob/master/easy_stable_diffusion.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## <font color=\"red\">실행하는 방법</font>\n",
        "\n",
        "- [버그 보고 및 질답용 디스코드 서버](https://discord.gg/6wQeA2QXgM)\n",
        "- [원클릭 다운로더 코랩 노트북](https://colab.research.google.com/drive/1cDsP8Ofgd7xtA_kMSdx-fJyI1WPRjFR4?usp=sharing)\n",
        "- [원클릭 코랩 노트북 (최신 버전)](https://colab.research.google.com/drive/1nBaePtwcW_ds7OQdFebcxB91n_aORQY5)\n",
        "- [레포지토리 (최신 버전)](https://github.com/toriato/easy-stable-diffusion)\n",
        "- [아카라이브 게시글 (**<font color=\"red\">자주 물어보는 질문, 오류</font>)**](https://arca.live/b/aiart/60472214)\n",
        "\n",
        "1. 스마트폰으로 접속했다면 **<font color=\"red\">데스크탑 모드</font>**로 페이지 열기 ([여는 방법](https://www.google.com/search?q=%EB%8D%B0%EC%8A%A4%ED%81%AC%ED%83%91+%EB%B3%B4%EA%B8%B0))\n",
        "1. 아래에 있는 **<font color=\"orange\">실행</font>** 셀 좌측에 있는 있는 <font color=\"red\">**재생(▶) 아이콘**</font> 클릭 ([예시 이미지](https://cdn.discordapp.com/attachments/872959812407816235/1029512475781103757/2022-10-12_06-55-01_02b9_librewolf.png))  \n",
        "    <font color=\"orange\" size=\"2\">오류가 발생하면 빨간색 경고 메세지가 나옴 ([예시 이미지](https://cdn.discordapp.com/attachments/872959812407816235/1029468903216267294/2022-10-12_04-01-54_02b5_librewolf.png))</font>  \n",
        "    <font color=\"orange\" size=\"2\">정상적으로 완료되면 초록색 메세지가 나옴 ([예시 이미지](https://media.discordapp.net/attachments/872959812407816235/1029468507328499784/2022-10-12_04-00-19_02b4_librewolf.png))</font>  \n",
        "    <font color=\"red\" size=\"2\">처음 실행할 때 약 10분 정도 소요됨</font>\n",
        "1. **<font color=\"red\">사용 중 이 페이지 절대 끄면 안됨</font>**"
      ],
      "metadata": {
        "id": "83yYK4Gh5AKI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import io\n",
        "import json\n",
        "import os\n",
        "import shlex\n",
        "import shutil\n",
        "import signal\n",
        "import subprocess\n",
        "import sys\n",
        "import time\n",
        "from datetime import datetime\n",
        "from distutils.spawn import find_executable\n",
        "from importlib.util import find_spec\n",
        "from pathlib import Path\n",
        "from typing import Callable, List, Optional, Tuple, Union\n",
        "\n",
        "import requests\n",
        "\n",
        "OPTIONS = {}\n",
        "\n",
        "# fmt: off\n",
        "#####################################################\n",
        "# 코랩 노트북에선 #@param 문법으로 사용자로부터 설정 값을 가져올 수 있음\n",
        "# 다른 환경일 땐 override.json 파일 등을 사용해야함\n",
        "#####################################################\n",
        "#@title\n",
        "\n",
        "#@markdown ### <font color=\"orange\">***작업 디렉터리 경로***</font>\n",
        "#@markdown 임베딩, 모델, 결과와 설정 파일 등이 영구적으로 보관될 디렉터리 경로\n",
        "WORKSPACE = 'SD' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown ##### <font color=\"orange\">***자동으로 코랩 런타임을 종료할지?***</font>\n",
        "DISCONNECT_RUNTIME = True  #@param {type:\"boolean\"}\n",
        "OPTIONS['DISCONNECT_RUNTIME'] = DISCONNECT_RUNTIME\n",
        "\n",
        "#@markdown ##### <font color=\"orange\">***구글 드라이브와 동기화할지?***</font>\n",
        "#@markdown <font color=\"red\">**주의**</font>: 동기화 전 남은 용량이 충분한지 확인 필수 (5GB 이상)\n",
        "USE_GOOGLE_DRIVE = True  #@param {type:\"boolean\"}\n",
        "OPTIONS['USE_GOOGLE_DRIVE'] = USE_GOOGLE_DRIVE\n",
        "\n",
        "#@markdown ##### <font color=\"orange\">***xformers 를 사용할지?***</font>\n",
        "#@markdown - <font color=\"green\">장점</font>: 이미지 생성 속도 개선 가능성 있음\n",
        "#@markdown - <font color=\"red\">단점</font>: 출력한 그림의 질이 조금 떨어질 수 있음\n",
        "USE_XFORMERS = True  #@param {type:\"boolean\"}\n",
        "OPTIONS['USE_XFORMERS'] = USE_XFORMERS\n",
        "\n",
        "#@markdown ##### <font color=\"orange\">***인증 정보***</font>\n",
        "#@markdown 접속 시 사용할 사용자 아이디와 비밀번호\n",
        "#@markdown <br>`GRADIO_USERNAME` 입력 란에 `user1:pass1,user,pass2`처럼 입력하면 여러 사용자 추가 가능\n",
        "#@markdown <br>`GRADIO_USERNAME` 입력 란을 <font color=\"red\">비워두면</font> 인증 과정을 사용하지 않음\n",
        "GRADIO_USERNAME = '' #@param {type:\"string\"}\n",
        "GRADIO_PASSWORD = '' #@param {type:\"string\"}\n",
        "OPTIONS['GRADIO_USERNAME'] = GRADIO_USERNAME\n",
        "OPTIONS['GRADIO_PASSWORD'] = GRADIO_PASSWORD\n",
        "\n",
        "#@markdown ##### <font color=\"orange\">***터널링 서비스***</font>\n",
        "TUNNEL = 'gradio' #@param [\"none\", \"gradio\", \"cloudflared\", \"ngrok\"]\n",
        "TUNNEL_URL: Optional[str] = None\n",
        "OPTIONS['TUNNEL'] = TUNNEL\n",
        "\n",
        "#@markdown ##### <font color=\"orange\">***ngrok API 키***</font>\n",
        "#@markdown ngrok 터널에 사용할 API 토큰\n",
        "#@markdown <br>[설정하는 방법은 여기를 클릭해 확인](https://arca.live/b/aiart/60683088), [API 토큰은 여기를 눌러 계정을 만든 뒤 얻을 수 있음](https://dashboard.ngrok.com/get-started/your-authtoken)\n",
        "#@markdown <br>입력 란을 <font color=\"red\">비워두면</font> ngrok 터널을 비활성화함\n",
        "#@markdown - <font color=\"green\">장점</font>: 접속이 빠른 편이고 타임아웃이 거의 발생하지 않음\n",
        "#@markdown - <font color=\"red\">**단점**</font>: 계정을 만들고 API 토큰을 직접 입력해줘야함\n",
        "NGROK_API_TOKEN = '' #@param {type:\"string\"}\n",
        "OPTIONS['NGROK_API_TOKEN'] = NGROK_API_TOKEN\n",
        "\n",
        "#@markdown ##### <font color=\"orange\">***레포지토리 주소***</font>\n",
        "REPO_URL = 'https://github.com/nicarinai/nicarin.git' #@param {type:\"string\"}\n",
        "OPTIONS['REPO_URL'] = REPO_URL\n",
        "\n",
        "#@markdown ##### <font color=\"orange\">***레포지토리 커밋 해시***</font>\n",
        "#@markdown <br>과거 커밋 해시 값[(영문과 숫자로된 난수 값; 예시 이미지)](https://vmm.pw/MzMy)을 아래에 붙여넣은 뒤 실행하면 과거 버전을 사용할 수 있음\n",
        "#@markdown <br>입력 란을 <font color=\"red\">비워두면</font> 가장 최신 커밋을 가져옴\n",
        "REPO_COMMIT = '' #@param {type:\"string\"}\n",
        "OPTIONS['REPO_COMMIT'] = REPO_COMMIT\n",
        "\n",
        "#@markdown ##### <font color=\"orange\">***Python 바이너리 이름***</font>\n",
        "#@markdown 입력 란을 <font color=\"red\">비워두면</font> 시스템에 설치된 Python 을 사용함\n",
        "PYTHON_EXECUTABLE = '' #@param {type:\"string\"}\n",
        "OPTIONS['PYTHON_EXECUTABLE'] = PYTHON_EXECUTABLE\n",
        "\n",
        "#@markdown ##### <font color=\"orange\">***인자***</font>\n",
        "#@markdown <font color=\"red\">**주의**</font>: 비어있지 않으면 실행에 필요한 인자가 자동으로 생성되지 않음\n",
        "ARGS = '' #@param {type:\"string\"}\n",
        "OPTIONS['ARGS'] = shlex.split(ARGS)\n",
        "\n",
        "#@markdown ##### <font color=\"orange\">***추가 인자***</font>\n",
        "EXTRA_ARGS = '' #@param {type:\"string\"}\n",
        "OPTIONS['EXTRA_ARGS'] = shlex.split(EXTRA_ARGS)\n",
        "\n",
        "#####################################################\n",
        "# 사용자 설정 값 끝\n",
        "#####################################################\n",
        "# fmt: on\n",
        "\n",
        "# 로그 변수\n",
        "LOG_FILE: Optional[io.TextIOWrapper] = None\n",
        "LOG_WIDGET = None\n",
        "LOG_BLOCKS = []\n",
        "\n",
        "# 로그 HTML 위젯 스타일\n",
        "LOG_WIDGET_STYLES = {\n",
        "    'wrapper': {\n",
        "        'overflow-x': 'auto',\n",
        "        'max-width': '100%',\n",
        "        'padding': '1em',\n",
        "        'background-color': 'black',\n",
        "        'white-space': 'pre',\n",
        "        'font-family': 'monospace',\n",
        "        'font-size': '1em',\n",
        "        'line-height': '1.1em',\n",
        "        'color': 'white'\n",
        "    },\n",
        "    'dialog': {\n",
        "        'display': 'block',\n",
        "        'margin-top': '.5em',\n",
        "        'padding': '.5em',\n",
        "        'font-weight': 'bold',\n",
        "        'font-size': '1.5em',\n",
        "        'line-height': '1em',\n",
        "        'color': 'black'\n",
        "    }\n",
        "}\n",
        "LOG_WIDGET_STYLES['dialog_success'] = {\n",
        "    **LOG_WIDGET_STYLES['dialog'],\n",
        "    'border': '3px dashed darkgreen',\n",
        "    'background-color': 'green',\n",
        "}\n",
        "LOG_WIDGET_STYLES['dialog_error'] = {\n",
        "    **LOG_WIDGET_STYLES['dialog'],\n",
        "    'border': '3px dashed darkred',\n",
        "    'background-color': 'red',\n",
        "}\n",
        "\n",
        "IN_INTERACTIVE = hasattr(sys, 'ps1')\n",
        "IN_COLAB = False\n",
        "\n",
        "try:\n",
        "    from IPython import get_ipython\n",
        "    IN_COLAB = 'google.colab' in str(get_ipython())\n",
        "except ImportError:\n",
        "    pass\n",
        "\n",
        "\n",
        "def hook_runtime_disconnect():\n",
        "    \"\"\"\n",
        "    셀이 종료됐을 때 자동으로 런타임을 해제하도록 asyncio 스레드를 생성합니다\n",
        "    \"\"\"\n",
        "    if not IN_COLAB:\n",
        "        return\n",
        "\n",
        "    from google.colab import runtime\n",
        "\n",
        "    # asyncio 는 여러 겹으로 사용할 수 없게끔 설계됐기 때문에\n",
        "    # 주피터 노트북 등 이미 루프가 돌고 있는 곳에선 사용할 수 없음\n",
        "    # 이는 nest-asyncio 패키지를 통해 어느정도 우회하여 사용할 수 있음\n",
        "    # https://pypi.org/project/nest-asyncio/\n",
        "    if not has_python_package('nest_asyncio'):\n",
        "        execute(['pip', 'install', 'nest-asyncio'])\n",
        "\n",
        "    import nest_asyncio\n",
        "    nest_asyncio.apply()\n",
        "\n",
        "    import asyncio\n",
        "\n",
        "    async def unassign():\n",
        "        time.sleep(1)\n",
        "        runtime.unassign()\n",
        "\n",
        "    # 평범한 환경에선 비동기로 동작하여 바로 실행되나\n",
        "    # 코랩 런타임에선 순차적으로 실행되기 때문에 현재 셀 종료 후 즉시 실행됨\n",
        "    asyncio.create_task(unassign())\n",
        "\n",
        "\n",
        "def setup_tunnels():\n",
        "    global TUNNEL_URL\n",
        "\n",
        "    tunnel = OPTIONS['TUNNEL']\n",
        "\n",
        "    if tunnel == 'none':\n",
        "        pass\n",
        "\n",
        "    elif tunnel == 'gradio':\n",
        "        if not has_python_package('gradio'):\n",
        "            # https://fastapi.tiangolo.com/release-notes/#0910\n",
        "            execute(['pip', 'install', 'gradio', 'fastapi==0.90.1'])\n",
        "\n",
        "        import secrets\n",
        "\n",
        "        from gradio.networking import setup_tunnel\n",
        "        TUNNEL_URL = setup_tunnel('localhost', 7860, secrets.token_urlsafe(32))\n",
        "\n",
        "    elif tunnel == 'cloudflared':\n",
        "        if not has_python_package('pycloudflared'):\n",
        "            execute(['pip', 'install', 'pycloudflared'])\n",
        "\n",
        "        from pycloudflared import try_cloudflare\n",
        "        TUNNEL_URL = try_cloudflare(port=7860).tunnel\n",
        "\n",
        "    elif tunnel == 'ngrok':\n",
        "        if not has_python_package('pyngrok'):\n",
        "            execute(['pip', 'install', 'pyngrok'])\n",
        "\n",
        "        auth = None\n",
        "        token = OPTIONS['NGROK_API_TOKEN']\n",
        "\n",
        "        if ':' in token:\n",
        "            parts = token.split(':')\n",
        "            auth = parts[1] + ':' + parts[-1]\n",
        "            token = parts[0]\n",
        "\n",
        "        from pyngrok import conf, exception, ngrok, process\n",
        "\n",
        "        # 로컬 포트가 닫혀있으면 경고 메세지가 스팸마냥 출력되므로 오류만 표시되게 수정함\n",
        "        process.ngrok_logger.setLevel('ERROR')\n",
        "\n",
        "        try:\n",
        "            tunnel = ngrok.connect(\n",
        "                7860,\n",
        "                pyngrok_config=conf.PyngrokConfig(\n",
        "                    auth_token=token,\n",
        "                    region='jp'\n",
        "                ),\n",
        "                auth=auth,\n",
        "                bind_tls=True\n",
        "            )\n",
        "        except exception.PyngrokNgrokError:\n",
        "            alert('ngrok 연결에 실패했습니다, 토큰을 확인해주세요!', True)\n",
        "        else:\n",
        "            assert isinstance(tunnel, ngrok.NgrokTunnel)\n",
        "            TUNNEL_URL = tunnel.public_url\n",
        "\n",
        "    else:\n",
        "        raise ValueError(f'{tunnel} 에 대응하는 터널 서비스가 존재하지 않습니다')\n",
        "\n",
        "\n",
        "def setup_environment():\n",
        "    # 노트북 환경이라면 로그 표시를 위한 HTML 요소 만들기\n",
        "    if IN_INTERACTIVE:\n",
        "        try:\n",
        "            from IPython.display import display\n",
        "            from ipywidgets import widgets\n",
        "\n",
        "            global LOG_WIDGET\n",
        "            LOG_WIDGET = widgets.HTML()\n",
        "            display(LOG_WIDGET)\n",
        "\n",
        "        except ImportError:\n",
        "            pass\n",
        "\n",
        "    # 구글 드라이브 마운트하기\n",
        "    if IN_COLAB and OPTIONS['USE_GOOGLE_DRIVE']:\n",
        "        from google.colab import drive\n",
        "        drive.mount('/content/drive')\n",
        "\n",
        "        global WORKSPACE\n",
        "        WORKSPACE = str(\n",
        "            Path('drive', 'MyDrive', WORKSPACE).resolve()\n",
        "        )\n",
        "\n",
        "    # 로그 파일 만들기\n",
        "    global LOG_FILE\n",
        "    workspace = Path(WORKSPACE).resolve()\n",
        "    log_path = workspace.joinpath(\n",
        "        'logs',\n",
        "        datetime.strftime(datetime.now(), '%Y-%m-%d_%H-%M-%S.log')\n",
        "    )\n",
        "\n",
        "    log_path.parent.mkdir(0o777, True, True)\n",
        "\n",
        "    LOG_FILE = log_path.open('a')\n",
        "\n",
        "    # 현재 환경 출력\n",
        "    import platform\n",
        "    log(' '.join(os.uname()))\n",
        "    log(f'Python {platform.python_version()}')\n",
        "    log(str(Path().resolve()))\n",
        "\n",
        "    # 덮어쓸 설정 파일 가져오기\n",
        "    override_path = workspace.joinpath('override.json')\n",
        "    if override_path.exists():\n",
        "        with override_path.open('r') as file:\n",
        "            override_options = json.loads(file.read())\n",
        "            for key, value in override_options.items():\n",
        "                if key not in OPTIONS:\n",
        "                    log(f'{key} 키는 존재하지 않는 설정입니다', styles={'color': 'red'})\n",
        "                    continue\n",
        "\n",
        "                if type(value) != type(OPTIONS[key]):\n",
        "                    log(f'{key} 키는 {type(OPTIONS[key]).__name__} 자료형이여만 합니다', styles={\n",
        "                        'color': 'red'})\n",
        "                    continue\n",
        "\n",
        "                OPTIONS[key] = value\n",
        "\n",
        "                log(f'override.json: {key} = {json.dumps(value)}')\n",
        "\n",
        "    if IN_COLAB:\n",
        "        # 다른 Python 버전 설치\n",
        "        if OPTIONS['PYTHON_EXECUTABLE'] and not find_executable(OPTIONS['PYTHON_EXECUTABLE']):\n",
        "            execute(['apt', 'install', OPTIONS['PYTHON_EXECUTABLE']])\n",
        "            execute(\n",
        "                f\"curl -sS https://bootstrap.pypa.io/get-pip.py | {OPTIONS['PYTHON_EXECUTABLE']}\"\n",
        "            )\n",
        "\n",
        "        # 런타임이 정상적으로 초기화 됐는지 확인하기\n",
        "        try:\n",
        "            import torch\n",
        "        except:\n",
        "            alert('torch 패키지가 잘못됐습니다, 런타임을 다시 실행해주세요!', True)\n",
        "        else:\n",
        "            if not torch.cuda.is_available():\n",
        "                alert('GPU 런타임이 아닙니다, 할당량이 초과 됐을 수도 있습니다!')\n",
        "\n",
        "                OPTIONS['EXTRA_ARGS'] += [\n",
        "                    '--skip-torch-cuda-test',\n",
        "                    '--no-half',\n",
        "                    '--opt-sub-quad-attention'\n",
        "                ]\n",
        "\n",
        "        # 코랩 tcmalloc 관련 이슈 우회\n",
        "        # https://github.com/googlecolab/colabtools/issues/3412\n",
        "        try:\n",
        "            # 패키지가 이미 다운그레이드 됐는지 확인하기\n",
        "            execute('dpkg -l libunwind8-dev', hide_summary=True)\n",
        "        except subprocess.CalledProcessError:\n",
        "            for url in (\n",
        "                'http://launchpadlibrarian.net/367274644/libgoogle-perftools-dev_2.5-2.2ubuntu3_amd64.deb',\n",
        "                'https://launchpad.net/ubuntu/+source/google-perftools/2.5-2.2ubuntu3/+build/14795286/+files/google-perftools_2.5-2.2ubuntu3_all.deb',\n",
        "                'https://launchpad.net/ubuntu/+source/google-perftools/2.5-2.2ubuntu3/+build/14795286/+files/libtcmalloc-minimal4_2.5-2.2ubuntu3_amd64.deb',\n",
        "                'https://launchpad.net/ubuntu/+source/google-perftools/2.5-2.2ubuntu3/+build/14795286/+files/libgoogle-perftools4_2.5-2.2ubuntu3_amd64.deb'\n",
        "            ):\n",
        "                download(url, ignore_aria2=True)\n",
        "            execute('apt install -qq libunwind8-dev')\n",
        "            execute('dpkg -i *.deb')\n",
        "            execute('rm *.deb')\n",
        "\n",
        "    # 외부 터널링 초기화\n",
        "    setup_tunnels()\n",
        "\n",
        "    # 체크포인트 모델이 존재하지 않는다면 기본 모델 받아오기\n",
        "    if not has_checkpoint():\n",
        "        for file in [\n",
        "            {\n",
        "                'url': 'https://huggingface.co/gsdf/Counterfeit-V2.5/resolve/main/Counterfeit-V2.5_fp16.safetensors',\n",
        "                'target': str(workspace.joinpath('models/Stable-diffusion/Counterfeit-V2.5_fp16.safetensors')),\n",
        "                'summary': '기본 체크포인트 파일을 받아옵니다'\n",
        "            },\n",
        "            {\n",
        "                'url': 'https://huggingface.co/saltacc/wd-1-4-anime/resolve/main/VAE/kl-f8-anime2.ckpt',\n",
        "                'target': str(workspace.joinpath(WORKSPACE, 'models/VAE/kl-f8-anime2.ckpt')),\n",
        "                'summary': '기본 VAE 파일을 받아옵니다'\n",
        "            }\n",
        "        ]:\n",
        "            download(**file)\n",
        "\n",
        "\n",
        "# ==============================\n",
        "# 로그\n",
        "# ==============================\n",
        "\n",
        "\n",
        "def format_styles(styles: dict) -> str:\n",
        "    return ';'.join(map(lambda kv: ':'.join(kv), styles.items()))\n",
        "\n",
        "\n",
        "def render_log() -> None:\n",
        "    try:\n",
        "        from ipywidgets import widgets\n",
        "    except ImportError:\n",
        "        return\n",
        "\n",
        "    if not isinstance(LOG_WIDGET, widgets.HTML):\n",
        "        return\n",
        "\n",
        "    html = f'''<div style=\"{format_styles(LOG_WIDGET_STYLES['wrapper'])}\">'''\n",
        "\n",
        "    for block in LOG_BLOCKS:\n",
        "        styles = {\n",
        "            'display': 'inline-block',\n",
        "            **block['styles']\n",
        "        }\n",
        "        child_styles = {\n",
        "            'display': 'inline-block',\n",
        "            **block['child_styles']\n",
        "        }\n",
        "\n",
        "        html += f'<span style=\"{format_styles(styles)}\">{block[\"msg\"]}</span>\\n'\n",
        "\n",
        "        if block['max_childs'] is not None and len(block['childs']) > 0:\n",
        "            html += f'<div style=\"{format_styles(child_styles)}\">'\n",
        "            html += ''.join(block['childs'][-block['max_childs']:])\n",
        "            html += '</div>'\n",
        "\n",
        "    html += '</div>'\n",
        "\n",
        "    LOG_WIDGET.value = html\n",
        "\n",
        "\n",
        "def log(\n",
        "    msg: str,\n",
        "    styles={},\n",
        "    newline=True,\n",
        "\n",
        "    parent=False,\n",
        "    parent_index: Optional[int] = None,\n",
        "    child_styles={\n",
        "        'padding-left': '1em',\n",
        "        'color': 'gray'\n",
        "    },\n",
        "    max_childs=0,\n",
        "\n",
        "    print_to_file=True,\n",
        "    print_to_widget=True\n",
        ") -> Optional[int]:\n",
        "    # 기록할 내용이 ngrok API 키와 일치한다면 숨기기\n",
        "    # TODO: 더 나은 문자열 검사, 원치 않은 내용이 가려질 수도 있음\n",
        "    if OPTIONS['NGROK_API_TOKEN'] != '':\n",
        "        msg = msg.replace(OPTIONS['NGROK_API_TOKEN'], '**REDACTED**')\n",
        "\n",
        "    if newline:\n",
        "        msg += '\\n'\n",
        "\n",
        "    # 파일에 기록하기\n",
        "    if print_to_file and LOG_FILE:\n",
        "        if parent_index and msg.endswith('\\n'):\n",
        "            LOG_FILE.write('\\t')\n",
        "        LOG_FILE.write(msg)\n",
        "        LOG_FILE.flush()\n",
        "\n",
        "    # 로그 위젯에 기록하기\n",
        "    if print_to_widget and LOG_WIDGET:\n",
        "        # 부모 로그가 없다면 새 블록 만들기\n",
        "        if parent or parent_index is None:\n",
        "            LOG_BLOCKS.append({\n",
        "                'msg': msg,\n",
        "                'styles': styles,\n",
        "                'childs': [],\n",
        "                'child_styles': child_styles,\n",
        "                'max_childs': max_childs\n",
        "            })\n",
        "            render_log()\n",
        "            return len(LOG_BLOCKS) - 1\n",
        "\n",
        "        # 부모 로그가 존재한다면 추가하기\n",
        "        if len(LOG_BLOCKS[parent_index]['childs']) > 100:\n",
        "            LOG_BLOCKS[parent_index]['childs'].pop(0)\n",
        "\n",
        "        LOG_BLOCKS[parent_index]['childs'].append(msg)\n",
        "        render_log()\n",
        "\n",
        "    print('\\t' if parent_index else '' + msg, end='')\n",
        "\n",
        "\n",
        "def log_trace() -> None:\n",
        "    import sys\n",
        "    import traceback\n",
        "\n",
        "    # 스택 가져오기\n",
        "    ex_type, ex_value, ex_traceback = sys.exc_info()\n",
        "\n",
        "    styles = {}\n",
        "\n",
        "    # 오류가 존재한다면 메세지 빨간색으로 출력하기\n",
        "    # https://docs.python.org/3/library/sys.html#sys.exc_info\n",
        "    # TODO: 오류 유무 이렇게 확인하면 안될거 같은데 일단 귀찮아서 대충 써둠\n",
        "    if ex_type is not None:\n",
        "        styles = LOG_WIDGET_STYLES['dialog_error']\n",
        "\n",
        "    parent_index = log(\n",
        "        '오류가 발생했습니다, <a href=\"https://discord.gg/6wQeA2QXgM\">디스코드 서버</a>에 보고해주세요',\n",
        "        styles)\n",
        "    assert parent_index\n",
        "\n",
        "    # 오류가 존재한다면 오류 정보와 스택 트레이스 출력하기\n",
        "    if ex_type is not None:\n",
        "        log(f'{ex_type.__name__}: {ex_value}', parent_index=parent_index)\n",
        "        log(\n",
        "            '\\n'.join(traceback.format_tb(ex_traceback)),\n",
        "            parent_index=parent_index\n",
        "        )\n",
        "\n",
        "    # 로그 파일이 없으면 보고하지 않기\n",
        "    # TODO: 로그 파일이 존재하지 않을 수가 있나...?\n",
        "    if not LOG_FILE:\n",
        "        log('로그 파일이 존재하지 않습니다, 보고서를 만들지 않습니다')\n",
        "        return\n",
        "\n",
        "\n",
        "def alert(message: str, unassign=False):\n",
        "    log(message)\n",
        "\n",
        "    if IN_INTERACTIVE:\n",
        "        from IPython.display import display\n",
        "        from ipywidgets import widgets\n",
        "\n",
        "        display(\n",
        "            widgets.HTML(f'<script>alert({json.dumps(message)})</script>')\n",
        "        )\n",
        "\n",
        "    if IN_COLAB and unassign:\n",
        "        from google.colab import runtime\n",
        "\n",
        "        time.sleep(1)\n",
        "        runtime.unassign()\n",
        "\n",
        "\n",
        "# ==============================\n",
        "# 서브 프로세스\n",
        "# ==============================\n",
        "def execute(\n",
        "    args: Union[str, List[str]],\n",
        "    parser: Optional[\n",
        "        Callable[[str], None]\n",
        "    ] = None,\n",
        "    summary: Optional[str] = None,\n",
        "    hide_summary=False,\n",
        "    print_to_file=True,\n",
        "    print_to_widget=True,\n",
        "    **kwargs\n",
        ") -> Tuple[str, int]:\n",
        "    if isinstance(args, str) and 'shell' not in kwargs:\n",
        "        kwargs['shell'] = True\n",
        "\n",
        "    # 서브 프로세스 만들기\n",
        "    p = subprocess.Popen(\n",
        "        args,\n",
        "        stdout=subprocess.PIPE,\n",
        "        stderr=subprocess.STDOUT,\n",
        "        encoding='utf-8',\n",
        "        **kwargs)\n",
        "\n",
        "    # 로그에 시작한 프로세스 정보 출력하기\n",
        "    formatted_args = args if isinstance(args, str) else ' '.join(args)\n",
        "    summary = formatted_args if summary is None else f'{summary}\\n   {formatted_args}'\n",
        "\n",
        "    log_index = log(\n",
        "        f'=> {summary}',\n",
        "        styles={'color': 'yellow'},\n",
        "        max_childs=10)\n",
        "\n",
        "    output = ''\n",
        "\n",
        "    # 프로세스 출력 위젯에 리다이렉션하기\n",
        "    while p.poll() is None:\n",
        "        # 출력이 비어있다면 넘어가기\n",
        "        assert p.stdout\n",
        "        line = p.stdout.readline()\n",
        "        if not line:\n",
        "            continue\n",
        "\n",
        "        # 프로세스 출력 버퍼에 추가하기\n",
        "        output += line\n",
        "\n",
        "        # 파서 함수 실행하기\n",
        "        if callable(parser):\n",
        "            parser(line)\n",
        "\n",
        "        # 프로세스 출력 로그하기\n",
        "        log(\n",
        "            line,\n",
        "            newline=False,\n",
        "            parent_index=log_index,\n",
        "            print_to_file=print_to_file,\n",
        "            print_to_widget=print_to_widget)\n",
        "\n",
        "    # 변수 정리하기\n",
        "    rc = p.poll()\n",
        "    assert rc is not None\n",
        "\n",
        "    # 로그 블록 업데이트\n",
        "    if LOG_WIDGET:\n",
        "        assert log_index\n",
        "\n",
        "        if rc == 0:\n",
        "            # 현재 로그 텍스트 초록색으로 변경하고 프로세스 출력 숨기기\n",
        "            LOG_BLOCKS[log_index]['styles']['color'] = 'green'\n",
        "            LOG_BLOCKS[log_index]['max_childs'] = None\n",
        "        else:\n",
        "            # 현재 로그 텍스트 빨간색으로 변경하고 프로세스 출력 모두 표시하기\n",
        "            LOG_BLOCKS[log_index]['styles']['color'] = 'red'\n",
        "            LOG_BLOCKS[log_index]['max_childs'] = 0\n",
        "\n",
        "        if hide_summary:\n",
        "            # 현재 로그 블록 숨기기 (제거하기)\n",
        "            del LOG_BLOCKS[log_index]\n",
        "\n",
        "        # 로그 블록 렌더링\n",
        "        render_log()\n",
        "\n",
        "    # 오류 코드를 반환했다면\n",
        "    if rc != 0:\n",
        "        if isinstance(rc, signal.Signals):\n",
        "            rc = rc.value\n",
        "\n",
        "        raise subprocess.CalledProcessError(rc, args)\n",
        "\n",
        "    return output, rc\n",
        "\n",
        "# ==============================\n",
        "# 작업 경로\n",
        "# ==============================\n",
        "\n",
        "\n",
        "def delete(path: os.PathLike) -> None:\n",
        "    path = Path(path)\n",
        "\n",
        "    if path.is_file() or path.is_symlink():\n",
        "        path.unlink()\n",
        "    else:\n",
        "        shutil.rmtree(path, ignore_errors=True)\n",
        "\n",
        "\n",
        "def has_python_package(pkg: str, executable: Optional[str] = None) -> bool:\n",
        "    if not executable:\n",
        "        return find_spec(pkg) is not None\n",
        "\n",
        "    _, rc = execute(\n",
        "        [\n",
        "            executable, '-c',\n",
        "            f'''\n",
        "            import importlib\n",
        "            import sys\n",
        "            sys.exit(0 if importlib.find_loader({shlex.quote(pkg)}) else 0)\n",
        "            '''\n",
        "        ])\n",
        "\n",
        "    return True if rc == 0 else False\n",
        "\n",
        "\n",
        "# ==============================\n",
        "# 파일 다운로드\n",
        "# ==============================\n",
        "def download(url: str, target: Optional[str] = None, ignore_aria2=False, **kwargs):\n",
        "    if not target:\n",
        "        # TODO: 경로 중 params 제거하기\n",
        "        target = url.split('/')[-1]\n",
        "\n",
        "    # 파일을 받을 디렉터리 만들기\n",
        "    Path(target).parent.mkdir(0o777, True, True)\n",
        "\n",
        "    # 빠른 다운로드를 위해 aria2 패키지 설치 시도하기\n",
        "    if not ignore_aria2:\n",
        "        if not find_executable('aria2c') and find_executable('apt'):\n",
        "            execute(['apt', 'install', 'aria2'])\n",
        "\n",
        "        if find_executable('aria2c'):\n",
        "            p = Path(target)\n",
        "            execute(\n",
        "                [\n",
        "                    'aria2c',\n",
        "                    '--continue',\n",
        "                    '--always-resume',\n",
        "                    '--summary-interval', '10',\n",
        "                    '--disk-cache', '64M',\n",
        "                    '--min-split-size', '8M',\n",
        "                    '--max-concurrent-downloads', '8',\n",
        "                    '--max-connection-per-server', '8',\n",
        "                    '--max-overall-download-limit', '0',\n",
        "                    '--max-download-limit', '0',\n",
        "                    '--split', '8',\n",
        "                    '--dir', str(p.parent),\n",
        "                    '--out', p.name,\n",
        "                    url\n",
        "                ],\n",
        "                **kwargs)\n",
        "\n",
        "    elif find_executable('curl'):\n",
        "        execute(\n",
        "            [\n",
        "                'curl',\n",
        "                '--location',\n",
        "                '--output', target,\n",
        "                url\n",
        "            ],\n",
        "            **kwargs)\n",
        "\n",
        "    else:\n",
        "        if 'summary' in kwargs.keys():\n",
        "            log(kwargs.pop('summary'), **kwargs)\n",
        "\n",
        "        with requests.get(url, stream=True) as res:\n",
        "            res.raise_for_status()\n",
        "\n",
        "            with open(target, 'wb') as file:\n",
        "                # 받아온 파일 디코딩하기\n",
        "                # https://github.com/psf/requests/issues/2155#issuecomment-50771010\n",
        "                import functools\n",
        "                res.raw.read = functools.partial(\n",
        "                    res.raw.read,\n",
        "                    decode_content=True)\n",
        "\n",
        "                # TODO: 파일 길이가 적합한지?\n",
        "                shutil.copyfileobj(res.raw, file, length=16*1024*1024)\n",
        "\n",
        "\n",
        "def has_checkpoint() -> bool:\n",
        "    workspace = Path(WORKSPACE)\n",
        "    for p in workspace.joinpath('models', 'Stable-diffusion').glob('**/*'):\n",
        "        if p.suffix != '.ckpt' and p.suffix != '.safetensors':\n",
        "            continue\n",
        "\n",
        "        # aria2 로 받다만 파일이면 무시하기\n",
        "        if p.with_suffix(p.suffix + '.aria2c').exists():\n",
        "            continue\n",
        "\n",
        "        return True\n",
        "    return False\n",
        "\n",
        "\n",
        "def parse_output(line: str) -> None:\n",
        "    # 첫 시작에 한해서 웹 서버 열렸을 때 다이어로그 표시하기\n",
        "    if line.startswith('Running on local URL:'):\n",
        "        log(\n",
        "            '\\n'.join([\n",
        "                '성공적으로 터널이 열렸습니다',\n",
        "                f'<a target=\"_blank\" href=\"{TUNNEL_URL}\">{TUNNEL_URL}</a>',\n",
        "            ]),\n",
        "            LOG_WIDGET_STYLES['dialog_success']\n",
        "        )\n",
        "        return\n",
        "\n",
        "\n",
        "def setup() -> None:\n",
        "    repo_dir = Path('repository')\n",
        "    need_clone = True\n",
        "\n",
        "    # 이미 디렉터리가 존재한다면 정상적인 레포인지 확인하기\n",
        "    if repo_dir.is_dir():\n",
        "        try:\n",
        "            # 사용자 파일만 남겨두고 레포지토리 초기화하기\n",
        "            # https://stackoverflow.com/a/12096327\n",
        "            execute(\n",
        "                'git stash && git pull',\n",
        "                cwd=repo_dir\n",
        "            )\n",
        "        except subprocess.CalledProcessError:\n",
        "            log('레포지토리가 잘못됐습니다, 디렉터리를 제거합니다')\n",
        "        else:\n",
        "            need_clone = False\n",
        "\n",
        "    # 레포지토리 클론이 필요하다면 기존 디렉터리 지우고 클론하기\n",
        "    if need_clone:\n",
        "        shutil.rmtree(repo_dir, ignore_errors=True)\n",
        "        execute(['git', 'clone', OPTIONS['REPO_URL'], str(repo_dir)])\n",
        "\n",
        "    # 특정 커밋이 지정됐다면 체크아웃하기\n",
        "    if OPTIONS['REPO_COMMIT']:\n",
        "        execute(\n",
        "            ['git', 'checkout', OPTIONS['REPO_COMMIT']],\n",
        "            cwd=repo_dir\n",
        "        )\n",
        "\n",
        "    if IN_COLAB:\n",
        "        patch_path = repo_dir.joinpath('scripts', 'patches.py')\n",
        "\n",
        "        if not patch_path.exists():\n",
        "            download(\n",
        "                'https://raw.githubusercontent.com/toriato/easy-stable-diffusion/main/scripts/patches.py',\n",
        "                str(patch_path),\n",
        "                ignore_aria2=True)\n",
        "\n",
        "\n",
        "def start(args: List[str] = OPTIONS['ARGS']) -> None:\n",
        "    workspace = Path(WORKSPACE).resolve()\n",
        "    repository = Path('repository').resolve()\n",
        "\n",
        "    # 기본 인자 만들기\n",
        "    if len(args) < 1:\n",
        "        args += ['--data-dir', str(workspace)]\n",
        "\n",
        "        # xformers\n",
        "        if OPTIONS['USE_XFORMERS']:\n",
        "            try:\n",
        "                import torch\n",
        "            except ImportError:\n",
        "                pass\n",
        "            else:\n",
        "                if torch.cuda.is_available():\n",
        "                    args += [\n",
        "                        '--xformers',\n",
        "                        '--xformers-flash-attention'\n",
        "                    ]\n",
        "\n",
        "        # Gradio 인증 정보\n",
        "        if OPTIONS['GRADIO_USERNAME'] != '':\n",
        "            args += [\n",
        "                f'--gradio-auth',\n",
        "                OPTIONS['GRADIO_USERNAME'] +\n",
        "                ('' if OPTIONS['GRADIO_PASSWORD'] ==\n",
        "                    '' else ':' + OPTIONS['GRADIO_PASSWORD'])\n",
        "            ]\n",
        "\n",
        "    # 추가 인자\n",
        "    args += OPTIONS['EXTRA_ARGS']\n",
        "\n",
        "    env = {\n",
        "        **os.environ,\n",
        "        # 'HF_HOME': str(workspace / 'cache' / 'huggingface'),\n",
        "    }\n",
        "\n",
        "    # https://github.com/googlecolab/colabtools/issues/3412\n",
        "    if IN_COLAB:\n",
        "        env['LD_PRELOAD'] = 'libtcmalloc.so'\n",
        "\n",
        "    try:\n",
        "        execute(\n",
        "            [\n",
        "                OPTIONS['PYTHON_EXECUTABLE'] or 'python',\n",
        "                '-u',\n",
        "                '-m', 'launch',\n",
        "                *args\n",
        "            ],\n",
        "            parser=parse_output,\n",
        "            cwd=str(repository),\n",
        "            env=env,\n",
        "            start_new_session=True,\n",
        "        )\n",
        "    except subprocess.CalledProcessError as e:\n",
        "        if IN_COLAB and e.returncode == signal.SIGINT.value:\n",
        "            raise RuntimeError(\n",
        "                '프로세스가 강제 종료됐습니다, 메모리가 부족해 발생한 문제일 수도 있습니다') from e\n",
        "\n",
        "\n",
        "try:\n",
        "    setup_environment()\n",
        "\n",
        "    # 3단 이상(?) 레벨에서 실행하면 nested 된 asyncio 이 문제를 일으킴\n",
        "    # 런타임을 종료해도 코랩 페이지에선 런타임이 실행 중(Busy)인 것으로 표시되므로 여기서 실행함\n",
        "    if OPTIONS['DISCONNECT_RUNTIME']:\n",
        "        hook_runtime_disconnect()\n",
        "\n",
        "    setup()\n",
        "    start()\n",
        "\n",
        "# ^c 종료 무시하기\n",
        "except KeyboardInterrupt:\n",
        "    pass\n",
        "\n",
        "except:\n",
        "    # 로그 위젯이 없다면 평범하게 오류 처리하기\n",
        "    if not LOG_WIDGET:\n",
        "        raise\n",
        "\n",
        "    log_trace()\n"
      ],
      "metadata": {
        "id": "UGSqtUJPJoOj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<br><br><br>\n",
        "<br><br><br>\n",
        "<!-- HOLY COW IS THIS A SECRET WORDS? CONGRATS. -->\n",
        "<br><br><br>\n",
        "<br><br><br>\n",
        "## 🔨 /dev/stuffs"
      ],
      "metadata": {
        "id": "68jj0NEN7w2b"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 좀비 프로세스 죽이기"
      ],
      "metadata": {
        "id": "YNylYhg2KzJl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title\n",
        "!kill -9 $(pgrep -f '\\-m launch')\n",
        "!free -h"
      ],
      "metadata": {
        "cellView": "form",
        "id": "VXz1XdVh1wtC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 모델 다운로드\n"
      ],
      "metadata": {
        "id": "cfFYv1jw-L2b"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "!apt -qq install -y aria2"
      ],
      "metadata": {
        "id": "tHBZa2MA-4sz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!aria2c --out wd15-beta1-fp16.safetensors https://huggingface.co/waifu-diffusion/wd-1-5-beta/resolve/main/checkpoints/wd15-beta1-fp16.safetensors\n",
        "!aria2c --out wd15-beta1-fp16.yaml https://huggingface.co/waifu-diffusion/wd-1-5-beta/resolve/main/checkpoints/wd15-beta1-fp16.yaml"
      ],
      "metadata": {
        "id": "EdT0sp8R-Qu8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!mv wd15-beta1-fp16* /content/drive/MyDrive/SD/models/Stable-diffusion"
      ],
      "metadata": {
        "id": "bXKRvS9--3Xp"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [
        "YNylYhg2KzJl",
        "YLYoxXP8K4jb",
        "cfFYv1jw-L2b"
      ],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard",
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}