Refactor CPU Offloading Backend Pattern (#1223)

Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Samit <285365963@qq.com>

examples/offline_inference/image_to_image/image_edit.py

@@ -300,12 +300,6 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
-    parser.add_argument(
-        "--layerwise-num-gpu-layers",
-        type=int,
-        default=1,
-        help="Number of ready layers (blocks) to keep on GPU during generation.",
-    )
     return parser.parse_args()
 
 
@@ -362,7 +356,6 @@ def main():
     omni = Omni(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
-        layerwise_num_gpu_layers=args.layerwise_num_gpu_layers,
         vae_use_slicing=args.vae_use_slicing,
         vae_use_tiling=args.vae_use_tiling,
         cache_backend=args.cache_backend,

examples/offline_inference/image_to_video/image_to_video.py

@@ -80,12 +80,6 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
-    parser.add_argument(
-        "--layerwise-num-gpu-layers",
-        type=int,
-        default=1,
-        help="Number of ready layers (blocks) to keep on GPU during generation.",
-    )
     parser.add_argument(
         "--cfg-parallel-size",
         type=int,
@@ -139,7 +133,6 @@ def main():
     omni = Omni(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
-        layerwise_num_gpu_layers=args.layerwise_num_gpu_layers,
         vae_use_slicing=args.vae_use_slicing,
         vae_use_tiling=args.vae_use_tiling,
         boundary_ratio=args.boundary_ratio,

examples/offline_inference/text_to_image/text_to_image.py

@@ -113,12 +113,6 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
-    parser.add_argument(
-        "--layerwise-num-gpu-layers",
-        type=int,
-        default=1,
-        help="Number of ready layers (blocks) to keep on GPU during generation.",
-    )
     parser.add_argument(
         "--quantization",
         type=str,
@@ -222,7 +216,6 @@ def main():
     omni = Omni(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
-        layerwise_num_gpu_layers=args.layerwise_num_gpu_layers,
         vae_use_slicing=args.vae_use_slicing,
         vae_use_tiling=args.vae_use_tiling,
         cache_backend=args.cache_backend,

examples/offline_inference/text_to_video/text_to_video.py

@@ -84,12 +84,6 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
-    parser.add_argument(
-        "--layerwise-num-gpu-layers",
-        type=int,
-        default=1,
-        help="Number of ready layers (blocks) to keep on GPU during generation.",
-    )
     parser.add_argument(
         "--ulysses-degree",
         type=int,
@@ -155,7 +149,6 @@ def main():
     omni = Omni(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
-        layerwise_num_gpu_layers=args.layerwise_num_gpu_layers,
         vae_use_slicing=args.vae_use_slicing,
         vae_use_tiling=args.vae_use_tiling,
         boundary_ratio=args.boundary_ratio,