diff --git a/ldm/invoke/textual_inversion_training.py b/ldm/invoke/textual_inversion_training.py index ab53f0801d..3c7ffd5663 100644 --- a/ldm/invoke/textual_inversion_training.py +++ b/ldm/invoke/textual_inversion_training.py @@ -518,10 +518,10 @@ def do_textual_inversion_training( pretrained_model_name_or_path = model_conf.get('repo_id',None) or Path(model_conf.get('path')) assert pretrained_model_name_or_path, f"models.yaml error: neither 'repo_id' nor 'path' is defined for {model}" pipeline_args = dict(cache_dir=global_cache_dir('diffusers')) - + # Load tokenizer if tokenizer_name: - tokenizer = CLIPTokenizer.from_pretrained(tokenizer_name,cache_dir=global_cache_dir('transformers')) + tokenizer = CLIPTokenizer.from_pretrained(tokenizer_name,**pipeline_args) else: tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer", **pipeline_args) @@ -670,24 +670,28 @@ def do_textual_inversion_training( logger.info(f" Total optimization steps = {max_train_steps}") global_step = 0 first_epoch = 0 + resume_step = None # Potentially load in the weights and states from a previous save if resume_from_checkpoint: - if resume_from_checkpoint != "latest": - path = os.path.basename(resume_from_checkpoint) - else: - # Get the most recent checkpoint - dirs = os.listdir(output_dir) - dirs = [d for d in dirs if d.startswith("checkpoint")] - dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) - path = dirs[-1] - accelerator.print(f"Resuming from checkpoint {path}") - accelerator.load_state(os.path.join(output_dir, path)) - global_step = int(path.split("-")[1]) + try: + if resume_from_checkpoint != "latest": + path = os.path.basename(resume_from_checkpoint) + else: + # Get the most recent checkpoint + dirs = os.listdir(output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] + accelerator.print(f"Resuming from checkpoint {path}") + accelerator.load_state(os.path.join(output_dir, path)) + global_step = int(path.split("-")[1]) - resume_global_step = global_step * gradient_accumulation_steps - first_epoch = resume_global_step // num_update_steps_per_epoch - resume_step = resume_global_step % num_update_steps_per_epoch + resume_global_step = global_step * gradient_accumulation_steps + first_epoch = resume_global_step // num_update_steps_per_epoch + resume_step = resume_global_step % num_update_steps_per_epoch + except: + logger.warn("No checkpoint available to resume from") # Only show the progress bar once on each machine. progress_bar = tqdm(range(global_step, max_train_steps), disable=not accelerator.is_local_main_process) @@ -700,7 +704,7 @@ def do_textual_inversion_training( text_encoder.train() for step, batch in enumerate(train_dataloader): # Skip steps until we reach the resumed step - if resume_from_checkpoint and epoch == first_epoch and step < resume_step: + if resume_step and resume_from_checkpoint and epoch == first_epoch and step < resume_step: if step % gradient_accumulation_steps == 0: progress_bar.update(1) continue diff --git a/scripts/textual_inversion_fe.py b/scripts/textual_inversion_fe.py index 941afcf613..4e31812f09 100755 --- a/scripts/textual_inversion_fe.py +++ b/scripts/textual_inversion_fe.py @@ -6,6 +6,7 @@ import sys import re import shutil import traceback +import curses from ldm.invoke.globals import Globals, global_set_root from omegaconf import OmegaConf from pathlib import Path @@ -43,6 +44,11 @@ class textualInversionForm(npyscreen.FormMultiPageAction): except: pass + self.add_widget_intelligent( + npyscreen.FixedText, + value='Use ctrl-N and ctrl-P to move to the ext and

revious fields, cursor arrows to make a selection, and space to toggle checkboxes.' + ) + self.model = self.add_widget_intelligent( npyscreen.TitleSelectOne, name='Model Name:', @@ -82,18 +88,18 @@ class textualInversionForm(npyscreen.FormMultiPageAction): max_height=4, ) self.train_data_dir = self.add_widget_intelligent( - npyscreen.TitleFilenameCombo, + npyscreen.TitleFilename, name='Data Training Directory:', select_dir=True, - must_exist=True, - value=saved_args.get('train_data_dir',Path(Globals.root) / TRAINING_DATA / default_placeholder_token) + must_exist=False, + value=str(saved_args.get('train_data_dir',Path(Globals.root) / TRAINING_DATA / default_placeholder_token)) ) self.output_dir = self.add_widget_intelligent( - npyscreen.TitleFilenameCombo, + npyscreen.TitleFilename, name='Output Destination Directory:', select_dir=True, must_exist=False, - value=saved_args.get('output_dir',Path(Globals.root) / TRAINING_DIR / default_placeholder_token) + value=str(saved_args.get('output_dir',Path(Globals.root) / TRAINING_DIR / default_placeholder_token)) ) self.resolution = self.add_widget_intelligent( npyscreen.TitleSelectOne, @@ -174,8 +180,8 @@ class textualInversionForm(npyscreen.FormMultiPageAction): def initializer_changed(self): placeholder = self.placeholder_token.value self.prompt_token.value = f'(Trigger by using <{placeholder}> in your prompts)' - self.train_data_dir.value = Path(Globals.root) / TRAINING_DATA / placeholder - self.output_dir.value = Path(Globals.root) / TRAINING_DIR / placeholder + self.train_data_dir.value = str(Path(Globals.root) / TRAINING_DATA / placeholder) + self.output_dir.value = str(Path(Globals.root) / TRAINING_DIR / placeholder) self.resume_from_checkpoint.value = Path(self.output_dir.value).exists() def on_ok(self): @@ -280,7 +286,9 @@ def save_args(args:dict): ''' Save the current argument values to an omegaconf file ''' - conf_file = Path(Globals.root) / TRAINING_DIR / CONF_FILE + dest_dir = Path(Globals.root) / TRAINING_DIR + os.makedirs(dest_dir, exist_ok=True) + conf_file = dest_dir / CONF_FILE conf = OmegaConf.create(args) OmegaConf.save(config=conf, f=conf_file)