diff options
Diffstat (limited to 'training/strategy/ti.py')
-rw-r--r-- | training/strategy/ti.py | 14 |
1 files changed, 8 insertions, 6 deletions
diff --git a/training/strategy/ti.py b/training/strategy/ti.py index ba78b98..e922954 100644 --- a/training/strategy/ti.py +++ b/training/strategy/ti.py | |||
@@ -117,14 +117,15 @@ def textual_inversion_strategy_callbacks( | |||
117 | with ema_context(): | 117 | with ema_context(): |
118 | yield | 118 | yield |
119 | 119 | ||
120 | def on_after_optimize(lr: float): | 120 | @torch.no_grad() |
121 | def on_before_optimize(lr: float, epoch: int): | ||
121 | if use_emb_decay: | 122 | if use_emb_decay: |
122 | with torch.no_grad(): | 123 | text_encoder.text_model.embeddings.normalize( |
123 | text_encoder.text_model.embeddings.normalize( | 124 | emb_decay_target, |
124 | emb_decay_target, | 125 | min(1.0, emb_decay * lr) |
125 | min(1.0, emb_decay * lr) | 126 | ) |
126 | ) | ||
127 | 127 | ||
128 | def on_after_optimize(lr: float): | ||
128 | if ema_embeddings is not None: | 129 | if ema_embeddings is not None: |
129 | ema_embeddings.step(text_encoder.text_model.embeddings.temp_token_embedding.parameters()) | 130 | ema_embeddings.step(text_encoder.text_model.embeddings.temp_token_embedding.parameters()) |
130 | 131 | ||
@@ -154,6 +155,7 @@ def textual_inversion_strategy_callbacks( | |||
154 | on_model=on_model, | 155 | on_model=on_model, |
155 | on_train=on_train, | 156 | on_train=on_train, |
156 | on_eval=on_eval, | 157 | on_eval=on_eval, |
158 | on_before_optimize=on_before_optimize, | ||
157 | on_after_optimize=on_after_optimize, | 159 | on_after_optimize=on_after_optimize, |
158 | on_log=on_log, | 160 | on_log=on_log, |
159 | on_checkpoint=on_checkpoint, | 161 | on_checkpoint=on_checkpoint, |