setup.sh

#!/bin/bash
set -e
cd "$(dirname "$0")"

echo "========================================"
echo " CSM-1B on RTX 5090 — Setup"
echo "========================================"

# Preflight
nvidia-smi --query-gpu=name,driver_version,compute_cap --format=csv,noheader || {
    echo "ERROR: nvidia-smi not found. Install NVIDIA drivers."; exit 1; }

PYTHON=${PYTHON:-python3}
VENV=${1:-.venv}

# Create venv
if [ ! -d "$VENV" ]; then
    echo "[1/5] Creating venv at $VENV..."
    $PYTHON -m venv "$VENV"
else
    echo "[1/5] Reusing venv at $VENV"
fi
source "$VENV/bin/activate"
pip install --upgrade pip -q

# PyTorch nightly (sm_120 kernels)
echo "[2/5] Installing PyTorch nightly cu128..."
pip install --pre torch torchaudio \
    --index-url https://download.pytorch.org/whl/nightly/cu128

# Fix missing nightly deps
pip install packaging numpy -q

# Verify CUDA
echo "[3/5] Verifying sm_120..."
python3 -c "
import torch
assert torch.cuda.is_available(), 'CUDA not available'
t = torch.randn(10, device='cuda') * 2
print(f'  torch={torch.__version__} cuda={torch.version.cuda} sm_120=OK')
"

# Install dependencies
echo "[4/5] Installing dependencies..."
pip install -r requirements.txt -q

# Patch transformers
echo "[5/5] Patching transformers for CUDA graphs..."
python3 patch_transformers.py

# ptxas hint
PTXAS=$(which ptxas 2>/dev/null || find /usr/local/cuda* -name ptxas 2>/dev/null | head -1)
if [ -n "$PTXAS" ]; then
    echo ""
    echo "ptxas found: $PTXAS"
    echo "Add to your shell: export TRITON_PTXAS_PATH=$PTXAS"
fi

echo ""
echo "========================================"
echo " Setup complete!"
echo "========================================"
echo ""
echo "  source $VENV/bin/activate"
echo "  python csm_pipeline.py"
echo ""
echo "  Accept sesame/csm-1b license at:"
echo "  https://huggingface.co/sesame/csm-1b"