refactored infile handling

This commit is contained in:
yun saki 2022-08-26 13:10:37 +02:00
parent 0f28663805
commit cf750f62db

View File

@ -64,13 +64,16 @@ def main():
# gets rid of annoying messages about random seed # gets rid of annoying messages about random seed
logging.getLogger('pytorch_lightning').setLevel(logging.ERROR) logging.getLogger('pytorch_lightning').setLevel(logging.ERROR)
# load the infile as a list of lines
infile = None infile = None
try:
if opt.infile: if opt.infile:
infile = open(opt.infile, 'r') if os.path.isfile(opt.infile):
except FileNotFoundError as e: with open(opt.infile, "r") as file:
print(e) infile = file.read()
exit(-1) infile = infile.split("\n")
else:
print(f"WARNING: '{opt.infile}' not found. Aborting.")
sys.exit(-1) # exit does not work on every os, sys.exit does afaik
# preload the model # preload the model
t2i.load_model() t2i.load_model()
@ -119,8 +122,6 @@ def main():
cmd_parser = create_cmd_parser() cmd_parser = create_cmd_parser()
main_loop(t2i, opt.outdir, cmd_parser, log, infile) main_loop(t2i, opt.outdir, cmd_parser, log, infile)
log.close() log.close()
if infile:
infile.close()
def main_loop(t2i, outdir, parser, log, infile): def main_loop(t2i, outdir, parser, log, infile):
@ -129,15 +130,19 @@ def main_loop(t2i, outdir, parser, log, infile):
last_seeds = [] last_seeds = []
while not done: while not done:
if not infile:
command = input("dream> ")
else:
try: try:
command = infile.readline() if infile else input('dream> ') # get the next line of the infile
except EOFError: command = infile.pop(0)
except IndexError:
done = True done = True
break break
if infile and len(command) == 0: # skip empty lines
done = True if not command.strip():
break continue
if command.startswith(('#', '//')): if command.startswith(('#', '//')):
continue continue
@ -152,9 +157,6 @@ def main_loop(t2i, outdir, parser, log, infile):
print(str(e)) print(str(e))
continue continue
if len(elements) == 0:
continue
if elements[0] == 'q': if elements[0] == 'q':
done = True done = True
break break