chore:enable model cpu offloading in all tasks pipelines

livepeer · Dec 29, 2024 · 441e48c · 441e48c
1 parent 7cd8af6
commit 441e48c
Showing 1 changed file with 3 additions and 0 deletions.
diff --git 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)