#!/usr/bin/env bash
# Place required weight files under MODELS_DIR with Power Pack names.
# SEALED CONFIG (2026-08-13): STOCK int8-convrot TE + bf16 pruned DiT for finals.
# The heretic TE is RETIRED (Heretic's author does not recommend heretic models
# for H3; stock int8 TE measured equal-or-better). int8-convrot DiT = draft tier.
set -euo pipefail
MODELS_DIR="${MODELS_DIR:-$HOME/keys-2k-minimax-h3-parallel/models}"
mkdir -p "$MODELS_DIR"/{diffusion_models,text_encoders,vae,upscale_models,loras}

echo "=== fetch / place weights → $MODELS_DIR ==="

hf_get() {  # hf_get <repo> <remote-path> <local-path>
  local repo="$1" remote="$2" local_p="$MODELS_DIR/$3"
  if [[ -f "$local_p" ]]; then echo "  present: $3"; return 0; fi
  echo "--> $3  ($repo)"
  if command -v hf >/dev/null; then
    hf download "$repo" "$remote" --local-dir /tmp/h3-dl >/dev/null || {
      echo "    download failed — place $3 manually"; return 0; }
    mkdir -p "$(dirname "$local_p")"
    mv -n "/tmp/h3-dl/$remote" "$local_p" 2>/dev/null || cp -n "/tmp/h3-dl/$remote" "$local_p"
  else
    echo "    install the hf CLI (pip install -U huggingface_hub) or place manually"
  fi
}

REPO="Comfy-Org/MiniMax-H3"

# Text encoder — STOCK int8-convrot (the sealed choice; heretic TE retired)
hf_get "$REPO" text_encoders/qwen3vl_32b_minimax_h3_int8_convrot.safetensors \
               text_encoders/qwen3vl_32b_minimax_h3_int8_convrot.safetensors

# DiT — FINALS tier: pruned bf16 (~40 GB each). Keyframes AND spans must use bf16;
# int8 keyframes poison a bf16 span (FLF spans inherit their anchors' look).
hf_get "$REPO" diffusion_models/minimax_h3_fl2va_pruned_bf16.safetensors \
               diffusion_models/minimax_h3_fl2va_pruned_bf16.safetensors
hf_get "$REPO" diffusion_models/minimax_h3_ref2va_pruned_bf16.safetensors \
               diffusion_models/minimax_h3_ref2va_pruned_bf16.safetensors

# DiT — DRAFT tier: int8-convrot (fast; visibly stylized — never for deliverables)
hf_get "$REPO" diffusion_models/minimax_h3_fl2va_pruned_int8_convrot.safetensors \
               diffusion_models/minimax_h3_fl2va_pruned_int8_convrot.safetensors
hf_get "$REPO" diffusion_models/minimax_h3_ref2va_pruned_int8_convrot.safetensors \
               diffusion_models/minimax_h3_ref2va_pruned_int8_convrot.safetensors

# RealESRGAN x2
if [[ ! -f "$MODELS_DIR/upscale_models/RealESRGAN_x2plus.pth" ]]; then
  echo "--> RealESRGAN_x2plus.pth"
  mkdir -p "$MODELS_DIR/upscale_models"
  curl -fL --retry 3 -o "$MODELS_DIR/upscale_models/RealESRGAN_x2plus.pth" \
    "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth" \
    || echo "    download failed — place RealESRGAN_x2plus.pth manually"
else
  echo "  ESRGAN present"
fi

cat <<EOF

Required (place or symlink into $MODELS_DIR):

  diffusion_models/minimax_h3_fl2va_pruned_bf16.safetensors        # FINALS (sealed)
  diffusion_models/minimax_h3_ref2va_pruned_bf16.safetensors       # FINALS talking-head/ref2va
  diffusion_models/minimax_h3_fl2va_pruned_int8_convrot.safetensors   # draft tier only
  diffusion_models/minimax_h3_ref2va_pruned_int8_convrot.safetensors  # draft tier only
  vae/minimax_h3_video_vae_fp16.safetensors
  vae/minimax_h3_audio_vae_fp32.safetensors
  text_encoders/qwen3vl_32b_minimax_h3_int8_convrot.safetensors    # STOCK TE (heretic retired)
  upscale_models/RealESRGAN_x2plus.pth

Optional:
  loras/h3-realism-people-t2v-i2v-r2v.safetensors

Upstream packs (layout may need rename):
  https://huggingface.co/Comfy-Org/MiniMax-H3
  https://huggingface.co/Kijai/MiniMax-H3_comfy

Lab shortcut (keyspark):
  rsync -a ~/h3-cotenancy/ComfyUI/models/ \$MODELS_DIR/
EOF

# Lab auto-fill if local tree exists
if [[ -d "$HOME/h3-cotenancy/ComfyUI/models/diffusion_models" ]]; then
  echo "--> detected local h3-cotenancy models — rsync (no overwrite of existing)"
  rsync -a --ignore-existing "$HOME/h3-cotenancy/ComfyUI/models/" "$MODELS_DIR/" || true
fi