diff options
| author | Volpeon <git@volpeon.ink> | 2023-06-25 08:40:05 +0200 |
|---|---|---|
| committer | Volpeon <git@volpeon.ink> | 2023-06-25 08:40:05 +0200 |
| commit | 4f6ecc2fd38dc894ac7fcb4f130d3ab8af49d132 (patch) | |
| tree | c6ab59d6726b818638fe90a3ea8bb8403a8b0c30 /training/strategy/dreambooth.py | |
| parent | Update (diff) | |
| download | textual-inversion-diff-4f6ecc2fd38dc894ac7fcb4f130d3ab8af49d132.tar.gz textual-inversion-diff-4f6ecc2fd38dc894ac7fcb4f130d3ab8af49d132.tar.bz2 textual-inversion-diff-4f6ecc2fd38dc894ac7fcb4f130d3ab8af49d132.zip | |
Update
Diffstat (limited to 'training/strategy/dreambooth.py')
| -rw-r--r-- | training/strategy/dreambooth.py | 3 |
1 files changed, 3 insertions, 0 deletions
diff --git a/training/strategy/dreambooth.py b/training/strategy/dreambooth.py index 3d1abf7..7e67589 100644 --- a/training/strategy/dreambooth.py +++ b/training/strategy/dreambooth.py | |||
| @@ -154,6 +154,9 @@ def dreambooth_strategy_callbacks( | |||
| 154 | unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=False) | 154 | unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=False) |
| 155 | text_encoder_ = accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=False) | 155 | text_encoder_ = accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=False) |
| 156 | 156 | ||
| 157 | text_encoder_.forward = MethodType(text_encoder_.forward, text_encoder_) | ||
| 158 | unet_.forward = MethodType(unet_.forward, unet_) | ||
| 159 | |||
| 157 | text_encoder_.text_model.embeddings.persist(False) | 160 | text_encoder_.text_model.embeddings.persist(False) |
| 158 | 161 | ||
| 159 | with ema_context(): | 162 | with ema_context(): |