apptainer files for CryoZeta

.gitattributes

@@ -36,3 +36,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 components.v20240608.cif filter=lfs diff=lfs merge=lfs -text
 msa/prot/24ee4eb8f30911b98ae173d19225440514a110ba50da1d89327c9c74ca0ac7fe/uniref100_hits.a3m filter=lfs diff=lfs merge=lfs -text
 examples/msa/prot/24ee4eb8f30911b98ae173d19225440514a110ba50da1d89327c9c74ca0ac7fe/uniref100_hits.a3m filter=lfs diff=lfs merge=lfs -text
+apptainer/CryoZeta.sif filter=lfs diff=lfs merge=lfs -text

apptainer/CryoZeta.sif

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bff5f716227db0b7da90f29ae9526216d6ef595a66d729f4e4d98b73eb9b14d3
+size 9623916544

apptainer/run_cryozeta.sh

@@ -0,0 +1,293 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# ════════════════════════════════════════════════════════════════════
+# CONFIGURE BEFORE RUNNING (assets_dir is a required CLI argument)
+SIF=""      # e.g. /path/to/CryoZeta.sif
+# ════════════════════════════════════════════════════════════════════
+
+GPU=""
+
+usage() {
+    cat <<EOF
+Usage: $(basename "$0") --gpu ID [OPTIONS] <input_json> <output_dir> <assets_dir>
+       $(basename "$0") prepare-ui [--sif PATH]
+
+Run the CryoZeta pipeline (detection → inference → combine) via Apptainer.
+
+Environment (prepare-ui):
+  CRYOZETA_PREP_UI_PORT   HTTP port when the image supports --port (default: 8501). Ignored for legacy images.
+
+Subcommands:
+  prepare-ui          Launch the web UI for preparing input JSON files
+
+Arguments:
+  input_json    Path to input JSON file (e.g. examples/example.json)
+  output_dir    Directory for output results
+  assets_dir    Host directory with .safetensors weights (mounted as /app/CryoZeta/assets)
+
+Required:
+  --gpu ID            CUDA device ID (e.g. 0, 1, 2)
+
+Options:
+  --sif PATH          Path to CryoZeta.sif    (default: $SIF)
+  --interpolation     Use interpolation model in Stage 2
+  --seeds SEEDS       Random seeds           (default: 101)
+  --n-sample N        Number of diffusion samples (default: 5)
+  --n-step N          Number of diffusion steps   (default: 20)
+  --n-cycle N         Number of model cycles      (default: 10)
+  --num-select N      Number to select in combine (default: 5)
+  --skip-detection    Skip Stage 1
+  --skip-inference    Skip Stage 2
+  --skip-combine      Skip Stage 3
+  -h, --help          Show this help
+
+Weights (not in the SIF): download into a directory and pass that path as <assets_dir>, e.g.
+  huggingface-cli download KiharaLab/CryoZeta --repo-type model --local-dir /path/to/assets
+  (pip install -U huggingface_hub). Stage 1: cryozeta-detection-v0.0.1.safetensors. Stage 2:
+  cryozeta-v0.0.1.safetensors, or cryozeta-interpolate-v0.0.1.safetensors if --interpolation.
+EOF
+    exit 0
+}
+
+require_weight_file() {
+    local path=$1
+    local label=$2
+    if [[ ! -f "$path" ]]; then
+        echo "Error: missing $label: $path" >&2
+        echo "Download model files with Hugging Face CLI:" >&2
+        echo "  huggingface-cli download KiharaLab/CryoZeta --repo-type model --local-dir <assets-dir>" >&2
+        echo "  (Install CLI: pip install -U huggingface_hub)" >&2
+        echo "Then pass that directory as the <assets_dir> argument (same path as --local-dir)." >&2
+        exit 1
+    fi
+}
+
+# ── Defaults ─────────────────────────────────────────────────────────
+USE_INTERPOLATION=false
+SEEDS=101
+N_SAMPLE=5
+N_STEP=20
+N_CYCLE=10
+NUM_SELECT=5
+SKIP_DETECTION=false
+SKIP_INFERENCE=false
+SKIP_COMBINE=false
+
+# ── Handle prepare-ui subcommand ─────────────────────────────────────
+if [[ "${1:-}" == "prepare-ui" ]]; then
+    shift
+    while [[ $# -gt 0 ]]; do
+        case $1 in
+            --sif) SIF="$2"; shift 2 ;;
+            -h|--help) usage ;;
+            *)     echo "Unknown option: $1" >&2; exit 1 ;;
+        esac
+    done
+    [[ -f "$SIF" ]] || { echo "Error: SIF not found: $SIF" >&2; exit 1; }
+
+    set +e
+    _ui_help_msg="$(apptainer exec --bind "$(pwd):$(pwd)" "$SIF" bash -c \
+        "cd /app/CryoZeta && source /opt/conda/etc/profile.d/conda.sh && conda activate cryozeta && cryozeta-prepare ui --help" 2>&1)"
+    _ui_help_ec=$?
+    set -e
+
+    UI_SUPPORTS_PORT=false
+    if [[ "$_ui_help_ec" -eq 0 ]] && printf '%s\n' "$_ui_help_msg" | grep -qE '\-\-port([[:space:]=]|$)'; then
+        UI_SUPPORTS_PORT=true
+    else
+        echo "prepare-ui: could not confirm --port from 'cryozeta-prepare ui --help' (exit ${_ui_help_ec}); using legacy invocation without --port." >&2
+    fi
+
+    if [[ "$UI_SUPPORTS_PORT" != true ]]; then
+        echo "prepare-ui: http://localhost:8501 (default)" >&2
+        echo "prepare-ui: SSH tunnel: ssh -L 8501:localhost:8501 <user>@<host>" >&2
+        exec apptainer exec --bind "$(pwd):$(pwd)" "$SIF" bash -c \
+            "cd /app/CryoZeta && source /opt/conda/etc/profile.d/conda.sh && conda activate cryozeta && cryozeta-prepare ui"
+    fi
+
+    prep_ui_port="${CRYOZETA_PREP_UI_PORT:-8501}"
+    if ! [[ "$prep_ui_port" =~ ^[0-9]+$ ]] || (( prep_ui_port > 65535 )); then
+        echo "Error: CRYOZETA_PREP_UI_PORT must be an integer 0–65535" >&2
+        exit 1
+    fi
+
+    echo "prepare-ui: http://localhost:${prep_ui_port}" >&2
+    echo "prepare-ui: SSH tunnel: ssh -L ${prep_ui_port}:localhost:${prep_ui_port} <user>@<host>" >&2
+    exec apptainer exec --bind "$(pwd):$(pwd)" "$SIF" bash -c \
+        "cd /app/CryoZeta && source /opt/conda/etc/profile.d/conda.sh && conda activate cryozeta && cryozeta-prepare ui --port ${prep_ui_port}"
+fi
+
+# ── Parse arguments ──────────────────────────────────────────────────
+POSITIONAL=()
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --sif)            SIF="$2";            shift 2 ;;
+        --gpu)            GPU="$2";            shift 2 ;;
+        --interpolation)  USE_INTERPOLATION=true; shift ;;
+        --seeds)          SEEDS="$2";          shift 2 ;;
+        --n-sample)       N_SAMPLE="$2";       shift 2 ;;
+        --n-step)         N_STEP="$2";         shift 2 ;;
+        --n-cycle)        N_CYCLE="$2";        shift 2 ;;
+        --num-select)     NUM_SELECT="$2";     shift 2 ;;
+        --skip-detection) SKIP_DETECTION=true; shift ;;
+        --skip-inference) SKIP_INFERENCE=true; shift ;;
+        --skip-combine)   SKIP_COMBINE=true;   shift ;;
+        -h|--help)        usage ;;
+        -*)               echo "Unknown option: $1" >&2; exit 1 ;;
+        *)                POSITIONAL+=("$1");  shift ;;
+    esac
+done
+
+if [[ ${#POSITIONAL[@]} -ne 3 ]]; then
+    echo "Error: expected 3 positional arguments (<input_json> <output_dir> <assets_dir>), got ${#POSITIONAL[@]}" >&2
+    usage
+fi
+
+INPUT_JSON="$(realpath "${POSITIONAL[0]}")"
+OUTPUT_DIR="$(realpath -m "${POSITIONAL[1]}")"
+ASSETS="$(realpath "${POSITIONAL[2]}")"
+
+# ── Validate ─────────────────────────────────────────────────────────
+[[ -n "$GPU" ]]        || { echo "Error: --gpu is required" >&2; usage; }
+[[ -n "$SIF" ]]        || { echo "Error: SIF not set — edit top of script or pass --sif PATH" >&2; exit 1; }
+[[ -f "$SIF" ]]        || { echo "Error: SIF not found: $SIF" >&2; exit 1; }
+[[ -f "$INPUT_JSON" ]] || { echo "Error: input JSON not found: $INPUT_JSON" >&2; exit 1; }
+[[ -d "$ASSETS" ]]     || { echo "Error: assets dir not found: $ASSETS" >&2; exit 1; }
+mkdir -p "$OUTPUT_DIR"
+
+BIND_FLAGS=("--bind" "$(pwd):$(pwd)" "--bind" "$ASSETS:/app/CryoZeta/assets")
+
+if [[ "$SKIP_DETECTION" == false ]]; then
+    require_weight_file "$ASSETS/cryozeta-detection-v0.0.1.safetensors" "detection weights (Stage 1)"
+fi
+if [[ "$SKIP_INFERENCE" == false ]]; then
+    if [[ "$USE_INTERPOLATION" == true ]]; then
+        require_weight_file "$ASSETS/cryozeta-interpolate-v0.0.1.safetensors" "interpolation checkpoint (Stage 2)"
+    else
+        require_weight_file "$ASSETS/cryozeta-v0.0.1.safetensors" "structure-prediction checkpoint (Stage 2)"
+    fi
+fi
+
+echo "Bind mounts: ${BIND_FLAGS[*]}"
+
+# ── GPU capability detection ──────────────────────────────────────────
+# nvidia-smi -i GPU uses PCI bus order, matching CUDA_DEVICE_ORDER=PCI_BUS_ID
+DRIVER_MAJOR=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader -i "$GPU" 2>/dev/null \
+    | head -1 | cut -d. -f1)
+GPU_CC_RAW=$(nvidia-smi --query-gpu=compute_cap   --format=csv,noheader -i "$GPU" 2>/dev/null \
+    | head -1 | tr -d '.')
+DRIVER_MAJOR="${DRIVER_MAJOR:-0}"
+GPU_CC_RAW="${GPU_CC_RAW:-75}"
+echo "GPU $GPU: sm_${GPU_CC_RAW}, driver ${DRIVER_MAJOR}.x"
+
+# DS4Sci EvoformerAttention requires sm_80+.
+# CUDA 12.9 JIT custom kernels (fast_layernorm, DS4Sci) produce PTX ISA 8.9,
+# which drivers older than 555 (CUDA 12.9 minimum) cannot assemble.
+# Fall back to standard PyTorch ops on older hardware.
+if [[ "$DRIVER_MAJOR" -ge 555 ]] && [[ "$GPU_CC_RAW" -ge 80 ]]; then
+    USE_DS_EVO_ATTN=true
+    USE_CUEQ_ATTN=true
+    USE_LMA=false
+    ATTN_CHUNK_SIZE=256
+    LAYERNORM_PREFIX="export LAYERNORM_TYPE=fast_layernorm && "
+    echo "GPU compat: sm_${GPU_CC_RAW} + driver ${DRIVER_MAJOR} — JIT custom kernels enabled"
+else
+    # DS4Sci EvoformerAttention requires sm_80+.
+    # cuEquivariance falls back to DS4Sci when shape limits are hit, so must
+    # also be disabled — otherwise primitives.py forces use_deepspeed_evo_attention=True.
+    # CUDA 12.9 JIT (fast_layernorm, DS4Sci) produces PTX ISA 8.9 which
+    # drivers older than 555 cannot assemble.
+    # LMA (Low Memory Attention) is used instead — chunks Q in 1024-token steps to
+    # avoid the O(N²) memory spike that OOMs with standard attention on N=1313.
+    # DS4Sci EvoformerAttention requires sm_80+.
+    # cuEquivariance falls back to DS4Sci when shape limits are hit, so must
+    # also be disabled — otherwise primitives.py forces use_deepspeed_evo_attention=True.
+    # CUDA 12.9 JIT (fast_layernorm, DS4Sci) produces PTX ISA 8.9 which
+    # drivers older than 555 cannot assemble.
+    # LMA (Low Memory Attention) is used instead — chunks Q in 1024-token steps.
+    # Default chunk_size=256 gives a=[1,256,8,1024,1313]=5.4GB tensor → OOM.
+    # Reduce to 32: a=[1,32,8,1024,1313]=~684MB, safely fits in reserved pool.
+    USE_DS_EVO_ATTN=false
+    USE_CUEQ_ATTN=false
+    USE_LMA=true
+    ATTN_CHUNK_SIZE=32
+    LAYERNORM_PREFIX=""
+    echo "GPU compat: sm_${GPU_CC_RAW} + driver ${DRIVER_MAJOR} — JIT custom kernels disabled (LMA, chunk_size=${ATTN_CHUNK_SIZE})"
+fi
+
+# Helper: run a command inside the container
+run_in_container() {
+    local need_gpu="$1"; shift
+    # Cache + Triton ptxas: match CryoZeta/env_setup.sh; \${…} expands inside the container.
+    # ptxas-blackwell shim lives under CRYOZETA_CACHE_DIR (same tree as Triton/Torch caches).
+    local cmd
+    cmd="cd /app/CryoZeta && source /opt/conda/etc/profile.d/conda.sh && conda activate cryozeta && "
+    cmd+="_CRYOZETA_CACHE=\"\${CRYOZETA_CACHE_DIR:-\${HOME}/.cache/cryozeta}\" && "
+    cmd+="export CRYOZETA_CACHE_DIR=\"\${_CRYOZETA_CACHE}\" && "
+    cmd+="export TORCH_EXTENSIONS_DIR=\"\${TORCH_EXTENSIONS_DIR:-\${_CRYOZETA_CACHE}/torch_extensions}\" && "
+    cmd+="export TORCHINDUCTOR_CACHE_DIR=\"\${TORCHINDUCTOR_CACHE_DIR:-\${_CRYOZETA_CACHE}/inductor}\" && "
+    cmd+="export TRITON_CACHE_DIR=\"\${TRITON_CACHE_DIR:-\${_CRYOZETA_CACHE}/triton}\" && "
+    cmd+="unset _CRYOZETA_CACHE && "
+    cmd+="[ -x \"\${CONDA_PREFIX}/bin/ptxas\" ] && export TRITON_PTXAS_PATH=\"\${CONDA_PREFIX}/bin/ptxas\" TRITON_PTXAS_BLACKWELL_PATH=\"\${CONDA_PREFIX}/bin/ptxas\" ; "
+    cmd+="mkdir -p \"\${CRYOZETA_CACHE_DIR}/triton-bin\" && ln -sf \"\${CONDA_PREFIX}/bin/ptxas\" \"\${CRYOZETA_CACHE_DIR}/triton-bin/ptxas-blackwell\" && export PATH=\"\${CRYOZETA_CACHE_DIR}/triton-bin:\$PATH\" && $*"
+
+    if [[ "$need_gpu" == "gpu" ]]; then
+        CUDA_VISIBLE_DEVICES="$GPU" apptainer exec --nv \
+            "${BIND_FLAGS[@]}" \
+            "$SIF" bash -c "$cmd"
+    else
+        apptainer exec \
+            "${BIND_FLAGS[@]}" \
+            "$SIF" bash -c "$cmd"
+    fi
+}
+
+# Pick checkpoint based on interpolation flag
+if [[ "$USE_INTERPOLATION" == true ]]; then
+    CHECKPOINT="/app/CryoZeta/assets/cryozeta-interpolate-v0.0.1.safetensors"
+else
+    CHECKPOINT="/app/CryoZeta/assets/cryozeta-v0.0.1.safetensors"
+fi
+
+# ── Stage 1: Atom Detection ─────────────────────────────────────────
+if [[ "$SKIP_DETECTION" == false ]]; then
+    echo "=== Stage 1: Atom Detection ==="
+    run_in_container gpu \
+        cryozeta-detection json-run "$INPUT_JSON" "$OUTPUT_DIR" --device cuda
+fi
+
+# ── Stage 2: Structure Prediction ────────────────────────────────────
+if [[ "$SKIP_INFERENCE" == false ]]; then
+    echo "=== Stage 2: Structure Prediction (interpolation=$USE_INTERPOLATION) ==="
+    run_in_container gpu \
+        "${LAYERNORM_PREFIX}cryozeta-inference \
+            --seeds $SEEDS \
+            --load_checkpoint_path $CHECKPOINT \
+            --em_file_dir $OUTPUT_DIR \
+            --dump_dir $OUTPUT_DIR \
+            --input_json_path $INPUT_JSON \
+            --use_deepspeed_evo_attention $USE_DS_EVO_ATTN \
+            --use_cuequivariance_attention $USE_CUEQ_ATTN \
+            --use_lma $USE_LMA \
+            --infer_setting.chunk_size $ATTN_CHUNK_SIZE \
+            --model.N_cycle $N_CYCLE \
+            --sample_diffusion.N_sample $N_SAMPLE \
+            --sample_diffusion.N_step $N_STEP \
+            --data.num_dl_workers 1 \
+            --use_interpolation $USE_INTERPOLATION \
+            --overwrite false"
+fi
+
+# ── Stage 3: Combine Results ─────────────────────────────────────────
+if [[ "$SKIP_COMBINE" == false ]]; then
+    echo "=== Stage 3: Combine Results ==="
+    run_in_container cpu \
+        cryozeta-combine \
+            --dump-dir "$OUTPUT_DIR" \
+            --input-json-path "$INPUT_JSON" \
+            --seeds "$SEEDS" \
+            --num-select "$NUM_SELECT"
+fi
+
+echo "=== Done. Results in: $OUTPUT_DIR ==="

apptainer/run_cryozeta_large.sh

@@ -0,0 +1,388 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# ════════════════════════════════════════════════════════════════════
+# CONFIGURE BEFORE RUNNING (assets_dir is a required CLI argument)
+SIF=""      # e.g. /path/to/CryoZeta.sif
+# ════════════════════════════════════════════════════════════════════
+
+GPU=""
+
+usage() {
+    cat <<EOF
+Usage: $(basename "$0") --gpu IDS [OPTIONS] <input_json> <output_dir> <assets_dir>
+
+Run the CryoZeta large-complex cycle pipeline via Apptainer.
+
+Arguments:
+  input_json    Path to a large-example JSON list
+  output_dir    Directory for output results
+  assets_dir    Host directory with .safetensors weights (mounted as /app/CryoZeta/assets)
+
+Required:
+  --gpu IDS            CUDA device IDs (e.g. 0 or 0,1)
+
+Options:
+  --sif PATH           Path to CryoZeta.sif    (default: $SIF)
+  --example SEL        Select one JSON entry by 0-based index or name (default: 0)
+  --registration M     Registration method: auto, teaser, svd, vesper (default: auto)
+  --threshold X        EM point cropping threshold in Angstrom (default: 5.0)
+  --seeds SEEDS        Random seeds (default: 101)
+  --n-sample N         Number of diffusion samples (default: 5)
+  --n-step N           Number of diffusion steps (default: 20)
+  --n-cycle N          Number of model cycles (default: 10)
+  --use-opm-tilelang   Enable TileLang OPM kernel
+  --skip-detection     Skip detection and reuse existing .pt output
+  --skip-combine       Skip combine-stages
+  --vesper-map PATH    Original MRC map path for VESPER mode
+  --vesper-contour X   Contour level for VESPER mode
+  --vesper-resolution X
+                       Simulated-map resolution for VESPER mode (default: 3.0)
+  --vesper-angle-spacing X
+                       Angular sampling interval for VESPER mode (default: 5.0)
+  -h, --help           Show this help
+
+Weights (not in the SIF): download into a directory and pass that path as <assets_dir>, e.g.
+  huggingface-cli download KiharaLab/CryoZeta --repo-type model --local-dir /path/to/assets
+  (pip install -U huggingface_hub). Required here:
+  cryozeta-detection-v0.0.1.safetensors and cryozeta-v0.0.1.safetensors
+EOF
+    exit 0
+}
+
+require_weight_file() {
+    local path=$1
+    local label=$2
+    if [[ ! -f "$path" ]]; then
+        echo "Error: missing $label: $path" >&2
+        echo "Download model files with Hugging Face CLI:" >&2
+        echo "  huggingface-cli download KiharaLab/CryoZeta --repo-type model --local-dir <assets-dir>" >&2
+        echo "  (Install CLI: pip install -U huggingface_hub)" >&2
+        echo "Then pass that directory as the <assets_dir> argument (same path as --local-dir)." >&2
+        exit 1
+    fi
+}
+
+EXAMPLE_SEL="0"
+REGISTRATION_METHOD="auto"
+THRESHOLD="5.0"
+SEEDS=101
+N_SAMPLE=5
+N_STEP=20
+N_CYCLE=10
+USE_OPM_TILELANG=false
+SKIP_DETECTION=false
+SKIP_COMBINE=false
+VESPER_MAP_PATH=""
+VESPER_CONTOUR_LEVEL=""
+VESPER_RESOLUTION="3.0"
+VESPER_ANGLE_SPACING="5.0"
+
+POSITIONAL=()
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --sif)                  SIF="$2"; shift 2 ;;
+        --gpu)                  GPU="$2"; shift 2 ;;
+        --example)              EXAMPLE_SEL="$2"; shift 2 ;;
+        --registration)         REGISTRATION_METHOD="$2"; shift 2 ;;
+        --threshold)            THRESHOLD="$2"; shift 2 ;;
+        --seeds)                SEEDS="$2"; shift 2 ;;
+        --n-sample)             N_SAMPLE="$2"; shift 2 ;;
+        --n-step)               N_STEP="$2"; shift 2 ;;
+        --n-cycle)              N_CYCLE="$2"; shift 2 ;;
+        --use-opm-tilelang)     USE_OPM_TILELANG=true; shift ;;
+        --skip-detection)       SKIP_DETECTION=true; shift ;;
+        --skip-combine)         SKIP_COMBINE=true; shift ;;
+        --vesper-map)           VESPER_MAP_PATH="$2"; shift 2 ;;
+        --vesper-contour)       VESPER_CONTOUR_LEVEL="$2"; shift 2 ;;
+        --vesper-resolution)    VESPER_RESOLUTION="$2"; shift 2 ;;
+        --vesper-angle-spacing) VESPER_ANGLE_SPACING="$2"; shift 2 ;;
+        -h|--help)              usage ;;
+        -*)                     echo "Unknown option: $1" >&2; exit 1 ;;
+        *)                      POSITIONAL+=("$1"); shift ;;
+    esac
+done
+
+if [[ ${#POSITIONAL[@]} -ne 3 ]]; then
+    echo "Error: expected 3 positional arguments (<input_json> <output_dir> <assets_dir>), got ${#POSITIONAL[@]}" >&2
+    usage
+fi
+
+INPUT_JSON="$(realpath "${POSITIONAL[0]}")"
+OUTPUT_DIR="$(realpath -m "${POSITIONAL[1]}")"
+ASSETS="$(realpath "${POSITIONAL[2]}")"
+SELECTED_INPUT_JSON="$(realpath -m "${OUTPUT_DIR}/selected_input_entry.json")"
+CYCLE_INPUT_JSON="$(realpath -m "${OUTPUT_DIR}/selected_cycle_input_entry.json")"
+TRITON_BIN_DIR="$(realpath -m "$(pwd)/tmp/triton-bin")"
+
+[[ -n "$GPU" ]]        || { echo "Error: --gpu is required" >&2; usage; }
+[[ -n "$SIF" ]]        || { echo "Error: SIF not set — edit top of script or pass --sif PATH" >&2; exit 1; }
+[[ -f "$SIF" ]]        || { echo "Error: SIF not found: $SIF" >&2; exit 1; }
+[[ -f "$INPUT_JSON" ]] || { echo "Error: input JSON not found: $INPUT_JSON" >&2; exit 1; }
+[[ -d "$ASSETS" ]]     || { echo "Error: assets dir not found: $ASSETS" >&2; exit 1; }
+
+if [[ "$SKIP_DETECTION" == false ]]; then
+    require_weight_file "$ASSETS/cryozeta-detection-v0.0.1.safetensors" "detection weights (Stage 1)"
+fi
+require_weight_file "$ASSETS/cryozeta-v0.0.1.safetensors" "cycle-prediction checkpoint"
+
+mkdir -p "$OUTPUT_DIR" "$TRITON_BIN_DIR"
+
+python3 - "$INPUT_JSON" "$SELECTED_INPUT_JSON" "$CYCLE_INPUT_JSON" "$OUTPUT_DIR" "$EXAMPLE_SEL" <<'PY'
+import json
+import os
+import sys
+from pathlib import Path
+
+input_json = Path(sys.argv[1]).resolve()
+selected_input_json = Path(sys.argv[2]).resolve()
+cycle_input_json = Path(sys.argv[3]).resolve()
+dump_dir = Path(sys.argv[4]).resolve()
+selection = sys.argv[5]
+repo_root = Path.cwd().resolve()
+
+data = json.loads(input_json.read_text())
+if not isinstance(data, list) or len(data) == 0:
+    raise ValueError(f"Input JSON must be a non-empty list: {input_json}")
+
+if selection.isdigit():
+    entry_index = int(selection)
+    if not (0 <= entry_index < len(data)):
+        raise IndexError(
+            f"example index {entry_index} is out of range for {len(data)} entries"
+        )
+else:
+    matches = [i for i, item in enumerate(data) if item.get("name") == selection]
+    if not matches:
+        names = ", ".join(item.get("name", "<missing>") for item in data)
+        raise ValueError(
+            f"example '{selection}' not found in {input_json}. Available names: {names}"
+        )
+    entry_index = matches[0]
+
+json_dir = input_json.parent
+
+def resolve_path(value: str) -> str:
+    if not value:
+        return value
+    p = Path(value)
+    if p.is_absolute():
+        return str(p.resolve())
+
+    # First support repo-root-relative paths such as server/8zui/60484.map,
+    # which are common in server-generated JSON files.
+    repo_candidate = (repo_root / p).resolve()
+    if repo_candidate.exists():
+        return str(repo_candidate)
+
+    # Otherwise fall back to JSON-local relative resolution.
+    json_candidate = (json_dir / p).resolve()
+    return str(json_candidate)
+
+def normalize_entry_paths(entry):
+    if "map_path" in entry and entry["map_path"]:
+        entry["map_path"] = resolve_path(entry["map_path"])
+    if "em_file" in entry and entry["em_file"]:
+        entry["em_file"] = resolve_path(entry["em_file"])
+    for seq in entry.get("sequences", []):
+        for kind in ("proteinChain", "dnaSequence", "rnaSequence"):
+            node = seq.get(kind, {})
+            msa = node.get("msa", {})
+            if msa.get("precomputed_msa_dir"):
+                msa["precomputed_msa_dir"] = resolve_path(msa["precomputed_msa_dir"])
+        lig = seq.get("ligand", {})
+        lig_value = lig.get("ligand")
+        if isinstance(lig_value, str) and lig_value.startswith("FILE_"):
+            lig["ligand"] = "FILE_" + resolve_path(lig_value[5:])
+    return entry
+
+entry = normalize_entry_paths(dict(data[entry_index]))
+name = entry.get("name")
+if not name:
+    raise ValueError("Selected entry must contain 'name'")
+
+selected_input_json.parent.mkdir(parents=True, exist_ok=True)
+selected_input_json.write_text(json.dumps([entry], indent=2))
+
+cycle_entry = dict(entry)
+cycle_entry["em_file"] = str(
+    (dump_dir / name / "CryoZeta-Detection" / f"{name}.pt").resolve()
+)
+cycle_input_json.parent.mkdir(parents=True, exist_ok=True)
+cycle_input_json.write_text(json.dumps([cycle_entry], indent=2))
+
+print(f"Prepared sample: {name} (index {entry_index})")
+print(f"  detection input: {selected_input_json}")
+print(f"  cycle input:     {cycle_input_json}")
+PY
+
+SELECTED_NAME="$(python3 - "$SELECTED_INPUT_JSON" <<'PY'
+import json
+import sys
+with open(sys.argv[1]) as fh:
+    data = json.load(fh)
+print(data[0]["name"])
+PY
+)"
+
+BIND_FLAGS=(
+    "--bind" "$INPUT_JSON:$INPUT_JSON"
+    "--bind" "$SELECTED_INPUT_JSON:$SELECTED_INPUT_JSON"
+    "--bind" "$CYCLE_INPUT_JSON:$CYCLE_INPUT_JSON"
+    "--bind" "$OUTPUT_DIR:$OUTPUT_DIR"
+    "--bind" "$ASSETS:/app/CryoZeta/assets"
+    "--bind" "$(pwd):$(pwd)"
+)
+
+while IFS= read -r p; do
+    [[ -n "$p" ]] && [[ -e "$p" ]] && BIND_FLAGS+=("--bind" "$p:$p")
+done < <(python3 - "$SELECTED_INPUT_JSON" "$VESPER_MAP_PATH" <<'PY'
+import json
+import os
+import sys
+
+selected_json = sys.argv[1]
+vesper_map = sys.argv[2]
+
+with open(selected_json) as f:
+    entries = json.load(f)
+
+dirs = set()
+for entry in entries:
+    for key in ("map_path", "em_file"):
+        path = entry.get(key, "")
+        if path:
+            dirs.add(os.path.dirname(os.path.abspath(path)))
+    for seq in entry.get("sequences", []):
+        for kind in ("proteinChain", "dnaSequence", "rnaSequence"):
+            msa_dir = seq.get(kind, {}).get("msa", {}).get("precomputed_msa_dir", "")
+            if msa_dir:
+                dirs.add(os.path.abspath(msa_dir))
+        lig = seq.get("ligand", {}).get("ligand", "")
+        if isinstance(lig, str) and lig.startswith("FILE_"):
+            dirs.add(os.path.dirname(os.path.abspath(lig[5:])))
+
+if vesper_map:
+    dirs.add(os.path.dirname(os.path.abspath(vesper_map)))
+
+for d in sorted(dirs):
+    print(d)
+PY
+)
+
+echo "Bind mounts: ${BIND_FLAGS[*]}"
+
+DRIVER_MAJOR=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader -i "${GPU%%,*}" 2>/dev/null \
+    | head -1 | cut -d. -f1)
+GPU_CC_RAW=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader -i "${GPU%%,*}" 2>/dev/null \
+    | head -1 | tr -d '.')
+DRIVER_MAJOR="${DRIVER_MAJOR:-0}"
+GPU_CC_RAW="${GPU_CC_RAW:-75}"
+echo "GPU ${GPU}: sm_${GPU_CC_RAW}, driver ${DRIVER_MAJOR}.x"
+
+if [[ "$DRIVER_MAJOR" -ge 555 ]] && [[ "$GPU_CC_RAW" -ge 80 ]]; then
+    USE_DS_EVO_ATTN=true
+    USE_CUEQ_ATTN=true
+    LAYERNORM_PREFIX="export LAYERNORM_TYPE=fast_layernorm && "
+    echo "GPU compat: sm_${GPU_CC_RAW} + driver ${DRIVER_MAJOR} — JIT custom kernels enabled"
+else
+    USE_DS_EVO_ATTN=false
+    USE_CUEQ_ATTN=false
+    LAYERNORM_PREFIX=""
+    echo "GPU compat: sm_${GPU_CC_RAW} + driver ${DRIVER_MAJOR} — JIT custom kernels disabled"
+fi
+
+run_in_container() {
+    local need_gpu="$1"; shift
+    local cmd
+    cmd="cd /app/CryoZeta && source /opt/conda/etc/profile.d/conda.sh && conda activate cryozeta && "
+    cmd+="_CRYOZETA_CACHE=\"\${CRYOZETA_CACHE_DIR:-\${HOME}/.cache/cryozeta}\" && "
+    cmd+="export CRYOZETA_CACHE_DIR=\"\${_CRYOZETA_CACHE}\" && "
+    cmd+="export TORCH_EXTENSIONS_DIR=\"\${TORCH_EXTENSIONS_DIR:-\${_CRYOZETA_CACHE}/torch_extensions}\" && "
+    cmd+="export TORCHINDUCTOR_CACHE_DIR=\"\${TORCHINDUCTOR_CACHE_DIR:-\${_CRYOZETA_CACHE}/inductor}\" && "
+    cmd+="export TRITON_CACHE_DIR=\"\${TRITON_CACHE_DIR:-\${_CRYOZETA_CACHE}/triton}\" && "
+    cmd+="export TMPDIR=\"\${TMPDIR:-\${_CRYOZETA_CACHE}/tmp}\" && "
+    cmd+="unset _CRYOZETA_CACHE && "
+    cmd+="mkdir -p \"\${CRYOZETA_CACHE_DIR}/triton-bin\" \"\${TORCH_EXTENSIONS_DIR}\" \"\${TORCHINDUCTOR_CACHE_DIR}\" \"\${TRITON_CACHE_DIR}\" \"\${TMPDIR}\" && "
+    cmd+="[ -x \"\${CONDA_PREFIX}/bin/ptxas\" ] && export TRITON_PTXAS_PATH=\"\${CONDA_PREFIX}/bin/ptxas\" TRITON_PTXAS_BLACKWELL_PATH=\"\${CONDA_PREFIX}/bin/ptxas\" ; "
+    cmd+="ln -sf \"\${CONDA_PREFIX}/bin/ptxas\" \"\${CRYOZETA_CACHE_DIR}/triton-bin/ptxas-blackwell\" && export PATH=\"\${CRYOZETA_CACHE_DIR}/triton-bin:\$PATH\" && $*"
+
+    if [[ "$need_gpu" == "gpu" ]]; then
+        CUDA_VISIBLE_DEVICES="$GPU" apptainer exec --nv \
+            "${BIND_FLAGS[@]}" \
+            "$SIF" bash -c "$cmd"
+    else
+        apptainer exec \
+            "${BIND_FLAGS[@]}" \
+            "$SIF" bash -c "$cmd"
+    fi
+}
+
+REGISTRATION_FLAG=""
+case "$REGISTRATION_METHOD" in
+    auto)
+        echo "Registration method: auto (compare TEASER++ and SVD, or VESPER if requested via flags)"
+        ;;
+    teaser)
+        REGISTRATION_FLAG="--teaser_only"
+        echo "Registration method: TEASER++ only"
+        ;;
+    svd)
+        REGISTRATION_FLAG="--svd_only"
+        echo "Registration method: SVD only"
+        ;;
+    vesper)
+        REGISTRATION_FLAG="--vesper_only"
+        echo "Registration method: VESPER only"
+        ;;
+    *)
+        echo "Error: unknown registration method '$REGISTRATION_METHOD' (expected auto|teaser|svd|vesper)" >&2
+        exit 1
+        ;;
+esac
+
+VESPER_ARGS=()
+if [[ -n "$VESPER_MAP_PATH" ]]; then
+    VESPER_MAP_PATH="$(realpath "$VESPER_MAP_PATH")"
+    VESPER_ARGS+=("--vesper_map_path" "$VESPER_MAP_PATH")
+fi
+if [[ -n "$VESPER_CONTOUR_LEVEL" ]]; then
+    VESPER_ARGS+=("--vesper_contour_level" "$VESPER_CONTOUR_LEVEL")
+fi
+VESPER_ARGS+=("--vesper_resolution" "$VESPER_RESOLUTION" "--vesper_angle_spacing" "$VESPER_ANGLE_SPACING")
+
+echo "Selected sample: $SELECTED_NAME"
+echo "Output directory: $OUTPUT_DIR"
+
+if [[ "$SKIP_DETECTION" == false ]]; then
+    echo "=== Stage 1: Atom Detection ==="
+    run_in_container gpu \
+        cryozeta-detection json-run "$SELECTED_INPUT_JSON" "$OUTPUT_DIR" --device cuda --overwrite
+fi
+
+echo "=== Stage 2: Large-Complex Cycle Prediction ==="
+run_in_container gpu \
+    "${LAYERNORM_PREFIX}cryozeta-cycle-predict \
+        --base_json $CYCLE_INPUT_JSON \
+        --load_checkpoint_path /app/CryoZeta/assets/cryozeta-v0.0.1.safetensors \
+        --dump_root $OUTPUT_DIR \
+        --threshold $THRESHOLD \
+        --seeds $SEEDS \
+        --sample.N_sample $N_SAMPLE \
+        --sample.N_step $N_STEP \
+        --model.N_cycle $N_CYCLE \
+        --use_deepspeed_evo_attention $USE_DS_EVO_ATTN \
+        --use_cuequivariance_attention $USE_CUEQ_ATTN \
+        --use_cuequivariance_multiplicative_update $USE_CUEQ_ATTN \
+        --use_cuequivariance_attention_pair_bias false \
+        --use_opm_tilelang $USE_OPM_TILELANG \
+        --num_dl_workers 1 \
+        ${REGISTRATION_FLAG} ${VESPER_ARGS[*]}"
+
+if [[ "$SKIP_COMBINE" == false ]]; then
+    echo "=== Stage 3: Combine Stages ==="
+    run_in_container cpu \
+        "cryozeta-combine-stages --dump_root $OUTPUT_DIR --output $OUTPUT_DIR/combined.cif"
+fi
+
+echo "=== Done. Results in: $OUTPUT_DIR ==="