nodes/MinusZoneAI-ComfyUI-CogVideoX-MZ
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

'update'

Browse Source
This commit is contained in:
wailovet
2024-09-05 12:29:13 +08:00
parent c6eca73aaa
commit 0696faaa5f
6 changed files with 298 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+6
View File
@@ -2,6 +2,12 @@
# ComfyUI-CogVideoX-MZ
## Update - 2024-09-05
部分层4位量化模型, 采样方式选择DPM, weight_dtype选择GGUF, cpu_offload不启用分辨率480x720在32帧下占用显存小于8G
[CogVideoX_5b_GGUF_Q4_0.safetensors](https://modelscope.cn/models/wailovet/CogVideoX-5b/resolve/master/CogVideoX_5b_GGUF_Q4_0.safetensors) 下载到unet文件夹
4位量化工作流在examples\workflow_GGUF_Q4_0.png
## 使用
CogVideoX-5B的模型加载器
依赖 https://github.com/kijai/ComfyUI-CogVideoXWrapper
__init__.py
+1 -1
View File
@@ -26,7 +26,7 @@ class MZ_CogVideoXLoader:
"required": {
"unet_name": (folder_paths.get_filename_list("unet"), ),
"vae_name": (folder_paths.get_filename_list("vae"), ),
"weight_dtype": (["default", "fp8_e4m3fn", "fp8_e5m2"],),
"weight_dtype": (["default", "fp8_e4m3fn", "fp8_e5m2", "GGUF"],),
"fp8_fast_mode": ("BOOLEAN", {"default": False}),
# "dyn_offload_cpu_layer": ("INT", {"default": 0, "tooltip": "0-42,默认0不启用"}),
"enable_sequential_cpu_offload": ("BOOLEAN", {"default": False, "tooltip": "significantly reducing memory usage and slows down the inference"}),
examples/workflow_GGUF_Q4_0.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 348 KiB

mz_cogvideox_core.py
+19 -6
View File
@@ -263,20 +263,33 @@ def MZ_CogVideoXLoader_call(args={}):
os.path.dirname(__file__),
"configs5b",
)
transformer = CogVideoXTransformer3DModel.from_config(
transformer_config)
transformer.load_state_dict(unet_sd)
dtype = None
weight_dtype = args.get("weight_dtype")
transformer = None
if weight_dtype not in ["GGUF"]:
transformer = CogVideoXTransformer3DModel.from_config(
transformer_config)
transformer.load_state_dict(unet_sd)
if weight_dtype == "fp8_e4m3fn":
dtype = torch.float8_e4m3fn
transformer.to(dtype)
elif weight_dtype == "fp8_e5m2":
dtype = torch.float8_e5m2
transformer.to(dtype)
elif weight_dtype == "GGUF":
dtype = torch.float8_e4m3fn
from . import mz_gguf_loader
import importlib
importlib.reload(mz_gguf_loader)
with mz_gguf_loader.quantize_lazy_load():
transformer = CogVideoXTransformer3DModel.from_config(
transformer_config)
transformer.to(dtype)
transformer = mz_gguf_loader.quantize_load_state_dict(
transformer, unet_sd, device="cpu")
transformer.to(device)
else:
dtype = transformer.parameters().__next__().dtype
mz_gguf_loader.py
+185
View File
@@ -0,0 +1,185 @@
import torch
import torch.nn as nn
import gc
class quantize_lazy_load():
def __init__(self):
self.device = None
def __enter__(self):
self.device = torch.device("meta")
self.device.__enter__()
return self
def __exit__(self, exc_type, exc_value, traceback):
self.device.__exit__(exc_type, exc_value, traceback)
def quantize_load_state_dict(model, state_dict, device="cpu"):
Q4_0_qkey = []
for key in state_dict.keys():
if key.endswith(".Q4_0_qweight"):
Q4_0_qkey.append(key.replace(".Q4_0_qweight", ""))
for name, module in model.named_modules():
if name in Q4_0_qkey:
q_linear = WQLinear_GGUF.from_linear(
linear=module,
device=device,
qtype="Q4_0",
)
set_op_by_name(model, name, q_linear)
model.to_empty(device=device)
model.load_state_dict(state_dict, strict=False)
model.to(device)
return model
def set_op_by_name(layer, name, new_module):
levels = name.split(".")
if len(levels) > 1:
mod_ = layer
for l_idx in range(len(levels) - 1):
if levels[l_idx].isdigit():
mod_ = mod_[int(levels[l_idx])]
else:
mod_ = getattr(mod_, levels[l_idx])
setattr(mod_, levels[-1], new_module)
else:
setattr(layer, name, new_module)
import torch.nn.functional as F
class WQLinear_GGUF(nn.Module):
def __init__(
self, in_features, out_features, bias, dev, qtype="Q4_0"
):
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.qtype = qtype
qweight_shape = quant_shape_to_byte_shape(
(out_features, in_features), qtype
)
self.register_buffer(
f"{qtype}_qweight",
torch.zeros(
qweight_shape,
dtype=torch.uint8,
device=dev,
),
)
if bias:
self.register_buffer(
"bias",
torch.zeros(
(out_features),
dtype=torch.float16,
device=dev,
),
)
else:
self.bias = None
@classmethod
def from_linear(
cls, linear,
device="cpu",
qtype="Q4_0",
):
q_linear = cls(
linear.in_features,
linear.out_features,
linear.bias is not None,
device,
qtype=qtype,
)
return q_linear
def extra_repr(self) -> str:
return (
"in_features={}, out_features={}, bias={}, w_bit={}, group_size={}".format(
self.in_features,
self.out_features,
self.bias is not None,
self.w_bit,
self.group_size,
)
)
@torch.no_grad()
def forward(self, x):
# x = torch.matmul(x, dequantize_blocks_Q4_0(self.qweight))
if self.qtype == "Q4_0":
x = F.linear(x, dequantize_blocks_Q4_0(
self.Q4_0_qweight, x.dtype), self.bias.to(x.dtype) if self.bias is not None else None)
else:
raise ValueError(f"Unknown qtype: {self.qtype}")
return x
def split_block_dims(blocks, *args):
n_max = blocks.shape[1]
dims = list(args) + [n_max - sum(args)]
return torch.split(blocks, dims, dim=1)
def quant_shape_to_byte_shape(shape, qtype) -> tuple[int, ...]:
# shape = shape[::-1]
block_size, type_size = GGML_QUANT_SIZES[qtype]
if shape[-1] % block_size != 0:
raise ValueError(
f"Quantized tensor row size ({shape[-1]}) is not a multiple of Q4_0 block size ({block_size})")
return (*shape[:-1], shape[-1] // block_size * type_size)
def quant_shape_from_byte_shape(shape, qtype) -> tuple[int, ...]:
# shape = shape[::-1]
block_size, type_size = GGML_QUANT_SIZES[qtype]
if shape[-1] % type_size != 0:
raise ValueError(
f"Quantized tensor bytes per row ({shape[-1]}) is not a multiple of Q4_0 type size ({type_size})")
return (*shape[:-1], shape[-1] // type_size * block_size)
GGML_QUANT_SIZES = {
"Q4_0": (32, 2 + 16),
}
def dequantize_blocks_Q4_0(data, dtype=torch.float16):
block_size, type_size = GGML_QUANT_SIZES["Q4_0"]
data = data.to(torch.uint8)
shape = data.shape
rows = data.reshape(
(-1, data.shape[-1])
).view(torch.uint8)
n_blocks = rows.numel() // type_size
blocks = data.reshape((n_blocks, type_size))
n_blocks = blocks.shape[0]
d, qs = split_block_dims(blocks, 2)
d = d.view(torch.float16)
qs = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor(
[0, 4], device=d.device, dtype=torch.uint8).reshape((1, 1, 2, 1))
qs = (qs & 0x0F).reshape((n_blocks, -1)).to(torch.int8) - 8
out = (d * qs)
out = out.reshape(quant_shape_from_byte_shape(
shape,
qtype="Q4_0",
)).to(dtype)
return out
pipeline_cogvideox.py
+87 -35
View File
@@ -17,6 +17,7 @@ import inspect
from typing import Callable, Dict, List, Optional, Tuple, Union
import torch
import torch.nn.functional as F
import math
from diffusers.models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel
@@ -138,6 +139,7 @@ class CogVideoXPipeline(DiffusionPipeline):
vae: AutoencoderKLCogVideoX,
transformer: CogVideoXTransformer3DModel,
scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler],
original_mask = None
):
super().__init__()
@@ -150,7 +152,7 @@ class CogVideoXPipeline(DiffusionPipeline):
self.vae_scale_factor_temporal = (
self.vae.config.temporal_compression_ratio if hasattr(self, "vae") and self.vae is not None else 4
)
self.original_mask = original_mask
self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
def prepare_latents(
@@ -168,9 +170,9 @@ class CogVideoXPipeline(DiffusionPipeline):
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
noise = randn_tensor(shape, generator=generator, device=device, dtype=self.vae.dtype)
if latents is None:
latents = randn_tensor(shape, generator=generator, device=device, dtype=self.vae.dtype)
latents = noise
else:
latents = latents.to(device)
timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, denoise_strength, device)
@@ -189,7 +191,7 @@ class CogVideoXPipeline(DiffusionPipeline):
latents = self.scheduler.add_noise(latents, noise, latent_timestep)
latents = latents * self.scheduler.init_noise_sigma # scale the initial noise by the standard deviation required by the scheduler
return latents, timesteps
return latents, timesteps, noise
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
def prepare_extra_step_kwargs(self, generator, eta):
@@ -330,10 +332,11 @@ class CogVideoXPipeline(DiffusionPipeline):
num_videos_per_prompt: int = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
device = torch.device("cuda"),
scheduler_name: str = "DPM",
):
"""
Function invoked when calling the pipeline for generation.
@@ -420,7 +423,11 @@ class CogVideoXPipeline(DiffusionPipeline):
if latents is None and num_frames == t_tile_length:
num_frames += 1
latents, timesteps = self.prepare_latents(
if self.original_mask is not None:
image_latents = latents
original_image_latents = image_latents
latents, timesteps, noise = self.prepare_latents(
batch_size * num_videos_per_prompt,
latent_channels,
num_frames,
@@ -435,15 +442,10 @@ class CogVideoXPipeline(DiffusionPipeline):
latents
)
latents = latents.to(self.transformer.dtype)
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
t_tile_weights = self._gaussian_weights(t_tile_length=t_tile_length, t_batch_size=1).to(latents.device).to(latents.dtype)
print("latents.shape", latents.shape)
print("latents.device", latents.device)
# 6.5. Create rotary embeds if required
image_rotary_emb = (
self._prepare_rotary_positional_embeddings(height, width, latents.size(1), device)
@@ -451,16 +453,39 @@ class CogVideoXPipeline(DiffusionPipeline):
else None
)
# masks
if self.original_mask is not None:
mask = self.original_mask.to(device)
print("self.original_mask: ", self.original_mask.shape)
mask = F.interpolate(self.original_mask.unsqueeze(1), size=(latents.shape[-2], latents.shape[-1]), mode='bilinear', align_corners=False)
if mask.shape[0] != latents.shape[1]:
mask = mask.unsqueeze(1).repeat(1, latents.shape[1], 16, 1, 1)
else:
mask = mask.unsqueeze(0).repeat(1, 1, 16, 1, 1)
print("latents: ", latents.shape)
print("mask: ", mask.shape)
# 7. Denoising loop
num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
comfy_pbar = ProgressBar(num_inference_steps)
with self.progress_bar(total=num_inference_steps) as progress_bar:
# for DPM-solver++
old_pred_original_sample = None
# 8. Temporal tiling prep
if "tiled" in scheduler_name:
t_tile_weights = self._gaussian_weights(t_tile_length=t_tile_length, t_batch_size=1).to(latents.device).to(self.vae.dtype)
temporal_tiling = True
print("Temporal tiling enabled")
else:
temporal_tiling = False
print("Temporal tiling disabled")
print("latents.shape", latents.shape)
with self.progress_bar(total=num_inference_steps) as progress_bar:
old_pred_original_sample = None # for DPM-solver++
for i, t in enumerate(timesteps):
if self.interrupt:
continue
if not isinstance(self.scheduler, CogVideoXDPMScheduler):
if temporal_tiling and isinstance(self.scheduler, CogVideoXDDIMScheduler):
#temporal tiling code based on https://github.com/mayuelala/FollowYourEmoji/blob/main/models/video_pipeline.py
# =====================================================
grid_ts = 0
@@ -513,12 +538,12 @@ class CogVideoXPipeline(DiffusionPipeline):
noise_pred = noise_pred_uncond + self._guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents_tile = self.scheduler.step(noise_pred, t, latents_tile, **extra_step_kwargs, return_dict=False)[0]
latents_tile = self.scheduler.step(noise_pred, t, latents_tile.to(self.vae.dtype), **extra_step_kwargs, return_dict=False)[0]
latents_all_list.append(latents_tile)
# ==========================================
latents_all = torch.zeros(latents.shape, device=latents.device, dtype=latents.dtype)
contributors = torch.zeros(latents.shape, device=latents.device, dtype=latents.dtype)
latents_all = torch.zeros(latents.shape, device=latents.device, dtype=self.vae.dtype)
contributors = torch.zeros(latents.shape, device=latents.device, dtype=self.vae.dtype)
# Add each tile contribution to overall latents
for t_i in range(grid_ts):
if t_i < grid_ts - 1:
@@ -535,13 +560,25 @@ class CogVideoXPipeline(DiffusionPipeline):
latents_all /= contributors
latents = latents_all
#print("latents",latents.shape)
# start diff diff
if i < len(timesteps) - 1 and self.original_mask is not None:
noise_timestep = timesteps[i + 1]
image_latent = self.scheduler.add_noise(original_image_latents, noise, torch.tensor([noise_timestep])
)
mask = mask.to(latents)
ts_from = timesteps[0]
ts_to = timesteps[-1]
threshold = (t - ts_to) / (ts_from - ts_to)
mask = torch.where(mask >= threshold, mask, torch.zeros_like(mask))
latents = image_latent * mask + latents * (1 - mask)
# end diff diff
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
progress_bar.update()
comfy_pbar.update(1)
# ==========================================
else:
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
@@ -558,26 +595,41 @@ class CogVideoXPipeline(DiffusionPipeline):
)[0]
noise_pred = noise_pred.float()
self._guidance_scale = 1 + guidance_scale * (
(1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2
)
if isinstance(self.scheduler, CogVideoXDPMScheduler):
self._guidance_scale = 1 + guidance_scale * (
(1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2
)
if do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + self._guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents, old_pred_original_sample = self.scheduler.step(
noise_pred,
old_pred_original_sample,
t,
timesteps[i - 1] if i > 0 else None,
latents.to(self.vae.dtype),
**extra_step_kwargs,
return_dict=False,
)
if not isinstance(self.scheduler, CogVideoXDPMScheduler):
latents = self.scheduler.step(noise_pred, t, latents.to(self.vae.dtype), **extra_step_kwargs, return_dict=False)[0]
else:
latents, old_pred_original_sample = self.scheduler.step(
noise_pred,
old_pred_original_sample,
t,
timesteps[i - 1] if i > 0 else None,
latents.to(self.vae.dtype),
**extra_step_kwargs,
return_dict=False,
)
latents = latents.to(prompt_embeds.dtype)
# start diff diff
if i < len(timesteps) - 1 and self.original_mask is not None:
noise_timestep = timesteps[i + 1]
image_latent = self.scheduler.add_noise(original_image_latents, noise, torch.tensor([noise_timestep])
)
mask = mask.to(latents)
ts_from = timesteps[0]
ts_to = timesteps[-1]
threshold = (t - ts_to) / (ts_from - ts_to)
mask = torch.where(mask >= threshold, mask, torch.zeros_like(mask))
latents = image_latent * mask + latents * (1 - mask)
# end diff diff
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
progress_bar.update()