From 3d9cd16df741c6de4e74ea9bd2511de661f6c1c0 Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Fri, 23 Jan 2026 22:43:00 +0800 Subject: [PATCH] [examples] add --enable-cpu-offload args (#930) Signed-off-by: David Chen <530634352@qq.com> --- examples/offline_inference/image_to_image/image_edit.py | 6 ++++++ examples/offline_inference/image_to_image/image_to_image.md | 1 + examples/offline_inference/image_to_video/README.md | 1 + examples/offline_inference/image_to_video/image_to_video.py | 6 ++++++ examples/offline_inference/text_to_image/README.md | 1 + examples/offline_inference/text_to_image/text_to_image.py | 6 ++++++ examples/offline_inference/text_to_video/text_to_video.md | 1 + examples/offline_inference/text_to_video/text_to_video.py | 6 ++++++ 8 files changed, 28 insertions(+) diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 2f1a8166d..405a0cafa 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -279,6 +279,11 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Disable torch.compile and force eager execution.", ) + parser.add_argument( + "--enable-cpu-offload", + action="store_true", + help="Enable CPU offloading for diffusion models.", + ) return parser.parse_args() @@ -344,6 +349,7 @@ def main(): cache_config=cache_config, parallel_config=parallel_config, enforce_eager=args.enforce_eager, + enable_cpu_offload=args.enable_cpu_offload, ) print("Pipeline loaded") diff --git a/examples/offline_inference/image_to_image/image_to_image.md b/examples/offline_inference/image_to_image/image_to_image.md index 87bb8b68d..d624b4684 100644 --- a/examples/offline_inference/image_to_image/image_to_image.md +++ b/examples/offline_inference/image_to_image/image_to_image.md @@ -47,3 +47,4 @@ Key arguments: - `--guidance_scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models. - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower). - `--output`: path to save the generated PNG. +- `--enable-cpu-offload`: enable CPU offloading for diffusion models. diff --git a/examples/offline_inference/image_to_video/README.md b/examples/offline_inference/image_to_video/README.md index b2c141683..afebb1963 100644 --- a/examples/offline_inference/image_to_video/README.md +++ b/examples/offline_inference/image_to_video/README.md @@ -54,3 +54,4 @@ Key arguments: - `--num_inference_steps`: Number of denoising steps (default 50). - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video). - `--output`: Path to save the generated video. +- `--enable-cpu-offload`: enable CPU offloading for diffusion models. diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index f40eebaef..2104ec0fd 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -58,6 +58,11 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).") parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.") + parser.add_argument( + "--enable-cpu-offload", + action="store_true", + help="Enable CPU offloading for diffusion models.", + ) return parser.parse_args() @@ -105,6 +110,7 @@ def main(): vae_use_tiling=vae_use_tiling, boundary_ratio=args.boundary_ratio, flow_shift=args.flow_shift, + enable_cpu_offload=args.enable_cpu_offload, ) if profiler_enabled: diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md index daf014041..0f4351ac2 100644 --- a/examples/offline_inference/text_to_image/README.md +++ b/examples/offline_inference/text_to_image/README.md @@ -61,6 +61,7 @@ Key arguments: - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower). - `--height/--width`: output resolution (defaults 1024x1024). - `--output`: path to save the generated PNG. +- `--enable-cpu-offload`: enable CPU offloading for diffusion models. > ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes. diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index b07dbb27a..428c2437e 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -96,6 +96,11 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Disable torch.compile and force eager execution.", ) + parser.add_argument( + "--enable-cpu-offload", + action="store_true", + help="Enable CPU offloading for diffusion models.", + ) parser.add_argument( "--tensor_parallel_size", type=int, @@ -162,6 +167,7 @@ def main(): cache_config=cache_config, parallel_config=parallel_config, enforce_eager=args.enforce_eager, + enable_cpu_offload=args.enable_cpu_offload, ) if profiler_enabled: diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md index da622dcd2..f0ce5495a 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/examples/offline_inference/text_to_video/text_to_video.md @@ -28,3 +28,4 @@ Key arguments: - `--boundary_ratio`: Boundary split ratio for low/high DiT. - `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video). - `--output`: path to save the generated video. +- `--enable-cpu-offload`: enable CPU offloading for diffusion models. diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index f50692ce6..066714c54 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -35,6 +35,11 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument("--output", type=str, default="wan22_output.mp4", help="Path to save the video (mp4).") parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.") + parser.add_argument( + "--enable-cpu-offload", + action="store_true", + help="Enable CPU offloading for diffusion models.", + ) return parser.parse_args() @@ -56,6 +61,7 @@ def main(): vae_use_tiling=vae_use_tiling, boundary_ratio=args.boundary_ratio, flow_shift=args.flow_shift, + enable_cpu_offload=args.enable_cpu_offload, ) if profiler_enabled: