Faster data downloads

.gitignore

@@ -0,0 +1,5 @@
+
+.DS_Store
+/checkpoints
+/data
+/stable_diffusion/models/ldm/stable-diffusion-v1

scripts/download_checkpoints.sh

@@ -4,4 +4,4 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 
 mkdir -p $SCRIPT_DIR/../checkpoints
 
-curl http://instruct-pix2pix.eecs.berkeley.edu/instruct-pix2pix-00-22000.ckpt -o $SCRIPT_DIR/../checkpoints/instruct-pix2pix-00-22000.ckpt
+wget -t15 -c --progress=bar -w15 --retry-connrefused http://instruct-pix2pix.eecs.berkeley.edu/instruct-pix2pix-00-22000.ckpt -O $SCRIPT_DIR/../checkpoints/instruct-pix2pix-00-22000.ckpt

scripts/download_data.sh

@@ -6,8 +6,8 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 mkdir -p $SCRIPT_DIR/../data
 
 # Copy text datasets
-wget -q --show-progress http://instruct-pix2pix.eecs.berkeley.edu/gpt-generated-prompts.jsonl -O $SCRIPT_DIR/../data/gpt-generated-prompts.jsonl
-wget -q --show-progress http://instruct-pix2pix.eecs.berkeley.edu/human-written-prompts.jsonl -O $SCRIPT_DIR/../data/human-written-prompts.jsonl
+wget -nc -c -q --show-progress http://instruct-pix2pix.eecs.berkeley.edu/gpt-generated-prompts.jsonl -O $SCRIPT_DIR/../data/gpt-generated-prompts.jsonl
+wget -nc -c -q --show-progress http://instruct-pix2pix.eecs.berkeley.edu/human-written-prompts.jsonl -O $SCRIPT_DIR/../data/human-written-prompts.jsonl
 
 # If dataset name isn't provided, exit. 
 if [ -z $1 ] 
@@ -17,7 +17,13 @@ fi
 
 # Copy dataset files
 mkdir $SCRIPT_DIR/../data/$1
-wget -A zip,json -R "index.html*" -q --show-progress -r --no-parent http://instruct-pix2pix.eecs.berkeley.edu/$1/ -nd -P $SCRIPT_DIR/../data/$1/
+wget -nc -A zip,json -R "index.html*" -q --show-progress -r --no-parent -c http://instruct-pix2pix.eecs.berkeley.edu/$1/ -nd -P $SCRIPT_DIR/../data/$1/
+
+if [ -z $2 -a ( $2="master" )]; then
+	True
+else
+	exit 2
+fi
 
 # Unzip to folders
 unzip $SCRIPT_DIR/../data/$1/\*.zip -d $SCRIPT_DIR/../data/$1/

scripts/download_pretrained_sd.sh

@@ -1,7 +1,20 @@
 #!/bin/bash
 
 SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+DATA_DIR="${SCRIPT_DIR}/../stable_diffusion/models/ldm/stable-diffusion-v1"
 
-mkdir -p $SCRIPT_DIR/../stable_diffusion/models/ldm/stable-diffusion-v1
-curl -L https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.ckpt -o $SCRIPT_DIR/../stable_diffusion/models/ldm/stable-diffusion-v1/v1-5-pruned-emaonly.ckpt
-curl -L https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt -o $SCRIPT_DIR/../stable_diffusion/models/ldm/stable-diffusion-v1/vae-ft-mse-840000-ema-pruned.ckpt
+mkdir -p ${DATA_DIR}
+
+if ( [ $HF_HUB_ENABLE_HF_TRANSFER ] ); then
+	echo "got into if"
+	# assumes that huggingface-cli and hf-transfer are installed and HF_HUB_ENABLE_HF_TRANSFER=1
+	# in environment globals. It also requires you have a HF_TOKEN set.
+	# NB although faster, hf-transfer doesn't support debugging or resuming partial downloads
+	# see https://huggingface.co/docs/huggingface_hub/en/guides/download#download-from-the-cli
+	huggingface-cli download --local-dir $DATA_DIR "runwayml/stable-diffusion-v1-5" v1-5-pruned-emaonly.ckpt
+	huggingface-cli download --local-dir $DATA_DIR "stabilityai/sd-vae-ft-mse-original" vae-ft-mse-840000-ema-pruned.ckpt
+else
+	echo "got into else"
+	curl -L https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.ckpt -o $DATA_DIR/v1-5-pruned-emaonly.ckpt
+	curl -L https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt -o $DATA_DIR/vae-ft-mse-840000-ema-pruned.ckpt
+fi