import os
import asyncio
import subprocess
from google.colab.output import eval_js

!pip install nest-asyncio
import nest_asyncio
nest_asyncio.apply()

async def async_system_command(cmd):
  '''异步执行系统命令'''
  print("run command: " + cmd)
  process = await asyncio.create_subprocess_shell(
      cmd,
      stdout=subprocess.PIPE,
      stderr=subprocess.PIPE
  )
  stdout, stderr = await process.communicate()
  outStr, errStr = stdout.decode(), stderr.decode()
  print(outStr, errStr)
  return outStr, errStr

async def async_command_group(urls):
  '''并发执行多个命令'''
  tasks = []
  for url in urls:
      task = asyncio.ensure_future(async_system_command(url))
      tasks.append(task)

  await asyncio.gather(*tasks)
  print("命令组执行完成...")

async def main():
  await async_command_group([
    "apt -y update -qq",
    "wget http://launchpadlibrarian.net/367274644/libgoogle-perftools-dev_2.5-2.2ubuntu3_amd64.deb",
    "wget https://launchpad.net/ubuntu/+source/google-perftools/2.5-2.2ubuntu3/+build/14795286/+files/google-perftools_2.5-2.2ubuntu3_all.deb",
    "wget https://launchpad.net/ubuntu/+source/google-perftools/2.5-2.2ubuntu3/+build/14795286/+files/libtcmalloc-minimal4_2.5-2.2ubuntu3_amd64.deb",
    "wget https://launchpad.net/ubuntu/+source/google-perftools/2.5-2.2ubuntu3/+build/14795286/+files/libgoogle-perftools4_2.5-2.2ubuntu3_amd64.deb",  
  ])
  !apt install -qq libunwind8-dev
  !dpkg -i *.deb
  %env LD_PRELOAD=libtcmalloc.so
  !rm *.deb
  !pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116 -U
  await async_command_group([
    "apt -y install -qq aria2",
    "pip install -q --pre xformers==0.0.17.dev476 -U",
    "pip install -q --pre triton",
  ])
  %cd /content
  !git clone -b v2.1 https://github.com/camenduru/stable-diffusion-webui
  await async_command_group([  
    "git clone https://github.com/camenduru/sd-webui-tunnels /content/stable-diffusion-webui/extensions/sd-webui-tunnels",
  ])
  %cd /content/stable-diffusion-webui
  !git reset --hard
  await async_command_group([
    "aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/datasets/gsdf/EasyNegative/resolve/main/EasyNegative.safetensors -d /content/stable-diffusion-webui/embeddings -o EasyNegative.safetensors",
    "aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/datasets/Nerfgun3/bad_prompt/resolve/main/bad_prompt_version2.pt -d /content/stable-diffusion-webui/embeddings -o bad_prompt_version2.pt",
    "aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/andite/anything-v4.0/resolve/main/anything-v4.5.ckpt -d /content/stable-diffusion-webui/models/Stable-diffusion -o anything-v4.5.ckpt",
    "aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/andite/anything-v4.0/resolve/main/anything-v4.0.vae.pt -d /content/stable-diffusion-webui/models/Stable-diffusion -o anything-v4.5.vae.pt",
])

os.environ['colab_url'] = eval_js("google.colab.kernel.proxyPort(7860, {'cache': false})")

asyncio.run(main())

!sed -i -e '''/    prepare_environment()/a\    os.system\(f\"""sed -i -e ''\"s/dict()))/dict())).cuda()/g\"'' /content/stable-diffusion-webui/repositories/stable-diffusion-stability-ai/ldm/util.py""")''' /content/stable-diffusion-webui/launch.py

!echo '{"sd_model_checkpoint": "anything-v4.5.ckpt"}' > config.json

!python launch.py --share --xformers --enable-insecure-extension-access --theme dark --gradio-queue --cloudflared --no-half-vae