From 1ce3983d686df7314e484d2217251a6eaaa3ec82 Mon Sep 17 00:00:00 2001 From: Yongqi Chen Date: Fri, 6 Dec 2024 12:45:08 -0500 Subject: [PATCH] add cpu offload --- fastvideo/demo/gradio_web_demo.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/fastvideo/demo/gradio_web_demo.py b/fastvideo/demo/gradio_web_demo.py index 173517012..bb4bd810d 100644 --- a/fastvideo/demo/gradio_web_demo.py +++ b/fastvideo/demo/gradio_web_demo.py @@ -26,6 +26,7 @@ def init_args(): parser.add_argument("--num_euler_timesteps", type=int, default=100) parser.add_argument("--linear_threshold", type=float, default=0.025) parser.add_argument("--linear_range", type=float, default=0.5) + parser.add_argument("--cpu_offload", action="store_true") return parser.parse_args() def load_model(args): @@ -43,7 +44,8 @@ def load_model(args): pipe = MochiPipeline.from_pretrained(args.model_path, transformer=transformer, scheduler=scheduler) pipe.enable_vae_tiling() pipe.to(device) - pipe.enable_model_cpu_offload() + if args.cpu_offload: + pipe.enable_model_cpu_offload() return pipe def generate_video(prompt, negative_prompt, use_negative_prompt, seed, guidance_scale,