446 lines
14 KiB
Python
446 lines
14 KiB
Python
|
|
|
|
from contextlib import nullcontext
|
|
import os
|
|
import torch
|
|
import comfy.supported_models
|
|
import comfy.model_base
|
|
import comfy.ldm.flux.model
|
|
import comfy.model_patcher
|
|
|
|
import comfy.model_management
|
|
import folder_paths
|
|
import safetensors.torch
|
|
|
|
from .pipeline_cogvideox import CogVideoXPipeline
|
|
from .cogvideox_fun.transformer_3d import CogVideoXTransformer3DModel as CogVideoXTransformer3DModelFun
|
|
from .cogvideox_fun.fun_pab_transformer_3d import CogVideoXTransformer3DModel as CogVideoXTransformer3DModelFunPAB
|
|
from .cogvideox_fun.autoencoder_magvit import AutoencoderKLCogVideoX as AutoencoderKLCogVideoXFun
|
|
from .cogvideox_fun.utils import get_image_to_video_latent, ASPECT_RATIO_512, get_closest_ratio, to_pil
|
|
from .cogvideox_fun.pipeline_cogvideox_inpaint import CogVideoX_Fun_Pipeline_Inpaint
|
|
# from diffusers.models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel
|
|
from diffusers.models import AutoencoderKLCogVideoX
|
|
from .custom_cogvideox_transformer_3d import CogVideoXTransformer3DModel
|
|
from diffusers.schedulers import CogVideoXDDIMScheduler
|
|
|
|
|
|
cogVideoXVaeConfig = {
|
|
"act_fn": "silu",
|
|
"block_out_channels": [
|
|
128,
|
|
256,
|
|
256,
|
|
512
|
|
],
|
|
"down_block_types": [
|
|
"CogVideoXDownBlock3D",
|
|
"CogVideoXDownBlock3D",
|
|
"CogVideoXDownBlock3D",
|
|
"CogVideoXDownBlock3D"
|
|
],
|
|
"force_upcast": True,
|
|
"in_channels": 3,
|
|
"latent_channels": 16,
|
|
"latents_mean": None,
|
|
"latents_std": None,
|
|
"layers_per_block": 3,
|
|
"mid_block_add_attention": True,
|
|
"norm_eps": 1e-06,
|
|
"norm_num_groups": 32,
|
|
"out_channels": 3,
|
|
"sample_size": 256,
|
|
"scaling_factor": 1.15258426,
|
|
"shift_factor": None,
|
|
"temporal_compression_ratio": 4,
|
|
"up_block_types": [
|
|
"CogVideoXUpBlock3D",
|
|
"CogVideoXUpBlock3D",
|
|
"CogVideoXUpBlock3D",
|
|
"CogVideoXUpBlock3D"
|
|
],
|
|
"use_post_quant_conv": False,
|
|
"use_quant_conv": False
|
|
}
|
|
|
|
cogVideoXVaeConfig5B = {
|
|
"act_fn": "silu",
|
|
"block_out_channels": [
|
|
128,
|
|
256,
|
|
256,
|
|
512
|
|
],
|
|
"down_block_types": [
|
|
"CogVideoXDownBlock3D",
|
|
"CogVideoXDownBlock3D",
|
|
"CogVideoXDownBlock3D",
|
|
"CogVideoXDownBlock3D"
|
|
],
|
|
"force_upcast": True,
|
|
"in_channels": 3,
|
|
"latent_channels": 16,
|
|
"latents_mean": None,
|
|
"latents_std": None,
|
|
"layers_per_block": 3,
|
|
"norm_eps": 1e-06,
|
|
"norm_num_groups": 32,
|
|
"out_channels": 3,
|
|
"sample_height": 480,
|
|
"sample_width": 720,
|
|
"scaling_factor": 0.7,
|
|
"shift_factor": None,
|
|
"temporal_compression_ratio": 4,
|
|
"up_block_types": [
|
|
"CogVideoXUpBlock3D",
|
|
"CogVideoXUpBlock3D",
|
|
"CogVideoXUpBlock3D",
|
|
"CogVideoXUpBlock3D"
|
|
],
|
|
"use_post_quant_conv": False,
|
|
"use_quant_conv": False
|
|
}
|
|
|
|
cogVideoXTransformerConfig = {
|
|
"activation_fn": "gelu-approximate",
|
|
"attention_bias": True,
|
|
"attention_head_dim": 64,
|
|
"dropout": 0.0,
|
|
"flip_sin_to_cos": True,
|
|
"freq_shift": 0,
|
|
"in_channels": 16,
|
|
"max_text_seq_length": 226,
|
|
"norm_elementwise_affine": True,
|
|
"norm_eps": 1e-05,
|
|
"num_attention_heads": 30,
|
|
"num_layers": 30,
|
|
"out_channels": 16,
|
|
"patch_size": 2,
|
|
"sample_frames": 49,
|
|
"sample_height": 60,
|
|
"sample_width": 90,
|
|
"spatial_interpolation_scale": 1.875,
|
|
"temporal_compression_ratio": 4,
|
|
"temporal_interpolation_scale": 1.0,
|
|
"text_embed_dim": 4096,
|
|
"time_embed_dim": 512,
|
|
"timestep_activation_fn": "silu",
|
|
"use_rotary_positional_embeddings": False
|
|
}
|
|
|
|
cogVideoXTransformerConfig5B = {
|
|
"activation_fn": "gelu-approximate",
|
|
"attention_bias": True,
|
|
"attention_head_dim": 64,
|
|
"dropout": 0.0,
|
|
"flip_sin_to_cos": True,
|
|
"freq_shift": 0,
|
|
"in_channels": 16,
|
|
"max_text_seq_length": 226,
|
|
"norm_elementwise_affine": True,
|
|
"norm_eps": 1e-05,
|
|
"num_attention_heads": 48,
|
|
"num_layers": 42,
|
|
"out_channels": 16,
|
|
"patch_size": 2,
|
|
"sample_frames": 49,
|
|
"sample_height": 60,
|
|
"sample_width": 90,
|
|
"spatial_interpolation_scale": 1.875,
|
|
"temporal_compression_ratio": 4,
|
|
"temporal_interpolation_scale": 1.0,
|
|
"text_embed_dim": 4096,
|
|
"time_embed_dim": 512,
|
|
"timestep_activation_fn": "silu",
|
|
"use_rotary_positional_embeddings": True
|
|
}
|
|
|
|
cogVideoXDDIMSchedulerConfig = {
|
|
"beta_end": 0.012,
|
|
"beta_schedule": "scaled_linear",
|
|
"beta_start": 0.00085,
|
|
"clip_sample": False,
|
|
"clip_sample_range": 1.0,
|
|
"num_train_timesteps": 1000,
|
|
"prediction_type": "v_prediction",
|
|
"rescale_betas_zero_snr": True,
|
|
"sample_max_value": 1.0,
|
|
"set_alpha_to_one": True,
|
|
"snr_shift_scale": 3.0,
|
|
"steps_offset": 0,
|
|
"timestep_spacing": "linspace",
|
|
"trained_betas": None,
|
|
}
|
|
|
|
cogVideoXDDIMSchedulerConfig5B = {
|
|
"beta_end": 0.012,
|
|
"beta_schedule": "scaled_linear",
|
|
"beta_start": 0.00085,
|
|
"clip_sample": False,
|
|
"clip_sample_range": 1.0,
|
|
"num_train_timesteps": 1000,
|
|
"prediction_type": "v_prediction",
|
|
"rescale_betas_zero_snr": True,
|
|
"sample_max_value": 1.0,
|
|
"set_alpha_to_one": True,
|
|
"snr_shift_scale": 1.0,
|
|
"steps_offset": 0,
|
|
"timestep_spacing": "trailing",
|
|
"trained_betas": None,
|
|
}
|
|
|
|
|
|
def gen_fp8_linear_forward(cast_dtype):
|
|
def fp8_linear_forward(cls, x):
|
|
original_dtype = cls.weight.dtype
|
|
if original_dtype == torch.float8_e4m3fn or original_dtype == torch.float8_e5m2:
|
|
if len(x.shape) == 3:
|
|
with torch.no_grad():
|
|
if original_dtype == torch.float8_e4m3fn:
|
|
inn = x.reshape(-1, x.shape[2]).to(torch.float8_e5m2)
|
|
else:
|
|
inn = x.reshape(-1, x.shape[2]).to(torch.float8_e4m3fn)
|
|
w = cls.weight.t()
|
|
|
|
scale_weight = torch.ones(
|
|
(1), device=x.device, dtype=torch.float32)
|
|
scale_input = scale_weight
|
|
|
|
bias = cls.bias.to(
|
|
cast_dtype) if cls.bias is not None else None
|
|
|
|
if bias is not None:
|
|
o = torch._scaled_mm(
|
|
inn, w, out_dtype=cast_dtype, bias=bias, scale_a=scale_input, scale_b=scale_weight)
|
|
else:
|
|
o = torch._scaled_mm(
|
|
inn, w, out_dtype=cast_dtype, scale_a=scale_input, scale_b=scale_weight)
|
|
|
|
if isinstance(o, tuple):
|
|
o = o[0]
|
|
|
|
return o.reshape((-1, x.shape[1], cls.weight.shape[0]))
|
|
else:
|
|
cls.to(cast_dtype)
|
|
out = cls.original_forward(x.to(
|
|
cast_dtype
|
|
))
|
|
cls.to(original_dtype)
|
|
return out
|
|
else:
|
|
return cls.original_forward(x)
|
|
return fp8_linear_forward
|
|
|
|
|
|
import torch.nn as nn
|
|
from types import MethodType
|
|
|
|
|
|
def convert_fp8_linear(module, dtype, cast_dtype):
|
|
for name, module in module.named_modules():
|
|
if isinstance(module, nn.Linear):
|
|
module.to(dtype)
|
|
original_forward = module.forward
|
|
setattr(module, "original_forward", original_forward)
|
|
setattr(module, "forward", MethodType(
|
|
gen_fp8_linear_forward(cast_dtype), module))
|
|
|
|
|
|
def MZ_CogVideoXLoader_call(args={}):
|
|
unet_name = args.get("unet_name")
|
|
|
|
unet_path = folder_paths.get_full_path("unet", unet_name)
|
|
|
|
enable_sequential_cpu_offload = args.get(
|
|
"enable_sequential_cpu_offload", False)
|
|
|
|
device = comfy.model_management.get_torch_device()
|
|
|
|
comfy.model_management.soft_empty_cache()
|
|
|
|
unet_sd = safetensors.torch.load_file(unet_path)
|
|
unet_sd_keys = list(unet_sd.keys())
|
|
|
|
transformer_type = ""
|
|
if unet_sd["patch_embed.proj.weight"].shape == (3072, 33, 2, 2):
|
|
transformer_type = "fun_5b"
|
|
elif unet_sd["patch_embed.proj.weight"].shape == (3072, 16, 2, 2):
|
|
transformer_type = "5b"
|
|
elif unet_sd["patch_embed.proj.weight"].shape == (1920, 33, 2, 2):
|
|
transformer_type = "fun_2b"
|
|
elif unet_sd["patch_embed.proj.weight"].shape == (1920, 16, 2, 2):
|
|
transformer_type = "2b"
|
|
elif unet_sd["patch_embed.proj.weight"].shape == (3072, 32, 2, 2):
|
|
transformer_type = "i2v_5b"
|
|
else:
|
|
raise Exception("This model is not supported")
|
|
|
|
is_GGUF = False
|
|
if len([k for k in unet_sd_keys if "Q4_0_qweight" in k]) > 0:
|
|
is_GGUF = True
|
|
|
|
print(f"transformer type: {transformer_type}")
|
|
print(f"GGUF: {is_GGUF}")
|
|
|
|
transformer_config = None
|
|
vae_config = None
|
|
scheduler_config = None
|
|
base_path = None
|
|
|
|
if transformer_type.endswith("5b"):
|
|
transformer_config = cogVideoXTransformerConfig5B
|
|
vae_config = cogVideoXVaeConfig5B
|
|
scheduler_config = cogVideoXDDIMSchedulerConfig5B
|
|
base_path = os.path.join(
|
|
os.path.dirname(__file__),
|
|
"configs5b",
|
|
)
|
|
|
|
if transformer_type == "fun_5b":
|
|
transformer_config["in_channels"] = 33
|
|
base_path = os.path.join(
|
|
os.path.dirname(__file__),
|
|
"configs5b-Fun",
|
|
)
|
|
elif transformer_type == "i2v_5b":
|
|
transformer_config["in_channels"] = 32
|
|
transformer_config["use_learned_positional_embeddings"] = True
|
|
base_path = os.path.join(
|
|
os.path.dirname(__file__),
|
|
"configs5b-i2v",
|
|
)
|
|
|
|
if transformer_type.endswith("2b"):
|
|
transformer_config = cogVideoXTransformerConfig
|
|
vae_config = cogVideoXVaeConfig
|
|
scheduler_config = cogVideoXDDIMSchedulerConfig
|
|
base_path = os.path.join(
|
|
os.path.dirname(__file__),
|
|
"configs",
|
|
)
|
|
if transformer_type == "fun_2b":
|
|
transformer_config["in_channels"] = 33
|
|
base_path = os.path.join(
|
|
os.path.dirname(__file__),
|
|
"configs2b-Fun",
|
|
)
|
|
|
|
weight_dtype = None
|
|
manual_cast_dtype = None
|
|
_weight_dtype = args.get("weight_dtype")
|
|
|
|
if _weight_dtype == "fp8_e4m3fn":
|
|
weight_dtype = torch.float8_e4m3fn
|
|
manual_cast_dtype = torch.float16
|
|
elif _weight_dtype == "fp8_e5m2":
|
|
weight_dtype = torch.float8_e5m2
|
|
manual_cast_dtype = torch.bfloat16
|
|
elif _weight_dtype == "fp16":
|
|
weight_dtype = torch.float16
|
|
manual_cast_dtype = torch.float16
|
|
elif _weight_dtype == "bf16":
|
|
weight_dtype = torch.bfloat16
|
|
manual_cast_dtype = torch.bfloat16
|
|
else:
|
|
weight_dtype = torch.float32
|
|
manual_cast_dtype = torch.float32
|
|
|
|
print(
|
|
f"model weight dtype: {weight_dtype} manual cast dtype: {manual_cast_dtype}")
|
|
|
|
pab_config = args.get("pab_config", None)
|
|
|
|
transformer = None
|
|
CogVideoXTransformer3DModelImp = None
|
|
if pab_config is not None:
|
|
CogVideoXTransformer3DModelImp = CogVideoXTransformer3DModelFunPAB
|
|
elif transformer_type.startswith("fun"):
|
|
CogVideoXTransformer3DModelImp = CogVideoXTransformer3DModelFun
|
|
else:
|
|
CogVideoXTransformer3DModelImp = CogVideoXTransformer3DModel
|
|
|
|
from . import mz_gguf_loader
|
|
import importlib
|
|
importlib.reload(mz_gguf_loader)
|
|
with mz_gguf_loader.quantize_lazy_load() if is_GGUF else nullcontext():
|
|
transformer = CogVideoXTransformer3DModelImp.from_config(
|
|
transformer_config)
|
|
transformer.to(weight_dtype)
|
|
block_edit = args.get("block_edit", None)
|
|
|
|
if block_edit is not None:
|
|
transformer = remove_specific_blocks(transformer, block_edit)
|
|
|
|
if is_GGUF:
|
|
transformer = mz_gguf_loader.quantize_load_state_dict(
|
|
transformer, unet_sd, device="cpu", cast_dtype=manual_cast_dtype)
|
|
transformer.to(device)
|
|
else:
|
|
transformer.load_state_dict(unet_sd)
|
|
|
|
if weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e5m2:
|
|
fp8_fast_mode = args.get("fp8_fast_mode", False)
|
|
if fp8_fast_mode:
|
|
print("convert to fp8 linear")
|
|
convert_fp8_linear(transformer, weight_dtype, manual_cast_dtype)
|
|
|
|
if transformer_type.endswith("2b") or transformer_type == "i2v_5b":
|
|
if hasattr(transformer, "pos_embedding"):
|
|
transformer.pos_embedding = transformer.pos_embedding.to(
|
|
manual_cast_dtype)
|
|
if hasattr(transformer, "patch_embed") and hasattr(transformer.patch_embed, "pos_embedding"):
|
|
transformer.patch_embed.pos_embedding = transformer.patch_embed.pos_embedding.to(
|
|
manual_cast_dtype)
|
|
|
|
transformer.to(device)
|
|
|
|
vae_name = args.get("vae_name")
|
|
vae_path = folder_paths.get_full_path("vae", vae_name)
|
|
if transformer_type.startswith("fun"):
|
|
vae = AutoencoderKLCogVideoXFun.from_config(vae_config)
|
|
else:
|
|
vae = AutoencoderKLCogVideoX.from_config(vae_config)
|
|
vae_sd = safetensors.torch.load_file(vae_path)
|
|
vae.load_state_dict(vae_sd)
|
|
vae.to(device)
|
|
|
|
enable_vae_encode_tiling = args.get("enable_vae_encode_tiling", False)
|
|
if enable_vae_encode_tiling:
|
|
from .mz_enable_vae_encode_tiling import enable_vae_encode_tiling
|
|
enable_vae_encode_tiling(vae)
|
|
|
|
# from .mz_dyn_cpu_offload import dyn_cpu_offload_model_vae
|
|
# vae = dyn_cpu_offload_model_vae(vae)
|
|
|
|
scheduler = CogVideoXDDIMScheduler.from_config(
|
|
scheduler_config)
|
|
|
|
if transformer_type.startswith("fun"):
|
|
pipe = CogVideoX_Fun_Pipeline_Inpaint(
|
|
vae, transformer, scheduler, pab_config=pab_config)
|
|
else:
|
|
pipe = CogVideoXPipeline(
|
|
vae, transformer, scheduler, pab_config=pab_config)
|
|
|
|
if enable_sequential_cpu_offload:
|
|
pipe.enable_sequential_cpu_offload()
|
|
|
|
pipeline = {
|
|
"pipe": pipe,
|
|
"dtype": manual_cast_dtype,
|
|
"base_path": base_path,
|
|
"onediff": False,
|
|
"cpu_offloading": enable_sequential_cpu_offload,
|
|
"scheduler_config": scheduler_config,
|
|
}
|
|
return (pipeline, )
|
|
|
|
|
|
def remove_specific_blocks(model, block_indices_to_remove):
|
|
import torch.nn as nn
|
|
transformer_blocks = model.transformer_blocks
|
|
new_blocks = [block for i, block in enumerate(
|
|
transformer_blocks) if i not in block_indices_to_remove]
|
|
model.transformer_blocks = nn.ModuleList(new_blocks)
|
|
|
|
return model
|