SongGeneration/generate.sh at main · tencent-ailab/SongGeneration · GitHub

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
74
75
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
export OMP_NUM_THREADS=1
export MKL_NUM_THREADS=1
export CUDA_LAUNCH_BLOCKING=0

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