intercept keyboard interrupt during processing and return to prompt;

remove "!dream" from beginning of prompt;
user can quit by typing <q>
This commit is contained in:
Lincoln Stein 2022-08-18 23:03:22 -04:00
parent ac440a1197
commit 80c0e30099

View File

@ -4,14 +4,14 @@ import shlex
import atexit import atexit
import os import os
# readline will be disabled on windows systems # readline unavailable on windows systems
try: try:
import readline import readline
readline_available = True readline_available = True
except: except:
readline_available = False readline_available = False
debugging = False debugging = True
def main(): def main():
''' Initialize command-line parsers and the diffusion model ''' ''' Initialize command-line parsers and the diffusion model '''
@ -60,7 +60,7 @@ def main():
# preload the model # preload the model
if not debugging: if not debugging:
t2i.load_model() t2i.load_model()
print("\n* Initialization done! Awaiting your command (-h for help)...") print("\n* Initialization done! Awaiting your command (-h for help, q to quit)...")
log_path = os.path.join(opt.outdir,"dream_log.txt") log_path = os.path.join(opt.outdir,"dream_log.txt")
with open(log_path,'a') as log: with open(log_path,'a') as log:
@ -70,15 +70,23 @@ def main():
def main_loop(t2i,parser,log): def main_loop(t2i,parser,log):
''' prompt/read/execute loop ''' ''' prompt/read/execute loop '''
while True: done = False
while not done:
try: try:
command = input("dream> ") command = input("dream> ")
except EOFError: except EOFError:
print("goodbye!") done = True
break break
# rearrange the arguments to mimic how it works in the Dream bot.
elements = shlex.split(command) elements = shlex.split(command)
if elements[0]=='q': #
done = True
break
if elements[0].startswith('!dream'): # in case a stored prompt still contains the !dream command
elements.pop(0)
# rearrange the arguments to mimic how it works in the Dream bot.
switches = [''] switches = ['']
switches_started = False switches_started = False
@ -101,12 +109,17 @@ def main_loop(t2i,parser,log):
print("Try again with a prompt!") print("Try again with a prompt!")
continue continue
try:
if opt.init_img is None: if opt.init_img is None:
results = t2i.txt2img(**vars(opt)) results = t2i.txt2img(**vars(opt))
else: else:
results = t2i.img2img(**vars(opt)) results = t2i.img2img(**vars(opt))
print("Outputs:") print("Outputs:")
write_log_message(opt,switches,results,log) write_log_message(opt,switches,results,log)
except KeyboardInterrupt:
continue
print("goodbye!")
def write_log_message(opt,switches,results,logfile): def write_log_message(opt,switches,results,logfile):
''' logs the name of the output image, its prompt and seed to both the terminal and the log file ''' ''' logs the name of the output image, its prompt and seed to both the terminal and the log file '''