diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index c1db6d6..dd395af 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -2254,6 +2254,16 @@ class BaseSDTrainProcess(BaseTrainProcess): pool_datasets=self.train_config.pool_datasets, ) + # a fresh text-embedding cache leaves the text encoder on the gpu (the + # datasets don't restore placement: that would bounce it per dataset). + # Park it once here, before accelerator.prepare brings the unet up. + if ( + self.is_caching_text_embeddings + and self.model_config.low_vram + and self.sd.text_encoder is not None + ): + self.sd.text_encoder_to('cpu') + flush() self.last_save_step = self.step_num ### HOOK ###