From 0767c7bc82645186159965c2a6be4278e33c6721 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Thu, 23 Mar 2023 11:07:57 +0100 Subject: Update --- training/strategy/lora.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) (limited to 'training/strategy/lora.py') diff --git a/training/strategy/lora.py b/training/strategy/lora.py index aa75bec..1c8fad6 100644 --- a/training/strategy/lora.py +++ b/training/strategy/lora.py @@ -47,7 +47,6 @@ def lora_strategy_callbacks( save_samples_ = partial( save_samples, accelerator=accelerator, - text_encoder=text_encoder, tokenizer=tokenizer, vae=vae, sample_scheduler=sample_scheduler, @@ -72,6 +71,7 @@ def lora_strategy_callbacks( @contextmanager def on_train(epoch: int): tokenizer.train() + text_encoder.train() yield @contextmanager @@ -89,8 +89,8 @@ def lora_strategy_callbacks( def on_checkpoint(step, postfix): print(f"Saving checkpoint for step {step}...") - unet_ = accelerator.unwrap_model(unet, False) - text_encoder_ = accelerator.unwrap_model(text_encoder, False) + unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=False) + text_encoder_ = accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=False) lora_config = {} state_dict = get_peft_model_state_dict(unet, state_dict=accelerator.get_state_dict(unet)) @@ -111,10 +111,10 @@ def lora_strategy_callbacks( @torch.no_grad() def on_sample(step): - unet_ = accelerator.unwrap_model(unet, False) - text_encoder_ = accelerator.unwrap_model(text_encoder, False) + unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=True) + text_encoder_ = accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=True) - save_samples_(step=step, unet=unet_) + save_samples_(step=step, unet=unet_, text_encoder=text_encoder_) del unet_ del text_encoder_ -- cgit v1.2.3-54-g00ecf