# Author: Theodore Zacharia
# V0.1: 12/10/2022 - Initial Release
# V1.1: 09/11/2022 - Modifed to support MY version of Demo to allow faster create of multiple images
#
# This script provides control to create an image using
# Stable Diffusion OpenVino demo.py application from
# https://github.com/bes-dev/stable_diffusion.openvino
# why bother ?  Just to make using makesetofimgs.sh simpler, which
# creates a whole bunch of images by varying some of the parameters
#
# Imperical research indicates:
# 1. seed has a fundamental variance on the image created
# 2. strength varies from 0.0 to 1.0, with each 0.05 making some difference
# 3. guidance starts at 1.1, with each 0.5 making some difference, upto about 15
#
# With start image + mask, seed makes the biggest difference if it will work out or not
#
# ENVIRONMET: You may need to run the following before this script works "conda activate base"
# Use "conda env list" to check, activate one has a "*"
#

# *** Globals
TRACE=0
strng=0.5
guida=7.5
steps=16
seed=2345678
srcimg=""
outimg=""
outdir="pics"
STARTIMAGE=""
MASKIMAGE=""
BATCHFILE=""
prompt="monkey"
LONGNAME=""
IPTC=""
USEMODIFIED=1  # default to use the modified python script for speed


# *** Functions
_usage()
{
	echo "usage: $0 [-h] [-t] [-l] [-X] [-z seed] [-s strength] [-g guidestrength] [-i numberofiterations] [-I startimage] [-M mask] [-o outimage] prompt"
	echo "where"
	cat <<- _EOF_
	  -h				Displays this help message
	  -t				Set trace on
	  -z seed			Sets the seed to use as part of the diffusion process, set -1 for random value
	  -s strength 		How much noise to put into image, more noise more variation from any supplied start image
	  -g guidestrength 	How closely to keep to any original image supplied
	  -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, defaults to a format which includes the above values plus starting with a date stamp
	  -O outputdir 		The output image dir
	  -l				Long file name, appends the prompt as well
	  -X                Add IPTC tags to EXIF info
	  -U				Use the original python script, default to the modified version
_EOF_
}


# *** Mainline

# process input parameters
while getopts thXUlz:s:g:i:I:M:O:o:B: AOPT
do
case $AOPT in
	t) TRACE=1 ;; # set TRACE mode
	l) LONGNAME="-l" ;;
	z) seed=$OPTARG ;;
	s) strng=$OPTARG ;;
	g) guida=$OPTARG ;;
	i) steps=$OPTARG ;;
	I) STARTIMAGE="--init-image $OPTARG" ;;
	M) MASKIMAGE="--mask $OPTARG" ;;
	o) outimg=$OPTARG ;;
	O) outdir=$OPTARG ;;
	B) BATCHFILE="--batchfile $OPTARG" ;;
	U) USEMODIFIED=0 ;;
	X) IPTC="--iptc" ;;
	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

# do some envronment checking, very basic check
if [ $TRACE -gt 0 ] ; then echo "checking conda environment" ; fi
CONDAENV=$(conda env list | grep "*" | cut -d' ' -f1 | grep -c base)
if [ $CONDAENV -eq 0 ]
then
	echo "***** WARNING: Did NOT detect the expected conda environment"
	echo "Will continue but may fail, if so please run the following"
	echo "(assuming base is env expected)"
	echo "  conda activate base"
	echo " "
	sleep 2
fi

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

if [ $seed -lt 0 ]
then
	seed=$RANDOM
fi

if [ ! -n "$outimg" ] && [ ! -n "$BATCHFILE" ]
then
	IMGNM=$(date +"%s")
	if [ "$LONGNAME" = "-l" ]
	then
		fprompt=$(echo $prompt | tr " " "_")
		outimg=img${IMGNM}_se${seed}_sr${strng}_gu${guida}_st${steps}_${fprompt}.png
	else
		outimg=img${IMGNM}_se${seed}_sr${strng}_gu${guida}_st${steps}.png
	fi
	OUTPUT="--output $outimg"
else
	OUTPUT=""
fi

echo "Prompt: $prompt"
if [ -n "$BATCHFILE" ]
then
	echo "  using values from $BATCHFILE"
else
	echo "  with: Seed $seed, Strength $strng, Guidance $guida, Steps $steps"
fi
if [ -n "$STARTIMAGE" ] ; then echo "  $STARTIMAGE" ; fi
if [ -n "$MASKIMAGE" ] ; then echo "  $MASKIMAGE" ; fi
echo "  save to $outimg"

date
# using new TZ modified for batch
if [ $USEMODIFIED -eq 1 ]
then
	# new version, for best speed locally
	python demo_tz.py $IPTC $LONGNAME --seed $seed --outputdir $outdir $OUTPUT --strength $strng --guidance-scale $guida --num-inference-steps $steps $BATCHFILE $STARTIMAGE $MASKIMAGE --prompt "$prompt"
else
	# original version
	python demo.py --seed $seed --output "$outimg" --strength $strng --guidance-scale $guida --num-inference-steps $steps $BATCHFILE $STARTIMAGE $MASKIMAGE --prompt "$prompt"
fi
date
