#!/usr/bin/env bash
# RunPod ComfyUI LTX-2.5 pod fixed installer by Aitrepreneur


set -euo pipefail

# ───────────────────── Config ─────────────────────

WORKSPACE="${WORKSPACE:-/workspace}"
COMFY_ROOT="${COMFY_ROOT:-$WORKSPACE/ComfyUI}"
HF_BASE="${HF_BASE:-https://huggingface.co/Aitrepreneur/FLX/resolve/main}"

PYTHON_BIN="${PYTHON_BIN:-python3}"
VENV_DIR="${VENV_DIR:-venv}"

# RunPod CUDA 12.1 safe stack
TORCH_VERSION="${TORCH_VERSION:-2.4.0}"
TORCHVISION_VERSION="${TORCHVISION_VERSION:-0.19.0}"
TORCHAUDIO_VERSION="${TORCHAUDIO_VERSION:-2.4.0}"
CUDA_TAG="${CUDA_TAG:-cu121}"
TORCH_INDEX="${TORCH_INDEX:-https://download.pytorch.org/whl/${CUDA_TAG}}"
XFORMERS_VERSION="${XFORMERS_VERSION:-0.0.27.post2}"

# LTXVideo pin from the workflow-compatible version
LTXVIDEO_REF="${LTXVIDEO_REF:-cd5d371518afb07d6b3641be8012f644f25269fc}"
PIN_LTXVIDEO_NODE="${PIN_LTXVIDEO_NODE:-true}"

# Behavior switches
INSTALL_MODELS="${INSTALL_MODELS:-true}"
INSTALL_NODE_REQUIREMENTS="${INSTALL_NODE_REQUIREMENTS:-true}"
INSTALL_ALL_NODE_REQUIREMENTS="${INSTALL_ALL_NODE_REQUIREMENTS:-false}"
START_AFTER_INSTALL="${START_AFTER_INSTALL:-true}"

# Important: this installer is intentionally persistent.
# It will only rebuild the venv if the user explicitly asks for it.
FORCE_RECREATE_VENV="${FORCE_RECREATE_VENV:-false}"
FORCE_UPDATE_NODES="${FORCE_UPDATE_NODES:-false}"

# Python pins
PIN_TRANSFORMERS="${PIN_TRANSFORMERS:-4.51.3}"
PIN_TOKENIZERS_RANGE="${PIN_TOKENIZERS_RANGE:->=0.21,<0.22}"
PIN_HF_HUB_RANGE="${PIN_HF_HUB_RANGE:->=0.25.2,<1.0}"
PIN_TIMM="${PIN_TIMM:-1.0.15}"
PIN_OPENCV_HEADLESS="${PIN_OPENCV_HEADLESS:-4.12.0.88}"
PIN_PILLOW_MIN="${PIN_PILLOW_MIN:-11.0.0}"
PIN_NUMPY_RANGE="${PIN_NUMPY_RANGE:->=1.26,<3}"
PIN_LIBROSA_VERSION="${PIN_LIBROSA_VERSION:-}"

LOG_DIR="${LOG_DIR:-$WORKSPACE/logs}"
LOG_FILE="${LOG_FILE:-$LOG_DIR/comfyui.log}"

# Required LTX node set. Keep this list in sync with the cloned nodes below.
REQUIRED_NODES="${REQUIRED_NODES:-ComfyUI-Manager ComfyUI-GGUF rgthree-comfy ComfyUI-Easy-Use ComfyUI-KJNodes RES4LYF ComfyUI-LTXVideo ComfyUI-Custom-Scripts ComfyUI-VideoHelperSuite ComfyUI-WanVideoWrapper ComfyUI-Impact-Pack Comfyui_TTP_Toolset ComfyMath WhatDreamsCost-ComfyUI}"

export PIP_DISABLE_PIP_VERSION_CHECK=1
export PIP_ROOT_USER_ACTION=ignore
export PYTHONNOUSERSITE=1
unset PYTHONPATH || true

# ───────────────────── Helpers ─────────────────────

[[ "$(id -u)" -eq 0 ]] && SUDO="" || SUDO="sudo"

log() {
  echo
  echo "============================================================"
  echo "$1"
  echo "============================================================"
}

warn() {
  echo "[WARN] $*"
}

die() {
  echo
  echo "[ERROR] $*"
  exit 1
}

need_pkg() {
  local pkg="$1"
  local cmd="${2:-$1}"

  if command -v "$cmd" >/dev/null 2>&1; then
    return 0
  fi

  echo "[INFO] Installing system package: $pkg"
  $SUDO apt-get update -y
  $SUDO apt-get install -y "$pkg"
}

grab() {
  local target="$1"
  local url="$2"

  if [[ -f "$target" ]]; then
    echo " [SKIP] $(basename "$target") already exists"
    return 0
  fi

  echo " [DL] $(basename "$target")"
  mkdir -p "$(dirname "$target")"
  curl -L --fail --progress-bar --show-error -o "$target" "$url"
}

get_node() {
  local dir="$1"
  local url="$2"
  local ref="${3:-}"
  local target="$COMFY_ROOT/custom_nodes/$dir"

  mkdir -p "$COMFY_ROOT/custom_nodes"

  if [[ -d "$target/.git" ]]; then
    echo " [SKIP] $dir already present"

    if [[ "$FORCE_UPDATE_NODES" == "true" ]]; then
      echo " [GIT] Updating $dir"
      git -C "$target" fetch --all --tags || true
      git -C "$target" pull --ff-only || true
    fi
  elif [[ -d "$target" ]]; then
    echo " [SKIP] $dir already exists but is not a git repo"
  else
    echo " [GIT] Cloning $dir"
    git clone "$url" "$target"
    git -C "$target" fetch --all --tags || true
  fi

  if [[ -n "$ref" && "$PIN_LTXVIDEO_NODE" == "true" && -d "$target/.git" ]]; then
    echo " [GIT] Checking out $dir -> $ref"
    git -C "$target" checkout "$ref" || warn "Could not checkout $dir to $ref"
  fi
}

write_constraints() {
  CONSTRAINT_FILE="/tmp/ait_ltx25_klein_style_constraints.txt"

  cat > "$CONSTRAINT_FILE" <<EOF
torch==${TORCH_VERSION}+${CUDA_TAG}
torchvision==${TORCHVISION_VERSION}+${CUDA_TAG}
torchaudio==${TORCHAUDIO_VERSION}+${CUDA_TAG}
xformers==${XFORMERS_VERSION}
numpy${PIN_NUMPY_RANGE}
transformers==${PIN_TRANSFORMERS}
tokenizers${PIN_TOKENIZERS_RANGE}
huggingface-hub${PIN_HF_HUB_RANGE}
timm==${PIN_TIMM}
opencv-python-headless==${PIN_OPENCV_HEADLESS}
Pillow>=${PIN_PILLOW_MIN}
SQLAlchemy>=2.0,<3
alembic>=1.13,<2
EOF

  echo "──────── Using pip constraints ────────"
  cat "$CONSTRAINT_FILE"
}

sanitize_requirements_file() {
  local input="$1"
  local output="$2"

  "$PYTHON" - "$input" "$output" <<'PY'
import re
import sys
from pathlib import Path

src = Path(sys.argv[1])
dst = Path(sys.argv[2])

blocked_prefixes = [
    "torch",
    "torchvision",
    "torchaudio",
    "xformers",
    "triton",
    "transformers",
    "tokenizers",
    "huggingface-hub",
    "huggingface_hub",
    "timm",
    "numpy",
    "opencv-python",
    "opencv-contrib-python",
    "opencv-python-headless",
    "opencv-contrib-python-headless",
    "cuda-toolkit",
    "cuda-bindings",
    "nvidia-",
    "sageattention",
    "flash-attn",
]

blocked_contains = [
    "github.com/facebookresearch/sam2",
]

def should_block(line: str) -> bool:
    stripped = line.strip()
    lower = stripped.lower()

    if not stripped or lower.startswith("#"):
        return False

    for item in blocked_contains:
        if item in lower:
            return True

    for prefix in blocked_prefixes:
        if re.match(rf"^(-e\s+)?{re.escape(prefix)}(\[|==|>=|<=|~=|!=|>|<|\s|$)", lower):
            return True

    return False

out = []
for raw in src.read_text(encoding="utf-8", errors="ignore").splitlines():
    if should_block(raw):
        out.append("# skipped by Aitrepreneur installer to protect Torch/CUDA stack: " + raw)
    else:
        out.append(raw)

dst.write_text("\n".join(out) + "\n", encoding="utf-8")
PY
}

install_req_sanitized() {
  local req="$1"
  local label="$2"

  [[ -f "$req" ]] || {
    echo " [SKIP] No requirements file for $label"
    return 0
  }

  local sanitized="/tmp/ait_ltx25_$(basename "$(dirname "$req")")_$(basename "$req").sanitized.txt"

  sanitize_requirements_file "$req" "$sanitized"

  echo " [REQ] Installing sanitized requirements for $label"
  echo "       Original:  $req"
  echo "       Sanitized: $sanitized"

  if [[ ! -s "$sanitized" ]]; then
    echo "       [SKIP] No safe packages left after filtering"
    return 0
  fi

  if ! "$PYTHON" -m pip install --no-input --prefer-binary \
    --upgrade-strategy only-if-needed \
    --constraint "$CONSTRAINT_FILE" \
    -r "$sanitized"; then

    warn "First attempt failed for $label. Retrying without --no-build-isolation."

    "$PYTHON" -m pip install --no-input --prefer-binary \
      --upgrade-strategy only-if-needed \
      --constraint "$CONSTRAINT_FILE" \
      -r "$sanitized" || warn "Failed requirements for $label. Continuing."
  fi
}

install_torch_stack() {
  log "Installing locked Torch cu121 stack"

  # Do not uninstall unless needed. This makes the installer safer after pod restarts.
  "$PYTHON" -m pip install --no-input --prefer-binary \
    --upgrade-strategy only-if-needed \
    --index-url "$TORCH_INDEX" \
    --extra-index-url https://pypi.org/simple \
    "torch==${TORCH_VERSION}+${CUDA_TAG}" \
    "torchvision==${TORCHVISION_VERSION}+${CUDA_TAG}" \
    "torchaudio==${TORCHAUDIO_VERSION}+${CUDA_TAG}"

  "$PYTHON" -m pip install --no-input --prefer-binary --no-deps \
    "xformers==${XFORMERS_VERSION}"
}

install_core_pins() {
  log "Installing core pinned packages"

  "$PYTHON" -m pip install --no-input --prefer-binary \
    --upgrade-strategy only-if-needed \
    --constraint "$CONSTRAINT_FILE" \
    "numpy${PIN_NUMPY_RANGE}" \
    "pillow>=${PIN_PILLOW_MIN}" \
    "opencv-python-headless==${PIN_OPENCV_HEADLESS}" \
    "safetensors>=0.4.3" \
    "huggingface-hub${PIN_HF_HUB_RANGE}" \
    "accelerate>=0.34.0" \
    "SQLAlchemy>=2.0,<3" \
    "alembic>=1.13,<2" \
    "filelock" \
    "aiohttp" \
    "aiofiles" \
    "psutil" \
    "packaging" \
    "pyyaml" \
    "regex" \
    "requests>=2.32.3,<3" \
    "charset-normalizer>=2,<4" \
    "chardet<6" \
    "fsspec<=2026.4.0,>=2023.1.0" \
    "tqdm" \
    "typing_extensions" \
    "einops" \
    "sentencepiece" \
    "protobuf" \
    "av" \
    "imageio" \
    "imageio-ffmpeg" \
    "soundfile" \
    "scipy"

  echo " [PIN] Installing Transformers without dependencies"
  "$PYTHON" -m pip install --no-input --no-deps \
    "transformers==${PIN_TRANSFORMERS}"

  echo " [PIN] Installing Tokenizers with constraints"
  "$PYTHON" -m pip install --no-input --prefer-binary \
    --constraint "$CONSTRAINT_FILE" \
    "tokenizers${PIN_TOKENIZERS_RANGE}"

  echo " [PIN] Installing timm without dependencies"
  "$PYTHON" -m pip install --no-input --no-deps \
    "timm==${PIN_TIMM}"
}

install_base_extras() {
  log "Installing useful LTX/ComfyUI extras"

  "$PYTHON" -m pip install --no-input --prefer-binary \
    --upgrade-strategy only-if-needed \
    --constraint "$CONSTRAINT_FILE" \
    boto3 \
    rotary-embedding-torch \
    deepdiff \
    py-cpuinfo \
    diffusers \
    gguf \
    piexif \
    lark \
    matplotlib

  if [[ -n "$PIN_LIBROSA_VERSION" ]]; then
    "$PYTHON" -m pip install --no-input --prefer-binary \
      --upgrade-strategy only-if-needed \
      --constraint "$CONSTRAINT_FILE" \
      "librosa==${PIN_LIBROSA_VERSION}"
  else
    "$PYTHON" -m pip install --no-input --prefer-binary \
      --upgrade-strategy only-if-needed \
      --constraint "$CONSTRAINT_FILE" \
      librosa
  fi
}

protect_comfyui_manager_pins() {
  log "Protecting pinned packages from ComfyUI-Manager"

  local manager_dir="$COMFY_ROOT/user/__manager"
  mkdir -p "$manager_dir"

  cat > "$manager_dir/pip_blacklist.list" <<EOF
torch
torchvision
torchaudio
xformers
transformers
tokenizers
huggingface-hub
huggingface_hub
numpy
opencv-python
opencv-contrib-python
opencv-python-headless
opencv-contrib-python-headless
nvidia-
cuda-toolkit
cuda-bindings
triton
torchcodec
sageattention
flash-attn
EOF

  cat > "$manager_dir/pip_auto_fix.list" <<EOF
torch==${TORCH_VERSION}+${CUDA_TAG}
torchvision==${TORCHVISION_VERSION}+${CUDA_TAG}
torchaudio==${TORCHAUDIO_VERSION}+${CUDA_TAG}
xformers==${XFORMERS_VERSION}
transformers==${PIN_TRANSFORMERS}
tokenizers${PIN_TOKENIZERS_RANGE}
huggingface-hub${PIN_HF_HUB_RANGE}
numpy${PIN_NUMPY_RANGE}
opencv-python-headless==${PIN_OPENCV_HEADLESS}
SQLAlchemy>=2.0,<3
alembic>=1.13,<2
filelock
EOF

  if [[ -f "$manager_dir/config.ini" ]]; then
    sed -i 's/^always_lazy_install *= *.*/always_lazy_install = False/' "$manager_dir/config.ini" || true
    sed -i 's/^use_uv *= *.*/use_uv = False/' "$manager_dir/config.ini" || true
    sed -i 's/^network_mode *= *.*/network_mode = public/' "$manager_dir/config.ini" || true
  else
    cat > "$manager_dir/config.ini" <<EOF
[default]
always_lazy_install = False
use_uv = False
network_mode = public
security_level = normal
EOF
  fi
}

verify_torch_stack() {
  log "Verifying Torch stack"

  "$PYTHON" - <<PY
import torch
import torchvision
import torchaudio

print("torch:", torch.__version__)
print("torch cuda:", torch.version.cuda)
print("cuda available:", torch.cuda.is_available())
print("torchvision:", torchvision.__version__)
print("torchaudio:", torchaudio.__version__)

if not torch.__version__.startswith("${TORCH_VERSION}"):
    raise SystemExit(f"[ERROR] Wrong torch version: {torch.__version__}")

if torch.version.cuda != "12.1":
    raise SystemExit(f"[ERROR] Wrong torch CUDA version: {torch.version.cuda}, expected 12.1")

if not torch.cuda.is_available():
    raise SystemExit("[ERROR] CUDA is not available. This usually means wrong Torch build or incompatible RunPod image.")

print("gpu:", torch.cuda.get_device_name(0))
print("[OK] Torch stack is good")
PY
}

verify_no_bad_cuda_packages() {
  log "Checking for bad CUDA 13 packages"

  if "$PYTHON" -m pip freeze | grep -Ei "nvidia-.*cu13|torch==.*cu13" >/tmp/ait_ltx25_bad_cuda.txt; then
    cat /tmp/ait_ltx25_bad_cuda.txt
    die "Detected CUDA 13 packages. This would break older RunPod drivers."
  fi

  echo "[OK] No CUDA 13 packages detected"
}

verify_core_imports() {
  log "Verifying ComfyUI/LTX core imports"

  cd "$COMFY_ROOT"

  "$PYTHON" - <<'PY'
import sys
import torch
import torchvision
import torchaudio
import sqlalchemy
import alembic
import filelock
import aiohttp
import yaml
import safetensors
import einops
import PIL
import cv2
import transformers
import tokenizers
import timm

print("Python:", sys.executable)
print("Torch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
print("TorchVision:", torchvision.__version__)
print("TorchAudio:", torchaudio.__version__)
print("SQLAlchemy:", sqlalchemy.__version__)
print("Alembic: OK")
print("FileLock: OK")
print("aiohttp: OK")
print("PyYAML: OK")
print("Safetensors: OK")
print("Einops: OK")
print("Pillow:", PIL.__version__)
print("OpenCV:", cv2.__version__)
print("Transformers:", transformers.__version__)
print("Tokenizers:", tokenizers.__version__)
print("timm:", timm.__version__)

# These are the exact modules that caused the broken port 3000 issue.
from app.database.db import create_session
import comfy.ldm.lightricks.vae.audio_vae

print("ComfyUI database import: OK")
print("Lightricks audio VAE import: OK")
print("[OK] Core imports are good")
PY
}

verify_ltx_node_import() {
  log "Verifying ComfyUI-LTXVideo node import"

  cd "$COMFY_ROOT"

  "$PYTHON" - <<'PY'
import importlib.util
import sys
from pathlib import Path

root = Path.cwd()
pkg_dir = root / "custom_nodes" / "ComfyUI-LTXVideo"
init_file = pkg_dir / "__init__.py"

if not init_file.exists():
    raise SystemExit("[ERROR] ComfyUI-LTXVideo __init__.py not found")

sys.path.insert(0, str(root))
sys.path.insert(0, str(root / "custom_nodes"))

spec = importlib.util.spec_from_file_location(
    "ComfyUI_LTXVideo",
    init_file,
    submodule_search_locations=[str(pkg_dir)],
)

module = importlib.util.module_from_spec(spec)
sys.modules["ComfyUI_LTXVideo"] = module
spec.loader.exec_module(module)

print("[OK] ComfyUI-LTXVideo imports successfully")
PY
}

write_launch_scripts() {
  log "Writing RunPod start/stop scripts"

  mkdir -p "$LOG_DIR"

  cat > "$WORKSPACE/START_COMFYUI_LTX25.sh" <<START
#!/usr/bin/env bash
set -euo pipefail

COMFY_ROOT="$COMFY_ROOT"
VENV_PY="\$COMFY_ROOT/$VENV_DIR/bin/python"
LOG_DIR="$LOG_DIR"
LOG_FILE="$LOG_FILE"

mkdir -p "\$LOG_DIR"

if [[ ! -x "\$VENV_PY" ]]; then
  echo "ERROR: ComfyUI venv python not found at \$VENV_PY"
  echo "Run the LTX installer again from inside /workspace/ComfyUI."
  exit 1
fi

cd "\$COMFY_ROOT"

if ! "\$VENV_PY" - <<'PY' >/tmp/ltx25_start_check.log 2>&1
import torch
import torchaudio
import sqlalchemy
import filelock
from app.database.db import create_session
import comfy.ldm.lightricks.vae.audio_vae
print("start check ok")
PY
then
  echo "ERROR: Python dependency check failed."
  cat /tmp/ltx25_start_check.log
  echo
  echo "Run the LTX installer again. It will repair the existing venv without deleting models."
  exit 1
fi

echo "Stopping old ComfyUI processes if any..."
pkill -f "$COMFY_ROOT/main.py" || true
pkill -f "ComfyUI/main.py" || true
sleep 2

echo "Starting ComfyUI on 0.0.0.0:3000"
nohup "\$VENV_PY" "\$COMFY_ROOT/main.py" \\
  --listen 0.0.0.0 \\
  --port 3000 \\
  > "\$LOG_FILE" 2>&1 &

echo "ComfyUI launched. Log: \$LOG_FILE"
echo "Open RunPod port 3000 after 10 to 20 seconds."
sleep 8
tail -n 160 "\$LOG_FILE" || true
START

  cat > "$WORKSPACE/STOP_COMFYUI_LTX25.sh" <<STOP
#!/usr/bin/env bash
set -euo pipefail

echo "Stopping ComfyUI..."
pkill -f "$COMFY_ROOT/main.py" || true
pkill -f "ComfyUI/main.py" || true
echo "Done."
STOP

  chmod +x "$WORKSPACE/START_COMFYUI_LTX25.sh" "$WORKSPACE/STOP_COMFYUI_LTX25.sh"
}

# ───────────────────── Main ─────────────────────

log "LTX-2.5 V2 RunPod installer, Klein-style persistent mode"

# Be more idiot-proof than Klein: if the user runs from /workspace, auto-enter ComfyUI.
if [[ ! -f "main.py" || ! -d "models" || ! -d "custom_nodes" ]]; then
  if [[ -f "$COMFY_ROOT/main.py" && -d "$COMFY_ROOT/models" && -d "$COMFY_ROOT/custom_nodes" ]]; then
    cd "$COMFY_ROOT"
  else
    die "This installer needs an existing RunPod ComfyUI template at $COMFY_ROOT. Deploy a ComfyUI RunPod template first, then run this installer."
  fi
fi

COMFY_ROOT="$(pwd)"
echo "ComfyUI root: $COMFY_ROOT"

action_text="This installer will NOT delete ComfyUI, models, custom nodes, or venv unless FORCE_RECREATE_VENV=true."
echo "$action_text"

log "Checking system packages"
need_pkg curl
need_pkg git
need_pkg git-lfs git-lfs
need_pkg ffmpeg
need_pkg python3-venv
need_pkg python3-pip
need_pkg build-essential gcc
need_pkg libgl1
need_pkg libglib2.0-0

git lfs install || true

log "Preparing persistent venv"

if [[ "$FORCE_RECREATE_VENV" == "true" ]]; then
  echo " [CLEAN] FORCE_RECREATE_VENV=true, deleting $COMFY_ROOT/$VENV_DIR"
  rm -rf "$COMFY_ROOT/$VENV_DIR"
fi

if [[ ! -d "$COMFY_ROOT/$VENV_DIR" ]]; then
  echo " [VENV] Creating $COMFY_ROOT/$VENV_DIR"
  "$PYTHON_BIN" -m venv "$COMFY_ROOT/$VENV_DIR"
fi

if [[ -f "$COMFY_ROOT/$VENV_DIR/pyvenv.cfg" ]]; then
  sed -i 's/^include-system-site-packages = .*/include-system-site-packages = false/' "$COMFY_ROOT/$VENV_DIR/pyvenv.cfg" || true
fi

# shellcheck disable=SC1091
source "$COMFY_ROOT/$VENV_DIR/bin/activate"

PYTHON="$(command -v python)"
PIP="$(command -v pip)"

echo "Python: $PYTHON"
echo "Pip:    $PIP"

"$PYTHON" -m ensurepip --upgrade || true
"$PYTHON" -m pip install --no-input --upgrade pip setuptools wheel

write_constraints

if [[ "$INSTALL_MODELS" == "true" ]]; then
  log "Downloading LTX-2.5 model files"

  mkdir -p \
    models/text_encoders \
    models/vae \
    models/diffusion_models \
    models/latent_upscale_models \
    models/loras

  grab "models/text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors" \
       "$HF_BASE/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors?download=true"

  grab "models/vae/ltx-2.5-audio-vae-bf16.safetensors" \
       "$HF_BASE/ltx-2.5-audio-vae-bf16.safetensors?download=true"

  grab "models/vae/ltx-2.5-video-vae-conv-bf16.safetensors" \
       "$HF_BASE/ltx-2.5-video-vae-conv-bf16.safetensors?download=true"

  grab "models/diffusion_models/ltx-2.5-22b-distilled-transformer-comfy-int8-convrot.safetensors" \
       "$HF_BASE/ltx-2.5-22b-distilled-transformer-comfy-int8-convrot.safetensors?download=true"

  grab "models/latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors" \
       "$HF_BASE/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors?download=true"

  grab "models/loras/ltx-2-19b-ic-lora-detailer.safetensors" \
       "$HF_BASE/ltx-2-19b-ic-lora-detailer.safetensors?download=true"
else
  echo " [SKIP] INSTALL_MODELS=false"
fi

log "Cloning missing custom nodes"

get_node "ComfyUI-Manager"          "https://github.com/ltdrdata/ComfyUI-Manager.git"
get_node "ComfyUI-GGUF"             "https://github.com/city96/ComfyUI-GGUF.git"
get_node "rgthree-comfy"            "https://github.com/rgthree/rgthree-comfy.git"
get_node "ComfyUI-Easy-Use"         "https://github.com/yolain/ComfyUI-Easy-Use.git"
get_node "ComfyUI-KJNodes"          "https://github.com/kijai/ComfyUI-KJNodes.git"
get_node "RES4LYF"                  "https://github.com/ClownsharkBatwing/RES4LYF.git"
get_node "ComfyUI-LTXVideo"         "https://github.com/Lightricks/ComfyUI-LTXVideo.git" "$LTXVIDEO_REF"
get_node "ComfyUI-Custom-Scripts"   "https://github.com/pythongosssss/ComfyUI-Custom-Scripts.git"
get_node "ComfyUI-VideoHelperSuite" "https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite.git"
get_node "ComfyUI-WanVideoWrapper"  "https://github.com/kijai/ComfyUI-WanVideoWrapper.git"
get_node "ComfyUI-Impact-Pack"      "https://github.com/ltdrdata/ComfyUI-Impact-Pack.git"
get_node "Comfyui_TTP_Toolset"      "https://github.com/TTPlanetPig/Comfyui_TTP_Toolset.git"
get_node "ComfyMath"                "https://github.com/evanspearman/ComfyMath.git"
get_node "WhatDreamsCost-ComfyUI"   "https://github.com/WhatDreamsCost/WhatDreamsCost-ComfyUI.git"

install_torch_stack
verify_torch_stack
verify_no_bad_cuda_packages

install_core_pins

log "Installing ComfyUI core requirements safely"
install_req_sanitized "$COMFY_ROOT/requirements.txt" "ComfyUI core"
install_req_sanitized "$COMFY_ROOT/manager_requirements.txt" "ComfyUI manager requirements"

install_core_pins
verify_torch_stack
verify_no_bad_cuda_packages

if [[ "$INSTALL_NODE_REQUIREMENTS" == "true" ]]; then
  log "Installing custom node requirements safely"

  if [[ "$INSTALL_ALL_NODE_REQUIREMENTS" == "true" ]]; then
    while IFS= read -r req; do
      install_req_sanitized "$req" "$req"
    done < <(find "$COMFY_ROOT/custom_nodes" -maxdepth 2 -name requirements.txt -print)
  else
    for node_dir in $REQUIRED_NODES; do
      install_req_sanitized "$COMFY_ROOT/custom_nodes/$node_dir/requirements.txt" "$node_dir"
    done
  fi
else
  echo " [SKIP] INSTALL_NODE_REQUIREMENTS=false"
fi

install_base_extras
install_core_pins

verify_torch_stack
verify_no_bad_cuda_packages
verify_core_imports
verify_ltx_node_import
protect_comfyui_manager_pins
write_launch_scripts

log "Install complete"

echo "✅ LTX-2.5 is installed."
echo
echo "Start ComfyUI later with:"
echo "bash $WORKSPACE/START_COMFYUI_LTX25.sh"
echo
echo "Stop ComfyUI before stopping the pod with:"
echo "bash $WORKSPACE/STOP_COMFYUI_LTX25.sh"
echo
echo "Important: after stopping and starting the same pod, do NOT reinstall. Just open port 3000 or run the START script."

if [[ "$START_AFTER_INSTALL" == "true" ]]; then
  log "Starting ComfyUI"
  bash "$WORKSPACE/START_COMFYUI_LTX25.sh"
else
  echo "START_AFTER_INSTALL=false, not starting ComfyUI automatically."
fi