diff --git a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py index 502239212..7781346fb 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py @@ -20,6 +20,10 @@ from PIL import Image from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.managers.component_manager import ( + ComponentResidencyStrategy, + get_global_component_residency_manager, +) from sglang.multimodal_gen.runtime.models.vision_utils import ( load_image as load_vision_image, ) @@ -371,6 +375,8 @@ class DiffusersPipeline(ComposedPipelineBase): self._stage_name_mapping: dict[str, PipelineStage] = {} self.modules: dict[str, Any] = {} self.memory_usages: dict[str, float] = {} + self.component_residency_strategies: dict[str, ComponentResidencyStrategy] = {} + self.component_residency_manager = None self.post_init_called = False self.executor = executor or SyncExecutor(server_args=server_args) self._cache_dit_enabled = False @@ -726,7 +732,13 @@ class DiffusersPipeline(ComposedPipelineBase): """Execute the pipeline on the given batch.""" if not self.post_init_called: self.post_init() - return self.executor.execute(self.stages, batch, server_args) + + self.component_residency_manager = get_global_component_residency_manager( + self, server_args + ) + self.executor.component_residency_manager = self.component_residency_manager + + return self.executor.execute_with_profiling(self.stages, batch, server_args) @classmethod def from_pretrained(