From a35d833abf044a74565a0c0bb60152348e2d17e2 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Tue, 22 Sep 2026 12:08:29 -0600 Subject: [PATCH] Fix text encoder unload skipping models whose pipeline does not expose the text encoder (Qwen-Image-2.1, Z-Image, LTX2), so cached-embedding runs actually free the TE. --- toolkit/unloader.py | 34 ++++++++++++++++++++++++---------- 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/toolkit/unloader.py b/toolkit/unloader.py index 3364755..ea9eb3f 100644 --- a/toolkit/unloader.py +++ b/toolkit/unloader.py @@ -40,26 +40,40 @@ def unload_text_encoder(model: "BaseModel"): if model.text_encoder is not None: if isinstance(model.text_encoder, list): - text_encoder_list = [] pipe = model.pipeline + freed = set() + + def _free(te): + if te is None or isinstance(te, FakeTextEncoder) or id(te) in freed: + return + MemoryManager.free(te) + freed.add(id(te)) + + def _new_fake(): + return FakeTextEncoder(device=model.device_torch, dtype=model.torch_dtype) + + # model.text_encoder is the source of truth; some pipelines (self-referencing + # or built with text_encoder=None) never expose the TE as an attribute + for real_te in model.text_encoder: + _free(real_te) + fakes = [_new_fake() for _ in model.text_encoder] + model.text_encoder = fakes + + def _fake(idx): + return fakes[idx] if idx < len(fakes) else _new_fake() # the pipeline stores text encoders like text_encoder, text_encoder_2, text_encoder_3, etc. if getattr(pipe, "text_encoder", None) is not None: - MemoryManager.free(pipe.text_encoder) - te = FakeTextEncoder(device=model.device_torch, dtype=model.torch_dtype) - text_encoder_list.append(te) - pipe.text_encoder = te + _free(pipe.text_encoder) + pipe.text_encoder = _fake(0) i = 2 while hasattr(pipe, f"text_encoder_{i}"): real_te = getattr(pipe, f"text_encoder_{i}") if real_te is not None: - MemoryManager.free(real_te) - te = FakeTextEncoder(device=model.device_torch, dtype=model.torch_dtype) - text_encoder_list.append(te) - setattr(pipe, f"text_encoder_{i}", te) + _free(real_te) + setattr(pipe, f"text_encoder_{i}", _fake(i - 1)) i += 1 - model.text_encoder = text_encoder_list else: # only has a single text encoder MemoryManager.free(model.text_encoder)