5 Commits
Author SHA1 Message Date
Jukka Seppänen 909d7026f3 Update model_loading.py 2024-11-09 20:20:15 +02:00
kijai 806a0fa1d6 Update model_loading.py 2024-11-08 21:31:31 +02:00
kijai f7a88cbd56 Update model_loading.py 2024-11-08 21:23:29 +02:00
kijai 4c2ce52f57 Update model_loading.py 2024-11-08 17:30:51 +02:00
Jukka Seppänen 4a597f1955 Update requirements.txt 2024-11-08 16:43:32 +02:00
2 changed files with 18 additions and 19 deletions
+16 -17
View File
@@ -4,6 +4,8 @@ import torch.nn as nn
import json import json
import folder_paths import folder_paths
import comfy.model_management as mm import comfy.model_management as mm
from .utils import check_diffusers_version, remove_specific_blocks, log
check_diffusers_version()
from diffusers.models import AutoencoderKLCogVideoX from diffusers.models import AutoencoderKLCogVideoX
from diffusers.schedulers import CogVideoXDDIMScheduler from diffusers.schedulers import CogVideoXDDIMScheduler
@@ -20,7 +22,6 @@ from .cogvideox_fun.pipeline_cogvideox_control import CogVideoX_Fun_Pipeline_Con
from .videosys.cogvideox_transformer_3d import CogVideoXTransformer3DModel as CogVideoXTransformer3DModelPAB from .videosys.cogvideox_transformer_3d import CogVideoXTransformer3DModel as CogVideoXTransformer3DModelPAB
from .utils import check_diffusers_version, remove_specific_blocks, log
from comfy.utils import load_torch_file from comfy.utils import load_torch_file
script_directory = os.path.dirname(os.path.abspath(__file__)) script_directory = os.path.dirname(os.path.abspath(__file__))
@@ -32,7 +33,7 @@ class CogVideoLoraSelect:
"required": { "required": {
"lora": (folder_paths.get_filename_list("cogvideox_loras"), "lora": (folder_paths.get_filename_list("cogvideox_loras"),
{"tooltip": "LORA models are expected to be in ComfyUI/models/CogVideo/loras with .safetensors extension"}), {"tooltip": "LORA models are expected to be in ComfyUI/models/CogVideo/loras with .safetensors extension"}),
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "LORA strength, set to 0.0 to unmerge the LORA"}), "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.0001, "tooltip": "LORA strength, set to 0.0 to unmerge the LORA"}),
}, },
"optional": { "optional": {
"prev_lora":("COGLORA", {"default": None, "tooltip": "For loading multiple LoRAs"}), "prev_lora":("COGLORA", {"default": None, "tooltip": "For loading multiple LoRAs"}),
@@ -182,16 +183,7 @@ class DownloadAndLoadCogVideoModel:
if block_edit is not None: if block_edit is not None:
transformer = remove_specific_blocks(transformer, block_edit) transformer = remove_specific_blocks(transformer, block_edit)
#fp8
if fp8_transformer == "enabled" or fp8_transformer == "fastmode":
for name, param in transformer.named_parameters():
params_to_keep = {"patch_embed", "lora", "pos_embedding"}
if not any(keyword in name for keyword in params_to_keep):
param.data = param.data.to(torch.float8_e4m3fn)
if fp8_transformer == "fastmode":
from .fp8_optimization import convert_fp8_linear
convert_fp8_linear(transformer, dtype)
with open(scheduler_path) as f: with open(scheduler_path) as f:
scheduler_config = json.load(f) scheduler_config = json.load(f)
@@ -221,8 +213,7 @@ class DownloadAndLoadCogVideoModel:
adapter_list = [] adapter_list = []
adapter_weights = [] adapter_weights = []
for l in lora: for l in lora:
if l["fuse_lora"]: fuse = True if l["fuse_lora"] else False
fuse = True
lora_sd = load_torch_file(l["path"]) lora_sd = load_torch_file(l["path"])
for key, val in lora_sd.items(): for key, val in lora_sd.items():
if "lora_B" in key: if "lora_B" in key:
@@ -239,9 +230,18 @@ class DownloadAndLoadCogVideoModel:
for l in lora: for l in lora:
pipe.set_adapters(adapter_list, adapter_weights=adapter_weights) pipe.set_adapters(adapter_list, adapter_weights=adapter_weights)
if fuse: if fuse:
pipe.fuse_lora(lora_scale=1 / lora_rank, components=["transformer"]) pipe.fuse_lora(lora_scale=lora[-1]["strength"] / lora_rank, components=["transformer"])
#fp8
if fp8_transformer == "enabled" or fp8_transformer == "fastmode":
for name, param in pipe.transformer.named_parameters():
params_to_keep = {"patch_embed", "lora", "pos_embedding"}
if not any(keyword in name for keyword in params_to_keep):
param.data = param.data.to(torch.float8_e4m3fn)
if fp8_transformer == "fastmode":
from .fp8_optimization import convert_fp8_linear
convert_fp8_linear(pipe.transformer, dtype)
if enable_sequential_cpu_offload: if enable_sequential_cpu_offload:
pipe.enable_sequential_cpu_offload() pipe.enable_sequential_cpu_offload()
@@ -312,7 +312,6 @@ class DownloadAndLoadCogVideoGGUFModel:
"pab_config": ("PAB_CONFIG", {"default": None}), "pab_config": ("PAB_CONFIG", {"default": None}),
"block_edit": ("TRANSFORMERBLOCKS", {"default": None}), "block_edit": ("TRANSFORMERBLOCKS", {"default": None}),
"compile": (["disabled","torch"], {"tooltip": "compile the model for faster inference, these are advanced options only available on Linux, see readme for more info"}), "compile": (["disabled","torch"], {"tooltip": "compile the model for faster inference, these are advanced options only available on Linux, see readme for more info"}),
} }
} }
@@ -633,4 +632,4 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"DownloadAndLoadCogVideoControlNet": "(Down)load CogVideo ControlNet", "DownloadAndLoadCogVideoControlNet": "(Down)load CogVideo ControlNet",
"DownloadAndLoadToraModel": "(Down)load Tora Model", "DownloadAndLoadToraModel": "(Down)load Tora Model",
"CogVideoLoraSelect": "CogVideo LoraSelect", "CogVideoLoraSelect": "CogVideo LoraSelect",
} }
+2 -2
View File
@@ -1,5 +1,5 @@
huggingface_hub huggingface_hub
diffusers>=0.30.3 diffusers>=0.31.0
accelerate>=0.33.0 accelerate>=0.33.0
einops einops
peft peft