From a1ddeeef13443825b2b548e1d635cd11ae7d5b53 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 12 Aug 2026 10:46:58 -0600 Subject: [PATCH] Fixed issue with offloading text encoder on ltx 2.5 --- extensions_built_in/diffusion_models/ltx2/ltx2.py | 7 ++++++- toolkit/memory_management/manager.py | 6 +++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/extensions_built_in/diffusion_models/ltx2/ltx2.py b/extensions_built_in/diffusion_models/ltx2/ltx2.py index ae0d4df6..bbf5e577 100644 --- a/extensions_built_in/diffusion_models/ltx2/ltx2.py +++ b/extensions_built_in/diffusion_models/ltx2/ltx2.py @@ -1536,11 +1536,16 @@ class LTX25Model(LTX2Model): self.model_config.layer_offloading and self.model_config.layer_offloading_text_encoder_percent > 0 ): + # layer_scalar is a bare tensor buffer on each decoder layer; the + # manager never enumerates it, so it must ride along explicitly + ignore_modules = [text_encoder.embed_tokens] + for layer in text_encoder.layers: + ignore_modules.append(layer.layer_scalar) MemoryManager.attach( text_encoder, self.device_torch, offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ignore_modules=[text_encoder.embed_tokens], + ignore_modules=ignore_modules, ) text_encoder.to(self.device_torch) diff --git a/toolkit/memory_management/manager.py b/toolkit/memory_management/manager.py index 7e8267a8..910c426b 100644 --- a/toolkit/memory_management/manager.py +++ b/toolkit/memory_management/manager.py @@ -56,8 +56,8 @@ class MemoryManager: def memory_managed_to(self, *args, **kwargs): # first move all the unmanaged modules for module in self.unmanaged_modules: - if isinstance(module, torch.nn.Parameter): - # Parameter cannot move this way + if isinstance(module, torch.Tensor): + # Parameters and bare tensor buffers cannot move this way module.data = module.data.to(*args, **kwargs) else: module.to(*args, **kwargs) @@ -181,7 +181,7 @@ class MemoryManager: for unmanaged in module._memory_manager.unmanaged_modules: try: - if isinstance(unmanaged, torch.nn.Parameter): + if isinstance(unmanaged, torch.Tensor): unmanaged.data = unmanaged.data.to('cpu') else: unmanaged.to('cpu')