summaryrefslogtreecommitdiffstats
path: root/training/strategy
diff options
context:
space:
mode:
authorVolpeon <git@volpeon.ink>2023-04-01 18:22:44 +0200
committerVolpeon <git@volpeon.ink>2023-04-01 18:22:44 +0200
commit67afa02ec4acf18fc815be034fc8673e0bdba4d9 (patch)
tree8904d63dd16fc2f3579fcea5999284f79d833647 /training/strategy
parentCombined TI with embedding and LoRA (diff)
downloadtextual-inversion-diff-67afa02ec4acf18fc815be034fc8673e0bdba4d9.tar.gz
textual-inversion-diff-67afa02ec4acf18fc815be034fc8673e0bdba4d9.tar.bz2
textual-inversion-diff-67afa02ec4acf18fc815be034fc8673e0bdba4d9.zip
Fix
Diffstat (limited to 'training/strategy')
-rw-r--r--training/strategy/ti.py4
1 files changed, 3 insertions, 1 deletions
diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index 33f5fb9..1b5adab 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -64,7 +64,7 @@ def textual_inversion_strategy_callbacks(
64 ) 64 )
65 65
66 def on_accum_model(): 66 def on_accum_model():
67 return text_encoder.text_model.embeddings.overlay 67 return text_encoder.text_model.embeddings
68 68
69 @contextmanager 69 @contextmanager
70 def on_train(epoch: int): 70 def on_train(epoch: int):
@@ -145,6 +145,8 @@ def textual_inversion_prepare(
145 145
146 text_encoder.text_model.encoder.requires_grad_(False) 146 text_encoder.text_model.encoder.requires_grad_(False)
147 text_encoder.text_model.final_layer_norm.requires_grad_(False) 147 text_encoder.text_model.final_layer_norm.requires_grad_(False)
148 text_encoder.text_model.embeddings.position_embedding.requires_grad_(False)
149 text_encoder.text_model.embeddings.token_embedding.requires_grad_(False)
148 text_encoder.eval() 150 text_encoder.eval()
149 151
150 return text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler, {} 152 return text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler, {}