!/usr/bin/env bash

DOCKER_EXP_DIR="/checkpoints/results/25FPS_inference"

### Best for nonstreaming Great for streaming
# CKPT="/checkpoints/streaming/magpie/jason/magpieTTS--val_loss=5.1255-epoch=89-last.ckpt"
# HPARAM="/checkpoints/streaming/magpie/jason/magpietts_en_jason_inference.yaml"

# Great for nonstreaming Best for streaming with exponential weight
# CKPT="/checkpoints/streaming/magpie/sugh_BIN_F2P1E0.0/magpieTTS--val_loss=5.1851-epoch=143-last.ckpt"
# HPARAM="/checkpoints/streaming/magpie/sugh_BIN_F2P1E0.0/magpietts_en_subhankarg_BIN_F2P1E0.0.yaml"

# Good for nonstreaming
# CKPT="/checkpoints/streaming/magpie/BIN_F5P5E0_DECF1.0P1.0_wait1_strictwindowTrue/magpieTTS--val_loss=5.2150-epoch=173-last.ckpt"
# HPARAM="/checkpoints/streaming/magpie/BIN_F5P5E0_DECF1.0P1.0_wait1_strictwindowTrue/magpietts_en_subhankarg_BIN_F0P2E0.0_DECF1.0P1.0.yaml"

# CKPT="/checkpoints/streaming/magpie/sugh_BIN_F0P2E0.0_DECF1.0P1.0/magpieTTS--val_loss=5.2076-epoch=144-last.ckpt"
# HPARAM="/checkpoints/streaming/magpie/sugh_BIN_F0P2E0.0_DECF1.0P1.0/magpietts_en_subhankarg_BIN_F0P2E0.0_DECF1.0P1.0.yaml"

# 25FPS
# Non streaming ok
# CKPT="/TB/magpie2503_CE_CA_BIN_BPE_F6P2E0_DECF1.0P1.0_C25FPS_Causal_8cb_HRLLM_lr1e-4_bs12_precbf16_wait3_strictwindowTrue/magpieTTS/checkpoints/magpieTTS--val_loss=5.2386-epoch=153-last.ckpt"

# CKPT="/TB/magpie2503_CE_CA_BIN_BPE_F6P2E0_DECF1.0P1.0_C25FPS_Causal_8cb_HRLLM_lr1e-4_bs12_precbf16_NOwaitK_strictwindowTrue/magpieTTS/checkpoints/magpieTTS--val_loss=5.2388-epoch=150-last.ckpt"

#
# CKPT="/TB/magpie2503_DC_CE_CA_BIN_BPE_F5P2E0_DECF1.0P1.0_C25FPS_Causal_8cb_HRLLM_lr1e-4_bs12_precbf16_wait3_strictwindowTrue_bpe/magpieTTS/checkpoints/magpieTTS--val_loss=5.2238-epoch=229.ckpt"
# HPARAM="/workspace/NeMo/examples/tts/conf/magpietts/magpietts_en_bpe_25fps_inference.yaml"

### SMALL MAGPIE
# CKPT="/TB/DC_CE_BIN_small_C21FPS_Causal_8cb_HRLLM_lr1e-4_bs12_precbf16_lhotse/magpieTTS/checkpoints/magpieTTS.nemo"
# CKPT="/TB/DC_CE_BIN_small_C21FPS_Causal_8cb_HRLLM_lr1e-4_bs12_precbf16_lhotse/magpieTTS/checkpoints/MagpieTTS-EN-Lhotse--val_loss=5.1882-step=349010-epoch=349-last.ckpt"
# HPARAM="/workspace/NeMo/examples/tts/conf/magpietts/magpietts_lhotse_dc_en_tiny_inference.yaml"

# SMALL MAGPIE BETA-BINOMIAL
# CKPT="/TB/DC_CE_small_C21FPS_Causal_8cb_HRLLM_lr1e-4_bs12_precbf16_lhotse/magpieTTS/checkpoints/magpieTTS.nemo"


### SMALL MAGPIE WITH GRPO
# CKPT="/TB/GRPO_DC_CE_BIN_small_C21FPS_Causal_8cb_lr2e-7_bs2_prec32/magpieTTS/checkpoints/Magpie-TTS-EN--val_cer_gt=0.0237-step=1373-epoch=1.ckpt"
# CKPT="/TB/GRPO_KLDiv_infcfg0.5_DC_CE_BIN_small_C21FPS_Causal_8cb_lr2e-7_bs2_prec32/magpieTTS/checkpoints/Magpie-TTS-EN--val_cer_gt=0.0461-step=1846-epoch=2.ckpt"
# HPARAM="/workspace/NeMo/examples/tts/conf/magpietts/magpietts_lhotse_dc_en_tiny_inference.yaml"

### SMALL MAGPIE BETA-BINOMIAL WITH GRPO
# CKPT="/TB/GRPO_DC_CE_small_C21FPS_Causal_8cb_lr5e-7_bs2_prec32/magpieTTS/checkpoints/Magpie-TTS-EN--val_cer_gt=0.0250-step=1373-epoch=1.ckpt"
# CKPT="/checkpoints/streaming/magpie/shehzeen_ckpt/cer_ssim_priornull_2e7_beta001_llmservice_wcfg_wthreholds_epoch0.ckpt"
# HPARAM="/workspace/NeMo/examples/tts/conf/magpietts/magpietts_lhotse_dc_en_tiny_betabinomial_inference.yaml"


#### CICD TESTING
# CKPT="/home/TestData/tts/2506_SeenSpeaker/T5TTS--val_loss=0.3125-epoch=8.ckpt" # multi_encoder_context_tts
# HPARAM="/home/TestData/tts/2506_SeenSpeaker/hparams.yaml"

CKPT="/home/TestData/tts/2506_ZeroShot/dpo-T5TTS--val_loss=0.4513-epoch=3.ckpt" #decoder_context_tts
HPARAM="/home/TestData/tts/2506_ZeroShot/lrhm_short_yt_prioralways_alignement_0.002_priorscale_0.1.yaml"

#### Jensen
# CKPT="/checkpoints/streaming/magpie/jensen/magpieTTS.nemo"
# decoder_context_tts
# CKPT="/TB/magpie2508_DC_CE_CE_CA_TextCond_BetaBinom_C21FPS_Causal_8cb_En-Mn-Fr-Es-De-Vi-Hi_IPA_lr1e-4_bs12_prec32_corrected/magpieTTS/checkpoints/magpieTTS.nemo"
# decoder_ce

# CODEC="/nemo_codec_checkpoints/AudioCodec_21Hz_no_eliz.nemo"
# CODEC="/nemo_codec_checkpoints/21fps_causal_codecmodel.nemo"
CODEC="/home/TestData/tts/AudioCodec_21Hz_no_eliz_without_wavlm_disc.nemo"
# CODEC="/nemo_codec_checkpoints/Low_Frame-rate_Speech_Codec++.nemo"
# CODEC="/nemo_codec_checkpoints/Low_Frame-rate_25FPS_Speech_Codec++.nemo"

OUT_DIR="/checkpoints/results/cicd_zs"

export CUDA_VISIBLE_DEVICES=0
export CUDA_LAUNCH_BLOCKING=1

# --nemo_files $CKPT \

# --checkpoint_files $CKPT \
# --hparams_files ${HPARAM} \

# python scripts/magpietts/infer_and_evaluate.py \
#     --codecmodel_path ${CODEC} \
#     --datasets an4_val_ci \
#     --out_dir ${OUT_DIR} \
#     --batch_size 4 \
#     --use_cfg \
#     --cfg_scale 2.5 \
#     --num_repeats 1 \
#     --temperature 0.6 \
#     --hparams_files ${HPARAM} \
#     --checkpoint_files ${CKPT} \
#     --legacy_codebooks \
#     --legacy_text_conditioning \
#     --apply_attention_prior \
#     --clean_up_disk \
#     --cer_target 0.1 \
#     --ssim_target 0.7

OUT_DIR="/checkpoints/results/cicd_streaming"

python scripts/magpietts/infer_and_evaluate_streaming.py \
    --codecmodel_path ${CODEC} \
    --datasets an4_val_ci \
    --out_dir ${OUT_DIR} \
    --batch_size 1 \
    --use_cfg \
    --cfg_scale 2.5 \
    --hparams_files ${HPARAM} \
    --checkpoint_files ${CKPT} \
    --legacy_codebooks \
    --apply_attention_prior \
    --clean_up_disk \
    --disable_fcd \
    --cer_target 0.8 \
    --ssim_target 0.3