From 1484c43183a18488c73618abe3e64354deb1acbd Mon Sep 17 00:00:00 2001 From: Daniel Gu Date: Tue, 23 Dec 2025 10:56:32 +0100 Subject: [PATCH] Improve CPU offload support --- scripts/ltx2_test_full_pipeline.py | 3 +++ src/diffusers/pipelines/ltx2/pipeline_ltx2.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/scripts/ltx2_test_full_pipeline.py b/scripts/ltx2_test_full_pipeline.py index 1907be2da8..63c6f0400f 100644 --- a/scripts/ltx2_test_full_pipeline.py +++ b/scripts/ltx2_test_full_pipeline.py @@ -153,6 +153,7 @@ def parse_args(): parser.add_argument("--device", type=str, default="cuda:0") parser.add_argument("--dtype", type=str, default="bf16") + parser.add_argument("--cpu_offload", action="store_true") parser.add_argument( "--output_dir", @@ -179,6 +180,8 @@ def main(args): torch_dtype=args.dtype, ) pipeline.to(device=args.device) + if args.cpu_offload: + pipeline.enable_model_cpu_offload() video, audio = pipeline( prompt=args.prompt, diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py index 5aed290aa0..45cfc8e3cd 100644 --- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py @@ -194,7 +194,7 @@ class LTX2Pipeline(DiffusionPipeline, FromSingleFileMixin, LTXVideoLoraLoaderMix [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast). """ - model_cpu_offload_seq = "text_encoder->transformer->vae" + model_cpu_offload_seq = "text_encoder->transformer->vae->audio_vae->vocoder" _optional_components = [] _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]