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

Remove to device to avoid memory allocation errors

Browse Source 
Remove to(device), so it will apply the settings first
This commit is contained in:
Rodrigo Antônio de Araújo 2024-08-29 14:54:19 -03:00 committed by GitHub
parent fa8c830a4e
commit e226010ffc
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()