#!/usr/bin/env bash
# Download pinned official weights and prepare the production worker stopped.
set -Eeuo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
ENV_FILE=${STACK_ENV:-/etc/mike-ai/stack.env}
MODEL_DIR=${QWEN_IMAGE_21_MODEL_DIR:-/data/models/Qwen-Image-2.1-ComfyUI}
PE_I2I_DIR=${QWEN_IMAGE_PE_I2I_MODEL_DIR:-/data/models/qwen-image-2.1-pe-i2i-q5}
PE_T2I_DIR=${QWEN_IMAGE_PE_T2I_MODEL_DIR:-/data/models/qwen-image-2.1-pe-t2i-q5}
BASE=https://huggingface.co/Comfy-Org/Qwen-Image-2.1/resolve/ace0edeb3791a594ddfa36ed5f41a178a394e921

download() {
  local relative=$1 expected=$2 target="$MODEL_DIR/$1"
  install -d -m 0755 "$(dirname "$target")"
  if [[ -f $target ]] && echo "$expected  $target" | sha256sum -c --status; then
    echo "OK: $relative"
    return
  fi
  curl -fL --retry 5 --retry-all-errors -C - -o "$target.part" "$BASE/$relative"
  echo "$expected  $target.part" | sha256sum -c --status
  mv "$target.part" "$target"
}

download_url() {
  local url=$1 target=$2 expected=$3
  install -d -m 0755 "$(dirname "$target")"
  if [[ -f $target ]] && echo "$expected  $target" | sha256sum -c --status; then
    echo "OK: $target"
    return
  fi
  curl -fL --retry 5 --retry-all-errors -C - -o "$target.part" "$url"
  echo "$expected  $target.part" | sha256sum -c --status
  mv "$target.part" "$target"
}

download diffusion_models/qwen_image_2.1_int8_convrot.safetensors cb74113cb03faecd79611b01fd7fd642f0aa60d6f0b95086abee214d75eaa57d
download text_encoders/qwen3vl_8b_int8_convrot.safetensors 8bfd0f6e12abf2d2d697ecc888e5e90b0d6741d6708f05799f53afa560452e8f
download vae/qwen_image_2.1_vae_bf16.safetensors bb21f7473051e1ac368515dd3f2e15cd44d7a11748ee8823e1ddca3e4876b7c9

# Official Qwen Image 2.1 Prompt Enhancers, quantized to Q5_K_M for the
# 12-GiB RTX 3060. I2I uses the separate multimodal projector; T2I does not.
download_url \
  https://huggingface.co/prithivMLmods/Qwen-Image-2.1-PE-I2I-GGUF/resolve/55b9c1a326599e142d59bcad8715d5601ccf8daa/Qwen-Image-2.1-PE-I2I.Q5_K_M.gguf \
  "$PE_I2I_DIR/Qwen-Image-2.1-PE-I2I.Q5_K_M.gguf" \
  cb71f71fe5fe5938570d65a7c403fbcb1a9065dff9dd9a021f58aec42785b84e
download_url \
  https://huggingface.co/prithivMLmods/Qwen-Image-2.1-PE-I2I-GGUF/resolve/55b9c1a326599e142d59bcad8715d5601ccf8daa/Qwen-Image-2.1-PE-I2I.mmproj-bf16.gguf \
  "$PE_I2I_DIR/Qwen-Image-2.1-PE-I2I.mmproj-bf16.gguf" \
  8dedb71dbc3092dc47de9108ad373d68a12854e2527d59bd9399601738f3bce1
download_url \
  https://huggingface.co/Qwen/Qwen-Image-2.1-PE-I2I/resolve/72927bc08afc99b7888ceb7d7d51a12db3700bbd/system_prompt.txt \
  "$PE_I2I_DIR/system_prompt.txt" \
  e378fea686a1431581ba4c654d332ae96adad633f144ae738ec8ce9c4fd66439
download_url \
  https://huggingface.co/prithivMLmods/Qwen-Image-2.1-PE-T2I-GGUF/resolve/e18d4a3e0830ab157770738b16830e6fcf5f57d4/Qwen-Image-2.1-PE-T2I.Q5_K_M.gguf \
  "$PE_T2I_DIR/Qwen-Image-2.1-PE-T2I.Q5_K_M.gguf" \
  749f5652fd6e8b760ca860091f7a5bffa254a7954203ae5c9bcdf5f93197702f
download_url \
  https://huggingface.co/Qwen/Qwen-Image-2.1-PE-T2I/resolve/f3ed7985c788ad75b3ab7223e0c4c51e2a43545b/system_prompt.txt \
  "$PE_T2I_DIR/system_prompt.txt" \
  a77c9a06c59b120741141d9514b95682bb8761d02bec49ca61def7b2b3d9fb99

compose=(docker compose --env-file "$ENV_FILE" -f "$ROOT_DIR/compose.yaml" --profile image --profile flux-standby)
"${compose[@]}" build image-worker
"${compose[@]}" up -d --build --no-deps profile-controller
"${compose[@]}" create image-worker image-prompt-enhancer-i2i \
  image-prompt-enhancer-t2i flux-image-worker
"${compose[@]}" stop image-worker image-prompt-enhancer-i2i \
  image-prompt-enhancer-t2i flux-image-worker
state=$(docker inspect -f '{{.State.Status}}' mike-ai-image-worker)
[[ $state == exited || $state == created ]]
echo "Qwen-Image-2.1, both prompt enhancers and the FLUX standby are prepared and stopped."