Update gradio_demo.py

使用huggingface下载模型路径
This commit is contained in:
learningpro 2024-08-06 15:16:34 +08:00 committed by GitHub
parent 239d912b6f
commit b158fbdadd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -17,7 +17,7 @@ import PIL
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = CogVideoXPipeline.from_pretrained("/share/home/zyx/Models/cogvideox-hf-0805", torch_dtype=dtype).to(device)
pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-2b", torch_dtype=dtype).to(device)
sys_prompt = """You are part of a team of bots that creates videos. You work with an assistant bot that will draw anything you say in square brackets.