JLNTV_IPD/CogVideo
4
0
Fork 0
mirror of https://github.com/THUDM/CogVideo.git synced 2025-04-06 03:57:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #210 from rodjjo/patch-1

Browse Source 
Remove to device to avoid memory allocation errors
This commit is contained in:
zR 2024-08-30 19:36:01 +08:00 committed by GitHub
commit 7ddd59a811
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
inference/gradio_web_demo.py
View File

@ -19,9 +19,8 @@ from openai import OpenAI
import moviepy.editor as mp import moviepy.editor as mp
dtype = torch.bfloat16 dtype = torch.bfloat16
device = "cuda" # Need to use cuda
pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=dtype).to(device) pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=dtype)
pipe.enable_model_cpu_offload() pipe.enable_model_cpu_offload()
pipe.enable_sequential_cpu_offload() pipe.enable_sequential_cpu_offload()
pipe.vae.enable_slicing() pipe.vae.enable_slicing()