nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add cpu offload

Browse Source
This commit is contained in:
Yongqi Chen
2024-12-06 12:45:08 -05:00
parent ee241cfa4d
commit 1ce3983d68
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/demo/gradio_web_demo.py
+3 -1
View File
@@ -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,