Merge pull request #210 from rodjjo/patch-1

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

View File

@ -19,9 +19,8 @@ from openai import OpenAI
import moviepy.editor as mp
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_sequential_cpu_offload()
pipe.vae.enable_slicing()