VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-16 17:02:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #846 from kohya-ss/dev

Browse Source 
Fix to work training U-Net only LoRA for SD1/2
This commit is contained in:
Kohya S
2023-10-01 16:44:13 +09:00
committed by GitHub
parent 6bd6cd9c51 4cc919607a
commit 93155242fa
1 changed files with 9 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
train_network.py
View File

@@ -12,10 +12,13 @@ import toml
from tqdm import tqdm
import torch
try:
import intel_extension_for_pytorch as ipex
if torch.xpu.is_available():
from library.ipex import ipex_init
ipex_init()
except Exception:
pass
@@ -426,7 +429,12 @@ class NetworkTrainer:
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)
# set top parameter requires_grad = True for gradient checkpointing works
if not train_text_encoder: # train U-Net only
unet.parameters().__next__().requires_grad_(True)
else:
unet.eval()
for t_enc in text_encoders: