Skip to content

Commit 4b4111a

Browse files
committed
fix invoke.py crash if no models.yaml file present
- Script will now offer the user the ability to create a minimal models.yaml and then gracefully exit. - Closes #1420
1 parent 832f183 commit 4b4111a

File tree

2 files changed

+39
-1
lines changed

2 files changed

+39
-1
lines changed

ldm/invoke/readline.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -284,6 +284,7 @@ def _path_completions(self, text, state, extensions, shortcut_ok=True):
284284
switch,partial_path = match.groups()
285285
partial_path = partial_path.lstrip()
286286

287+
287288
matches = list()
288289
path = os.path.expanduser(partial_path)
289290

@@ -321,6 +322,7 @@ def _path_completions(self, text, state, extensions, shortcut_ok=True):
321322
matches.append(
322323
switch+os.path.join(os.path.dirname(full_path), node)
323324
)
325+
324326
return matches
325327

326328
class DummyCompleter(Completer):

scripts/invoke.py

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,12 @@ def main():
9090
safety_checker=opt.safety_checker,
9191
max_loaded_models=opt.max_loaded_models,
9292
)
93-
except (FileNotFoundError, IOError, KeyError) as e:
93+
except FileNotFoundError:
94+
print('** You appear to be missing configs/models.yaml')
95+
print('** You can either exit this script and run scripts/preload_models.py, or fix the problem now.')
96+
emergency_model_create(opt)
97+
sys.exit(-1)
98+
except (IOError, KeyError) as e:
9499
print(f'{e}. Aborting.')
95100
sys.exit(-1)
96101

@@ -482,6 +487,7 @@ def do_command(command:str, gen, opt:Args, completer) -> tuple:
482487
command = '-h'
483488
return command, operation
484489

490+
485491
def add_weights_to_config(model_path:str, gen, opt, completer):
486492
print(f'>> Model import in process. Please enter the values needed to configure this model:')
487493
print()
@@ -887,6 +893,36 @@ def write_commands(opt, file_path:str, outfilepath:str):
887893
f.write('\n'.join(commands))
888894
print(f'>> File {outfilepath} with commands created')
889895

896+
def emergency_model_create(opt:Args):
897+
completer = get_completer(opt)
898+
completer.complete_extensions(('.yaml','.yml','.ckpt','.vae.pt'))
899+
completer.set_default_dir('.')
900+
valid_path = False
901+
while not valid_path:
902+
weights_file = input('Enter the path to a downloaded models file, or ^C to exit: ')
903+
valid_path = os.path.exists(weights_file)
904+
dir,basename = os.path.split(weights_file)
905+
906+
valid_name = False
907+
while not valid_name:
908+
name = input('Enter a short name for this model (no spaces): ')
909+
name = 'unnamed model' if len(name)==0 else name
910+
valid_name = ' ' not in name
911+
912+
description = input('Enter a description for this model: ')
913+
description = 'no description' if len(description)==0 else description
914+
915+
with open(opt.conf, 'w', encoding='utf-8') as f:
916+
f.write(f'{name}:\n')
917+
f.write(f' description: {description}\n')
918+
f.write(f' weights: {weights_file}\n')
919+
f.write(f' config: ./configs/stable-diffusion/v1-inference.yaml\n')
920+
f.write(f' width: 512\n')
921+
f.write(f' height: 512\n')
922+
f.write(f' default: true\n')
923+
print(f'Config file {opt.conf} is created. This script will now exit.')
924+
print(f'After restarting you may examine the entry with !models and edit it with !edit.')
925+
890926
######################################
891927

892928
if __name__ == '__main__':

0 commit comments

Comments
 (0)