diff options
Diffstat (limited to 'training/strategy/ti.py')
-rw-r--r-- | training/strategy/ti.py | 25 |
1 files changed, 12 insertions, 13 deletions
diff --git a/training/strategy/ti.py b/training/strategy/ti.py index 363c3f9..f00045f 100644 --- a/training/strategy/ti.py +++ b/training/strategy/ti.py | |||
@@ -142,25 +142,24 @@ def textual_inversion_strategy_callbacks( | |||
142 | checkpoint_output_dir / f"{slugify(token)}_{step}_{postfix}.bin" | 142 | checkpoint_output_dir / f"{slugify(token)}_{step}_{postfix}.bin" |
143 | ) | 143 | ) |
144 | 144 | ||
145 | @torch.no_grad() | 145 | @on_eval() |
146 | def on_sample(step): | 146 | def on_sample(step): |
147 | with ema_context(): | 147 | unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=True) |
148 | unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=True) | 148 | text_encoder_ = accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=True) |
149 | text_encoder_ = accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=True) | ||
150 | 149 | ||
151 | orig_unet_dtype = unet_.dtype | 150 | orig_unet_dtype = unet_.dtype |
152 | orig_text_encoder_dtype = text_encoder_.dtype | 151 | orig_text_encoder_dtype = text_encoder_.dtype |
153 | 152 | ||
154 | unet_.to(dtype=weight_dtype) | 153 | unet_.to(dtype=weight_dtype) |
155 | text_encoder_.to(dtype=weight_dtype) | 154 | text_encoder_.to(dtype=weight_dtype) |
156 | 155 | ||
157 | save_samples_(step=step, unet=unet_, text_encoder=text_encoder_) | 156 | save_samples_(step=step, unet=unet_, text_encoder=text_encoder_) |
158 | 157 | ||
159 | unet_.to(dtype=orig_unet_dtype) | 158 | unet_.to(dtype=orig_unet_dtype) |
160 | text_encoder_.to(dtype=orig_text_encoder_dtype) | 159 | text_encoder_.to(dtype=orig_text_encoder_dtype) |
161 | 160 | ||
162 | del unet_ | 161 | del unet_ |
163 | del text_encoder_ | 162 | del text_encoder_ |
164 | 163 | ||
165 | if torch.cuda.is_available(): | 164 | if torch.cuda.is_available(): |
166 | torch.cuda.empty_cache() | 165 | torch.cuda.empty_cache() |