init.sh

#!/usr/bin/env bash
# ==============================================================================
# init.sh — One-shot bootstrap for EditScore on this dev machine.
#
# Idempotent: re-running just verifies / fills in the missing pieces.
#
# Usage examples
# --------------
#   bash init.sh                                    # env + benchmark deps only
#   WITH_RL=1            bash init.sh               # + OmniGen2-RL python deps
#   WITH_RL=1 WITH_RL_DATA=1   bash init.sh         # + download/prepare data + models (~140GB)
#   WITH_RL=1 WITH_RL_DATA=1 WITH_RL_RUN=1 bash init.sh  # + actually launch reward server
#                                                   #   and a 50-step smoke-test training
#
# Optional knobs
# --------------
#   ENV_NAME=editscore   conda env name
#   PY_VERSION=3.12      python version
#   PROXY=10.140.15.68:3128   intranet HTTP proxy
#   SKIP_PROXY=1         do not export http_proxy / https_proxy
#   SKIP_VLLM=1          skip vLLM (saves ~3GB; *_vllm scripts won't work)
#
# Verified on: NVIDIA L20Y x8 (CUDA 13.0 driver, backwards compatible with cu126)
# ==============================================================================

set -euo pipefail

# ---- 0. Config ---------------------------------------------------------------
ENV_NAME="${ENV_NAME:-editscore}"
PY_VERSION="${PY_VERSION:-3.12}"
PROXY="${PROXY:-10.140.15.68:3128}"
SKIP_VLLM="${SKIP_VLLM:-0}"
SKIP_PROXY="${SKIP_PROXY:-0}"
WITH_RL="${WITH_RL:-0}"
WITH_RL_DATA="${WITH_RL_DATA:-0}"
WITH_RL_RUN="${WITH_RL_RUN:-0}"

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
RL_DIR="${SCRIPT_DIR}/examples/OmniGen2-RL"
DATA_DIR="${SCRIPT_DIR}/datasets/EditScore-RL-Data"
MODELS_DIR="${SCRIPT_DIR}/pretrained_models"

cd "${SCRIPT_DIR}"

echo "==> EditScore env bootstrap"
echo "    workdir       : ${SCRIPT_DIR}"
echo "    env name      : ${ENV_NAME}"
echo "    python        : ${PY_VERSION}"
echo "    skip vllm     : ${SKIP_VLLM}"
echo "    skip proxy    : ${SKIP_PROXY}"
echo "    with rl       : ${WITH_RL}"
echo "    with rl data  : ${WITH_RL_DATA}"
echo "    with rl run   : ${WITH_RL_RUN}"
echo

# WITH_RL_RUN / WITH_RL_DATA imply WITH_RL
if [[ "${WITH_RL_RUN}" == "1" || "${WITH_RL_DATA}" == "1" ]]; then
  WITH_RL=1
fi

# ---- 1. Network proxy (intranet only; needs to reach pypi/HF) ----------------
if [[ "${SKIP_PROXY}" != "1" ]]; then
  export http_proxy="${PROXY}"
  export https_proxy="${PROXY}"
  echo "==> Proxy set to ${PROXY}"
fi

# ---- 2. Conda activation hook ------------------------------------------------
if ! command -v conda >/dev/null 2>&1; then
  echo "ERROR: conda not found in PATH. Install Miniconda/Anaconda first." >&2
  exit 1
fi
# shellcheck disable=SC1091
source "$(conda info --base)/etc/profile.d/conda.sh"

# ---- 3. Create / reuse conda env --------------------------------------------
if conda env list | awk '{print $1}' | grep -qx "${ENV_NAME}"; then
  echo "==> Conda env '${ENV_NAME}' already exists; reusing"
else
  echo "==> Creating conda env '${ENV_NAME}' (python=${PY_VERSION})"
  conda create -n "${ENV_NAME}" "python=${PY_VERSION}" -y
fi
conda activate "${ENV_NAME}"
echo "    using python: $(which python)  ($(python --version 2>&1))"

# ---- 4. PyTorch (must be installed BEFORE editscore + vllm) -----------------
# vLLM (step 6) pins its own torch (currently 2.10 + cu128). To avoid the
# downgrade-then-upgrade churn, install README-pinned torch 2.7.1+cu126 only
# when vLLM is being skipped OR when no torch is installed.
need_torch_install=1
if python -c "import torch" 2>/dev/null; then
  echo "==> torch already installed: $(python -c 'import torch;print(torch.__version__)')"
  need_torch_install=0
fi
if [[ "${SKIP_VLLM}" == "1" && "${need_torch_install}" == "1" ]]; then
  echo "==> Installing PyTorch 2.7.1 (cu126) — README pinned version"
  pip install --quiet \
    torch==2.7.1 torchvision \
    --extra-index-url https://download.pytorch.org/whl/cu126
elif [[ "${SKIP_VLLM}" != "1" && "${need_torch_install}" == "1" ]]; then
  echo "==> Skipping explicit torch install; vLLM will pull its own pinned torch"
fi

# ---- 5. EditScore + benchmark deps ------------------------------------------
echo "==> Installing editscore (editable) + benchmark deps"
pip install --quiet -e .
pip install --quiet -r requirements.txt

# ---- 6. vLLM (optional, for *_vllm backends) --------------------------------
if [[ "${SKIP_VLLM}" != "1" ]]; then
  echo "==> Installing vLLM (this will UPGRADE torch to vLLM's required version)"
  pip install --quiet -U vllm
else
  echo "==> Skipping vLLM (SKIP_VLLM=1). HF backends still work; *_vllm scripts will not."
fi

# ---- 7. (Optional) OmniGen2-RL python deps ----------------------------------
# flash-attn cannot be built from source quickly here (no nvcc) — pin a prebuilt
# wheel matching the active torch / cuda / python.
if [[ "${WITH_RL}" == "1" ]]; then
  echo "==> Installing OmniGen2-RL extra deps"
  if [[ "${SKIP_VLLM}" == "1" ]]; then
    echo "WARN: WITH_RL=1 but SKIP_VLLM=1 — flash-attn wheel below assumes vLLM's torch 2.10+cu128." >&2
  fi
  pip install --quiet \
    omegaconf diffusers einops wandb flask matplotlib pyyaml ninja

  # Verify flash_attn loads against the active torch ABI; reinstall if not.
  if ! python -c "import flash_attn; from flash_attn import flash_attn_func" 2>/dev/null; then
    FA_INFO=$(python - <<'PY'
import torch, sys
ver = torch.__version__.split("+")[0]
major_minor = ".".join(ver.split(".")[:2])         # e.g. "2.10"
abi = "TRUE" if torch._C._GLIBCXX_USE_CXX11_ABI else "FALSE"
py  = f"cp{sys.version_info.major}{sys.version_info.minor}"
print(major_minor, abi, py)
PY
)
    read -r FA_TORCH FA_ABI FA_PY <<<"${FA_INFO}"
    FA_URL="https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu12torch${FA_TORCH}cxx11abi${FA_ABI}-${FA_PY}-${FA_PY}-linux_x86_64.whl"
    echo "    Installing flash-attn wheel: ${FA_URL}"
    pip install --quiet --force-reinstall --no-deps "${FA_URL}" || {
      echo "ERROR: prebuilt flash-attn wheel not found for torch=${FA_TORCH} abi=${FA_ABI} py=${FA_PY}." >&2
      echo "       See https://github.com/Dao-AILab/flash-attention/releases for matching wheels." >&2
      exit 1
    }
  else
    echo "    flash-attn already importable: $(python -c 'import flash_attn;print(flash_attn.__version__)')"
  fi
fi

# ---- 8. (Optional) Download + prepare RL data and base models ---------------
# Total payload: ~140 GB on disk.
#   - EditScore-RL-Data dataset:    94 GB   (gzip parts) + 94 GB extracted
#   - OmniGen2 (incl. its VAE):     30 GB
#   - Qwen2.5-VL-3B-Instruct:        7 GB
#
# We deliberately DO NOT download FLUX.1-dev: it's a gated repo and OmniGen2
# itself ships an architecturally-identical 16-channel VAE under `vae/`.
if [[ "${WITH_RL_DATA}" == "1" ]]; then
  echo "==> Preparing RL data + base models"
  mkdir -p "${DATA_DIR}" "${MODELS_DIR}" "${SCRIPT_DIR}/logs"

  # 8a. Download dataset
  if [[ ! -f "${DATA_DIR}/rl.jsonl" ]]; then
    echo "    [8a] Downloading EditScore-RL-Data (~94GB)..."
    hf download EditScore/EditScore-RL-Data --repo-type dataset \
      --local-dir "${DATA_DIR}"
  else
    echo "    [8a] EditScore-RL-Data already downloaded"
  fi

  # 8b. Stream-extract images (avoids creating an intermediate 88GB tar.gz)
  if [[ ! -f "${DATA_DIR}/images/0.png" ]]; then
    echo "    [8b] Extracting images (cat images_part_* | tar -xz, can take a while)..."
    ( cd "${DATA_DIR}" && cat images_part_* | tar -xzf - )
  else
    img_count=$(ls "${DATA_DIR}/images" 2>/dev/null | wc -l)
    echo "    [8b] Images already extracted (${img_count} files)"
  fi

  # 8c. Convert relative -> absolute paths in jsonl, then filter to 9 tasks
  if [[ ! -f "${DATA_DIR}/rl_abs_9tasks.jsonl" ]]; then
    echo "    [8c] Processing rl.jsonl (abs paths + 9 tasks)..."
    cd "${RL_DIR}"
    python scripts/data/process_jsonl.py \
      --input  "${DATA_DIR}/rl.jsonl" \
      --output "${DATA_DIR}/rl_abs.jsonl" \
      --base-path "${DATA_DIR}" \
      --max-records 200000
    python scripts/data/extract_9_tasks.py \
      --input_path  "${DATA_DIR}/rl_abs.jsonl" \
      --output_path "${DATA_DIR}/rl_abs_9tasks.jsonl"
    cd "${SCRIPT_DIR}"
  else
    echo "    [8c] Processed jsonl already exists ($(wc -l < ${DATA_DIR}/rl_abs_9tasks.jsonl) lines)"
  fi

  # 8d. Patch data_configs/.../edit/all.yml to point at our processed jsonl
  ALL_YML="${RL_DIR}/data_configs/train/example/edit/all.yml"
  if grep -q "/path/to/EditScore-RL-Data" "${ALL_YML}" 2>/dev/null; then
    echo "    [8d] Patching ${ALL_YML} with absolute jsonl path"
    sed -i "s|/path/to/EditScore-RL-Data/rl_abs_9tasks.jsonl|${DATA_DIR}/rl_abs_9tasks.jsonl|" "${ALL_YML}"
  fi

  # 8e. Download OmniGen2 (transformer + bundled VAE + scheduler + processor + mllm)
  if [[ ! -f "${MODELS_DIR}/OmniGen2/transformer/diffusion_pytorch_model.safetensors.index.json" ]]; then
    echo "    [8e] Downloading OmniGen2 (~30GB)..."
    hf download OmniGen2/OmniGen2 --local-dir "${MODELS_DIR}/OmniGen2"
  else
    echo "    [8e] OmniGen2 already downloaded"
  fi

  # 8f. Download Qwen2.5-VL-3B-Instruct (text encoder used by train.py)
  if [[ ! -f "${MODELS_DIR}/Qwen2.5-VL-3B-Instruct/config.json" ]]; then
    echo "    [8f] Downloading Qwen2.5-VL-3B-Instruct (~7GB)..."
    hf download Qwen/Qwen2.5-VL-3B-Instruct \
      --local-dir "${MODELS_DIR}/Qwen2.5-VL-3B-Instruct"
  else
    echo "    [8f] Qwen2.5-VL-3B-Instruct already downloaded"
  fi

  # 8g. Extract OmniGen2 transformer state_dict to a single .bin
  # (train.py loads via raw torch.load; reading the safetensors shards directly
  #  requires some on-the-fly conversion that the script doesn't do.)
  BIN_PATH="${MODELS_DIR}/OmniGen2/transformer/pytorch_model.bin"
  if [[ ! -f "${BIN_PATH}" ]]; then
    echo "    [8g] Extracting OmniGen2 transformer .bin (~16GB)..."
    cd "${RL_DIR}"
    python - <<PY
import sys, os, torch
sys.path.insert(0, '.')
from omnigen2.models.transformers.transformer_omnigen2 import OmniGen2Transformer2DModel
local = '${MODELS_DIR}/OmniGen2'
print(f"Loading transformer from {local}")
m = OmniGen2Transformer2DModel.from_pretrained(local, subfolder='transformer')
torch.save(m.state_dict(), '${BIN_PATH}')
print("Saved", '${BIN_PATH}')
PY
    cd "${SCRIPT_DIR}"
  else
    echo "    [8g] OmniGen2 transformer .bin already extracted"
  fi

  # 8h. Symlink so train.py's relative `pretrained_models/...` resolves correctly
  # (train.py runs with cwd = examples/OmniGen2-RL/).
  if [[ ! -e "${RL_DIR}/pretrained_models" ]]; then
    echo "    [8h] Creating examples/OmniGen2-RL/pretrained_models -> ${MODELS_DIR}"
    ln -sfn "${MODELS_DIR}" "${RL_DIR}/pretrained_models"
  elif [[ -L "${RL_DIR}/pretrained_models" ]]; then
    echo "    [8h] examples/OmniGen2-RL/pretrained_models symlink already exists"
  else
    # original repo ships an empty pretrained_models/ dir with .gitkeep
    rm -rf "${RL_DIR}/pretrained_models"
    ln -sfn "${MODELS_DIR}" "${RL_DIR}/pretrained_models"
    echo "    [8h] Replaced empty pretrained_models/ dir with symlink"
  fi

  echo "    RL data + models ready."
fi

# ---- 9. (Optional) Smoke-test the full RL pipeline (50-step training run) ---
# Spins up:
#   GPU 0:    EditScore-Qwen3-VL-4B reward server (vLLM) + reward proxy
#   GPU 1-7:  7-way FSDP RL training of OmniGen2 (LoRA)
#
# Configs / scripts used (created alongside the upstream ones, not modifying them):
#   options/omnigen2_edit_rl_smoketest_7gpu.yml
#   scripts/train/omnigen2_edit_rl_smoketest_7gpu.sh
#   reward_server/server_configs/editscore_qwen3_4B_localhost.yml
#   reward_server/reward_proxy_singleserver.py    (drop-in replacement for
#                                                  reward_proxy.py that respects
#                                                  server.num_servers from yaml
#                                                  instead of the hard-coded
#                                                  `8 // tensor_parallel_size`)
#
# A 50-step run takes ~2.2 hours on 8x L20Y. The first step takes ~2.5 min
# (model loads, NCCL init, first reward server warmup); subsequent steps ~70 s.
if [[ "${WITH_RL_RUN}" == "1" ]]; then
  echo "==> Launching RL smoke-test (50-step training)"

  # Sanity-check that prerequisite files from step 8 exist
  for f in \
      "${RL_DIR}/options/omnigen2_edit_rl_smoketest_7gpu.yml" \
      "${RL_DIR}/scripts/train/omnigen2_edit_rl_smoketest_7gpu.sh" \
      "${RL_DIR}/reward_server/server_configs/editscore_qwen3_4B_localhost.yml" \
      "${RL_DIR}/reward_server/reward_proxy_singleserver.py" \
      "${MODELS_DIR}/OmniGen2/transformer/pytorch_model.bin" \
      "${DATA_DIR}/rl_abs_9tasks.jsonl"; do
    [[ -e "$f" ]] || { echo "ERROR: missing $f — re-run with WITH_RL_DATA=1" >&2; exit 1; }
  done

  mkdir -p "${RL_DIR}/logs"
  cd "${RL_DIR}"

  # Cleanup any stale reward_server / proxy from previous runs
  pkill -9 -f "reward_server.py.*editscore_qwen3_4B" 2>/dev/null || true
  pkill -9 -f "reward_proxy_singleserver"            2>/dev/null || true
  sleep 3

  # 9a. Start the reward server on GPU 0 (background, setsid so it survives this shell).
  # http_proxy is needed for the FIRST run (HF LoRA download); on re-runs the
  # weights come from cache and the proxy isn't strictly required. We export
  # both http_proxy and NO_PROXY=127.0.0.1 so the server can fetch from HF but
  # any internal localhost calls stay direct.
  echo "    [9a] Starting reward server on GPU 0..."
  setsid bash -c "
    source '$(conda info --base)/etc/profile.d/conda.sh' && conda activate ${ENV_NAME}
    export http_proxy='${PROXY}' https_proxy='${PROXY}'
    export NO_PROXY='127.0.0.1,localhost' no_proxy='127.0.0.1,localhost'
    cd '${RL_DIR}'
    CUDA_VISIBLE_DEVICES=0 VLLM_USE_V1=1 \\
      python reward_server/reward_server.py \\
        --port 18888 \\
        --config_path reward_server/server_configs/editscore_qwen3_4B_localhost.yml \\
        > logs/reward_server.log 2>&1
  " < /dev/null > /dev/null 2>&1 &
  disown

  # Wait for vLLM warm up (look for the "Starting VLM reward server" line).
  echo "    [9a] Waiting for reward server warm-up (~40-90s)..."
  for i in $(seq 1 60); do
    sleep 5
    if grep -q "Starting VLM reward server" "${RL_DIR}/logs/reward_server.log" 2>/dev/null; then
      echo "         reward server READY after $((i*5))s"
      break
    fi
    if [[ $i -eq 60 ]]; then
      echo "ERROR: reward server didn't come up in 5 minutes. See ${RL_DIR}/logs/reward_server.log" >&2
      exit 1
    fi
  done

  # 9b. Start the patched proxy (CPU only, very lightweight).
  # IMPORTANT: NO_PROXY must include 127.0.0.1, otherwise the proxy's calls to
  # http://127.0.0.1:18888 (the local reward_server) get routed through the
  # squid in http_proxy and return 503.
  echo "    [9b] Starting reward proxy..."
  setsid bash -c "
    source '$(conda info --base)/etc/profile.d/conda.sh' && conda activate ${ENV_NAME}
    unset http_proxy https_proxy
    cd '${RL_DIR}'
    python reward_server/reward_proxy_singleserver.py \\
      --config_path reward_server/server_configs/editscore_qwen3_4B_localhost.yml \\
      > logs/reward_proxy.log 2>&1
  " < /dev/null > /dev/null 2>&1 &
  disown
  sleep 4

  # 9c. Launch the 7-GPU FSDP training
  # The launcher itself sets NO_PROXY so the trainer's localhost reward calls
  # don't get routed through the squid (which would 503).
  echo "    [9c] Launching 7-GPU FSDP training (50 steps)..."
  echo "         training log: ${RL_DIR}/logs/train.log"
  echo "         reward log  : ${RL_DIR}/logs/reward_server.log"
  echo "         proxy log   : ${RL_DIR}/logs/reward_proxy.log"
  echo
  echo "    Tail the log to watch:  tail -f ${RL_DIR}/logs/train.log | tr '\\r' '\\n'"
  setsid bash -c "
    source '$(conda info --base)/etc/profile.d/conda.sh' && conda activate ${ENV_NAME}
    export http_proxy='${PROXY}' https_proxy='${PROXY}'
    cd '${RL_DIR}'
    bash scripts/train/omnigen2_edit_rl_smoketest_7gpu.sh \\
      > logs/train.log 2>&1
  " < /dev/null > /dev/null 2>&1 &
  disown

  # Don't block init.sh on training completion; just verify it actually started.
  sleep 30
  if pgrep -f "accelerate launch.*omnigen2" > /dev/null; then
    echo "    [9c] Training process running in the background."
    echo "         Stop it with:  pkill -9 -f 'accelerate launch.*omnigen2'"
    echo "         Stop reward stack: pkill -9 -f 'reward_(server|proxy_singleserver)'"
  else
    echo "ERROR: training process didn't start. See ${RL_DIR}/logs/train.log" >&2
    exit 1
  fi
  cd "${SCRIPT_DIR}"
fi

# ---- 10. Smoke test ---------------------------------------------------------
echo "==> Smoke test"
SCRIPT_DIR="${SCRIPT_DIR}" python - <<'PY'
import sys, importlib

mods = ["torch", "transformers", "peft", "qwen_vl_utils",
        "datasets", "dotenv", "json_repair", "editscore"]
for m in mods:
    importlib.import_module(m)

import torch
assert torch.cuda.is_available(), "CUDA not available — check driver / torch install"

# Verify the VLM classes used by mllm_tools are present in transformers
from transformers import (
    Qwen2_5_VLForConditionalGeneration,
    Qwen3VLForConditionalGeneration,
    AutoProcessor,
)

# Verify all backend modules import without side effects
from editscore.mllm_tools import qwen25vl, qwen3vl, openai
try:
    from editscore.mllm_tools import qwen25vl_vllm, qwen3vl_vllm  # needs vllm
    vllm_ok = True
except ImportError:
    vllm_ok = False

# Verify EditScore class wires the prompts
from editscore import EditScore
s = EditScore(backbone="openai", key="dummy",
              model_name_or_path="gpt-4.1", score_range=25)
assert s.SC_prompt and s.PQ_prompt

print(f"  torch       : {torch.__version__}  (cuda devices: {torch.cuda.device_count()})")
print(f"  vllm avail  : {vllm_ok}")
print("  EditScore   : OK")

# Optional: verify OmniGen2-RL stack if its deps are present
import os
rl_dir = os.path.join(os.environ.get("SCRIPT_DIR", "."), "examples", "OmniGen2-RL")
try:
    import flash_attn, omegaconf, diffusers, einops, flask, wandb  # noqa: F401
    sys.path.insert(0, rl_dir)
    from omnigen2.models.transformers.transformer_omnigen2 import OmniGen2Transformer2DModel  # noqa: F401
    from omnigen2.grpo.reward_client_edit import evaluate_images  # noqa: F401
    print(f"  OmniGen2-RL : OK  (flash_attn={flash_attn.__version__})")
except ImportError as e:
    print(f"  OmniGen2-RL : (deps not installed — re-run with WITH_RL=1)  [{type(e).__name__}: {e}]")
PY

echo
echo "=============================================================="
echo "✅ Environment '${ENV_NAME}' is ready."
echo
echo "To use it:"
echo "    conda activate ${ENV_NAME}"
if [[ "${SKIP_PROXY}" != "1" ]]; then
  echo "    export http_proxy=${PROXY} https_proxy=${PROXY}   # for HF model downloads"
fi
echo
echo "Then run e.g.:"
echo "    bash evaluate_qwen3_vl_4B.sh             # benchmark on EditReward-Bench"
if [[ "${WITH_RL}" == "1" ]]; then
  echo
  echo "OmniGen2-RL training (also automated above with WITH_RL_RUN=1):"
  echo "    cd ${RL_DIR}"
  echo "    # GPU 0 reward stack:"
  echo "    CUDA_VISIBLE_DEVICES=0 VLLM_USE_V1=1 \\"
  echo "      python reward_server/reward_server.py --port 18888 \\"
  echo "        --config_path reward_server/server_configs/editscore_qwen3_4B_localhost.yml &"
  echo "    python reward_server/reward_proxy_singleserver.py \\"
  echo "      --config_path reward_server/server_configs/editscore_qwen3_4B_localhost.yml &"
  echo "    # GPU 1-7 trainer:"
  echo "    bash scripts/train/omnigen2_edit_rl_smoketest_7gpu.sh"
fi
echo "=============================================================="