394 lines
12 KiB
Bash
394 lines
12 KiB
Bash
#!/bin/bash
|
|
|
|
#set -euxo pipefail
|
|
|
|
arch
|
|
|
|
# -------------------- common init --------------------
|
|
|
|
if [ -z "$ESD_VERSION" ]; then
|
|
echo "ESD_VERSION is not set"
|
|
exit 1
|
|
fi
|
|
|
|
if [ -z "$S3_BUCKET_NAME" ]; then
|
|
echo "S3_BUCKET_NAME is not set"
|
|
exit 1
|
|
fi
|
|
|
|
if [ -z "$SERVICE_TYPE" ]; then
|
|
echo "SERVICE_TYPE is not set"
|
|
exit 1
|
|
fi
|
|
|
|
export TAR_FILE="esd.tar"
|
|
export CACHE_ENDPOINT="endpoint-$ESD_VERSION-$ENDPOINT_NAME"
|
|
|
|
# Use verified cache version file for production: v1.5.0-fe21616
|
|
export CACHE_PUBLIC_SD="aws-gcr-solutions-$AWS_REGION/stable-diffusion-aws-extension-github-mainline/$ESD_VERSION/sd.tar"
|
|
|
|
# Use verified cache version file for production: v1.5.0-fe21616
|
|
export CACHE_PUBLIC_COMFY="aws-gcr-solutions-$AWS_REGION/stable-diffusion-aws-extension-github-mainline/$ESD_VERSION/comfy.tar"
|
|
|
|
export ENDPOINT_INSTANCE_ID=$(date +"%m%d%H%M%S")
|
|
|
|
cores=$(lscpu | grep "^Core(s) per socket:" | awk '{print $4}')
|
|
sockets=$(lscpu | grep "^Socket(s):" | awk '{print $2}')
|
|
export CUP_CORE_NUMS=$((cores * sockets))
|
|
|
|
echo "---------------------------------------------------------------------------------"
|
|
echo "whoami: $(whoami)"
|
|
echo "Current shell: $SHELL"
|
|
echo "Running in $(bash --version)"
|
|
echo "---------------------------------------------------------------------------------"
|
|
echo "CREATED_AT: $CREATED_AT"
|
|
created_time_seconds=$(date -d "$CREATED_AT" +%s)
|
|
current_time=$(date "+%Y-%m-%dT%H:%M:%S.%6N")
|
|
current_time_seconds=$(date -d "$current_time" +%s)
|
|
export INSTANCE_INIT_SECONDS=$(( current_time_seconds - created_time_seconds ))
|
|
echo "NOW_AT: $current_time"
|
|
echo "Init from Create: $INSTANCE_INIT_SECONDS seconds"
|
|
echo "---------------------------------------------------------------------------------"
|
|
printenv
|
|
echo "---------------------------------------------------------------------------------"
|
|
nvidia-smi
|
|
echo "---------------------------------------------------------------------------------"
|
|
df -h
|
|
echo "---------------------------------------------------------------------------------"
|
|
|
|
# -------------------- common functions --------------------
|
|
|
|
set_conda(){
|
|
echo "set conda environment..."
|
|
export LD_LIBRARY_PATH=/home/ubuntu/conda/lib:$LD_LIBRARY_PATH
|
|
}
|
|
|
|
download_conda(){
|
|
echo "---------------------------------------------------------------------------------"
|
|
mkdir -p /home/ubuntu/conda/lib/
|
|
wget -qO /home/ubuntu/conda/lib/libcufft.so.10 https://huggingface.co/elonniu/esd/resolve/main/libcufft.so.10
|
|
wget -qO /home/ubuntu/conda/lib/libcurand.so.10 https://huggingface.co/elonniu/esd/resolve/main/libcurand.so.10
|
|
set_conda
|
|
}
|
|
|
|
# -------------------- sd functions --------------------
|
|
|
|
sd_cache_endpoint() {
|
|
start_at=$(date +%s)
|
|
|
|
echo "collection big files..."
|
|
upload_files=$(mktemp)
|
|
big_files=$(find "/home/ubuntu/stable-diffusion-webui" -type f -size +2520k)
|
|
for file in $big_files; do
|
|
key=$(echo "$file" | cut -d'/' -f4-)
|
|
echo "sync $file s3://$S3_BUCKET_NAME/$CACHE_ENDPOINT/$key" >> "$upload_files"
|
|
done
|
|
|
|
echo "tar files..."
|
|
filelist=$(mktemp)
|
|
# shellcheck disable=SC2164
|
|
cd /home/ubuntu/stable-diffusion-webui
|
|
find "./" \( -type f -o -type l \) -size -2530k > "$filelist"
|
|
tar -cf $TAR_FILE -T "$filelist"
|
|
|
|
echo "sync $TAR_FILE s3://$S3_BUCKET_NAME/$CACHE_ENDPOINT/" >> "$upload_files"
|
|
echo "sync /home/ubuntu/conda/* s3://$S3_BUCKET_NAME/$CACHE_ENDPOINT/conda/" >> "$upload_files"
|
|
|
|
echo "upload files..."
|
|
s5cmd run "$upload_files"
|
|
|
|
end_at=$(date +%s)
|
|
export UPLOAD_ENDPOINT_CACHE_SECONDS=$((end_at-start_at))
|
|
echo "sync endpoint files: $UPLOAD_ENDPOINT_CACHE_SECONDS seconds"
|
|
}
|
|
|
|
sd_install_build(){
|
|
cd /home/ubuntu || exit 1
|
|
bash install_sd.sh
|
|
}
|
|
|
|
sd_launch_cmd(){
|
|
python launch.py --enable-insecure-extension-access --api --api-log --log-startup --listen --port 8080 --xformers --no-half-vae --no-download-sd-model --no-hashing --nowebui --skip-torch-cuda-test --skip-load-model-at-start --disable-safe-unpickle --skip-prepare-environment --skip-python-version-check --skip-install --skip-version-check --disable-nan-check
|
|
}
|
|
|
|
sd_launch(){
|
|
echo "---------------------------------------------------------------------------------"
|
|
echo "accelerate sd launch..."
|
|
|
|
set_conda
|
|
|
|
ls -la /home/ubuntu/
|
|
|
|
mkdir -p /root/.u2net/
|
|
mv /home/ubuntu/stable-diffusion-webui/u2net_human_seg.onnx /root/.u2net/
|
|
|
|
if [[ $AWS_REGION == *"cn-"* ]]; then
|
|
mkdir -p /root/.cache/huggingface/accelerate
|
|
mv /home/ubuntu/stable-diffusion-webui/default_config.yaml /root/.cache/huggingface/accelerate/
|
|
fi
|
|
|
|
echo "initiated_lock" > /initiated_lock
|
|
|
|
cd /home/ubuntu/stable-diffusion-webui || exit 1
|
|
source venv/bin/activate
|
|
python /metrics.py &
|
|
|
|
sd_launch_cmd
|
|
}
|
|
|
|
sd_launch_from_private_s3(){
|
|
start_at=$(date +%s)
|
|
s5cmd sync "s3://$S3_BUCKET_NAME/$CACHE_ENDPOINT/*" /home/ubuntu/
|
|
end_at=$(date +%s)
|
|
export DOWNLOAD_FILE_SECONDS=$((end_at-start_at))
|
|
echo "download file: $DOWNLOAD_FILE_SECONDS seconds"
|
|
|
|
export DOWNLOAD_FILE_SIZE=$(du -sm /home/ubuntu | awk '{print $1}' | grep -oE '[0-9]+')
|
|
|
|
start_at=$(date +%s)
|
|
tar --overwrite -xf "$TAR_FILE" -C /home/ubuntu/stable-diffusion-webui/
|
|
rm -rf $TAR_FILE
|
|
end_at=$(date +%s)
|
|
export DECOMPRESS_SECONDS=$((end_at-start_at))
|
|
echo "decompress file: $DECOMPRESS_SECONDS seconds"
|
|
|
|
cd /home/ubuntu/stable-diffusion-webui/ || exit 1
|
|
|
|
mkdir -p models/Lora || true
|
|
mkdir -p models/hypernetworks || true
|
|
mkdir -p models/ControlNet || true
|
|
|
|
sd_launch
|
|
}
|
|
|
|
sd_launch_from_public_s3(){
|
|
start_at=$(date +%s)
|
|
s5cmd cp "s3://$CACHE_PUBLIC_SD" /home/ubuntu/
|
|
end_at=$(date +%s)
|
|
export DOWNLOAD_FILE_SECONDS=$((end_at-start_at))
|
|
echo "download file: $DOWNLOAD_FILE_SECONDS seconds"
|
|
|
|
export DOWNLOAD_FILE_SIZE=$(du -sm /home/ubuntu | awk '{print $1}' | grep -oE '[0-9]+')
|
|
|
|
start_at=$(date +%s)
|
|
tar --overwrite -xf "$SERVICE_TYPE.tar" -C /home/ubuntu/
|
|
rm -rf "$SERVICE_TYPE.tar"
|
|
end_at=$(date +%s)
|
|
export DECOMPRESS_SECONDS=$((end_at-start_at))
|
|
echo "decompress file: $DECOMPRESS_SECONDS seconds"
|
|
|
|
cd /home/ubuntu/stable-diffusion-webui/ || exit 1
|
|
|
|
mkdir -p models/Lora
|
|
mkdir -p models/hypernetworks
|
|
|
|
# if $EXTENSIONS is not empty, it will be executed
|
|
if [ -n "$EXTENSIONS" ]; then
|
|
echo "---------------------------------------------------------------------------------"
|
|
echo "install extensions..."
|
|
|
|
cd extensions || exit 1
|
|
|
|
read -ra array <<< "$(echo "$EXTENSIONS" | tr "," " ")"
|
|
|
|
for git_repo in "${array[@]}"; do
|
|
IFS='#' read -r -a repo <<< "$git_repo"
|
|
|
|
git_repo=${repo[0]}
|
|
repo_name=$(basename -s .git "$git_repo")
|
|
repo_branch=${repo[1]}
|
|
commit_sha=${repo[2]}
|
|
|
|
echo "rm -rf $repo_name for install $git_repo"
|
|
rm -rf "$repo_name"
|
|
|
|
start_at=$(date +%s)
|
|
|
|
echo "git clone $git_repo"
|
|
git clone "$git_repo"
|
|
|
|
cd "$repo_name" || exit 1
|
|
|
|
echo "git checkout $repo_branch"
|
|
git checkout "$repo_branch"
|
|
|
|
echo "git reset --hard $commit_sha"
|
|
git reset --hard "$commit_sha"
|
|
cd ..
|
|
|
|
end_at=$(date +%s)
|
|
cost=$((end_at-start_at))
|
|
echo "git clone $git_repo: $cost seconds"
|
|
done
|
|
|
|
cd /home/ubuntu/stable-diffusion-webui/ || exit 1
|
|
echo "---------------------------------------------------------------------------------"
|
|
echo "build for launch..."
|
|
python launch.py --enable-insecure-extension-access --api --api-log --log-startup --listen --xformers --no-half-vae --no-download-sd-model --no-hashing --nowebui --skip-torch-cuda-test --skip-load-model-at-start --disable-safe-unpickle --disable-nan-check --exit
|
|
fi
|
|
|
|
cd /home/ubuntu/stable-diffusion-webui/ || exit 1
|
|
|
|
/serve trim_sd
|
|
sd_cache_endpoint
|
|
sd_launch
|
|
}
|
|
|
|
# -------------------- comfy functions --------------------
|
|
|
|
comfy_install_build(){
|
|
cd /home/ubuntu || exit 1
|
|
bash install_comfy.sh
|
|
}
|
|
|
|
comfy_cache_endpoint() {
|
|
start_at=$(date +%s)
|
|
|
|
upload_files=$(mktemp)
|
|
big_files=$(find "/home/ubuntu/ComfyUI" -type f -size +2520k)
|
|
for file in $big_files; do
|
|
key=$(echo "$file" | cut -d'/' -f4-)
|
|
echo "sync $file s3://$S3_BUCKET_NAME/$CACHE_ENDPOINT/$key" >> "$upload_files"
|
|
done
|
|
|
|
echo "tar files..."
|
|
filelist=$(mktemp)
|
|
# shellcheck disable=SC2164
|
|
cd /home/ubuntu/ComfyUI
|
|
find "./" \( -type f -o -type l \) -size -2530k > "$filelist"
|
|
tar -cf $TAR_FILE -T "$filelist"
|
|
|
|
echo "sync $TAR_FILE s3://$S3_BUCKET_NAME/$CACHE_ENDPOINT/" >> "$upload_files"
|
|
echo "sync /home/ubuntu/conda/* s3://$S3_BUCKET_NAME/$CACHE_ENDPOINT/conda/" >> "$upload_files"
|
|
|
|
s5cmd run "$upload_files"
|
|
|
|
end_at=$(date +%s)
|
|
export UPLOAD_ENDPOINT_CACHE_SECONDS=$((end_at-start_at))
|
|
echo "sync endpoint files: $UPLOAD_ENDPOINT_CACHE_SECONDS seconds"
|
|
}
|
|
|
|
comfy_launch(){
|
|
echo "---------------------------------------------------------------------------------"
|
|
echo "accelerate comfy launch..."
|
|
|
|
set_conda
|
|
|
|
cd /home/ubuntu/ComfyUI || exit 1
|
|
rm -rf /home/ubuntu/ComfyUI/custom_nodes/ComfyUI-AWS-Extension
|
|
rm /home/ubuntu/ComfyUI/custom_nodes/comfy_local_proxy.py
|
|
source venv/bin/activate
|
|
python /metrics.py &
|
|
|
|
echo "initiated_lock" > /initiated_lock
|
|
|
|
python serve.py
|
|
}
|
|
|
|
comfy_launch_from_private_s3(){
|
|
start_at=$(date +%s)
|
|
s5cmd sync "s3://$S3_BUCKET_NAME/$CACHE_ENDPOINT/*" /home/ubuntu/
|
|
end_at=$(date +%s)
|
|
export DOWNLOAD_FILE_SECONDS=$((end_at-start_at))
|
|
echo "download file: $DOWNLOAD_FILE_SECONDS seconds"
|
|
|
|
export DOWNLOAD_FILE_SIZE=$(du -sm /home/ubuntu | awk '{print $1}' | grep -oE '[0-9]+')
|
|
|
|
start_at=$(date +%s)
|
|
tar --overwrite -xf "$TAR_FILE" -C /home/ubuntu/ComfyUI/
|
|
rm -rf $TAR_FILE
|
|
end_at=$(date +%s)
|
|
export DECOMPRESS_SECONDS=$((end_at-start_at))
|
|
echo "decompress file: $DECOMPRESS_SECONDS seconds"
|
|
|
|
comfy_launch
|
|
}
|
|
|
|
comfy_launch_from_public_s3(){
|
|
start_at=$(date +%s)
|
|
s5cmd cp "s3://$CACHE_PUBLIC_COMFY" /home/ubuntu/
|
|
end_at=$(date +%s)
|
|
export DOWNLOAD_FILE_SECONDS=$((end_at-start_at))
|
|
echo "download file: $DOWNLOAD_FILE_SECONDS seconds"
|
|
|
|
export DOWNLOAD_FILE_SIZE=$(du -sm /home/ubuntu | awk '{print $1}' | grep -oE '[0-9]+')
|
|
|
|
start_at=$(date +%s)
|
|
tar --overwrite -xf "$SERVICE_TYPE.tar" -C /home/ubuntu/
|
|
rm -rf "$SERVICE_TYPE.tar"
|
|
end_at=$(date +%s)
|
|
export DECOMPRESS_SECONDS=$((end_at-start_at))
|
|
echo "decompress file: $DECOMPRESS_SECONDS seconds"
|
|
|
|
/serve trim_comfy
|
|
comfy_cache_endpoint
|
|
comfy_launch
|
|
}
|
|
|
|
# -------------------- startup --------------------
|
|
|
|
# if pipeline finished, it will be executed
|
|
#if [[ $IMAGE_URL == *"dev"* ]]; then
|
|
# download_conda
|
|
# if [ "$SERVICE_TYPE" == "sd" ]; then
|
|
# sd_install_build
|
|
# /serve trim_sd.sh
|
|
# sd_cache_endpoint
|
|
# sd_launch
|
|
# exit 1
|
|
# else
|
|
# comfy_install_build
|
|
# /serve trim_comfy
|
|
# comfy_cache_endpoint
|
|
# comfy_launch
|
|
# exit 1
|
|
# fi
|
|
#fi
|
|
|
|
if [ -n "$APP_SOURCE" ]; then
|
|
if [ -n "$APP_CWD" ]; then
|
|
start_at=$(date +%s)
|
|
s5cmd --log=error sync "s3://$S3_BUCKET_NAME/$APP_SOURCE/*" "$APP_CWD"
|
|
end_at=$(date +%s)
|
|
export DOWNLOAD_FILE_SECONDS=$((end_at-start_at))
|
|
echo "download file: $DOWNLOAD_FILE_SECONDS seconds"
|
|
|
|
bash "$APP_CWD/start.sh"
|
|
exit 1
|
|
fi
|
|
fi
|
|
|
|
if [ -f "/initiated_lock" ]; then
|
|
echo "already initiated, start service directly..."
|
|
if [ "$SERVICE_TYPE" == "sd" ]; then
|
|
cd /home/ubuntu/stable-diffusion-webui || exit 1
|
|
source venv/bin/activate
|
|
sd_launch_cmd
|
|
exit 1
|
|
else
|
|
cd /home/ubuntu/ComfyUI || exit 1
|
|
source venv/bin/activate
|
|
python serve.py
|
|
exit 1
|
|
fi
|
|
fi
|
|
|
|
output=$(s5cmd ls "s3://$S3_BUCKET_NAME/")
|
|
if echo "$output" | grep -q "$CACHE_ENDPOINT"; then
|
|
if [ "$SERVICE_TYPE" == "sd" ]; then
|
|
sd_launch_from_private_s3
|
|
exit 1
|
|
else
|
|
comfy_launch_from_private_s3
|
|
exit 1
|
|
fi
|
|
fi
|
|
|
|
if [ "$SERVICE_TYPE" == "sd" ]; then
|
|
sd_launch_from_public_s3
|
|
exit 1
|
|
else
|
|
comfy_launch_from_public_s3
|
|
exit 1
|
|
fi
|