From 543a0d74d06fad88a155d74b041dda693c919c0d Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Tue, 22 Sep 2026 12:03:09 -0600 Subject: [PATCH] Fix issue with saving device state when a text encoder is unloaded. --- toolkit/models/base_model.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py index 8468595..35a28d5 100644 --- a/toolkit/models/base_model.py +++ b/toolkit/models/base_model.py @@ -26,6 +26,7 @@ from toolkit.prompt_utils import inject_trigger_into_prompt, PromptEmbeds, conca from toolkit.reference_adapter import ReferenceAdapter from toolkit.sd_device_states_presets import empty_preset from toolkit.train_tools import get_torch_dtype, apply_noise_offset +from toolkit.unloader import FakeTextEncoder import torch from toolkit.pipelines import CustomStableDiffusionXLPipeline from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, T2IAdapter, DDPMScheduler, \ @@ -1499,8 +1500,10 @@ class BaseModel: } if isinstance(self.text_encoder, list): self.device_state['text_encoder']: List[dict] = [] + # unloaded TEs are FakeTextEncoder stubs; arch probes into TE internals would raise + any_fake = any(isinstance(e, FakeTextEncoder) for e in self.text_encoder) for encoder in self.text_encoder: - te_has_grad = self.get_te_has_grad() + te_has_grad = False if any_fake else self.get_te_has_grad() self.device_state['text_encoder'].append({ 'training': encoder.training, 'device': encoder.device, @@ -1508,7 +1511,10 @@ class BaseModel: 'requires_grad': te_has_grad }) elif self.text_encoder is not None: - te_has_grad = self.get_te_has_grad() + if isinstance(self.text_encoder, FakeTextEncoder): + te_has_grad = False + else: + te_has_grad = self.get_te_has_grad() self.device_state['text_encoder'] = { 'training': self.text_encoder.training,