# Author: Theodore Zacharia
# V0.1: 12/10/2022 - Initial Release
# V0.2: 02/11/2022 - Add batch file mode
#
# This script provides control to create multiple images by
# calling another script which generates one image.
#
# Create a bunch of images based on variations of settings
# There is a batch file option, whereby, as things take so long, you can
# specify a whole bunch of different settings t otry and then walk away and le
# the machine get on with it, over night if necessary.
#
TRACE=0
startstrng=0.1
stepstrng=0.1
startguida=1.5
stepguida=1.5
maxguida=14
steps=16
seed=2345678
seedend=0
srcimg=""
outimg=""
outdir=""
STARTIMAGE=""
MASKIMAGE=""
prompt="monkey"
LONGNAME=""
BATCHFILE=""
CREATEBATCHFILE=0
IPTC=""
USEMODIFIED=""  # default to use the modified python script for speed
PROCESSFILE=set.$$.tmp

# Testing
#MASKIMAGE="-M /home/theodore/Documents/20160331_124836_blank02.jpg"
#STARTIMAGE="-I /home/theodore/Documents/20160331_124836.jpg"

# *** Functions
_beginswith() { case $2 in "$1"*) true ;; *) false ;; esac ; }
_usage()
{
	echo "usage: $0 [-h] [-t] [-X] [-l] [-z seed] [-Z seedend] [-s strength] [-g guidestrength] [-i numberofiterations] [-I startimage] [-M mask] [-o outimage] [-b batchfile] [-B] prompt"
	echo "where"
	cat <<- _EOF_
	  -h				Displays this help message
	  -t				Set trace on
	  -z seed			Sets the random seed to use as part of the diffusion process
	  -Z seedend		Sets the last value for the random seed, if not present only 1 seed value is used from -z
	  -s strength 		How much noise to put into image, more noise more variation from any supplied start image
	  -S stepstrng		How much to increase the strength by each loop
	  -g guidestrength 	How closely to keep to any original image supplied
	  -G stepguide		How much to increase the guide strength by each loop
	  -H maxguide		The max guide number to go to, after 16 it gets a bit crazy
	  -i numberofiterations	Number of steps to run create the final image
	  -I startimage		Provide a start image
	  -M mask 			Provide a mask to keep parts of the original image the same, use black and white where black is not modified by the AI
	  -o outimage 		Generated image name, including path, defaults to a format which includes the above values plus a date stamp
	  -b batchfile		Provide set of ids to operate on
	  -B				Create the batch, will run using a created batch file using the optimised version of the demo.py
	  -l				Long file name, appends the prompt as well
	  -X                Add IPTC tags to EXIF info
	  -U				Use the original python script parameter set, default to the modified version

	  Format of the batch file:
	  seed,strength,stepstrng,guidestrength,stepguide,numberofiterations[,prompt]
	  You may comment out a line by starting it with a #

	  Prompt is optional you do NOT need one on every line and you can use it from the command line only or
	  each time a prompt is present, it replaces the previous one and all images from that one follow that prompt.

	  Examples:

	  Use existing batch file
	    ./makesetofimgs.sh -b batch.csv -I ~/Documents/20160424_172243d.jpg -M ~/Documents/20160424_172243d_blank02.jpg -l starship enterprise in orbit over red planet

	  Run in full batch mode, create the batch file from parameters passed, then runs, will produce a LOT of image variants of the prompt + seed + strength + guide
	    ./makesetofimgs.sh -B -z 13460 -Z 13470 -s 0.7 -S 0.1 -g 5.5 -G 1.5 -i 16 -I ~/Documents/20160424_172243d.jpg -M ~/Documents/20160424_172243d_blank02.jpg -l starship enterprise in orbit over red planet

	  To make a batch file for a specific set of values to apply to different prompts
		./makesetofimgs.sh -B -z 0 -Z 10 -s 0.8 -S 0.1 -g 4.5 -G 2.0  -i 16
	  then when you see \"ready to process...\" do a CTRL-C and look for the latest .tmp file, rename that to something you prefer and
	  you can then use this as your -b inputfile.

_EOF_
}


# *** Mainline

# process input parameters
while getopts thlXUBz:Z:s:S:g:G:i:I:M:O:o:b:H: AOPT
do
case $AOPT in
	t) TRACE=1 ;; # set TRACE mode
	l) LONGNAME="-l" ;;
	z) seed=$OPTARG ;;
	Z) seedend=$OPTARG ;;
	s) startstrng=$OPTARG ;;
	S) stepstrng=$OPTARG ;;
	g) startguida=$OPTARG ;;
	G) stepguida=$OPTARG ;;
	H) maxguida=$OPTARG ;;
	i) steps=$OPTARG ;;
	I) STARTIMAGE="-I $OPTARG" ;;
	M) MASKIMAGE="-M $OPTARG" ;;
	o) outimg="-o $OPTARG" ;;
	O) outdir="-O $OPTARG" ;;
	b) BATCHFILE=$OPTARG ;;
	B) CREATEBATCHFILE=1 ;;
	X) IPTC="-X" ;;
	U) USEMODIFIED="-U" ;;
	h) _usage
		exit 1 ;;
	*) echo "$AOPT is an invalid option" >&2
		exit 2 ;;
esac
done

shift $((OPTIND-1))

if [ $TRACE -gt 0 ]
then
	echo "Starting at $(date)"
	echo "positional_args=$@" >&2
fi

if [ $# -gt 0 ]
then
	prompt="$@"
fi


echo "Starting set at $(date)"
echo "Seed $seed ($seedend), Strength $stepstrng ($stepstrng), Guidance $startguida ($stepguida), steps $steps"
echo "STARTIMAGE $STARTIMAGE , MASKIMAGE $MASKIMAGE"
echo "prompt $prompt"


if [ -f "$BATCHFILE" ] && [ $CREATEBATCHFILE -eq 0 ]
then
	echo "Processing existing batch file $BATCHFILE"
	grep -v "^#" $BATCHFILE > $PROCESSFILE
elif [ $CREATEBATCHFILE -eq 1 ]
then
	if [ $seedend -eq 0 ] ; then seedend=$seed ; fi
	echo "Creating new batch file $PROCESSFILE"
	rm -f $PROCESSFILE
	for aseed in $(seq $seed 1 $seedend)
	do
		for strng in $(seq $startstrng $stepstrng 1.0)	# 1.0 is the max for strength
		do
			for guida in $(seq $startguida $stepguida $maxguida) # after 16 it starts to get crazy
			do
				echo "$aseed,$strng,$stepstrng,$guida,$stepguida,$steps" >> $PROCESSFILE
			done
		done
	done
	if [ $TRACE -gt 0 ] ; then cat $PROCESSFILE ; fi
else
	echo "$seed,$startstrng,$stepstrng,$startguida,$stepguida,$steps" > $PROCESSFILE
fi

cat $PROCESSFILE
echo "ready to process... (CTRL-C now if you want the csv file of $PROCESSFILE)"
sleep 5
echo "Starting set at $(date)"


if [ ! -n "$USEMODIFIED" ]
then
	./makeimg.sh -B $PROCESSFILE $outimg $outdir $STARTIMAGE $MASKIMAGE $LONGNAME $IPTC "$prompt"
else
# ORIGINAL VERSION, we did all the work of looping outside of the model, so was slow
	# old style, before we did the loop inside
	while read ALINE
	do
		if [ $TRACE -gt 0 ] ; then echo "ALINE=$ALINE" ; fi
		if _beginswith "#" $ALINE ; then continue ; fi

		seed=$(echo $ALINE | cut -d, -f1)
		startstrng=$(echo $ALINE | cut -d, -f2)
		stepstrng=$(echo $ALINE | cut -d, -f3)
		startguida=$(echo $ALINE | cut -d, -f4)
		stepguida=$(echo $ALINE | cut -d, -f5)
		steps=$(echo $ALINE | cut -d, -f6)
		newprompt=$(echo $ALINE | cut -d, -f7-)
		if [ -n "$newprompt" ] ; then echo "detected new prompt: $newprompt" ; prompt=$newprompt ; fi

		for strng in $(seq $startstrng $stepstrng 1.0)	# 1.0 is the max for strength
		do
			for guida in $(seq $startguida $stepguida $maxguida) # after 16 it starts to get crazy
			do
				echo "Seed $seed, Strength $strng, Guidance $guida, Steps $steps"
				./makeimg.sh $USEMODIFIED -z $seed -s $strng -g $guida -i $steps $STARTIMAGE $MASKIMAGE $IPTC $LONGNAME "$prompt"
				sleep 3
			done
		done
	done<$PROCESSFILE
fi

rm -f $PROCESSFILE

echo "Finished set at $(date)"