xinyun/ComfyUI-CogVideoXWrapper
1
0
Fork 0
mirror of https://git.datalinker.icu/kijai/ComfyUI-CogVideoXWrapper.git synced 2026-05-18 13:16:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update nodes.py

Browse Source
This commit is contained in:
Brett Meirhofer 2024-08-27 08:58:52 -05:00
parent 8457fa7a4d
commit cea711c84d
1 changed files with 9 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
nodes.py
View File

@ -16,7 +16,13 @@ class DownloadAndLoadCogVideoModel:
def INPUT_TYPES(s): def INPUT_TYPES(s):
return { return {
"required": { "required": {
"model": (
[
"THUDM/CogVideoX-2b",
"THUDM/CogVideoX-5b",
],
{"default": "THUDM/CogVideoX-2b"},
),
}, },
"optional": { "optional": {
"precision": ( "precision": (
@ -35,7 +41,7 @@ class DownloadAndLoadCogVideoModel:
FUNCTION = "loadmodel" FUNCTION = "loadmodel"
CATEGORY = "CogVideoWrapper" CATEGORY = "CogVideoWrapper"
def loadmodel(self, precision): def loadmodel(self, model, precision):
device = mm.get_torch_device() device = mm.get_torch_device()
offload_device = mm.unet_offload_device() offload_device = mm.unet_offload_device()
mm.soft_empty_cache() mm.soft_empty_cache()
@ -49,7 +55,7 @@ class DownloadAndLoadCogVideoModel:
from huggingface_hub import snapshot_download from huggingface_hub import snapshot_download
snapshot_download( snapshot_download(
repo_id="THUDM/CogVideoX-2b", repo_id=model,
ignore_patterns=["*text_encoder*"], ignore_patterns=["*text_encoder*"],
local_dir=base_path, local_dir=base_path,
local_dir_use_symlinks=False, local_dir_use_symlinks=False,