#!/bin/bash

# Generate images using Illustrious model
# This script provides easy-to-use commands for running the image generator

set -e

# Default paths (relative to project root)
PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
MODEL_PATH="${PROJECT_ROOT}/models/waiNSFWIllustrious_v140.safetensors"
JSONL_PATH="${PROJECT_ROOT}/augmented_prompts.jsonl"
OUTPUT_DIR="${PROJECT_ROOT}/illustrious_generated"

# Default generation parameters
NUM_INFERENCE_STEPS=35
GUIDANCE_SCALE=7.5
MAX_IMAGES=""

# Function to display usage
usage() {
    echo "Usage: $0 [OPTIONS]"
    echo ""
    echo "Options:"
    echo "  --model-path PATH         Path to Illustrious model (default: models/waiNSFWIllustrious_v140.safetensors)"
    echo "  --jsonl-path PATH         Path to prompts JSONL file (default: augmented_prompts.jsonl)"
    echo "  --output-dir PATH         Output directory for images (default: illustrious_generated)"
    echo "  --max-images N            Maximum number of images to generate (for testing)"
    echo "  --num-inference-steps N   Number of inference steps (default: 35)"
    echo "  --guidance-scale N        Guidance scale (default: 7.5)"
    echo "  --help                    Show this help message"
    echo ""
    echo "Examples:"
    echo "  $0                        # Generate all images with default settings"
    echo "  $0 --max-images 10        # Generate only 10 images for testing"
    echo "  $0 --num-inference-steps 50 --guidance-scale 8.0  # Custom generation params"
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --model-path)
            MODEL_PATH="$2"
            shift 2
            ;;
        --jsonl-path)
            JSONL_PATH="$2"
            shift 2
            ;;
        --output-dir)
            OUTPUT_DIR="$2"
            shift 2
            ;;
        --max-images)
            MAX_IMAGES="$2"
            shift 2
            ;;
        --num-inference-steps)
            NUM_INFERENCE_STEPS="$2"
            shift 2
            ;;
        --guidance-scale)
            GUIDANCE_SCALE="$2"
            shift 2
            ;;
        --help)
            usage
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            usage
            exit 1
            ;;
    esac
done

# Check if required files exist
if [[ ! -f "$MODEL_PATH" ]]; then
    echo "Error: Model file not found at $MODEL_PATH"
    echo "Please make sure the Illustrious model is downloaded to the correct location."
    exit 1
fi

if [[ ! -f "$JSONL_PATH" ]]; then
    echo "Error: JSONL file not found at $JSONL_PATH"
    echo "Please make sure the augmented_prompts.jsonl file exists."
    exit 1
fi

# Check GPU availability
if ! nvidia-smi &>/dev/null; then
    echo "Warning: nvidia-smi not found. Make sure CUDA is available."
fi

# Create output directory if it doesn't exist
mkdir -p "$OUTPUT_DIR"

# Build the command
CMD="python ${PROJECT_ROOT}/data_tool/generate_alignment_data/generate_illustrious_images.py"
CMD="$CMD --model-path \"$MODEL_PATH\""
CMD="$CMD --jsonl-path \"$JSONL_PATH\""
CMD="$CMD --output-dir \"$OUTPUT_DIR\""
CMD="$CMD --num-inference-steps $NUM_INFERENCE_STEPS"
CMD="$CMD --guidance-scale $GUIDANCE_SCALE"

if [[ -n "$MAX_IMAGES" ]]; then
    CMD="$CMD --max-images $MAX_IMAGES"
fi

echo "Starting Illustrious image generation..."
echo "Model: $MODEL_PATH"
echo "Prompts: $JSONL_PATH"
echo "Output: $OUTPUT_DIR"
echo "Inference steps: $NUM_INFERENCE_STEPS"
echo "Guidance scale: $GUIDANCE_SCALE"
if [[ -n "$MAX_IMAGES" ]]; then
    echo "Max images: $MAX_IMAGES"
fi
echo ""

# Run the command
eval "$CMD"

echo ""
echo "Image generation completed!"
echo "Generated images are saved in: $OUTPUT_DIR"
echo "Metadata files are saved in: $OUTPUT_DIR/metadata"