waifu-diffusion-original/train.sh

16 lines
365 B
Bash

#!/bin/bash
ARGS=""
if [ ! -z "$NUM_GPU" ]; then
ARGS="--gpu="
for i in $(seq 0 $((NUM_GPU-1)))
do
ARGS="$ARGS$i,"
done
sed -i "s/batch_size: 4/batch_size: $NUM_GPU/g" ./configs/stable-diffusion/v1-finetune-4gpu.yaml
sed -i "s/num_workers: 4/num_workers: $NUM_GPU/g" ./configs/stable-diffusion/v1-finetune-4gpu.yaml
fi
python3 main.py $ARGS "$@"