#!/bin/bash
#SBATCH --job-name=sfp4-val-gen
#SBATCH --account=nvr_elm_llm
#SBATCH --partition=interactive
#SBATCH --nodes=1
#SBATCH --gres=gpu:8
#SBATCH --cpus-per-task=128
#SBATCH --mem=1440G
#SBATCH --time=02:00:00
#SBATCH --output=slurm_logs/sfp4_val_gen_%j.out
#SBATCH --error=slurm_logs/sfp4_val_gen_%j.err

set -ex

REPO_ROOT="/lustre/fsw/portfolios/nvr/projects/nvr_elm_llm/users/yitongl/code/FastVideo"
KERNEL_ROOT="${REPO_ROOT}/fastvideo-kernel"

mkdir -p "${REPO_ROOT}/slurm_logs"
cd "${REPO_ROOT}"
source .venv/bin/activate
export PYTHONPATH="${KERNEL_ROOT}/python:${KERNEL_ROOT}:${PYTHONPATH}"

echo "=== Environment ==="
nvidia-smi -L | head -1
python -c "import torch; print(f'torch={torch.__version__}, cuda={torch.cuda.is_available()}, gpus={torch.cuda.device_count()}')"
python -c "import triton; print(f'triton={triton.__version__}')"

echo ""
echo "######################################################################"
echo "# Generate 8 videos with sparse FP4 attention                       #"
echo "######################################################################"
cd "${REPO_ROOT}"

MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
PROMPT="Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting."
NEGATIVE_PROMPT="Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
SEED=1024

SPARSITY_LIST=(0.0 0.1 0.2 0.4 0.5 0.7 0.8 0.9)
OUTPUT_BASE="${REPO_ROOT}/outputs_sparse_fp4_sweep"
mkdir -p "${OUTPUT_BASE}"

echo "Sparsity levels: ${SPARSITY_LIST[*]}"

PIDS=()
for i in $(seq 0 7); do
    SPARSITY=${SPARSITY_LIST[$i]}
    OUT_DIR="${OUTPUT_BASE}/sparsity_${SPARSITY}"
    mkdir -p "${OUT_DIR}"
    echo "[GPU ${i}] sparsity=${SPARSITY}"

    (
        export CUDA_VISIBLE_DEVICES=${i}
        export FASTVIDEO_ATTENTION_BACKEND=SPARSE_FP4_ATTN

        fastvideo generate \
            --model-path "${MODEL_PATH}" \
            --sp-size 1 --tp-size 1 --num-gpus 1 \
            --dit-cpu-offload False \
            --vae-cpu-offload False \
            --text-encoder-cpu-offload True \
            --pin-cpu-memory False \
            --height 480 --width 832 --num-frames 81 \
            --num-inference-steps 50 --fps 16 \
            --guidance-scale 6.0 --flow-shift 8.0 \
            --prompt "${PROMPT}" \
            --negative-prompt "${NEGATIVE_PROMPT}" \
            --seed ${SEED} \
            --VSA-sparsity ${SPARSITY} \
            --output-path "${OUT_DIR}/" \
            2>&1 | tee "${OUT_DIR}/log.txt"

        echo "[GPU ${i}] sparsity=${SPARSITY} DONE"
    ) &
    PIDS+=($!)
done

echo "=== Waiting for all 8 jobs ==="
FAIL=0
for i in $(seq 0 7); do
    wait ${PIDS[$i]} || { echo "[GPU ${i}] FAILED"; FAIL=1; }
done

echo ""
if [ $FAIL -eq 0 ]; then
    echo "=== All 8 videos generated ==="
else
    echo "=== Some failed ==="
fi
find "${OUTPUT_BASE}" -name "*.mp4" | sort
echo "=== Done ==="