#!/bin/bash -l

# SLURM SUBMIT SCRIPT
#SBATCH --job-name=boost
#SBATCH --nodes=1             # This needs to match Trainer(num_nodes=...)
#SBATCH --gres=gpu:8
#SBATCH --ntasks-per-node=8   # This needs to match Trainer(devices=...)

# ACTIVATE ANACONDA
eval "$(conda shell.bash hook)"

# activate conda env
conda activate stable-audio-tools-env

# debugging flags (optional)
export NCCL_DEBUG=INFO
export PYTHONFAULTHANDLER=1

# on your cluster you might need these:
# set the network interface
# export NCCL_SOCKET_IFNAME=^docker0,lo

# might need the latest CUDA
#module load NCCL/2.4.7-1-cuda.10.0

# unwrap autoencoder
#python3 ./unwrap_model.py \
#--model-config stable_audio_tools/configs/model_configs/autoencoders/stable_audio_2_0_vae_48khz.json \
#--ckpt-path /home/christian/christian/stable-audio-tools/harmonai_train/usj037d6/checkpoints/epoch=61-step=1350000.ckpt \
#--name vae_model_unwrap-epoch=61-step=1350000

# unwrap upsample model (semantic)
#python3 ./unwrap_model.py \
#--model-config stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_semantic_48khz.json \
#--ckpt-path /home/christian/christian/stable-audio-tools/harmonai_train/mj07h6b3/checkpoints/epoch=0-step=10000.ckpt \
#--name upsample_model_semantic_unwrap-epoch=0-step=10000

# unwrap upsample model (codec)
#python3 ./unwrap_model.py \
#--model-config stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_codec_48khz.json \
#--ckpt-path /home/christian/christian/stable-audio-tools/harmonai_train/2zp4baao/checkpoints/epoch=45-step=370000.ckpt \
#--name upsample_model_codec_unwrap-epoch=45-step=370000


# unwrap upsample model (semantic+codec)
python3 ./unwrap_model.py \
--model-config stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_semantic+codec_48khz.json \
--ckpt-path /home/christian/christian/stable-audio-tools/harmonai_train/mj07h6b3/checkpoints/epoch=0-step=10000.ckpt \
--name upsample_model_semantic+codec_unwrap-epoch=0-step=10000