diff options
Diffstat (limited to 'training/strategy/dreambooth.py')
-rw-r--r-- | training/strategy/dreambooth.py | 27 |
1 files changed, 13 insertions, 14 deletions
diff --git a/training/strategy/dreambooth.py b/training/strategy/dreambooth.py index 42624cd..7cdfc7f 100644 --- a/training/strategy/dreambooth.py +++ b/training/strategy/dreambooth.py | |||
@@ -113,7 +113,7 @@ def dreambooth_strategy_callbacks( | |||
113 | accelerator.clip_grad_norm_(itertools.chain(*params_to_clip), max_grad_norm) | 113 | accelerator.clip_grad_norm_(itertools.chain(*params_to_clip), max_grad_norm) |
114 | 114 | ||
115 | @torch.no_grad() | 115 | @torch.no_grad() |
116 | def on_after_optimize(_, lr: float): | 116 | def on_after_optimize(_, lrs: dict[str, float]): |
117 | if ema_unet is not None: | 117 | if ema_unet is not None: |
118 | ema_unet.step(unet.parameters()) | 118 | ema_unet.step(unet.parameters()) |
119 | 119 | ||
@@ -149,25 +149,24 @@ def dreambooth_strategy_callbacks( | |||
149 | if torch.cuda.is_available(): | 149 | if torch.cuda.is_available(): |
150 | torch.cuda.empty_cache() | 150 | torch.cuda.empty_cache() |
151 | 151 | ||
152 | @torch.no_grad() | 152 | @on_eval() |
153 | def on_sample(step): | 153 | def on_sample(step): |
154 | with ema_context(): | 154 | unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=True) |
155 | unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=True) | 155 | text_encoder_ = accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=True) |
156 | text_encoder_ = accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=True) | ||
157 | 156 | ||
158 | orig_unet_dtype = unet_.dtype | 157 | orig_unet_dtype = unet_.dtype |
159 | orig_text_encoder_dtype = text_encoder_.dtype | 158 | orig_text_encoder_dtype = text_encoder_.dtype |
160 | 159 | ||
161 | unet_.to(dtype=weight_dtype) | 160 | unet_.to(dtype=weight_dtype) |
162 | text_encoder_.to(dtype=weight_dtype) | 161 | text_encoder_.to(dtype=weight_dtype) |
163 | 162 | ||
164 | save_samples_(step=step, unet=unet_, text_encoder=text_encoder_) | 163 | save_samples_(step=step, unet=unet_, text_encoder=text_encoder_) |
165 | 164 | ||
166 | unet_.to(dtype=orig_unet_dtype) | 165 | unet_.to(dtype=orig_unet_dtype) |
167 | text_encoder_.to(dtype=orig_text_encoder_dtype) | 166 | text_encoder_.to(dtype=orig_text_encoder_dtype) |
168 | 167 | ||
169 | del unet_ | 168 | del unet_ |
170 | del text_encoder_ | 169 | del text_encoder_ |
171 | 170 | ||
172 | if torch.cuda.is_available(): | 171 | if torch.cuda.is_available(): |
173 | torch.cuda.empty_cache() | 172 | torch.cuda.empty_cache() |