Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

improve UI of textual inversion frontend

Browse Source 
- File selection box now accepts directories that don't exist yet.
- Fixed crash when resume is selected and no files available to resume from.
This commit is contained in:
Lincoln Stein 2023-01-15 17:04:14 -05:00
parent 251a409087
commit bfeafa8d5e
2 changed files with 37 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
ldm/invoke/textual_inversion_training.py
View File

@ -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

24
scripts/textual_inversion_fe.py
View File

@ -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 <N>ext and <P>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)