From 81419f7f320930a0e3deb330f374fca6dffc671f Mon Sep 17 00:00:00 2001 From: Kohya S Date: Sun, 1 Oct 2023 16:37:23 +0900 Subject: [PATCH] Fix to work training U-Net only LoRA for SD1/2 --- train_network.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/train_network.py b/train_network.py index 0e2e0fa9f..6d98037ca 100644 --- a/train_network.py +++ b/train_network.py @@ -426,7 +426,10 @@ def train(self, args): t_enc.train() # set top parameter requires_grad = True for gradient checkpointing works - t_enc.text_model.embeddings.requires_grad_(True) + if train_text_encoder: + t_enc.text_model.embeddings.requires_grad_(True) + else: + unet.parameters().__next__().requires_grad_(True) else: unet.eval() for t_enc in text_encoders: