From e3277b3be2225c3438850be7b19f022f34b3e08e Mon Sep 17 00:00:00 2001 From: Ratish P <114130421+Ratish1@users.noreply.github.com> Date: Tue, 17 Mar 2026 12:44:00 +0530 Subject: [PATCH] [diffusion]: remove stale offload-manager in LTX2 AV denoising (#20624) --- .../runtime/pipelines_core/stages/denoising_av.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py index 18b2d8dbe..55deb10f4 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py @@ -27,6 +27,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler from sglang.multimodal_gen.utils import PRECISION_TO_TYPE @@ -704,10 +705,9 @@ class LTX2AVDenoisingStage(DenoisingStage): batch.latents = latents batch.audio_latents = audio_latents - # 4. Cleanup - offload_mgr = getattr(self.transformer, "_layerwise_offload_manager", None) - if offload_mgr is not None and getattr(offload_mgr, "enabled", False): - offload_mgr.release_all() + if isinstance(self.transformer, OffloadableDiTMixin): + for manager in self.transformer.layerwise_offload_managers: + manager.release_all() def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify denoising stage inputs.