#!/usr/bin/env bash
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
# Forked: replaces the upstream `kimodo_demo` Viser launch with a Gradio API
# (server.py) used programmatically by the GengaMachines webapp.
set -euo pipefail

cd /workspace

# pre-download checkpoints. We use SOMA-RP-v1.1 (per the integration plan); v1
# is downloaded too for parity with upstream.
python - <<'PY'
from huggingface_hub import snapshot_download
for repo in ("nvidia/Kimodo-SOMA-RP-v1", "nvidia/Kimodo-SOMA-RP-v1.1"):
    print(f"snapshot_download({repo}) ...")
    snapshot_download(repo)
print("Checkpoint download complete.")
PY

# launch text encoder (internal-only, port 9550)
echo "Starting text-encoder on :9550 ..."
kimodo_textencoder &
TEXT_ENCODER_PID=$!

cleanup() {
  echo "Shutting down text-encoder (pid=${TEXT_ENCODER_PID}) ..."
  kill "${TEXT_ENCODER_PID}" >/dev/null 2>&1 || true
}
trap cleanup EXIT

# wait for the text encoder to be healthy
echo "Waiting for text-encoder health ..."
for i in $(seq 1 1200); do
  if curl -fsS "http://127.0.0.1:9550/" >/dev/null 2>&1; then
    echo "Text-encoder is up."
    break
  fi
  sleep 1
  if [[ $i -eq 1200 ]]; then
    echo "ERROR: text-encoder did not become healthy on http://127.0.0.1:9550/ within 1200s" >&2
    exit 1
  fi
done

# launch our Gradio API in place of kimodo_demo
echo "Starting Gradio API on :7860 ..."
exec python -u /workspace/server.py