scripts/setup_nunchaku.sh

#!/usr/bin/env bash
# Re-install the Nunchaku FLUX.2 loader files into the installed nunchaku package.
# The dev wheel (v1.3.0dev20260306) does NOT yet ship NunchakuFlux2Transformer2DModel
# (PR nunchaku-ai/nunchaku#926, still in review) — the checkpoint repo bundles the
# modeling code, which must be copied into site-packages. A box restart wipes
# site-packages (ephemeral root), so re-run this after reinstalling the wheel.
#
# Prereq: nunchaku wheel installed, models/klein-9b-nunchaku downloaded.
set -euo pipefail
cd /workspace
SRC=models/klein-9b-nunchaku
NP=$(python3 -c "import nunchaku,os;print(os.path.dirname(nunchaku.__file__))" 2>/dev/null | tail -1)
echo "nunchaku pkg: $NP"
cp "$SRC/torch_transfer_utils.py" "$NP/torch_transfer_utils.py"
cp "$SRC/transformer_flux2.py"    "$NP/models/transformers/transformer_flux2.py"
cp -r "$SRC/common"               "$NP/lora/common"
python3 -c "from nunchaku.models.transformers.transformer_flux2 import NunchakuFlux2Transformer2DModel; from nunchaku.utils import get_precision; print('OK: FLUX.2 loader installed, precision=', get_precision())" 2>&1 | grep -vE "_pytree|torchao Tensor" | tail -1