Spaces:
Running
on
L40S
Running
on
L40S
File size: 2,103 Bytes
258fd02 d658154 98a0e3b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
export USER=root
export PYTHONDONTWRITEBYTECODE=1
export TRANSFORMERS_CACHE="$(pwd)/third_party/hub"
export NCCL_HOME=/usr/local/tccl
export PYTHONPATH="$(pwd)/codeclm/tokenizer/":"$(pwd)":"$(pwd)/codeclm/tokenizer/Flow1dVAE/":"$(pwd)/codeclm/tokenizer/":$PYTHONPATH
CKPT_PATH=$1
JSONL=$2
SAVE_DIR=$3
USE_FLASH_ATTN="True"
LOW_MEM="False"
GENERATE_TYPE="mixed"
for arg in "$@"; do
if [[ $arg == "--not_use_flash_attn" ]]; then
USE_FLASH_ATTN="False"
fi
done
for arg in "$@"; do
if [[ $arg == "--low_mem" ]]; then
LOW_MEM="True"
fi
done
for arg in "$@"; do
if [[ $arg == "--separate" ]]; then
GENERATE_TYPE="separate"
fi
done
for arg in "$@"; do
if [[ $arg == "--bgm" ]]; then
GENERATE_TYPE="bgm"
fi
done
for arg in "$@"; do
if [[ $arg == "--vocal" ]]; then
GENERATE_TYPE="vocal"
fi
done
if [ "$USE_FLASH_ATTN" == "True" ] && [ "$LOW_MEM" == "True" ]; then
echo "Use Flash Attention + Low Memory Mode"
python3 generate.py \
--ckpt_path $CKPT_PATH \
--input_jsonl $JSONL \
--save_dir $SAVE_DIR \
--generate_type $GENERATE_TYPE \
--use_flash_attn \
--low_mem
elif [ "$USE_FLASH_ATTN" == "True" ] && [ "$LOW_MEM" == "False" ]; then
echo "Use Flash Attention + Auto Memory Mode"
python3 generate.py \
--ckpt_path $CKPT_PATH \
--input_jsonl $JSONL \
--save_dir $SAVE_DIR \
--generate_type $GENERATE_TYPE \
--use_flash_attn
elif [ "$USE_FLASH_ATTN" == "False" ] && [ "$LOW_MEM" == "False" ]; then
echo "Not Use Flash Attention + Auto Memory Mode"
python3 generate.py \
--ckpt_path $CKPT_PATH \
--input_jsonl $JSONL \
--generate_type $GENERATE_TYPE \
--save_dir $SAVE_DIR
elif [ "$USE_FLASH_ATTN" == "False" ] && [ "$LOW_MEM" == "True" ]; then
echo "Not Use Flash Attention + Low Memory Mode"
python3 generate.py \
--ckpt_path $CKPT_PATH \
--input_jsonl $JSONL \
--save_dir $SAVE_DIR \
--generate_type $GENERATE_TYPE \
--low_mem
fi
|