From a630bb33140a7a90f55c818f9c7d9638be90e188 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Sat, 9 Nov 2024 22:17:10 +0200 Subject: [PATCH] update --- model_loading.py | 89 +++++++++++++++++++++++++++++++++---------- nodes.py | 1 + pipeline_cogvideox.py | 4 +- 3 files changed, 71 insertions(+), 23 deletions(-) diff --git a/model_loading.py b/model_loading.py index 896c8a4..ad0dfc4 100644 --- a/model_loading.py +++ b/model_loading.py @@ -157,7 +157,7 @@ class DownloadAndLoadCogVideoModel: base_path = os.path.join(download_path, "CogVideoX-5b-1.5") download_path = base_path subfolder = "transformer_T2V" if "1.5-T2V" in model else "transformer_I2V" - allow_patterns = [f"*{subfolder}*"] + allow_patterns = [f"*{subfolder}*", "*vae*", "*scheduler*"] repo_id = "kijai/CogVideoX-5b-1.5" else: base_path = os.path.join(download_path, (model.split("/")[-1])) @@ -204,15 +204,17 @@ class DownloadAndLoadCogVideoModel: #fp8 if fp8_transformer == "enabled" or fp8_transformer == "fastmode": + params_to_keep = {"patch_embed", "lora", "pos_embedding", "time_embedding"} + if "1.5" in model: + params_to_keep = {"patch_embed", "lora", "pos_embedding", "time_embedding", "norm","ofs_embedding", "norm_final", "norm_out", "proj_out"} for name, param in transformer.named_parameters(): - params_to_keep = {"patch_embed", "lora", "pos_embedding", "time_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 if "1.5" in model: - params_to_keep = {"norm","ff"} + params_to_keep.update({"ff"}) convert_fp8_linear(transformer, dtype, params_to_keep=params_to_keep) with open(scheduler_path) as f: @@ -311,12 +313,12 @@ class DownloadAndLoadCogVideoGGUFModel: [ "CogVideoX_5b_GGUF_Q4_0.safetensors", "CogVideoX_5b_I2V_GGUF_Q4_0.safetensors", + "CogVideoX_5b_1_5_I2V_GGUF_Q4_0.safetensors", "CogVideoX_5b_fun_GGUF_Q4_0.safetensors", "CogVideoX_5b_fun_1_1_GGUF_Q4_0.safetensors", "CogVideoX_5b_fun_1_1_Pose_GGUF_Q4_0.safetensors", "CogVideoX_5b_Interpolation_GGUF_Q4_0.safetensors", "CogVideoX_5b_Tora_GGUF_Q4_0.safetensors", - ], ), "vae_precision": (["fp16", "fp32", "bf16"], {"default": "bf16", "tooltip": "VAE dtype"}), @@ -327,8 +329,9 @@ class DownloadAndLoadCogVideoGGUFModel: "optional": { "pab_config": ("PAB_CONFIG", {"default": None}), "block_edit": ("TRANSFORMERBLOCKS", {"default": None}), + #"lora": ("COGLORA", {"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"}), - + "attention_mode": (["sdpa", "sageattn"], {"default": "sdpa"}), } } @@ -337,7 +340,8 @@ class DownloadAndLoadCogVideoGGUFModel: FUNCTION = "loadmodel" CATEGORY = "CogVideoWrapper" - def loadmodel(self, model, vae_precision, fp8_fastmode, load_device, enable_sequential_cpu_offload, pab_config=None, block_edit=None, compile="disabled"): + def loadmodel(self, model, vae_precision, fp8_fastmode, load_device, enable_sequential_cpu_offload, + pab_config=None, block_edit=None, compile="disabled", attention_mode="sdpa"): check_diffusers_version() @@ -375,7 +379,7 @@ class DownloadAndLoadCogVideoGGUFModel: with open(transformer_path) as f: transformer_config = json.load(f) - sd = load_torch_file(gguf_path) + from . import mz_gguf_loader import importlib @@ -393,6 +397,13 @@ class DownloadAndLoadCogVideoGGUFModel: transformer = CogVideoXTransformer3DModelFun.from_config(transformer_config) elif "I2V" in model or "Interpolation" in model: transformer_config["in_channels"] = 32 + if "1_5" in model: + transformer_config["ofs_embed_dim"] = 512 + transformer_config["use_learned_positional_embeddings"] = False + transformer_config["patch_size_t"] = 2 + transformer_config["patch_bias"] = False + transformer_config["sample_height"] = 96 + transformer_config["sample_width"] = 170 if pab_config is not None: transformer = CogVideoXTransformer3DModelPAB.from_config(transformer_config) else: @@ -405,23 +416,23 @@ class DownloadAndLoadCogVideoGGUFModel: transformer = CogVideoXTransformer3DModel.from_config(transformer_config) if "2b" in model: - for name, param in transformer.named_parameters(): - if name != "pos_embedding": - param.data = param.data.to(torch.float8_e4m3fn) - else: - param.data = param.data.to(torch.float16) - else: - transformer.to(torch.float8_e4m3fn) - + params_to_keep = {"patch_embed", "pos_embedding", "time_embedding"} + cast_dtype = torch.float16 + elif "1_5" in model: + params_to_keep = {"patch_embed", "time_embedding", "ofs_embedding", "norm_final", "norm_out", "proj_out", "norm"} + cast_dtype = torch.bfloat16 + for name, param in transformer.named_parameters(): + if not any(keyword in name for keyword in params_to_keep): + param.data = param.data.to(torch.bfloat16) + else: + param.data = param.data.to(cast_dtype) + #for name, param in transformer.named_parameters(): + # print(name, param.data.dtype) + if block_edit is not None: transformer = remove_specific_blocks(transformer, block_edit) - transformer = mz_gguf_loader.quantize_load_state_dict(transformer, sd, device="cpu") - if load_device == "offload_device": - transformer.to(offload_device) - else: - transformer.to(device) - + transformer.attention_mode = attention_mode if fp8_fastmode: from .fp8_optimization import convert_fp8_linear @@ -468,6 +479,42 @@ class DownloadAndLoadCogVideoGGUFModel: if enable_sequential_cpu_offload: pipe.enable_sequential_cpu_offload() + sd = load_torch_file(gguf_path) + + # #LoRAs + # if lora is not None: + # if "fun" in model.lower(): + # raise NotImplementedError("LoRA with GGUF is not supported for Fun models") + # from .lora_utils import merge_lora#, load_lora_into_transformer + # #for l in lora: + # # log.info(f"Merging LoRA weights from {l['path']} with strength {l['strength']}") + # # pipe.transformer = merge_lora(pipe.transformer, l["path"], l["strength"]) + # else: + # adapter_list = [] + # adapter_weights = [] + # for l in lora: + # lora_sd = load_torch_file(l["path"]) + # for key, val in lora_sd.items(): + # if "lora_B" in key: + # lora_rank = val.shape[1] + # break + # log.info(f"Loading rank {lora_rank} LoRA weights from {l['path']} with strength {l['strength']}") + # adapter_name = l['path'].split("/")[-1].split(".")[0] + # adapter_weight = l['strength'] + # pipe.load_lora_weights(l['path'], weight_name=l['path'].split("/")[-1], lora_rank=lora_rank, adapter_name=adapter_name) + + # #transformer = load_lora_into_transformer(lora, transformer) + # adapter_list.append(adapter_name) + # adapter_weights.append(adapter_weight) + # for l in lora: + # pipe.set_adapters(adapter_list, adapter_weights=adapter_weights) + # #pipe.fuse_lora(lora_scale=1 / lora_rank, components=["transformer"]) + + pipe.transformer = mz_gguf_loader.quantize_load_state_dict(pipe.transformer, sd, device="cpu") + if load_device == "offload_device": + pipe.transformer.to(offload_device) + else: + pipe.transformer.to(device) pipeline = { diff --git a/nodes.py b/nodes.py index 4603799..7fa9606 100644 --- a/nodes.py +++ b/nodes.py @@ -821,6 +821,7 @@ class CogVideoSampler: assert ( "I2V" not in pipeline.get("model_name", "") or "1.5" in pipeline.get("model_name", "") or + "1_5" in pipeline.get("model_name", "") or num_frames == 49 or context_options is not None ), "1.0 I2V model can only do 49 frames" diff --git a/pipeline_cogvideox.py b/pipeline_cogvideox.py index 4704255..22bb8e7 100644 --- a/pipeline_cogvideox.py +++ b/pipeline_cogvideox.py @@ -21,7 +21,7 @@ import torch.nn.functional as F import math from diffusers.models import AutoencoderKLCogVideoX#, CogVideoXTransformer3DModel -from diffusers.pipelines.pipeline_utils import DiffusionPipeline +#from diffusers.pipelines.pipeline_utils import DiffusionPipeline from diffusers.schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler from diffusers.utils import logging from diffusers.utils.torch_utils import randn_tensor @@ -115,7 +115,7 @@ def retrieve_timesteps( timesteps = scheduler.timesteps return timesteps, num_inference_steps -class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin): +class CogVideoXPipeline(VideoSysPipeline, CogVideoXLoraLoaderMixin): r""" Pipeline for text-to-video generation using CogVideoX.