diff --git a/README.md b/README.md index 687ab36..ce3ea62 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/__init__.py b/__init__.py index 94d3fdb..4c9e248 100644 --- a/__init__.py +++ b/__init__.py @@ -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"}), diff --git a/examples/workflow_GGUF_Q4_0.png b/examples/workflow_GGUF_Q4_0.png new file mode 100644 index 0000000..b28651c Binary files /dev/null and b/examples/workflow_GGUF_Q4_0.png differ diff --git a/mz_cogvideox_core.py b/mz_cogvideox_core.py index 227dd16..541f926 100644 --- a/mz_cogvideox_core.py +++ b/mz_cogvideox_core.py @@ -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 diff --git a/mz_gguf_loader.py b/mz_gguf_loader.py new file mode 100644 index 0000000..b53f1e7 --- /dev/null +++ b/mz_gguf_loader.py @@ -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 diff --git a/pipeline_cogvideox.py b/pipeline_cogvideox.py index fbd3927..09eeeec 100644 --- a/pipeline_cogvideox.py +++ b/pipeline_cogvideox.py @@ -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()