From 441e48c78b1e490d4f28ec661cbfaae516804019 Mon Sep 17 00:00:00 2001 From: RUFFY-369 Date: Mon, 30 Dec 2024 01:28:30 +0530 Subject: [PATCH] chore:enable model cpu offloading in all tasks pipelines --- runner/app/pipelines/image_to_image_generic.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/runner/app/pipelines/image_to_image_generic.py b/runner/app/pipelines/image_to_image_generic.py index 5986737dd..881d9d7df 100644 --- a/runner/app/pipelines/image_to_image_generic.py +++ b/runner/app/pipelines/image_to_image_generic.py @@ -93,9 +93,11 @@ def __init__(self, model_id: str, task: str): safety_checker=None, **kwargs, ).to(torch_device) + self.pipeline_stage1.enable_model_cpu_offload() self.pipeline_stage2 = StableDiffusionXLInpaintPipeline.from_pretrained( "OzzyGT/RealVisXL_V4.0_inpainting", vae=self.vae, **kwargs ).to(torch_device) + self.pipeline_stage1.enable_model_cpu_offload() elif self.task == TaskType.SKETCH_TO_IMAGE.value: self.controlnet = ControlNetModel.from_pretrained(model_id, **kwargs).to( @@ -115,6 +117,7 @@ def __init__(self, model_id: str, task: str): scheduler=eulera_scheduler, **kwargs, ).to(torch_device) + self.pipeline.enable_model_cpu_offload() self._lora_loader = LoraLoader(self.pipeline)