From 8f6f8413b12ea8e8ae072b4264c3459a38c7ad28 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 27 Jul 2023 09:30:41 -0400 Subject: [PATCH] fix ti training --- modules/lycoris | 2 +- modules/processing_diffusers.py | 4 ++-- modules/textual_inversion/dataset.py | 19 ++++--------------- .../textual_inversion/textual_inversion.py | 9 +++++---- modules/textual_inversion/ui.py | 10 +--------- modules/ui.py | 2 +- scripts/xyz_grid.py | 2 +- 7 files changed, 15 insertions(+), 33 deletions(-) diff --git a/modules/lycoris b/modules/lycoris index c8a662d61..6dc0dbf47 160000 --- a/modules/lycoris +++ b/modules/lycoris @@ -1 +1 @@ -Subproject commit c8a662d61271a14968bc6e5ff6ce4461a3488772 +Subproject commit 6dc0dbf477d5c71a44da25014a7b3d7146b08b52 diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 955e697c1..52175889c 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -139,8 +139,8 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro negative_prompt_2=[p.refiner_negative] if len(p.refiner_negative) > 0 else negative_prompts, eta=shared.opts.eta_ddim, guidance_rescale=p.diffusers_guidance_rescale, + denoising_start=0 if refiner_enabled and p.refiner_start > 0 and p.refiner_start < 1 else None, denoising_end=p.refiner_start if refiner_enabled and p.refiner_start > 0 and p.refiner_start < 1 else None, - # aesthetic_score=shared.opts.diffusers_aesthetics_score, output_type='latent' if hasattr(shared.sd_model, 'vae') else 'np', **task_specific_kwargs ) @@ -188,8 +188,8 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro strength=p.denoising_strength, guidance_scale=p.image_cfg_scale if p.image_cfg_scale is not None else p.cfg_scale, guidance_rescale=p.diffusers_guidance_rescale, - # aesthetic_score=shared.opts.diffusers_aesthetics_score, denoising_start=p.refiner_start if p.refiner_start > 0 and p.refiner_start < 1 else None, + denoising_end=1 if p.refiner_start > 0 and p.refiner_start < 1 else None, image=output.images[i], output_type='latent' if hasattr(shared.sd_refiner, 'vae') else 'np', ) diff --git a/modules/textual_inversion/dataset.py b/modules/textual_inversion/dataset.py index f53a73b89..912736f71 100644 --- a/modules/textual_inversion/dataset.py +++ b/modules/textual_inversion/dataset.py @@ -32,14 +32,10 @@ class PersonalizedBase(Dataset): re_word = re.compile(shared.opts.dataset_filename_word_regex) if len(shared.opts.dataset_filename_word_regex) > 0 else None self.placeholder_token = placeholder_token - self.flip = transforms.RandomHorizontalFlip(p=flip_p) - self.dataset = [] - - with open(template_file, "r") as file: + with open(template_file, "r", encoding="utf8") as file: lines = [x.strip() for x in file.readlines()] - self.lines = lines assert data_root, 'dataset directory not specified' @@ -47,20 +43,16 @@ class PersonalizedBase(Dataset): assert os.listdir(data_root), "Dataset directory is empty" self.image_paths = [os.path.join(data_root, file_path) for file_path in os.listdir(data_root)] - self.shuffle_tags = shuffle_tags self.tag_drop_out = tag_drop_out groups = defaultdict(list) - - print("Preparing dataset...") + shared.log.info(f"TI Training: Preparing dataset: {data_root}") for path in tqdm.tqdm(self.image_paths): alpha_channel = None if shared.state.interrupted: - raise Exception("interrupted") + raise RuntimeError("interrupted") try: image = Image.open(path) - #Currently does not work for single color transparency - #We would need to read image.info['transparency'] for that if use_weight and 'A' in image.getbands(): alpha_channel = image.getchannel('A') image = image.convert('RGB') @@ -91,12 +83,9 @@ class PersonalizedBase(Dataset): with devices.autocast(): latent_dist = model.encode_first_stage(torchdata.unsqueeze(dim=0)) - #Perform latent sampling, even for random sampling. - #We need the sample dimensions for the weights if latent_sampling_method == "deterministic": if isinstance(latent_dist, DiagonalGaussianDistribution): - # Works only for DiagonalGaussianDistribution - latent_dist.std = 0 + latent_dist.std = torch.exp(0 * latent_dist.logvar) else: latent_sampling_method = "once" latent_sample = model.get_first_stage_encoding(latent_dist).squeeze().to(devices.cpu) diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py index 41c2f8de9..25345a35a 100644 --- a/modules/textual_inversion/textual_inversion.py +++ b/modules/textual_inversion/textual_inversion.py @@ -338,7 +338,7 @@ def tensorboard_add_image(tensorboard_writer, tag, pil_image, step): tensorboard_writer.add_image(tag, img_tensor, global_step=step) -def validate_train_inputs(model_name, learn_rate, batch_size, gradient_step, data_root, template_file, template_filename, steps, save_model_every, create_image_every, log_directory, name="embedding"): +def validate_train_inputs(model_name, learn_rate, batch_size, gradient_step, data_root, template_file, template_filename, steps, save_model_every, create_image_every, name="embedding"): assert model_name, f"{name} not selected" assert learn_rate, "Learning rate is empty or 0" assert isinstance(batch_size, int), "Batch size must be integer" @@ -366,7 +366,7 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st save_embedding_every = save_embedding_every or 0 create_image_every = create_image_every or 0 template_file = textual_inversion_templates.get(template_filename, None) - validate_train_inputs(embedding_name, learn_rate, batch_size, gradient_step, data_root, template_file, template_filename, steps, save_embedding_every, create_image_every, log_directory, name="embedding") + validate_train_inputs(embedding_name, learn_rate, batch_size, gradient_step, data_root, template_file, template_filename, steps, save_embedding_every, create_image_every, name="embedding") if log_directory is None or log_directory == '': log_directory = f"{os.path.join(shared.cmd_opts.data_dir, 'train/log/embeddings')}" template_file = template_file.path @@ -421,10 +421,13 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st tensorboard_writer = tensorboard_setup(log_directory) pin_memory = shared.opts.pin_memory + # init dataset ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=embedding_name, model=shared.sd_model, cond_model=shared.sd_model.cond_stage_model, device=devices.device, template_file=template_file, batch_size=batch_size, gradient_step=gradient_step, shuffle_tags=shuffle_tags, tag_drop_out=tag_drop_out, latent_sampling_method=latent_sampling_method, varsize=varsize, use_weight=use_weight) + if shared.opts.save_training_settings_to_txt: save_settings_to_file(log_directory, {**dict(model_name=checkpoint.model_name, model_hash=checkpoint.shorthash, num_of_dataset_images=len(ds), num_vectors_per_token=len(embedding.vec)), **locals()}) latent_sampling_method = ds.latent_sampling_method + # init dataloader dl = modules.textual_inversion.dataset.PersonalizedDataLoader(ds, latent_sampling_method=latent_sampling_method, batch_size=ds.batch_size, pin_memory=pin_memory) if unload: shared.parallel_processing_allowed = False @@ -463,7 +466,6 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st pbar = tqdm(total=steps - initial_step) try: sd_hijack_checkpoint.add() - for _i in range((steps-initial_step) * gradient_step): if scheduler.finished: break @@ -480,7 +482,6 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st break if clip_grad: clip_grad_sched.step(embedding.step) - with devices.autocast(): x = batch.latent_sample.to(devices.device, non_blocking=pin_memory) if use_weight: diff --git a/modules/textual_inversion/ui.py b/modules/textual_inversion/ui.py index 8afd8ef04..76367b772 100644 --- a/modules/textual_inversion/ui.py +++ b/modules/textual_inversion/ui.py @@ -13,25 +13,17 @@ def create_embedding(name, initialization_text, nvpt, overwrite_old): def preprocess(*args): modules.textual_inversion.preprocess.preprocess(*args) - return f"Preprocessing {'interrupted' if shared.state.interrupted else 'finished'}.", "" def train_embedding(*args): - assert not shared.cmd_opts.lowvram, 'Training models with lowvram not possible' - apply_optimizations = False try: if not apply_optimizations: sd_hijack.undo_optimizations() - embedding, filename = modules.textual_inversion.textual_inversion.train_embedding(*args) - - res = f""" -Training {'interrupted' if shared.state.interrupted else 'finished'} at {embedding.step} steps. -Embedding saved to {html.escape(filename)} -""" + res = f"Training {'interrupted' if shared.state.interrupted else 'finished'} at {embedding.step} steps. Embedding saved to {html.escape(filename)}" return res, "" except Exception: raise diff --git a/modules/ui.py b/modules/ui.py index 924584019..dec77a27c 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -384,7 +384,7 @@ def create_ui(startup_timer = None): hr_second_pass_steps, latent_index = create_sampler_and_steps_selection(modules.sd_samplers.samplers, "txt2img", False) with FormRow(elem_id="txt2img_hires_fix_row1", variant="compact"): denoising_strength = gr.Slider(minimum=0.05, maximum=1.0, step=0.01, label='Denoising strength', value=0.3, elem_id="txt2img_denoising_strength") - refiner_start = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Denoise start', value=1.0, elem_id="txt2img_refiner_start") + refiner_start = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Denoise start', value=0.8, elem_id="txt2img_refiner_start") with FormRow(elem_id="txt2img_hires_finalres", variant="compact"): hr_final_resolution = FormHTML(value="", elem_id="txtimg_hr_finalres", label="Upscaled resolution", interactive=False) with FormRow(elem_id="txt2img_hires_fix_row2", variant="compact"): diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py index 9448f5ed0..d74b2b7a0 100644 --- a/scripts/xyz_grid.py +++ b/scripts/xyz_grid.py @@ -205,7 +205,7 @@ class AxisOptionTxt2Img(AxisOption): axis_options = [ AxisOption("Nothing", str, do_nothing, fmt=format_nothing), - AxisOption("Checkpoint name", str, apply_checkpoint, fmt=format_value, cost=1.0, choices=lambda: list(sd_models.checkpoints_list)), + AxisOption("Checkpoint name", str, apply_checkpoint, fmt=format_value, cost=1.0, choices=lambda: sorted(sd_models.checkpoints_list)), AxisOption("VAE", str, apply_vae, cost=0.7, choices=lambda: ['None'] + list(sd_vae.vae_dict)), AxisOption("Dict name", str, apply_dict, fmt=format_value, cost=1.0, choices=lambda: ['None'] + list(sd_models.checkpoints_list)), AxisOption("Prompt S/R", str, apply_prompt, fmt=format_value),