[examples] add --enable-cpu-offload args (#930)

Signed-off-by: David Chen <530634352@qq.com>

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")
 

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.

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.

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:

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.
 

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:

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.

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: