#!/usr/bin/env bash

# ==============================
# Generate metadata.jsonl before training
# Configure variables directly in this file.
# Environment variable override: QWEN_IMAGE_TYPE (edit-2509, edit-2511, layered)
# ==============================

echo "[QIE] Torch version check"
python - <<'PY'
try:
    import torch
    print(f"[QIE] torch: {torch.__version__}")
except Exception as e:
    print(f"[QIE] torch: not available ({e})")
PY

DATA_ROOT="/workspace/data"
DATASET_NAME=""

# Required inputs
CAPTION=""

IMAGE_FOLDER="image"

# CONTROL_FOLDER_0=""
# CONTROL_FOLDER_1=""
# CONTROL_FOLDER_2=""
# CONTROL_FOLDER_3=""
# CONTROL_FOLDER_4=""
# CONTROL_FOLDER_5=""
# CONTROL_FOLDER_6=""
# CONTROL_FOLDER_7=""

RUN_NAME="${DATASET_NAME%/}"
DATASET_DIR="${DATA_ROOT%/}/${DATASET_NAME}"

OUTPUT_DIR_BASE="/workspace/auto/train_LoRA"
DATASET_CONFIG="/workspace/auto/dataset_QIE.toml"
OUTPUT_JSON="${DATASET_DIR%/}/metadata.jsonl"

# Training hyperparameters (can be overridden by app)
LEARNING_RATE="1e-3"
NETWORK_DIM=4
SEED=42
MAX_TRAIN_EPOCHS=100
SAVE_EVERY_N_EPOCHS=10

# Model type (edit-2509, edit-2511, layered)
MODEL_VERSION="${QWEN_IMAGE_TYPE:-edit-2509}"
if [[ "$MODEL_VERSION" != "edit-2509" && "$MODEL_VERSION" != "edit-2511" && "$MODEL_VERSION" != "layered" ]]; then
  echo "[QIE] Unsupported QWEN_IMAGE_TYPE: $MODEL_VERSION, defaulting to edit-2509"
  MODEL_VERSION="edit-2509"
fi
if [[ "$MODEL_VERSION" == "layered" ]]; then
  DIT_FILENAME="qwen_image_layered_bf16.safetensors"
  VAE_FILENAME="qwen_image_layered_vae.safetensors"
else
  DIT_SUFFIX="${MODEL_VERSION#edit-}"
  DIT_FILENAME="qwen_image_edit_${DIT_SUFFIX}_bf16.safetensors"
  VAE_FILENAME="diffusion_pytorch_model.safetensors"
fi

# Build control args from folder names with auto-detect fallback
CONTROL_ARGS=()
for i in {0..7}; do
  var="CONTROL_FOLDER_${i}"
  folder_name=${!var}
  cpath=""
  if [[ -n "$folder_name" ]]; then
    cpath="${DATASET_DIR%/}/$folder_name"
  elif [[ -d "${DATASET_DIR%/}/control_${i}" ]]; then
    cpath="${DATASET_DIR%/}/control_${i}"
  elif [[ $i -eq 0 && -d "${DATASET_DIR%/}/control" ]]; then
    # Special fallback: allow single control folder named "control" for control_0
    cpath="${DATASET_DIR%/}/control"
  fi
  [[ -n "$cpath" ]] && CONTROL_ARGS+=("--control_dir_${i}" "$cpath")
done

# Sync dataset config's image_jsonl_file with OUTPUT_JSON if present
if [[ -f "$DATASET_CONFIG" ]]; then
  python - "$DATASET_CONFIG" "$OUTPUT_JSON" <<'PY'
import sys, re, os
path, out = sys.argv[1], sys.argv[2]
txt = open(path, 'r', encoding='utf-8').read()
base = os.path.dirname(path)
cache = os.path.join(base, 'cache').replace('\\\\', '/')

# Update image_jsonl_file
new = re.sub(r"(?m)^\s*image_jsonl_file\s*=.*$", f'image_jsonl_file = "{out}"', txt)
if new == txt and 'image_jsonl_file' not in txt:
    new = txt.rstrip('\n') + f"\nimage_jsonl_file = \"{out}\"\n"

# Update cache_directory to a writable folder under the config directory
if re.search(r"(?m)^\s*cache_directory\s*=", new):
    new = re.sub(r"(?m)^\s*cache_directory\s*=.*$", f'cache_directory = "{cache}"', new)
else:
    new = new.rstrip('\n') + f"\ncache_directory = \"{cache}\"\n"

open(path, 'w', encoding='utf-8').write(new)
print("[QIE] Updated {}: image_jsonl_file -> {}".format(path, out))
print("[QIE] Updated {}: cache_directory -> {}".format(path, cache))
PY
  mkdir -p "$(dirname "$DATASET_CONFIG")/cache"
else
  echo "[QIE] WARN: Dataset config not found at $DATASET_CONFIG. Ensure it points to $OUTPUT_JSON"
fi

cd /workspace/auto

if [[ "$MODEL_VERSION" == "layered" ]]; then
  echo "[QIE] Layered mode: skip metadata generation (expects prebuilt JSONL)."
else
  echo "[QIE] Generating metadata: $OUTPUT_JSON"
  python create_image_caption_json.py \
    -i "${DATASET_DIR%/}/${IMAGE_FOLDER}" \
    -c "$CAPTION" \
    -o "$OUTPUT_JSON" \
    --image-dir "${DATASET_DIR%/}/${IMAGE_FOLDER}" \
    "${CONTROL_ARGS[@]}"
fi

cd /musubi-tuner

python qwen_image_cache_latents.py  \
--dataset_config "$DATASET_CONFIG" \
--vae "/workspace/Qwen-Image_models/vae/${VAE_FILENAME}" \
--model_version "$MODEL_VERSION" \
--vae_spatial_tile_sample_min_size 16384

python qwen_image_cache_text_encoder_outputs.py  \
--dataset_config "$DATASET_CONFIG"  \
--text_encoder "/workspace/Qwen-Image_models/text_encoder/qwen_2.5_vl_7b.safetensors" \
--model_version "$MODEL_VERSION" \
--batch_size 16

accelerate launch src/musubi_tuner/qwen_image_train_network.py \
--model_version "$MODEL_VERSION" \
--dit "/workspace/Qwen-Image_models/dit/${DIT_FILENAME}" \
--vae "/workspace/Qwen-Image_models/vae/${VAE_FILENAME}" \
--text_encoder "/workspace/Qwen-Image_models/text_encoder/qwen_2.5_vl_7b.safetensors" \
--dataset_config "$DATASET_CONFIG"  \
--mixed_precision bf16 \
--sdpa \
--timestep_sampling shift \
--weighting_scheme none \
--discrete_flow_shift 2.0  \
--optimizer_type adamw8bit \
--learning_rate "$LEARNING_RATE" \
--gradient_checkpointing \
--max_data_loader_n_workers 2 \
--persistent_data_loader_workers \
--network_module networks.lora_qwen_image \
--network_dim "$NETWORK_DIM" \
--max_train_epochs "$MAX_TRAIN_EPOCHS" \
--save_every_n_epochs "$SAVE_EVERY_N_EPOCHS" \
--seed "$SEED" \
--output_dir "${OUTPUT_DIR_BASE}/${RUN_NAME}" \
--output_name "${RUN_NAME}" \
--ddp_gradient_as_bucket_view \
--ddp_static_graph