#!/bin/bash
# usage: run_yue.sh [segments] [seed]
SEG=${1:-1}; SEED=${2:-42}
cd /marimo/ai_studio/vendor/YuE/inference || exit 1
export PYTHONPATH=/marimo/ai_studio/vendor/YuE/inference
export PYTHONSAFEPATH=
export YUES_NO_COMPILE=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
exec /marimo/ai_studio/envs/music/bin/python infer.py \
  --stage1_model m-a-p/YuE-s1-7B-anneal-en-cot \
  --stage2_model m-a-p/YuE-s2-1B-general \
  --genre_txt /marimo/ai_studio/characters/lyra/genre_canto.txt \
  --lyrics_txt /marimo/ai_studio/characters/lyra/lyrics_canto.txt \
  --run_n_segments "$SEG" --stage2_batch_size 4 \
  --output_dir /marimo/ai_studio/outputs/song_canto_1 \
  --cuda_idx 0 --seed "$SEED"