File size: 3,796 Bytes
3f9fa87 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
#!/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"
|