#!/bin/bash

set -e
set -o nounset
set -o pipefail

apt update
apt install -y less vim unzip

cd /workspace

git clone https://github.com/vladmandic/automatic
mkdir -p automatic/models/Stable-diffusion \
    automatic/models/LyCORIS \
    automatic/models/Lora \
    automatic/models/VAE \
    automatic/models/embeddings \
    automatic/extensions/

# Download models

pushd /workspace/automatic/models/Stable-diffusion/
curl -# -O -L -J "https://huggingface.co/bluepen5805/blue_pencil/resolve/main/blue_pencil-v10.safetensors"
curl -# -O -L -J "https://www.heavensinferno.net/test/defbeacf3arMix_defbeacf3ar2dMixV2.safetensors"
popd
pushd /workspace/automatic/models/VAE
curl -# -O -L -J "https://www.heavensinferno.net/test/clearvae_main.safetensors"
popd
pushd /workspace/automatic/models/LyCORIS
# Better hands
curl -# -O -L -J "https://civitai.com/api/download/models/55199"

popd
pushd /workspace/automatic/models/embeddings
curl -# -O -L -J "https://civitai.com/api/download/models/76712"
popd

pushd /workspace/automatic/models/Lora/
curl -# -O -L -J "https://huggingface.co/ashen-sensored/mzpikas_tmnd_enhanced/resolve/main/Silicon-landscape-isolation.safetensors"
curl -# -O -L -J "https://civitai.com/api/download/models/19075"
curl -# -O -L -J "https://civitai.com/api/download/models/62833"
popd

sed -i 's#/notebooks/#/workspace/#g' ~/*.json
cp ~/*.json /workspace/automatic/
cp ~/*.toml /workspace
cp ~/styles.csv /workspace/automatic/

# Extensions
pushd /workspace/automatic/extensions/
git clone https://github.com/hako-mikan/sd-webui-regional-prompter
git clone https://github.com/hnmr293/sd-webui-cutoff
git clone https://github.com/ctwrs/a1111-sd-webui-tagcomplete
git clone https://github.com/ashen-sensored/sd_webui_SAG
git clone https://github.com/Bing-su/adetailer
git clone https://github.com/adieyal/sd-dynamic-prompts
git clone https://github.com/huchenlei/sd-webui-openpose-editor
git clone https://github.com/hako-mikan/sd-webui-lora-block-weight

popd

cp ~/lbwpresets.txt /workspace/automatic/extensions/sd-webui-lora-block-weight/

# ControlNet

pushd /workspace/automatic/extensions-builtin/sd-webui-controlnet/annotator/midas/
!curl -\# -O -L -J -C - https://huggingface.co/lllyasviel/Annotators/resolve/main/dpt_hybrid-midas-501f0c75.pt
popd

pushd /workspace/automatic/extensions-builtin/sd-webui-controlnet/annotator/openpose/
!curl -\# -O -L -J -C - https://huggingface.co/lllyasviel/Annotators/resolve/main/body_pose_model.pth
!curl -\# -O -L -J -C - https://huggingface.co/lllyasviel/Annotators/resolve/main/hand_pose_model.pth
!curl -\# -O -L -J -C - https://huggingface.co/lllyasviel/Annotators/resolve/main/facenet.pth
popd

pushd /workspace/automatic/extensions-builtin/sd-webui-controlnet/annotator/uniformer/
!curl -\# -O -L -J -C - https://huggingface.co/lllyasviel/Annotators/resolve/main/upernet_global_small.pth
popd

pushd /workspace/automatic/extensions-builtin/sd-webui-controlnet/annotator/zoedepth/
!curl -\# -O -L -J -C - https://huggingface.co/lllyasviel/Annotators/resolve/main/ZoeD_M12_N.pt
popd

pushd /workspace/automatic/extensions-builtin/sd-webui-controlnet/annotator/oneformer

!curl -\# -O -L -J -C - https://huggingface.co/lllyasviel/Annotators/resolve/main/150_16_swin_l_oneformer_coco_100ep.pth
!curl -\# -O -L -J -C - https://huggingface.co/lllyasviel/Annotators/resolve/main/250_16_swin_l_oneformer_ade20k_160k.pth
!curl -\# -O -L -J -C - https://huggingface.co/lllyasviel/Annotators/resolve/main/res101.pth
!curl -\# -O -L -J -C - https://huggingface.co/lllyasviel/Annotators/resolve/main/latest_net_G.pth
!curl -\# -O -L -J -C - https://huggingface.co/lllyasviel/Annotators/resolve/main/ControlNetLama.pth


cd /workspace
curl -L -O -J https://github.com/rapiz1/rathole/releases/download/v0.4.8/rathole-x86_64-unknown-linux-gnu.zip
unzip rathole-x86_64-unknown-linux-gnu.zip

nohup /workspace/rathole -c client-local.toml &

cd /workspace/automatic
python launch.py