From bc50936745e1a349afdc28cf1540109ba20bc71a Mon Sep 17 00:00:00 2001 From: Shondoit Date: Thu, 12 Jan 2023 15:34:11 +0100 Subject: Call weighted_forward during training --- modules/textual_inversion/textual_inversion.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) (limited to 'modules/textual_inversion') diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py index a1a406c2..8853c868 100644 --- a/modules/textual_inversion/textual_inversion.py +++ b/modules/textual_inversion/textual_inversion.py @@ -480,6 +480,7 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st with devices.autocast(): x = batch.latent_sample.to(devices.device, non_blocking=pin_memory) + w = batch.weight.to(devices.device, non_blocking=pin_memory) c = shared.sd_model.cond_stage_model(batch.cond_text) if is_training_inpainting_model: @@ -490,7 +491,7 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st else: cond = c - loss = shared.sd_model(x, cond)[0] / gradient_step + loss = shared.sd_model.weighted_forward(x, cond, w)[0] / gradient_step del x _loss_step += loss.item() -- cgit v1.2.3