fix the issue with training on SD2.0

2 files changed, 3 insertions, 2 deletions

diff --git a/modules/sd_models.py b/modules/sd_models.py
index ebd4dff7..bff8d6c9 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -228,6 +228,8 @@ def load_model_weights(model, checkpoint_info, vae_file="auto"):
     model.sd_model_checkpoint = checkpoint_file

     model.sd_checkpoint_info = checkpoint_info

 

+    model.logvar = model.logvar.to(devices.device)  # fix for training

+

     sd_vae.delete_base_vae()

     sd_vae.clear_loaded_vae()

     vae_file = sd_vae.resolve_vae(checkpoint_file, vae_file=vae_file)

diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index 66f40367..1e5722e7 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -282,7 +282,7 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
         return embedding, filename

     scheduler = LearnRateScheduler(learn_rate, steps, initial_step)

 

-   # dataset loading may take a while, so input validations and early returns should be done before this

+    # dataset loading may take a while, so input validations and early returns should be done before this

     shared.state.textinfo = f"Preparing dataset from {html.escape(data_root)}..."

     old_parallel_processing_allowed = shared.parallel_processing_allowed

 

@@ -310,7 +310,6 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
     loss_step = 0

     _loss_step = 0 #internal

 

-

     last_saved_file = "<none>"

     last_saved_image = "<none>"

     forced_filename = "<none>"