diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..1f152aa --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS +NODE_DISPLAY_NAME_MAPPINGS = {k:v.TITLE for k,v in NODE_CLASS_MAPPINGS.items()} +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/diffusers_helper/bucket_tools.py b/diffusers_helper/bucket_tools.py new file mode 100644 index 0000000..e499c8f --- /dev/null +++ b/diffusers_helper/bucket_tools.py @@ -0,0 +1,69 @@ +bucket_options = { + 480: [(416, 960), + (448, 864), + (480, 832), + (512, 768), + (544, 704), + (576, 672), + (608, 640), + (640, 608), + (672, 576), + (704, 544), + (768, 512), + (832, 480), + (864, 448), + (960, 416)], + 720: [(624, 1456), + (672, 1312), + (720, 1280), + (768, 1152), + (816, 1056), + (864, 1008), + (912, 960), + (960, 912), + (1008, 864), + (1056, 816), + (1152, 768), + (1280, 720), + (1312, 672), + (1456, 624)], + 540: [(464, 1088), + (496, 976), + (544, 960), + (576, 864), + (608, 800), + (656, 752), + (688, 720), + (720, 688), + (752, 656), + (800, 608), + (864, 576), + (960, 544), + (976, 496), + (1088, 464)], + 360: [(304, 720), + (336, 656), + (352, 640), + (384, 576), + (400, 528), + (432, 496), + (448, 480), + (480, 448), + (496, 432), + (528, 400), + (576, 384), + (640, 352), + (656, 336), + (720, 304)], +} + +def find_nearest_bucket(h, w, resolution=640): + min_metric = float('inf') + best_bucket = None + for (bucket_h, bucket_w) in bucket_options[resolution]: + metric = abs(h * bucket_w - w * bucket_h) + if metric <= min_metric: + min_metric = metric + best_bucket = (bucket_h, bucket_w) + return best_bucket + diff --git a/diffusers_helper/clip_vision.py b/diffusers_helper/clip_vision.py new file mode 100644 index 0000000..c834b7e --- /dev/null +++ b/diffusers_helper/clip_vision.py @@ -0,0 +1,12 @@ +import numpy as np + + +def hf_clip_vision_encode(image, feature_extractor, image_encoder): + assert isinstance(image, np.ndarray) + assert image.ndim == 3 and image.shape[2] == 3 + assert image.dtype == np.uint8 + + preprocessed = feature_extractor.preprocess(images=image, return_tensors="pt").to(device=image_encoder.device, dtype=image_encoder.dtype) + image_encoder_output = image_encoder(**preprocessed) + + return image_encoder_output diff --git a/diffusers_helper/dit_common.py b/diffusers_helper/dit_common.py new file mode 100644 index 0000000..d4983a9 --- /dev/null +++ b/diffusers_helper/dit_common.py @@ -0,0 +1,53 @@ +import torch +import accelerate.accelerator + +from diffusers.models.normalization import RMSNorm, LayerNorm, FP32LayerNorm, AdaLayerNormContinuous + + +accelerate.accelerator.convert_outputs_to_fp32 = lambda x: x + + +def LayerNorm_forward(self, x): + return torch.nn.functional.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps).to(x) + + +LayerNorm.forward = LayerNorm_forward +torch.nn.LayerNorm.forward = LayerNorm_forward + + +def FP32LayerNorm_forward(self, x): + origin_dtype = x.dtype + return torch.nn.functional.layer_norm( + x.float(), + self.normalized_shape, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ).to(origin_dtype) + + +FP32LayerNorm.forward = FP32LayerNorm_forward + + +def RMSNorm_forward(self, hidden_states): + input_dtype = hidden_states.dtype + variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + + if self.weight is None: + return hidden_states.to(input_dtype) + + return hidden_states.to(input_dtype) * self.weight.to(input_dtype) + + +RMSNorm.forward = RMSNorm_forward + + +def AdaLayerNormContinuous_forward(self, x, conditioning_embedding): + emb = self.linear(self.silu(conditioning_embedding)) + scale, shift = emb.chunk(2, dim=1) + x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + return x + + +AdaLayerNormContinuous.forward = AdaLayerNormContinuous_forward diff --git a/diffusers_helper/gradio/progress_bar.py b/diffusers_helper/gradio/progress_bar.py new file mode 100644 index 0000000..d32954b --- /dev/null +++ b/diffusers_helper/gradio/progress_bar.py @@ -0,0 +1,86 @@ +progress_html = ''' +
+
+
+ +
+ *text* +
+''' + +css = ''' +.loader-container { + display: flex; /* Use flex to align items horizontally */ + align-items: center; /* Center items vertically within the container */ + white-space: nowrap; /* Prevent line breaks within the container */ +} + +.loader { + border: 8px solid #f3f3f3; /* Light grey */ + border-top: 8px solid #3498db; /* Blue */ + border-radius: 50%; + width: 30px; + height: 30px; + animation: spin 2s linear infinite; +} + +@keyframes spin { + 0% { transform: rotate(0deg); } + 100% { transform: rotate(360deg); } +} + +/* Style the progress bar */ +progress { + appearance: none; /* Remove default styling */ + height: 20px; /* Set the height of the progress bar */ + border-radius: 5px; /* Round the corners of the progress bar */ + background-color: #f3f3f3; /* Light grey background */ + width: 100%; + vertical-align: middle !important; +} + +/* Style the progress bar container */ +.progress-container { + margin-left: 20px; + margin-right: 20px; + flex-grow: 1; /* Allow the progress container to take up remaining space */ +} + +/* Set the color of the progress bar fill */ +progress::-webkit-progress-value { + background-color: #3498db; /* Blue color for the fill */ +} + +progress::-moz-progress-bar { + background-color: #3498db; /* Blue color for the fill in Firefox */ +} + +/* Style the text on the progress bar */ +progress::after { + content: attr(value '%'); /* Display the progress value followed by '%' */ + position: absolute; + top: 50%; + left: 50%; + transform: translate(-50%, -50%); + color: white; /* Set text color */ + font-size: 14px; /* Set font size */ +} + +/* Style other texts */ +.loader-container > span { + margin-left: 5px; /* Add spacing between the progress bar and the text */ +} + +.no-generating-animation > .generating { + display: none !important; +} + +''' + + +def make_progress_bar_html(number, text): + return progress_html.replace('*number*', str(number)).replace('*text*', text) + + +def make_progress_bar_css(): + return css diff --git a/diffusers_helper/hf_login.py b/diffusers_helper/hf_login.py new file mode 100644 index 0000000..ccbe3cc --- /dev/null +++ b/diffusers_helper/hf_login.py @@ -0,0 +1,21 @@ +import os + + +def login(token): + from huggingface_hub import login + import time + + while True: + try: + login(token) + print('HF login ok.') + break + except Exception as e: + print(f'HF login failed: {e}. Retrying') + time.sleep(0.5) + + +hf_token = os.environ.get('HF_TOKEN', None) + +if hf_token is not None: + login(hf_token) diff --git a/diffusers_helper/hunyuan.py b/diffusers_helper/hunyuan.py new file mode 100644 index 0000000..2fd6ebf --- /dev/null +++ b/diffusers_helper/hunyuan.py @@ -0,0 +1,111 @@ +import torch + +from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import DEFAULT_PROMPT_TEMPLATE +from diffusers_helper.utils import crop_or_pad_yield_mask + + +@torch.no_grad() +def encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2, max_length=256): + assert isinstance(prompt, str) + + prompt = [prompt] + + # LLAMA + + prompt_llama = [DEFAULT_PROMPT_TEMPLATE["template"].format(p) for p in prompt] + crop_start = DEFAULT_PROMPT_TEMPLATE["crop_start"] + + llama_inputs = tokenizer( + prompt_llama, + padding="max_length", + max_length=max_length + crop_start, + truncation=True, + return_tensors="pt", + return_length=False, + return_overflowing_tokens=False, + return_attention_mask=True, + ) + + llama_input_ids = llama_inputs.input_ids.to(text_encoder.device) + llama_attention_mask = llama_inputs.attention_mask.to(text_encoder.device) + llama_attention_length = int(llama_attention_mask.sum()) + + llama_outputs = text_encoder( + input_ids=llama_input_ids, + attention_mask=llama_attention_mask, + output_hidden_states=True, + ) + + llama_vec = llama_outputs.hidden_states[-3][:, crop_start:llama_attention_length] + # llama_vec_remaining = llama_outputs.hidden_states[-3][:, llama_attention_length:] + llama_attention_mask = llama_attention_mask[:, crop_start:llama_attention_length] + + assert torch.all(llama_attention_mask.bool()) + + # CLIP + + clip_l_input_ids = tokenizer_2( + prompt, + padding="max_length", + max_length=77, + truncation=True, + return_overflowing_tokens=False, + return_length=False, + return_tensors="pt", + ).input_ids + clip_l_pooler = text_encoder_2(clip_l_input_ids.to(text_encoder_2.device), output_hidden_states=False).pooler_output + + return llama_vec, clip_l_pooler + + +@torch.no_grad() +def vae_decode_fake(latents): + latent_rgb_factors = [ + [-0.0395, -0.0331, 0.0445], + [0.0696, 0.0795, 0.0518], + [0.0135, -0.0945, -0.0282], + [0.0108, -0.0250, -0.0765], + [-0.0209, 0.0032, 0.0224], + [-0.0804, -0.0254, -0.0639], + [-0.0991, 0.0271, -0.0669], + [-0.0646, -0.0422, -0.0400], + [-0.0696, -0.0595, -0.0894], + [-0.0799, -0.0208, -0.0375], + [0.1166, 0.1627, 0.0962], + [0.1165, 0.0432, 0.0407], + [-0.2315, -0.1920, -0.1355], + [-0.0270, 0.0401, -0.0821], + [-0.0616, -0.0997, -0.0727], + [0.0249, -0.0469, -0.1703] + ] # From comfyui + + latent_rgb_factors_bias = [0.0259, -0.0192, -0.0761] + + weight = torch.tensor(latent_rgb_factors, device=latents.device, dtype=latents.dtype).transpose(0, 1)[:, :, None, None, None] + bias = torch.tensor(latent_rgb_factors_bias, device=latents.device, dtype=latents.dtype) + + images = torch.nn.functional.conv3d(latents, weight, bias=bias, stride=1, padding=0, dilation=1, groups=1) + images = images.clamp(0.0, 1.0) + + return images + + +@torch.no_grad() +def vae_decode(latents, vae, image_mode=False): + latents = latents / vae.config.scaling_factor + + if not image_mode: + image = vae.decode(latents.to(device=vae.device, dtype=vae.dtype)).sample + else: + latents = latents.to(device=vae.device, dtype=vae.dtype).unbind(2) + image = [vae.decode(l.unsqueeze(2)).sample for l in latents] + image = torch.cat(image, dim=2) + + return image + + +@torch.no_grad() +def vae_encode(image, vae): + latents = vae.encode(image.to(device=vae.device, dtype=vae.dtype)).latent_dist.sample() + latents = latents * vae.config.scaling_factor + return latents diff --git a/diffusers_helper/k_diffusion/uni_pc_fm.py b/diffusers_helper/k_diffusion/uni_pc_fm.py new file mode 100644 index 0000000..85e501a --- /dev/null +++ b/diffusers_helper/k_diffusion/uni_pc_fm.py @@ -0,0 +1,141 @@ +# Better Flow Matching UniPC by Lvmin Zhang +# (c) 2025 +# CC BY-SA 4.0 +# Attribution-ShareAlike 4.0 International Licence + + +import torch + +from tqdm.auto import trange + + +def expand_dims(v, dims): + return v[(...,) + (None,) * (dims - 1)] + + +class FlowMatchUniPC: + def __init__(self, model, extra_args, variant='bh1'): + self.model = model + self.variant = variant + self.extra_args = extra_args + + def model_fn(self, x, t): + return self.model(x, t, **self.extra_args) + + def update_fn(self, x, model_prev_list, t_prev_list, t, order): + assert order <= len(model_prev_list) + dims = x.dim() + + t_prev_0 = t_prev_list[-1] + lambda_prev_0 = - torch.log(t_prev_0) + lambda_t = - torch.log(t) + model_prev_0 = model_prev_list[-1] + + h = lambda_t - lambda_prev_0 + + rks = [] + D1s = [] + for i in range(1, order): + t_prev_i = t_prev_list[-(i + 1)] + model_prev_i = model_prev_list[-(i + 1)] + lambda_prev_i = - torch.log(t_prev_i) + rk = ((lambda_prev_i - lambda_prev_0) / h)[0] + rks.append(rk) + D1s.append((model_prev_i - model_prev_0) / rk) + + rks.append(1.) + rks = torch.tensor(rks, device=x.device) + + R = [] + b = [] + + hh = -h[0] + h_phi_1 = torch.expm1(hh) + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.variant == 'bh1': + B_h = hh + elif self.variant == 'bh2': + B_h = torch.expm1(hh) + else: + raise NotImplementedError('Bad variant!') + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= (i + 1) + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R) + b = torch.tensor(b, device=x.device) + + use_predictor = len(D1s) > 0 + + if use_predictor: + D1s = torch.stack(D1s, dim=1) + if order == 2: + rhos_p = torch.tensor([0.5], device=b.device) + else: + rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]) + else: + D1s = None + rhos_p = None + + if order == 1: + rhos_c = torch.tensor([0.5], device=b.device) + else: + rhos_c = torch.linalg.solve(R, b) + + x_t_ = expand_dims(t / t_prev_0, dims) * x - expand_dims(h_phi_1, dims) * model_prev_0 + + if use_predictor: + pred_res = torch.tensordot(D1s, rhos_p, dims=([1], [0])) + else: + pred_res = 0 + + x_t = x_t_ - expand_dims(B_h, dims) * pred_res + model_t = self.model_fn(x_t, t) + + if D1s is not None: + corr_res = torch.tensordot(D1s, rhos_c[:-1], dims=([1], [0])) + else: + corr_res = 0 + + D1_t = (model_t - model_prev_0) + x_t = x_t_ - expand_dims(B_h, dims) * (corr_res + rhos_c[-1] * D1_t) + + return x_t, model_t + + def sample(self, x, sigmas, callback=None, disable_pbar=False): + order = min(3, len(sigmas) - 2) + model_prev_list, t_prev_list = [], [] + for i in trange(len(sigmas) - 1, disable=disable_pbar): + vec_t = sigmas[i].expand(x.shape[0]) + + if i == 0: + model_prev_list = [self.model_fn(x, vec_t)] + t_prev_list = [vec_t] + elif i < order: + init_order = i + x, model_x = self.update_fn(x, model_prev_list, t_prev_list, vec_t, init_order) + model_prev_list.append(model_x) + t_prev_list.append(vec_t) + else: + x, model_x = self.update_fn(x, model_prev_list, t_prev_list, vec_t, order) + model_prev_list.append(model_x) + t_prev_list.append(vec_t) + + model_prev_list = model_prev_list[-order:] + t_prev_list = t_prev_list[-order:] + + if callback is not None: + callback({'x': x, 'i': i, 'denoised': model_prev_list[-1]}) + + return model_prev_list[-1] + + +def sample_unipc(model, noise, sigmas, extra_args=None, callback=None, disable=False, variant='bh1'): + assert variant in ['bh1', 'bh2'] + return FlowMatchUniPC(model, extra_args=extra_args, variant=variant).sample(noise, sigmas=sigmas, callback=callback, disable_pbar=disable) diff --git a/diffusers_helper/k_diffusion/wrapper.py b/diffusers_helper/k_diffusion/wrapper.py new file mode 100644 index 0000000..c2c05af --- /dev/null +++ b/diffusers_helper/k_diffusion/wrapper.py @@ -0,0 +1,51 @@ +import torch + + +def append_dims(x, target_dims): + return x[(...,) + (None,) * (target_dims - x.ndim)] + + +def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0): + if guidance_rescale == 0: + return noise_cfg + + std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) + std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) + noise_pred_rescaled = noise_cfg * (std_text / std_cfg) + noise_cfg = guidance_rescale * noise_pred_rescaled + (1.0 - guidance_rescale) * noise_cfg + return noise_cfg + + +def fm_wrapper(transformer, t_scale=1000.0): + def k_model(x, sigma, **extra_args): + dtype = extra_args['dtype'] + cfg_scale = extra_args['cfg_scale'] + cfg_rescale = extra_args['cfg_rescale'] + concat_latent = extra_args['concat_latent'] + + original_dtype = x.dtype + sigma = sigma.float() + + x = x.to(dtype) + timestep = (sigma * t_scale).to(dtype) + + if concat_latent is None: + hidden_states = x + else: + hidden_states = torch.cat([x, concat_latent.to(x)], dim=1) + + pred_positive = transformer(hidden_states=hidden_states, timestep=timestep, return_dict=False, **extra_args['positive'])[0].float() + + if cfg_scale == 1.0: + pred_negative = torch.zeros_like(pred_positive) + else: + pred_negative = transformer(hidden_states=hidden_states, timestep=timestep, return_dict=False, **extra_args['negative'])[0].float() + + pred_cfg = pred_negative + cfg_scale * (pred_positive - pred_negative) + pred = rescale_noise_cfg(pred_cfg, pred_positive, guidance_rescale=cfg_rescale) + + x0 = x.float() - pred.float() * append_dims(sigma, x.ndim) + + return x0.to(dtype=original_dtype) + + return k_model diff --git a/diffusers_helper/memory.py b/diffusers_helper/memory.py new file mode 100644 index 0000000..19f611a --- /dev/null +++ b/diffusers_helper/memory.py @@ -0,0 +1,136 @@ +# By lllyasviel + + +import torch + + +#cpu = torch.device('cpu') +#gpu = torch.device(f'cuda:{torch.cuda.current_device()}') +cpu = 'cpu' +gpu = 'cuda:0' +gpu_complete_modules = [] + + +class DynamicSwapInstaller: + @staticmethod + def _install_module(module: torch.nn.Module, **kwargs): + original_class = module.__class__ + module.__dict__['forge_backup_original_class'] = original_class + + def hacked_get_attr(self, name: str): + if '_parameters' in self.__dict__: + _parameters = self.__dict__['_parameters'] + if name in _parameters: + p = _parameters[name] + if p is None: + return None + if p.__class__ == torch.nn.Parameter: + return torch.nn.Parameter(p.to(**kwargs), requires_grad=p.requires_grad) + else: + return p.to(**kwargs) + if '_buffers' in self.__dict__: + _buffers = self.__dict__['_buffers'] + if name in _buffers: + return _buffers[name].to(**kwargs) + return super(original_class, self).__getattr__(name) + + module.__class__ = type('DynamicSwap_' + original_class.__name__, (original_class,), { + '__getattr__': hacked_get_attr, + }) + + return + + @staticmethod + def _uninstall_module(module: torch.nn.Module): + if 'forge_backup_original_class' in module.__dict__: + module.__class__ = module.__dict__.pop('forge_backup_original_class') + return + + @staticmethod + def install_model(model: torch.nn.Module, **kwargs): + for m in model.modules(): + DynamicSwapInstaller._install_module(m, **kwargs) + return + + @staticmethod + def uninstall_model(model: torch.nn.Module): + for m in model.modules(): + DynamicSwapInstaller._uninstall_module(m) + return + + +def fake_diffusers_current_device(model: torch.nn.Module, target_device: torch.device): + if hasattr(model, 'scale_shift_table'): + model.scale_shift_table.data = model.scale_shift_table.data.to(target_device) + return + + for k, p in model.named_modules(): + if hasattr(p, 'weight'): + p.to(target_device) + return + + +def get_cuda_free_memory_gb(device=None): + if device is None: + device = gpu + + memory_stats = torch.cuda.memory_stats(device) + bytes_active = memory_stats['active_bytes.all.current'] + bytes_reserved = memory_stats['reserved_bytes.all.current'] + bytes_free_cuda, _ = torch.cuda.mem_get_info(device) + bytes_inactive_reserved = bytes_reserved - bytes_active + bytes_total_available = bytes_free_cuda + bytes_inactive_reserved + return bytes_total_available / (1024 ** 3) + + +def move_model_to_device_with_memory_preservation(model, target_device, preserved_memory_gb=0): + print(f'Moving {model.__class__.__name__} to {target_device} with preserved memory: {preserved_memory_gb} GB') + + for m in model.modules(): + if get_cuda_free_memory_gb(target_device) <= preserved_memory_gb: + torch.cuda.empty_cache() + return + + if hasattr(m, 'weight'): + m.to(device=target_device) + + model.to(device=target_device) + torch.cuda.empty_cache() + return + + +def offload_model_from_device_for_memory_preservation(model, target_device, preserved_memory_gb=0): + print(f'Offloading {model.__class__.__name__} from {target_device} to preserve memory: {preserved_memory_gb} GB') + + for m in model.modules(): + if get_cuda_free_memory_gb(target_device) >= preserved_memory_gb: + torch.cuda.empty_cache() + return + + if hasattr(m, 'weight'): + m.to(device=cpu) + + model.to(device=cpu) + torch.cuda.empty_cache() + return + + +def unload_complete_models(*args): + for m in gpu_complete_modules + list(args): + m.to(device=cpu) + print(f'Unloaded {m.__class__.__name__} as complete.') + + gpu_complete_modules.clear() + torch.cuda.empty_cache() + return + + +def load_model_as_complete(model, target_device, unload=True): + if unload: + unload_complete_models() + + model.to(device=target_device) + print(f'Loaded {model.__class__.__name__} to {target_device} as complete.') + + gpu_complete_modules.append(model) + return diff --git a/diffusers_helper/models/hunyuan_video_packed.py b/diffusers_helper/models/hunyuan_video_packed.py new file mode 100644 index 0000000..9f1364e --- /dev/null +++ b/diffusers_helper/models/hunyuan_video_packed.py @@ -0,0 +1,1032 @@ +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import einops +import torch.nn as nn +import numpy as np + +from diffusers.loaders import FromOriginalModelMixin +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.loaders import PeftAdapterMixin +from diffusers.utils import logging +from diffusers.models.attention import FeedForward +from diffusers.models.attention_processor import Attention +from diffusers.models.embeddings import TimestepEmbedding, Timesteps, PixArtAlphaTextProjection +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers_helper.dit_common import LayerNorm +from diffusers_helper.utils import zero_module + + +enabled_backends = [] + +if torch.backends.cuda.flash_sdp_enabled(): + enabled_backends.append("flash") +if torch.backends.cuda.math_sdp_enabled(): + enabled_backends.append("math") +if torch.backends.cuda.mem_efficient_sdp_enabled(): + enabled_backends.append("mem_efficient") +if torch.backends.cuda.cudnn_sdp_enabled(): + enabled_backends.append("cudnn") + +print("Currently enabled native sdp backends:", enabled_backends) + +try: + # raise NotImplementedError + from xformers.ops import memory_efficient_attention as xformers_attn_func + print('Xformers is installed!') +except: + print('Xformers is not installed!') + xformers_attn_func = None + +try: + # raise NotImplementedError + from flash_attn import flash_attn_varlen_func, flash_attn_func + print('Flash Attn is installed!') +except: + print('Flash Attn is not installed!') + flash_attn_varlen_func = None + flash_attn_func = None + +try: + # raise NotImplementedError + from sageattention import sageattn_varlen, sageattn + print('Sage Attn is installed!') +except: + print('Sage Attn is not installed!') + sageattn_varlen = None + sageattn = None + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def pad_for_3d_conv(x, kernel_size): + b, c, t, h, w = x.shape + pt, ph, pw = kernel_size + pad_t = (pt - (t % pt)) % pt + pad_h = (ph - (h % ph)) % ph + pad_w = (pw - (w % pw)) % pw + return torch.nn.functional.pad(x, (0, pad_w, 0, pad_h, 0, pad_t), mode='replicate') + + +def center_down_sample_3d(x, kernel_size): + # pt, ph, pw = kernel_size + # cp = (pt * ph * pw) // 2 + # xp = einops.rearrange(x, 'b c (t pt) (h ph) (w pw) -> (pt ph pw) b c t h w', pt=pt, ph=ph, pw=pw) + # xc = xp[cp] + # return xc + return torch.nn.functional.avg_pool3d(x, kernel_size, stride=kernel_size) + + +def get_cu_seqlens(text_mask, img_len): + batch_size = text_mask.shape[0] + text_len = text_mask.sum(dim=1) + max_len = text_mask.shape[1] + img_len + + cu_seqlens = torch.zeros([2 * batch_size + 1], dtype=torch.int32, device="cuda") + + for i in range(batch_size): + s = text_len[i] + img_len + s1 = i * max_len + s + s2 = (i + 1) * max_len + cu_seqlens[2 * i + 1] = s1 + cu_seqlens[2 * i + 2] = s2 + + return cu_seqlens + + +def apply_rotary_emb_transposed(x, freqs_cis): + cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1) + x_real, x_imag = x.unflatten(-1, (-1, 2)).unbind(-1) + x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) + out = x.float() * cos + x_rotated.float() * sin + out = out.to(x) + return out + + +def attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv): + if cu_seqlens_q is None and cu_seqlens_kv is None and max_seqlen_q is None and max_seqlen_kv is None: + if sageattn is not None: + x = sageattn(q, k, v, tensor_layout='NHD') + return x + + if flash_attn_func is not None: + x = flash_attn_func(q, k, v) + return x + + if xformers_attn_func is not None: + x = xformers_attn_func(q, k, v) + return x + + x = torch.nn.functional.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2) + return x + + batch_size = q.shape[0] + q = q.view(q.shape[0] * q.shape[1], *q.shape[2:]) + k = k.view(k.shape[0] * k.shape[1], *k.shape[2:]) + v = v.view(v.shape[0] * v.shape[1], *v.shape[2:]) + if sageattn_varlen is not None: + x = sageattn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv) + elif flash_attn_varlen_func is not None: + x = flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv) + else: + raise NotImplementedError('No Attn Installed!') + x = x.view(batch_size, max_seqlen_q, *x.shape[2:]) + return x + + +class HunyuanAttnProcessorFlashAttnDouble: + def __call__(self, attn, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb): + cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv = attention_mask + + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + query = apply_rotary_emb_transposed(query, image_rotary_emb) + key = apply_rotary_emb_transposed(key, image_rotary_emb) + + encoder_query = attn.add_q_proj(encoder_hidden_states) + encoder_key = attn.add_k_proj(encoder_hidden_states) + encoder_value = attn.add_v_proj(encoder_hidden_states) + + encoder_query = encoder_query.unflatten(2, (attn.heads, -1)) + encoder_key = encoder_key.unflatten(2, (attn.heads, -1)) + encoder_value = encoder_value.unflatten(2, (attn.heads, -1)) + + encoder_query = attn.norm_added_q(encoder_query) + encoder_key = attn.norm_added_k(encoder_key) + + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + + hidden_states = attn_varlen_func(query, key, value, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv) + hidden_states = hidden_states.flatten(-2) + + txt_length = encoder_hidden_states.shape[1] + hidden_states, encoder_hidden_states = hidden_states[:, :-txt_length], hidden_states[:, -txt_length:] + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + encoder_hidden_states = attn.to_add_out(encoder_hidden_states) + + return hidden_states, encoder_hidden_states + + +class HunyuanAttnProcessorFlashAttnSingle: + def __call__(self, attn, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb): + cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv = attention_mask + + hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) + + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + txt_length = encoder_hidden_states.shape[1] + + query = torch.cat([apply_rotary_emb_transposed(query[:, :-txt_length], image_rotary_emb), query[:, -txt_length:]], dim=1) + key = torch.cat([apply_rotary_emb_transposed(key[:, :-txt_length], image_rotary_emb), key[:, -txt_length:]], dim=1) + + hidden_states = attn_varlen_func(query, key, value, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv) + hidden_states = hidden_states.flatten(-2) + + hidden_states, encoder_hidden_states = hidden_states[:, :-txt_length], hidden_states[:, -txt_length:] + + return hidden_states, encoder_hidden_states + + +class CombinedTimestepGuidanceTextProjEmbeddings(nn.Module): + def __init__(self, embedding_dim, pooled_projection_dim): + super().__init__() + + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + self.guidance_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu") + + def forward(self, timestep, guidance, pooled_projection): + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype)) + + guidance_proj = self.time_proj(guidance) + guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype)) + + time_guidance_emb = timesteps_emb + guidance_emb + + pooled_projections = self.text_embedder(pooled_projection) + conditioning = time_guidance_emb + pooled_projections + + return conditioning + + +class CombinedTimestepTextProjEmbeddings(nn.Module): + def __init__(self, embedding_dim, pooled_projection_dim): + super().__init__() + + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu") + + def forward(self, timestep, pooled_projection): + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype)) + + pooled_projections = self.text_embedder(pooled_projection) + + conditioning = timesteps_emb + pooled_projections + + return conditioning + + +class HunyuanVideoAdaNorm(nn.Module): + def __init__(self, in_features: int, out_features: Optional[int] = None) -> None: + super().__init__() + + out_features = out_features or 2 * in_features + self.linear = nn.Linear(in_features, out_features) + self.nonlinearity = nn.SiLU() + + def forward( + self, temb: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + temb = self.linear(self.nonlinearity(temb)) + gate_msa, gate_mlp = temb.chunk(2, dim=-1) + gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1) + return gate_msa, gate_mlp + + +class HunyuanVideoIndividualTokenRefinerBlock(nn.Module): + def __init__( + self, + num_attention_heads: int, + attention_head_dim: int, + mlp_width_ratio: str = 4.0, + mlp_drop_rate: float = 0.0, + attention_bias: bool = True, + ) -> None: + super().__init__() + + hidden_size = num_attention_heads * attention_head_dim + + self.norm1 = LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6) + self.attn = Attention( + query_dim=hidden_size, + cross_attention_dim=None, + heads=num_attention_heads, + dim_head=attention_head_dim, + bias=attention_bias, + ) + + self.norm2 = LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6) + self.ff = FeedForward(hidden_size, mult=mlp_width_ratio, activation_fn="linear-silu", dropout=mlp_drop_rate) + + self.norm_out = HunyuanVideoAdaNorm(hidden_size, 2 * hidden_size) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + norm_hidden_states = self.norm1(hidden_states) + + attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=None, + attention_mask=attention_mask, + ) + + gate_msa, gate_mlp = self.norm_out(temb) + hidden_states = hidden_states + attn_output * gate_msa + + ff_output = self.ff(self.norm2(hidden_states)) + hidden_states = hidden_states + ff_output * gate_mlp + + return hidden_states + + +class HunyuanVideoIndividualTokenRefiner(nn.Module): + def __init__( + self, + num_attention_heads: int, + attention_head_dim: int, + num_layers: int, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + attention_bias: bool = True, + ) -> None: + super().__init__() + + self.refiner_blocks = nn.ModuleList( + [ + HunyuanVideoIndividualTokenRefinerBlock( + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + attention_bias=attention_bias, + ) + for _ in range(num_layers) + ] + ) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + ) -> None: + self_attn_mask = None + if attention_mask is not None: + batch_size = attention_mask.shape[0] + seq_len = attention_mask.shape[1] + attention_mask = attention_mask.to(hidden_states.device).bool() + self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1) + self_attn_mask_2 = self_attn_mask_1.transpose(2, 3) + self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool() + self_attn_mask[:, :, :, 0] = True + + for block in self.refiner_blocks: + hidden_states = block(hidden_states, temb, self_attn_mask) + + return hidden_states + + +class HunyuanVideoTokenRefiner(nn.Module): + def __init__( + self, + in_channels: int, + num_attention_heads: int, + attention_head_dim: int, + num_layers: int, + mlp_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + attention_bias: bool = True, + ) -> None: + super().__init__() + + hidden_size = num_attention_heads * attention_head_dim + + self.time_text_embed = CombinedTimestepTextProjEmbeddings( + embedding_dim=hidden_size, pooled_projection_dim=in_channels + ) + self.proj_in = nn.Linear(in_channels, hidden_size, bias=True) + self.token_refiner = HunyuanVideoIndividualTokenRefiner( + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + num_layers=num_layers, + mlp_width_ratio=mlp_ratio, + mlp_drop_rate=mlp_drop_rate, + attention_bias=attention_bias, + ) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.LongTensor, + attention_mask: Optional[torch.LongTensor] = None, + ) -> torch.Tensor: + if attention_mask is None: + pooled_projections = hidden_states.mean(dim=1) + else: + original_dtype = hidden_states.dtype + mask_float = attention_mask.float().unsqueeze(-1) + pooled_projections = (hidden_states * mask_float).sum(dim=1) / mask_float.sum(dim=1) + pooled_projections = pooled_projections.to(original_dtype) + + temb = self.time_text_embed(timestep, pooled_projections) + hidden_states = self.proj_in(hidden_states) + hidden_states = self.token_refiner(hidden_states, temb, attention_mask) + + return hidden_states + + +class HunyuanVideoRotaryPosEmbed(nn.Module): + def __init__(self, rope_dim, theta): + super().__init__() + self.DT, self.DY, self.DX = rope_dim + self.theta = theta + + @torch.no_grad() + def get_frequency(self, dim, pos): + T, H, W = pos.shape + freqs = 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=pos.device)[: (dim // 2)] / dim)) + freqs = torch.outer(freqs, pos.reshape(-1)).unflatten(-1, (T, H, W)).repeat_interleave(2, dim=0) + return freqs.cos(), freqs.sin() + + @torch.no_grad() + def forward_inner(self, frame_indices, height, width, device): + GT, GY, GX = torch.meshgrid( + frame_indices.to(device=device, dtype=torch.float32), + torch.arange(0, height, device=device, dtype=torch.float32), + torch.arange(0, width, device=device, dtype=torch.float32), + indexing="ij" + ) + + FCT, FST = self.get_frequency(self.DT, GT) + FCY, FSY = self.get_frequency(self.DY, GY) + FCX, FSX = self.get_frequency(self.DX, GX) + + result = torch.cat([FCT, FCY, FCX, FST, FSY, FSX], dim=0) + + return result.to(device) + + @torch.no_grad() + def forward(self, frame_indices, height, width, device): + frame_indices = frame_indices.unbind(0) + results = [self.forward_inner(f, height, width, device) for f in frame_indices] + results = torch.stack(results, dim=0) + return results + + +class AdaLayerNormZero(nn.Module): + def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=bias) + if norm_type == "layer_norm": + self.norm = LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) + else: + raise ValueError(f"unknown norm_type {norm_type}") + + def forward( + self, + x: torch.Tensor, + emb: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + emb = emb.unsqueeze(-2) + emb = self.linear(self.silu(emb)) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1) + x = self.norm(x) * (1 + scale_msa) + shift_msa + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +class AdaLayerNormZeroSingle(nn.Module): + def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True): + super().__init__() + + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, 3 * embedding_dim, bias=bias) + if norm_type == "layer_norm": + self.norm = LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) + else: + raise ValueError(f"unknown norm_type {norm_type}") + + def forward( + self, + x: torch.Tensor, + emb: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + emb = emb.unsqueeze(-2) + emb = self.linear(self.silu(emb)) + shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=-1) + x = self.norm(x) * (1 + scale_msa) + shift_msa + return x, gate_msa + + +class AdaLayerNormContinuous(nn.Module): + def __init__( + self, + embedding_dim: int, + conditioning_embedding_dim: int, + elementwise_affine=True, + eps=1e-5, + bias=True, + norm_type="layer_norm", + ): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias) + if norm_type == "layer_norm": + self.norm = LayerNorm(embedding_dim, eps, elementwise_affine, bias) + else: + raise ValueError(f"unknown norm_type {norm_type}") + + def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: + emb = emb.unsqueeze(-2) + emb = self.linear(self.silu(emb)) + scale, shift = emb.chunk(2, dim=-1) + x = self.norm(x) * (1 + scale) + shift + return x + + +class HunyuanVideoSingleTransformerBlock(nn.Module): + def __init__( + self, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float = 4.0, + qk_norm: str = "rms_norm", + ) -> None: + super().__init__() + + hidden_size = num_attention_heads * attention_head_dim + mlp_dim = int(hidden_size * mlp_ratio) + + self.attn = Attention( + query_dim=hidden_size, + cross_attention_dim=None, + dim_head=attention_head_dim, + heads=num_attention_heads, + out_dim=hidden_size, + bias=True, + processor=HunyuanAttnProcessorFlashAttnSingle(), + qk_norm=qk_norm, + eps=1e-6, + pre_only=True, + ) + + self.norm = AdaLayerNormZeroSingle(hidden_size, norm_type="layer_norm") + self.proj_mlp = nn.Linear(hidden_size, mlp_dim) + self.act_mlp = nn.GELU(approximate="tanh") + self.proj_out = nn.Linear(hidden_size + mlp_dim, hidden_size) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> torch.Tensor: + text_seq_length = encoder_hidden_states.shape[1] + hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) + + residual = hidden_states + + # 1. Input normalization + norm_hidden_states, gate = self.norm(hidden_states, emb=temb) + mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)) + + norm_hidden_states, norm_encoder_hidden_states = ( + norm_hidden_states[:, :-text_seq_length, :], + norm_hidden_states[:, -text_seq_length:, :], + ) + + # 2. Attention + attn_output, context_attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + attn_output = torch.cat([attn_output, context_attn_output], dim=1) + + # 3. Modulation and residual connection + hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2) + hidden_states = gate * self.proj_out(hidden_states) + hidden_states = hidden_states + residual + + hidden_states, encoder_hidden_states = ( + hidden_states[:, :-text_seq_length, :], + hidden_states[:, -text_seq_length:, :], + ) + return hidden_states, encoder_hidden_states + + +class HunyuanVideoTransformerBlock(nn.Module): + def __init__( + self, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float, + qk_norm: str = "rms_norm", + ) -> None: + super().__init__() + + hidden_size = num_attention_heads * attention_head_dim + + self.norm1 = AdaLayerNormZero(hidden_size, norm_type="layer_norm") + self.norm1_context = AdaLayerNormZero(hidden_size, norm_type="layer_norm") + + self.attn = Attention( + query_dim=hidden_size, + cross_attention_dim=None, + added_kv_proj_dim=hidden_size, + dim_head=attention_head_dim, + heads=num_attention_heads, + out_dim=hidden_size, + context_pre_only=False, + bias=True, + processor=HunyuanAttnProcessorFlashAttnDouble(), + qk_norm=qk_norm, + eps=1e-6, + ) + + self.norm2 = LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate") + + self.norm2_context = LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.ff_context = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate") + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # 1. Input normalization + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) + norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(encoder_hidden_states, emb=temb) + + # 2. Joint attention + attn_output, context_attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + attention_mask=attention_mask, + image_rotary_emb=freqs_cis, + ) + + # 3. Modulation and residual connection + hidden_states = hidden_states + attn_output * gate_msa + encoder_hidden_states = encoder_hidden_states + context_attn_output * c_gate_msa + + norm_hidden_states = self.norm2(hidden_states) + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) + + norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp + norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp + + # 4. Feed-forward + ff_output = self.ff(norm_hidden_states) + context_ff_output = self.ff_context(norm_encoder_hidden_states) + + hidden_states = hidden_states + gate_mlp * ff_output + encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output + + return hidden_states, encoder_hidden_states + + +class ClipVisionProjection(nn.Module): + def __init__(self, in_channels, out_channels): + super().__init__() + self.up = nn.Linear(in_channels, out_channels * 3) + self.down = nn.Linear(out_channels * 3, out_channels) + + def forward(self, x): + projected_x = self.down(nn.functional.silu(self.up(x))) + return projected_x + + +class HunyuanVideoPatchEmbed(nn.Module): + def __init__(self, patch_size, in_chans, embed_dim): + super().__init__() + self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + + +class HunyuanVideoPatchEmbedForCleanLatents(nn.Module): + def __init__(self, inner_dim): + super().__init__() + self.proj = nn.Conv3d(16, inner_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2)) + self.proj_2x = nn.Conv3d(16, inner_dim, kernel_size=(2, 4, 4), stride=(2, 4, 4)) + self.proj_4x = nn.Conv3d(16, inner_dim, kernel_size=(4, 8, 8), stride=(4, 8, 8)) + + @torch.no_grad() + def initialize_weight_from_another_conv3d(self, another_layer): + weight = another_layer.weight.detach().clone() + bias = another_layer.bias.detach().clone() + + sd = { + 'proj.weight': weight.clone(), + 'proj.bias': bias.clone(), + 'proj_2x.weight': einops.repeat(weight, 'b c t h w -> b c (t tk) (h hk) (w wk)', tk=2, hk=2, wk=2) / 8.0, + 'proj_2x.bias': bias.clone(), + 'proj_4x.weight': einops.repeat(weight, 'b c t h w -> b c (t tk) (h hk) (w wk)', tk=4, hk=4, wk=4) / 64.0, + 'proj_4x.bias': bias.clone(), + } + + sd = {k: v.clone() for k, v in sd.items()} + + self.load_state_dict(sd) + return + + +class HunyuanVideoTransformer3DModelPacked(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): + @register_to_config + def __init__( + self, + in_channels: int = 16, + out_channels: int = 16, + num_attention_heads: int = 24, + attention_head_dim: int = 128, + num_layers: int = 20, + num_single_layers: int = 40, + num_refiner_layers: int = 2, + mlp_ratio: float = 4.0, + patch_size: int = 2, + patch_size_t: int = 1, + qk_norm: str = "rms_norm", + guidance_embeds: bool = True, + text_embed_dim: int = 4096, + pooled_projection_dim: int = 768, + rope_theta: float = 256.0, + rope_axes_dim: Tuple[int] = (16, 56, 56), + has_image_proj=False, + image_proj_dim=1152, + has_clean_x_embedder=False, + ) -> None: + super().__init__() + + inner_dim = num_attention_heads * attention_head_dim + out_channels = out_channels or in_channels + + # 1. Latent and condition embedders + self.x_embedder = HunyuanVideoPatchEmbed((patch_size_t, patch_size, patch_size), in_channels, inner_dim) + self.context_embedder = HunyuanVideoTokenRefiner( + text_embed_dim, num_attention_heads, attention_head_dim, num_layers=num_refiner_layers + ) + self.time_text_embed = CombinedTimestepGuidanceTextProjEmbeddings(inner_dim, pooled_projection_dim) + + self.clean_x_embedder = None + self.image_projection = None + + # 2. RoPE + self.rope = HunyuanVideoRotaryPosEmbed(rope_axes_dim, rope_theta) + + # 3. Dual stream transformer blocks + self.transformer_blocks = nn.ModuleList( + [ + HunyuanVideoTransformerBlock( + num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm + ) + for _ in range(num_layers) + ] + ) + + # 4. Single stream transformer blocks + self.single_transformer_blocks = nn.ModuleList( + [ + HunyuanVideoSingleTransformerBlock( + num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm + ) + for _ in range(num_single_layers) + ] + ) + + # 5. Output projection + self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out = nn.Linear(inner_dim, patch_size_t * patch_size * patch_size * out_channels) + + self.inner_dim = inner_dim + self.use_gradient_checkpointing = False + self.enable_teacache = False + + if has_image_proj: + self.install_image_projection(image_proj_dim) + + if has_clean_x_embedder: + self.install_clean_x_embedder() + + self.high_quality_fp32_output_for_inference = False + + def install_image_projection(self, in_channels): + self.image_projection = ClipVisionProjection(in_channels=in_channels, out_channels=self.inner_dim) + self.config['has_image_proj'] = True + self.config['image_proj_dim'] = in_channels + + def install_clean_x_embedder(self): + self.clean_x_embedder = HunyuanVideoPatchEmbedForCleanLatents(self.inner_dim) + self.config['has_clean_x_embedder'] = True + + def enable_gradient_checkpointing(self): + self.use_gradient_checkpointing = True + print('self.use_gradient_checkpointing = True') + + def disable_gradient_checkpointing(self): + self.use_gradient_checkpointing = False + print('self.use_gradient_checkpointing = False') + + def initialize_teacache(self, enable_teacache=True, num_steps=25, rel_l1_thresh=0.15): + self.enable_teacache = enable_teacache + self.cnt = 0 + self.num_steps = num_steps + self.rel_l1_thresh = rel_l1_thresh # 0.1 for 1.6x speedup, 0.15 for 2.1x speedup + self.accumulated_rel_l1_distance = 0 + self.previous_modulated_input = None + self.previous_residual = None + self.teacache_rescale_func = np.poly1d([7.33226126e+02, -4.01131952e+02, 6.75869174e+01, -3.14987800e+00, 9.61237896e-02]) + + def gradient_checkpointing_method(self, block, *args): + if self.use_gradient_checkpointing: + result = torch.utils.checkpoint.checkpoint(block, *args, use_reentrant=False) + else: + result = block(*args) + return result + + def process_input_hidden_states( + self, + latents, latent_indices=None, + clean_latents=None, clean_latent_indices=None, + clean_latents_2x=None, clean_latent_2x_indices=None, + clean_latents_4x=None, clean_latent_4x_indices=None + ): + hidden_states = self.gradient_checkpointing_method(self.x_embedder.proj, latents) + B, C, T, H, W = hidden_states.shape + + if latent_indices is None: + latent_indices = torch.arange(0, T).unsqueeze(0).expand(B, -1) + + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + rope_freqs = self.rope(frame_indices=latent_indices, height=H, width=W, device=hidden_states.device) + rope_freqs = rope_freqs.flatten(2).transpose(1, 2) + + if clean_latents is not None and clean_latent_indices is not None: + clean_latents = clean_latents.to(hidden_states) + clean_latents = self.gradient_checkpointing_method(self.clean_x_embedder.proj, clean_latents) + clean_latents = clean_latents.flatten(2).transpose(1, 2) + + clean_latent_rope_freqs = self.rope(frame_indices=clean_latent_indices, height=H, width=W, device=clean_latents.device) + clean_latent_rope_freqs = clean_latent_rope_freqs.flatten(2).transpose(1, 2) + + hidden_states = torch.cat([clean_latents, hidden_states], dim=1) + rope_freqs = torch.cat([clean_latent_rope_freqs, rope_freqs], dim=1) + + if clean_latents_2x is not None and clean_latent_2x_indices is not None: + clean_latents_2x = clean_latents_2x.to(hidden_states) + clean_latents_2x = pad_for_3d_conv(clean_latents_2x, (2, 4, 4)) + clean_latents_2x = self.gradient_checkpointing_method(self.clean_x_embedder.proj_2x, clean_latents_2x) + clean_latents_2x = clean_latents_2x.flatten(2).transpose(1, 2) + + clean_latent_2x_rope_freqs = self.rope(frame_indices=clean_latent_2x_indices, height=H, width=W, device=clean_latents_2x.device) + clean_latent_2x_rope_freqs = pad_for_3d_conv(clean_latent_2x_rope_freqs, (2, 2, 2)) + clean_latent_2x_rope_freqs = center_down_sample_3d(clean_latent_2x_rope_freqs, (2, 2, 2)) + clean_latent_2x_rope_freqs = clean_latent_2x_rope_freqs.flatten(2).transpose(1, 2) + + hidden_states = torch.cat([clean_latents_2x, hidden_states], dim=1) + rope_freqs = torch.cat([clean_latent_2x_rope_freqs, rope_freqs], dim=1) + + if clean_latents_4x is not None and clean_latent_4x_indices is not None: + clean_latents_4x = clean_latents_4x.to(hidden_states) + clean_latents_4x = pad_for_3d_conv(clean_latents_4x, (4, 8, 8)) + clean_latents_4x = self.gradient_checkpointing_method(self.clean_x_embedder.proj_4x, clean_latents_4x) + clean_latents_4x = clean_latents_4x.flatten(2).transpose(1, 2) + + clean_latent_4x_rope_freqs = self.rope(frame_indices=clean_latent_4x_indices, height=H, width=W, device=clean_latents_4x.device) + clean_latent_4x_rope_freqs = pad_for_3d_conv(clean_latent_4x_rope_freqs, (4, 4, 4)) + clean_latent_4x_rope_freqs = center_down_sample_3d(clean_latent_4x_rope_freqs, (4, 4, 4)) + clean_latent_4x_rope_freqs = clean_latent_4x_rope_freqs.flatten(2).transpose(1, 2) + + hidden_states = torch.cat([clean_latents_4x, hidden_states], dim=1) + rope_freqs = torch.cat([clean_latent_4x_rope_freqs, rope_freqs], dim=1) + + return hidden_states, rope_freqs + + def forward( + self, + hidden_states, timestep, encoder_hidden_states, encoder_attention_mask, pooled_projections, guidance, + latent_indices=None, + clean_latents=None, clean_latent_indices=None, + clean_latents_2x=None, clean_latent_2x_indices=None, + clean_latents_4x=None, clean_latent_4x_indices=None, + image_embeddings=None, + attention_kwargs=None, return_dict=True + ): + + if attention_kwargs is None: + attention_kwargs = {} + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p, p_t = self.config['patch_size'], self.config['patch_size_t'] + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p + post_patch_width = width // p + original_context_length = post_patch_num_frames * post_patch_height * post_patch_width + + hidden_states, rope_freqs = self.process_input_hidden_states(hidden_states, latent_indices, clean_latents, clean_latent_indices, clean_latents_2x, clean_latent_2x_indices, clean_latents_4x, clean_latent_4x_indices) + + temb = self.gradient_checkpointing_method(self.time_text_embed, timestep, guidance, pooled_projections) + encoder_hidden_states = self.gradient_checkpointing_method(self.context_embedder, encoder_hidden_states, timestep, encoder_attention_mask) + + if self.image_projection is not None: + assert image_embeddings is not None, 'You must use image embeddings!' + extra_encoder_hidden_states = self.gradient_checkpointing_method(self.image_projection, image_embeddings) + extra_attention_mask = torch.ones((batch_size, extra_encoder_hidden_states.shape[1]), dtype=encoder_attention_mask.dtype, device=encoder_attention_mask.device) + + # must cat before (not after) encoder_hidden_states, due to attn masking + encoder_hidden_states = torch.cat([extra_encoder_hidden_states, encoder_hidden_states], dim=1) + encoder_attention_mask = torch.cat([extra_attention_mask, encoder_attention_mask], dim=1) + + with torch.no_grad(): + if batch_size == 1: + # When batch size is 1, we do not need any masks or var-len funcs since cropping is mathematically same to what we want + # If they are not same, then their impls are wrong. Ours are always the correct one. + text_len = encoder_attention_mask.sum().item() + encoder_hidden_states = encoder_hidden_states[:, :text_len] + attention_mask = None, None, None, None + else: + img_seq_len = hidden_states.shape[1] + txt_seq_len = encoder_hidden_states.shape[1] + + cu_seqlens_q = get_cu_seqlens(encoder_attention_mask, img_seq_len) + cu_seqlens_kv = cu_seqlens_q + max_seqlen_q = img_seq_len + txt_seq_len + max_seqlen_kv = max_seqlen_q + + attention_mask = cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv + + if self.enable_teacache: + modulated_inp = self.transformer_blocks[0].norm1(hidden_states, emb=temb)[0] + + if self.cnt == 0 or self.cnt == self.num_steps-1: + should_calc = True + self.accumulated_rel_l1_distance = 0 + else: + curr_rel_l1 = ((modulated_inp - self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item() + self.accumulated_rel_l1_distance += self.teacache_rescale_func(curr_rel_l1) + should_calc = self.accumulated_rel_l1_distance >= self.rel_l1_thresh + + if should_calc: + self.accumulated_rel_l1_distance = 0 + + self.previous_modulated_input = modulated_inp + self.cnt += 1 + + if self.cnt == self.num_steps: + self.cnt = 0 + + if not should_calc: + hidden_states = hidden_states + self.previous_residual + else: + ori_hidden_states = hidden_states.clone() + + for block_id, block in enumerate(self.transformer_blocks): + hidden_states, encoder_hidden_states = self.gradient_checkpointing_method( + block, + hidden_states, + encoder_hidden_states, + temb, + attention_mask, + rope_freqs + ) + + for block_id, block in enumerate(self.single_transformer_blocks): + hidden_states, encoder_hidden_states = self.gradient_checkpointing_method( + block, + hidden_states, + encoder_hidden_states, + temb, + attention_mask, + rope_freqs + ) + + self.previous_residual = hidden_states - ori_hidden_states + else: + for block_id, block in enumerate(self.transformer_blocks): + hidden_states, encoder_hidden_states = self.gradient_checkpointing_method( + block, + hidden_states, + encoder_hidden_states, + temb, + attention_mask, + rope_freqs + ) + + for block_id, block in enumerate(self.single_transformer_blocks): + hidden_states, encoder_hidden_states = self.gradient_checkpointing_method( + block, + hidden_states, + encoder_hidden_states, + temb, + attention_mask, + rope_freqs + ) + + hidden_states = self.gradient_checkpointing_method(self.norm_out, hidden_states, temb) + + hidden_states = hidden_states[:, -original_context_length:, :] + + if self.high_quality_fp32_output_for_inference: + hidden_states = hidden_states.to(dtype=torch.float32) + if self.proj_out.weight.dtype != torch.float32: + self.proj_out.to(dtype=torch.float32) + + hidden_states = self.gradient_checkpointing_method(self.proj_out, hidden_states) + + hidden_states = einops.rearrange(hidden_states, 'b (t h w) (c pt ph pw) -> b c (t pt) (h ph) (w pw)', + t=post_patch_num_frames, h=post_patch_height, w=post_patch_width, + pt=p_t, ph=p, pw=p) + + if return_dict: + return Transformer2DModelOutput(sample=hidden_states) + + return hidden_states, diff --git a/diffusers_helper/pipelines/k_diffusion_hunyuan.py b/diffusers_helper/pipelines/k_diffusion_hunyuan.py new file mode 100644 index 0000000..85c0238 --- /dev/null +++ b/diffusers_helper/pipelines/k_diffusion_hunyuan.py @@ -0,0 +1,120 @@ +import torch +import math + +from diffusers_helper.k_diffusion.uni_pc_fm import sample_unipc +from diffusers_helper.k_diffusion.wrapper import fm_wrapper +from diffusers_helper.utils import repeat_to_batch_size + + +def flux_time_shift(t, mu=1.15, sigma=1.0): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + + +def calculate_flux_mu(context_length, x1=256, y1=0.5, x2=4096, y2=1.15, exp_max=7.0): + k = (y2 - y1) / (x2 - x1) + b = y1 - k * x1 + mu = k * context_length + b + mu = min(mu, math.log(exp_max)) + return mu + + +def get_flux_sigmas_from_mu(n, mu): + sigmas = torch.linspace(1, 0, steps=n + 1) + sigmas = flux_time_shift(sigmas, mu=mu) + return sigmas + + +@torch.inference_mode() +def sample_hunyuan( + transformer, + sampler='unipc', + initial_latent=None, + concat_latent=None, + strength=1.0, + width=512, + height=512, + frames=16, + real_guidance_scale=1.0, + distilled_guidance_scale=6.0, + guidance_rescale=0.0, + shift=None, + num_inference_steps=25, + batch_size=None, + generator=None, + prompt_embeds=None, + prompt_embeds_mask=None, + prompt_poolers=None, + negative_prompt_embeds=None, + negative_prompt_embeds_mask=None, + negative_prompt_poolers=None, + dtype=torch.bfloat16, + device=None, + negative_kwargs=None, + callback=None, + **kwargs, +): + device = device or transformer.device + + if batch_size is None: + batch_size = int(prompt_embeds.shape[0]) + + latents = torch.randn((batch_size, 16, (frames + 3) // 4, height // 8, width // 8), generator=generator, device=generator.device).to(device=device, dtype=torch.float32) + + B, C, T, H, W = latents.shape + seq_length = T * H * W // 4 + + if shift is None: + mu = calculate_flux_mu(seq_length, exp_max=7.0) + else: + mu = math.log(shift) + + sigmas = get_flux_sigmas_from_mu(num_inference_steps, mu).to(device) + + k_model = fm_wrapper(transformer) + + if initial_latent is not None: + sigmas = sigmas * strength + first_sigma = sigmas[0].to(device=device, dtype=torch.float32) + initial_latent = initial_latent.to(device=device, dtype=torch.float32) + latents = initial_latent.float() * (1.0 - first_sigma) + latents.float() * first_sigma + + if concat_latent is not None: + concat_latent = concat_latent.to(latents) + + distilled_guidance = torch.tensor([distilled_guidance_scale * 1000.0] * batch_size).to(device=device, dtype=dtype) + + prompt_embeds = repeat_to_batch_size(prompt_embeds, batch_size) + prompt_embeds_mask = repeat_to_batch_size(prompt_embeds_mask, batch_size) + prompt_poolers = repeat_to_batch_size(prompt_poolers, batch_size) + negative_prompt_embeds = repeat_to_batch_size(negative_prompt_embeds, batch_size) + negative_prompt_embeds_mask = repeat_to_batch_size(negative_prompt_embeds_mask, batch_size) + negative_prompt_poolers = repeat_to_batch_size(negative_prompt_poolers, batch_size) + concat_latent = repeat_to_batch_size(concat_latent, batch_size) + + sampler_kwargs = dict( + dtype=dtype, + cfg_scale=real_guidance_scale, + cfg_rescale=guidance_rescale, + concat_latent=concat_latent, + positive=dict( + pooled_projections=prompt_poolers, + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_embeds_mask, + guidance=distilled_guidance, + **kwargs, + ), + negative=dict( + pooled_projections=negative_prompt_poolers, + encoder_hidden_states=negative_prompt_embeds, + encoder_attention_mask=negative_prompt_embeds_mask, + guidance=distilled_guidance, + **(kwargs if negative_kwargs is None else {**kwargs, **negative_kwargs}), + ) + ) + + if sampler == 'unipc': + results = sample_unipc(k_model, latents, sigmas, extra_args=sampler_kwargs, disable=False, callback=callback) + else: + raise NotImplementedError(f'Sampler {sampler} is not supported.') + + return results diff --git a/diffusers_helper/thread_utils.py b/diffusers_helper/thread_utils.py new file mode 100644 index 0000000..c040426 --- /dev/null +++ b/diffusers_helper/thread_utils.py @@ -0,0 +1,76 @@ +import time + +from threading import Thread, Lock + + +class Listener: + task_queue = [] + lock = Lock() + thread = None + + @classmethod + def _process_tasks(cls): + while True: + task = None + with cls.lock: + if cls.task_queue: + task = cls.task_queue.pop(0) + + if task is None: + time.sleep(0.001) + continue + + func, args, kwargs = task + try: + func(*args, **kwargs) + except Exception as e: + print(f"Error in listener thread: {e}") + + @classmethod + def add_task(cls, func, *args, **kwargs): + with cls.lock: + cls.task_queue.append((func, args, kwargs)) + + if cls.thread is None: + cls.thread = Thread(target=cls._process_tasks, daemon=True) + cls.thread.start() + + +def async_run(func, *args, **kwargs): + Listener.add_task(func, *args, **kwargs) + + +class FIFOQueue: + def __init__(self): + self.queue = [] + self.lock = Lock() + + def push(self, item): + with self.lock: + self.queue.append(item) + + def pop(self): + with self.lock: + if self.queue: + return self.queue.pop(0) + return None + + def top(self): + with self.lock: + if self.queue: + return self.queue[0] + return None + + def next(self): + while True: + with self.lock: + if self.queue: + return self.queue.pop(0) + + time.sleep(0.001) + + +class AsyncStream: + def __init__(self): + self.input_queue = FIFOQueue() + self.output_queue = FIFOQueue() diff --git a/diffusers_helper/utils.py b/diffusers_helper/utils.py new file mode 100644 index 0000000..8ebda0f --- /dev/null +++ b/diffusers_helper/utils.py @@ -0,0 +1,613 @@ +import os +import cv2 +import json +import random +import glob +import torch +import einops +import numpy as np +import datetime +import torchvision + +import safetensors.torch as sf +from PIL import Image + + +def min_resize(x, m): + if x.shape[0] < x.shape[1]: + s0 = m + s1 = int(float(m) / float(x.shape[0]) * float(x.shape[1])) + else: + s0 = int(float(m) / float(x.shape[1]) * float(x.shape[0])) + s1 = m + new_max = max(s1, s0) + raw_max = max(x.shape[0], x.shape[1]) + if new_max < raw_max: + interpolation = cv2.INTER_AREA + else: + interpolation = cv2.INTER_LANCZOS4 + y = cv2.resize(x, (s1, s0), interpolation=interpolation) + return y + + +def d_resize(x, y): + H, W, C = y.shape + new_min = min(H, W) + raw_min = min(x.shape[0], x.shape[1]) + if new_min < raw_min: + interpolation = cv2.INTER_AREA + else: + interpolation = cv2.INTER_LANCZOS4 + y = cv2.resize(x, (W, H), interpolation=interpolation) + return y + + +def resize_and_center_crop(image, target_width, target_height): + if target_height == image.shape[0] and target_width == image.shape[1]: + return image + + pil_image = Image.fromarray(image) + original_width, original_height = pil_image.size + scale_factor = max(target_width / original_width, target_height / original_height) + resized_width = int(round(original_width * scale_factor)) + resized_height = int(round(original_height * scale_factor)) + resized_image = pil_image.resize((resized_width, resized_height), Image.LANCZOS) + left = (resized_width - target_width) / 2 + top = (resized_height - target_height) / 2 + right = (resized_width + target_width) / 2 + bottom = (resized_height + target_height) / 2 + cropped_image = resized_image.crop((left, top, right, bottom)) + return np.array(cropped_image) + + +def resize_and_center_crop_pytorch(image, target_width, target_height): + B, C, H, W = image.shape + + if H == target_height and W == target_width: + return image + + scale_factor = max(target_width / W, target_height / H) + resized_width = int(round(W * scale_factor)) + resized_height = int(round(H * scale_factor)) + + resized = torch.nn.functional.interpolate(image, size=(resized_height, resized_width), mode='bilinear', align_corners=False) + + top = (resized_height - target_height) // 2 + left = (resized_width - target_width) // 2 + cropped = resized[:, :, top:top + target_height, left:left + target_width] + + return cropped + + +def resize_without_crop(image, target_width, target_height): + if target_height == image.shape[0] and target_width == image.shape[1]: + return image + + pil_image = Image.fromarray(image) + resized_image = pil_image.resize((target_width, target_height), Image.LANCZOS) + return np.array(resized_image) + + +def just_crop(image, w, h): + if h == image.shape[0] and w == image.shape[1]: + return image + + original_height, original_width = image.shape[:2] + k = min(original_height / h, original_width / w) + new_width = int(round(w * k)) + new_height = int(round(h * k)) + x_start = (original_width - new_width) // 2 + y_start = (original_height - new_height) // 2 + cropped_image = image[y_start:y_start + new_height, x_start:x_start + new_width] + return cropped_image + + +def write_to_json(data, file_path): + temp_file_path = file_path + ".tmp" + with open(temp_file_path, 'wt', encoding='utf-8') as temp_file: + json.dump(data, temp_file, indent=4) + os.replace(temp_file_path, file_path) + return + + +def read_from_json(file_path): + with open(file_path, 'rt', encoding='utf-8') as file: + data = json.load(file) + return data + + +def get_active_parameters(m): + return {k: v for k, v in m.named_parameters() if v.requires_grad} + + +def cast_training_params(m, dtype=torch.float32): + result = {} + for n, param in m.named_parameters(): + if param.requires_grad: + param.data = param.to(dtype) + result[n] = param + return result + + +def separate_lora_AB(parameters, B_patterns=None): + parameters_normal = {} + parameters_B = {} + + if B_patterns is None: + B_patterns = ['.lora_B.', '__zero__'] + + for k, v in parameters.items(): + if any(B_pattern in k for B_pattern in B_patterns): + parameters_B[k] = v + else: + parameters_normal[k] = v + + return parameters_normal, parameters_B + + +def set_attr_recursive(obj, attr, value): + attrs = attr.split(".") + for name in attrs[:-1]: + obj = getattr(obj, name) + setattr(obj, attrs[-1], value) + return + + +def print_tensor_list_size(tensors): + total_size = 0 + total_elements = 0 + + if isinstance(tensors, dict): + tensors = tensors.values() + + for tensor in tensors: + total_size += tensor.nelement() * tensor.element_size() + total_elements += tensor.nelement() + + total_size_MB = total_size / (1024 ** 2) + total_elements_B = total_elements / 1e9 + + print(f"Total number of tensors: {len(tensors)}") + print(f"Total size of tensors: {total_size_MB:.2f} MB") + print(f"Total number of parameters: {total_elements_B:.3f} billion") + return + + +@torch.no_grad() +def batch_mixture(a, b=None, probability_a=0.5, mask_a=None): + batch_size = a.size(0) + + if b is None: + b = torch.zeros_like(a) + + if mask_a is None: + mask_a = torch.rand(batch_size) < probability_a + + mask_a = mask_a.to(a.device) + mask_a = mask_a.reshape((batch_size,) + (1,) * (a.dim() - 1)) + result = torch.where(mask_a, a, b) + return result + + +@torch.no_grad() +def zero_module(module): + for p in module.parameters(): + p.detach().zero_() + return module + + +@torch.no_grad() +def supress_lower_channels(m, k, alpha=0.01): + data = m.weight.data.clone() + + assert int(data.shape[1]) >= k + + data[:, :k] = data[:, :k] * alpha + m.weight.data = data.contiguous().clone() + return m + + +def freeze_module(m): + if not hasattr(m, '_forward_inside_frozen_module'): + m._forward_inside_frozen_module = m.forward + m.requires_grad_(False) + m.forward = torch.no_grad()(m.forward) + return m + + +def get_latest_safetensors(folder_path): + safetensors_files = glob.glob(os.path.join(folder_path, '*.safetensors')) + + if not safetensors_files: + raise ValueError('No file to resume!') + + latest_file = max(safetensors_files, key=os.path.getmtime) + latest_file = os.path.abspath(os.path.realpath(latest_file)) + return latest_file + + +def generate_random_prompt_from_tags(tags_str, min_length=3, max_length=32): + tags = tags_str.split(', ') + tags = random.sample(tags, k=min(random.randint(min_length, max_length), len(tags))) + prompt = ', '.join(tags) + return prompt + + +def interpolate_numbers(a, b, n, round_to_int=False, gamma=1.0): + numbers = a + (b - a) * (np.linspace(0, 1, n) ** gamma) + if round_to_int: + numbers = np.round(numbers).astype(int) + return numbers.tolist() + + +def uniform_random_by_intervals(inclusive, exclusive, n, round_to_int=False): + edges = np.linspace(0, 1, n + 1) + points = np.random.uniform(edges[:-1], edges[1:]) + numbers = inclusive + (exclusive - inclusive) * points + if round_to_int: + numbers = np.round(numbers).astype(int) + return numbers.tolist() + + +def soft_append_bcthw(history, current, overlap=0): + if overlap <= 0: + return torch.cat([history, current], dim=2) + + assert history.shape[2] >= overlap, f"History length ({history.shape[2]}) must be >= overlap ({overlap})" + assert current.shape[2] >= overlap, f"Current length ({current.shape[2]}) must be >= overlap ({overlap})" + + weights = torch.linspace(1, 0, overlap, dtype=history.dtype, device=history.device).view(1, 1, -1, 1, 1) + blended = weights * history[:, :, -overlap:] + (1 - weights) * current[:, :, :overlap] + output = torch.cat([history[:, :, :-overlap], blended, current[:, :, overlap:]], dim=2) + + return output.to(history) + + +def save_bcthw_as_mp4(x, output_filename, fps=10): + b, c, t, h, w = x.shape + + per_row = b + for p in [6, 5, 4, 3, 2]: + if b % p == 0: + per_row = p + break + + os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True) + x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5 + x = x.detach().cpu().to(torch.uint8) + x = einops.rearrange(x, '(m n) c t h w -> t (m h) (n w) c', n=per_row) + torchvision.io.write_video(output_filename, x, fps=fps, video_codec='h264', options={'crf': '0'}) + return x + + +def save_bcthw_as_png(x, output_filename): + os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True) + x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5 + x = x.detach().cpu().to(torch.uint8) + x = einops.rearrange(x, 'b c t h w -> c (b h) (t w)') + torchvision.io.write_png(x, output_filename) + return output_filename + + +def save_bchw_as_png(x, output_filename): + os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True) + x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5 + x = x.detach().cpu().to(torch.uint8) + x = einops.rearrange(x, 'b c h w -> c h (b w)') + torchvision.io.write_png(x, output_filename) + return output_filename + + +def add_tensors_with_padding(tensor1, tensor2): + if tensor1.shape == tensor2.shape: + return tensor1 + tensor2 + + shape1 = tensor1.shape + shape2 = tensor2.shape + + new_shape = tuple(max(s1, s2) for s1, s2 in zip(shape1, shape2)) + + padded_tensor1 = torch.zeros(new_shape) + padded_tensor2 = torch.zeros(new_shape) + + padded_tensor1[tuple(slice(0, s) for s in shape1)] = tensor1 + padded_tensor2[tuple(slice(0, s) for s in shape2)] = tensor2 + + result = padded_tensor1 + padded_tensor2 + return result + + +def print_free_mem(): + torch.cuda.empty_cache() + free_mem, total_mem = torch.cuda.mem_get_info(0) + free_mem_mb = free_mem / (1024 ** 2) + total_mem_mb = total_mem / (1024 ** 2) + print(f"Free memory: {free_mem_mb:.2f} MB") + print(f"Total memory: {total_mem_mb:.2f} MB") + return + + +def print_gpu_parameters(device, state_dict, log_count=1): + summary = {"device": device, "keys_count": len(state_dict)} + + logged_params = {} + for i, (key, tensor) in enumerate(state_dict.items()): + if i >= log_count: + break + logged_params[key] = tensor.flatten()[:3].tolist() + + summary["params"] = logged_params + + print(str(summary)) + return + + +def visualize_txt_as_img(width, height, text, font_path='font/DejaVuSans.ttf', size=18): + from PIL import Image, ImageDraw, ImageFont + + txt = Image.new("RGB", (width, height), color="white") + draw = ImageDraw.Draw(txt) + font = ImageFont.truetype(font_path, size=size) + + if text == '': + return np.array(txt) + + # Split text into lines that fit within the image width + lines = [] + words = text.split() + current_line = words[0] + + for word in words[1:]: + line_with_word = f"{current_line} {word}" + if draw.textbbox((0, 0), line_with_word, font=font)[2] <= width: + current_line = line_with_word + else: + lines.append(current_line) + current_line = word + + lines.append(current_line) + + # Draw the text line by line + y = 0 + line_height = draw.textbbox((0, 0), "A", font=font)[3] + + for line in lines: + if y + line_height > height: + break # stop drawing if the next line will be outside the image + draw.text((0, y), line, fill="black", font=font) + y += line_height + + return np.array(txt) + + +def blue_mark(x): + x = x.copy() + c = x[:, :, 2] + b = cv2.blur(c, (9, 9)) + x[:, :, 2] = ((c - b) * 16.0 + b).clip(-1, 1) + return x + + +def green_mark(x): + x = x.copy() + x[:, :, 2] = -1 + x[:, :, 0] = -1 + return x + + +def frame_mark(x): + x = x.copy() + x[:64] = -1 + x[-64:] = -1 + x[:, :8] = 1 + x[:, -8:] = 1 + return x + + +@torch.inference_mode() +def pytorch2numpy(imgs): + results = [] + for x in imgs: + y = x.movedim(0, -1) + y = y * 127.5 + 127.5 + y = y.detach().float().cpu().numpy().clip(0, 255).astype(np.uint8) + results.append(y) + return results + + +@torch.inference_mode() +def numpy2pytorch(imgs): + h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.5 - 1.0 + h = h.movedim(-1, 1) + return h + + +@torch.no_grad() +def duplicate_prefix_to_suffix(x, count, zero_out=False): + if zero_out: + return torch.cat([x, torch.zeros_like(x[:count])], dim=0) + else: + return torch.cat([x, x[:count]], dim=0) + + +def weighted_mse(a, b, weight): + return torch.mean(weight.float() * (a.float() - b.float()) ** 2) + + +def clamped_linear_interpolation(x, x_min, y_min, x_max, y_max, sigma=1.0): + x = (x - x_min) / (x_max - x_min) + x = max(0.0, min(x, 1.0)) + x = x ** sigma + return y_min + x * (y_max - y_min) + + +def expand_to_dims(x, target_dims): + return x.view(*x.shape, *([1] * max(0, target_dims - x.dim()))) + + +def repeat_to_batch_size(tensor: torch.Tensor, batch_size: int): + if tensor is None: + return None + + first_dim = tensor.shape[0] + + if first_dim == batch_size: + return tensor + + if batch_size % first_dim != 0: + raise ValueError(f"Cannot evenly repeat first dim {first_dim} to match batch_size {batch_size}.") + + repeat_times = batch_size // first_dim + + return tensor.repeat(repeat_times, *[1] * (tensor.dim() - 1)) + + +def dim5(x): + return expand_to_dims(x, 5) + + +def dim4(x): + return expand_to_dims(x, 4) + + +def dim3(x): + return expand_to_dims(x, 3) + + +def crop_or_pad_yield_mask(x, length): + B, F, C = x.shape + device = x.device + dtype = x.dtype + + if F < length: + y = torch.zeros((B, length, C), dtype=dtype, device=device) + mask = torch.zeros((B, length), dtype=torch.bool, device=device) + y[:, :F, :] = x + mask[:, :F] = True + return y, mask + + return x[:, :length, :], torch.ones((B, length), dtype=torch.bool, device=device) + + +def extend_dim(x, dim, minimal_length, zero_pad=False): + original_length = int(x.shape[dim]) + + if original_length >= minimal_length: + return x + + if zero_pad: + padding_shape = list(x.shape) + padding_shape[dim] = minimal_length - original_length + padding = torch.zeros(padding_shape, dtype=x.dtype, device=x.device) + else: + idx = (slice(None),) * dim + (slice(-1, None),) + (slice(None),) * (len(x.shape) - dim - 1) + last_element = x[idx] + padding = last_element.repeat_interleave(minimal_length - original_length, dim=dim) + + return torch.cat([x, padding], dim=dim) + + +def lazy_positional_encoding(t, repeats=None): + if not isinstance(t, list): + t = [t] + + from diffusers.models.embeddings import get_timestep_embedding + + te = torch.tensor(t) + te = get_timestep_embedding(timesteps=te, embedding_dim=256, flip_sin_to_cos=True, downscale_freq_shift=0.0, scale=1.0) + + if repeats is None: + return te + + te = te[:, None, :].expand(-1, repeats, -1) + + return te + + +def state_dict_offset_merge(A, B, C=None): + result = {} + keys = A.keys() + + for key in keys: + A_value = A[key] + B_value = B[key].to(A_value) + + if C is None: + result[key] = A_value + B_value + else: + C_value = C[key].to(A_value) + result[key] = A_value + B_value - C_value + + return result + + +def state_dict_weighted_merge(state_dicts, weights): + if len(state_dicts) != len(weights): + raise ValueError("Number of state dictionaries must match number of weights") + + if not state_dicts: + return {} + + total_weight = sum(weights) + + if total_weight == 0: + raise ValueError("Sum of weights cannot be zero") + + normalized_weights = [w / total_weight for w in weights] + + keys = state_dicts[0].keys() + result = {} + + for key in keys: + result[key] = state_dicts[0][key] * normalized_weights[0] + + for i in range(1, len(state_dicts)): + state_dict_value = state_dicts[i][key].to(result[key]) + result[key] += state_dict_value * normalized_weights[i] + + return result + + +def group_files_by_folder(all_files): + grouped_files = {} + + for file in all_files: + folder_name = os.path.basename(os.path.dirname(file)) + if folder_name not in grouped_files: + grouped_files[folder_name] = [] + grouped_files[folder_name].append(file) + + list_of_lists = list(grouped_files.values()) + return list_of_lists + + +def generate_timestamp(): + now = datetime.datetime.now() + timestamp = now.strftime('%y%m%d_%H%M%S') + milliseconds = f"{int(now.microsecond / 1000):03d}" + random_number = random.randint(0, 9999) + return f"{timestamp}_{milliseconds}_{random_number}" + + +def write_PIL_image_with_png_info(image, metadata, path): + from PIL.PngImagePlugin import PngInfo + + png_info = PngInfo() + for key, value in metadata.items(): + png_info.add_text(key, value) + + image.save(path, "PNG", pnginfo=png_info) + return image + + +def torch_safe_save(content, path): + torch.save(content, path + '_tmp') + os.replace(path + '_tmp', path) + return path + + +def move_optimizer_to_device(optimizer, device): + for state in optimizer.state.values(): + for k, v in state.items(): + if isinstance(v, torch.Tensor): + state[k] = v.to(device) diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..d5058c5 --- /dev/null +++ b/nodes.py @@ -0,0 +1,707 @@ +import os +import sys +current_dir = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, current_dir) + +import torch +import traceback +import einops +import safetensors.torch as sf +import numpy as np +import argparse +import math + +from PIL import Image +from diffusers import AutoencoderKLHunyuanVideo +from transformers import LlamaModel, CLIPTextModel, LlamaTokenizerFast, CLIPTokenizer +from diffusers_helper.hunyuan import encode_prompt_conds, vae_decode, vae_encode, vae_decode_fake +from diffusers_helper.utils import save_bcthw_as_mp4, crop_or_pad_yield_mask, soft_append_bcthw, resize_and_center_crop, state_dict_weighted_merge, state_dict_offset_merge, generate_timestamp +from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked +from diffusers_helper.pipelines.k_diffusion_hunyuan import sample_hunyuan +from diffusers_helper.memory import cpu, gpu, get_cuda_free_memory_gb, move_model_to_device_with_memory_preservation, offload_model_from_device_for_memory_preservation, fake_diffusers_current_device, DynamicSwapInstaller, unload_complete_models, load_model_as_complete +from transformers import SiglipImageProcessor, SiglipVisionModel +from diffusers_helper.clip_vision import hf_clip_vision_encode +from diffusers_helper.bucket_tools import find_nearest_bucket +import hashlib +import random +import string +import torchvision +from torchvision.transforms.functional import to_pil_image +import comfy.utils + +from PIL import Image +import folder_paths + +class Kiki_FramePack: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "ref_image": ("IMAGE", ), + "prompt": ("STRING", {"multiline": True}), + "total_second_length": ("INT", {"default": 5, "min": 1, "max": 120, "step": 1}), + "seed": ("INT", {"default": 3407, "min": 0, "max": 0xffffffffffffffff}), + "steps": ("INT", {"default": 25, "min": 1, "max": 100, "step": 1}), + "use_teacache": ("BOOLEAN", {"default": True}), + "resolution": (["360p", "480p", "540p", "720p"], {"default": "480p"}), + "padding_mode": (["optimized", "default (test)", "constant (test)", "waterfall (test)", "center_focus (test)"], {"default": "optimized"}), + "end_condition_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05}), + "enable_feature_fusion": ("BOOLEAN", {"default": True}), + "history_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "history_decay": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 0.9, "step": 0.01}), + "history_weight_min": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 0.8, "step": 0.01}), + "gpu_memory_preservation": ("FLOAT", {"default": 6, "min": 6, "max": 128, "step": 0.1, "description": "GPU内存保留(GB)，设置更大的值可避免OOM但会降低速度"}), + "use_flash_attention": ("BOOLEAN", {"default": False}), + "use_sage_attention": ("BOOLEAN", {"default": False}), + "overlap_frames": ("INT", {"default": 33, "min": 1, "max": 100, "step": 1, "description": "section之间的混合帧数，默认值33 (latent_window_size * 4 - 3)"}), + "blend_mode": (["linear", "cosine", "sigmoid"], {"default": "linear", "description": "混合曲线类型：linear(线性)、cosine(余弦)、sigmoid(S型)"}), + }, + "optional": { + "end_image": ("IMAGE", ), + }, + } + + RETURN_TYPES = ("IMAGE", "FLOAT") + RETURN_NAMES = ("frames", "fps") + CATEGORY = "TTPlanet/FramePack" + FUNCTION = "run" + + TITLE = 'TTPlanet FramePack' + OUTPUT_NODE = True + + def __init__(self): + self.high_vram = False + self.frames = None + self.fps = None + + hunyuan_root = os.path.join(folder_paths.models_dir, 'HunyuanVideo') + flux_redux_bfl_root = os.path.join(folder_paths.models_dir, 'flux_redux_bfl') + framePackI2V_root = os.path.join(folder_paths.models_dir, 'FramePackI2V_HY') + + self.text_encoder = LlamaModel.from_pretrained(hunyuan_root, subfolder='text_encoder', torch_dtype=torch.float16).cpu() + self.text_encoder_2 = CLIPTextModel.from_pretrained(hunyuan_root, subfolder='text_encoder_2', torch_dtype=torch.float16).cpu() + self.tokenizer = LlamaTokenizerFast.from_pretrained(hunyuan_root, subfolder='tokenizer') + self.tokenizer_2 = CLIPTokenizer.from_pretrained(hunyuan_root, subfolder='tokenizer_2') + self.vae = AutoencoderKLHunyuanVideo.from_pretrained(hunyuan_root, subfolder='vae', torch_dtype=torch.float16).cpu() + + self.feature_extractor = SiglipImageProcessor.from_pretrained(flux_redux_bfl_root, subfolder='feature_extractor') + self.image_encoder = SiglipVisionModel.from_pretrained(flux_redux_bfl_root, subfolder='image_encoder', torch_dtype=torch.float16).cpu() + + self.transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained(framePackI2V_root, torch_dtype=torch.bfloat16).cpu() + + self.vae.eval() + self.text_encoder.eval() + self.text_encoder_2.eval() + self.image_encoder.eval() + self.transformer.eval() + + if not self.high_vram: + self.vae.enable_slicing() + self.vae.enable_tiling() + + self.transformer.high_quality_fp32_output_for_inference = True + print('transformer.high_quality_fp32_output_for_inference = True') + + self.transformer.to(dtype=torch.bfloat16) + self.vae.to(dtype=torch.float16) + self.image_encoder.to(dtype=torch.float16) + self.text_encoder.to(dtype=torch.float16) + self.text_encoder_2.to(dtype=torch.float16) + + self.vae.requires_grad_(False) + self.text_encoder.requires_grad_(False) + self.text_encoder_2.requires_grad_(False) + self.image_encoder.requires_grad_(False) + self.transformer.requires_grad_(False) + + if not self.high_vram: + # DynamicSwapInstaller is same as huggingface's enable_sequential_offload but 3x faster + DynamicSwapInstaller.install_model(self.transformer, device=gpu) + DynamicSwapInstaller.install_model(self.text_encoder, device=gpu) + + def strict_align(self, h, w): + """确保高度和宽度是64的倍数，且潜在空间是8的倍数""" + aligned_h = int(round(h / 64)) * 64 + aligned_w = int(round(w / 64)) * 64 + + assert (aligned_h % 64 == 0) and (aligned_w % 64 == 0), "尺寸必须是64的倍数" + assert (aligned_h//8) % 8 == 0 and (aligned_w//8) % 8 == 0, "潜在空间需要8的倍数" + return aligned_h, aligned_w + + def preprocess_image(self, image): + if image is None: + return None + image_np = 255. * image[0].cpu().numpy() + image = Image.fromarray(np.clip(image_np, 0, 255).astype(np.uint8)).convert("RGB") + input_image = np.array(image) + return input_image + + def run(self, **kwargs): + try: + image = kwargs['ref_image'] + end_image = kwargs.get('end_image', None) # Use get with None as default + image_np = self.preprocess_image(image) + end_image_np = self.preprocess_image(end_image) if end_image is not None else None + prompt = kwargs['prompt'] + seed = kwargs['seed'] + total_second_length = kwargs['total_second_length'] + steps = kwargs['steps'] + use_teacache = kwargs['use_teacache'] + resolution = kwargs['resolution'] + padding_mode = kwargs['padding_mode'] + end_condition_strength = kwargs['end_condition_strength'] + enable_feature_fusion = kwargs['enable_feature_fusion'] + history_weight = kwargs['history_weight'] + history_decay = kwargs['history_decay'] + history_weight_min = kwargs['history_weight_min'] + gpu_memory_preservation = kwargs['gpu_memory_preservation'] + use_flash_attention = kwargs['use_flash_attention'] + use_sage_attention = kwargs['use_sage_attention'] + overlap_frames = kwargs['overlap_frames'] + blend_mode = kwargs['blend_mode'] + + # Remove "(test)" suffix from padding_mode + real_padding_mode = padding_mode.split(" ")[0] + + random_str = ''.join(random.choices(string.ascii_letters + string.digits, k=16)) + video_path = os.path.join(folder_paths.get_output_directory(), f'{random_str}.mp4') + + self.pbar = comfy.utils.ProgressBar(steps * total_second_length) + + self.exec( + input_image=image_np, + end_image=end_image_np, + prompt=prompt, + seed=seed, + total_second_length=total_second_length, + video_path=video_path, + steps=steps, + use_teacache=use_teacache, + resolution=resolution, + padding_mode=real_padding_mode, + end_condition_strength=end_condition_strength, + enable_feature_fusion=enable_feature_fusion, + history_weight=history_weight, + history_decay=history_decay, + history_weight_min=history_weight_min, + gpu_memory_preservation=gpu_memory_preservation, + use_flash_attention=use_flash_attention, + use_sage_attention=use_sage_attention, + overlap_frames=overlap_frames, + blend_mode=blend_mode + ) + + if os.path.exists(video_path): + self.fps = self.get_fps_with_torchvision(video_path) + self.frames = self.extract_frames_as_pil(video_path) + print(f'{video_path}:{self.fps} {len(self.frames)}') + else: + self.frames = [] + self.fps = 0.0 + except Exception as e: + print(f"Error in run: {str(e)}") + traceback.print_exc() + self.frames = [] + self.fps = 0.0 + + return (self.frames, self.fps) + + @torch.no_grad() + def exec(self, input_image, video_path, + end_image=None, + prompt="The girl dances gracefully, with clear movements, full of charm.", + n_prompt="", + seed=31337, + total_second_length=5, + latent_window_size=9, + steps=25, + cfg=1, + gs=32, + rs=0, + gpu_memory_preservation=6, + use_teacache=True, + resolution="480p", + padding_mode="optimized", + end_condition_strength=1.0, + enable_feature_fusion=True, + history_weight=1.0, + history_decay=0.0, + history_weight_min=0.0, + use_flash_attention=False, + use_sage_attention=False, + overlap_frames=33, + blend_mode="linear"): + + total_latent_sections = (total_second_length * 30) / (latent_window_size * 4) + total_latent_sections = int(max(round(total_latent_sections), 1)) + + # 启用Progressive Decay（按需求默认开启且不需要展示给用户） + progressive_decay = True + + try: + # 应用Flash Attention或Sage Attention + if use_flash_attention: + # 尝试导入并应用Flash Attention + try: + from flash_attn import flash_attn_func + print("Flash Attention enabled") + # 在此处应用Flash Attention的设置（如果需要的话） + except ImportError: + print("Flash Attention module not found, please install it first") + + if use_sage_attention: + # 尝试导入并应用Sage Attention + try: + print("Sage Attention enabled") + # 在此处应用Sage Attention的设置（如果需要的话） + except ImportError: + print("Sage Attention module not found or not applicable") + + # Clean GPU + if not self.high_vram: + unload_complete_models( + self.text_encoder, self.text_encoder_2, self.image_encoder, self.vae, self.transformer + ) + + # Text encoding + print('Text encoding') + + if not self.high_vram: + fake_diffusers_current_device(self.text_encoder, gpu) + load_model_as_complete(self.text_encoder_2, target_device=gpu) + + llama_vec, clip_l_pooler = encode_prompt_conds(prompt, self.text_encoder, self.text_encoder_2, self.tokenizer, self.tokenizer_2) + + if cfg == 1: + llama_vec_n, clip_l_pooler_n = torch.zeros_like(llama_vec), torch.zeros_like(clip_l_pooler) + else: + llama_vec_n, clip_l_pooler_n = encode_prompt_conds(n_prompt, self.text_encoder, self.text_encoder_2, self.tokenizer, self.tokenizer_2) + + llama_vec, llama_attention_mask = crop_or_pad_yield_mask(llama_vec, length=512) + llama_vec_n, llama_attention_mask_n = crop_or_pad_yield_mask(llama_vec_n, length=512) + + # Processing input image (start frame) + print('Processing start frame ...') + + H, W, C = input_image.shape + + # 根据所选分辨率确定目标高度 + if resolution == '360p': + target_height = 360 + elif resolution == '540p': + target_height = 540 + elif resolution == '720p': + target_height = 720 + else: # 默认480p + target_height = 480 + + # 保持宽高比计算宽度 + target_ratio = W / H + target_width = int(target_height * target_ratio) + if target_width % 2 != 0: # 确保是偶数 + target_width += 1 + + # 找到最接近的桶分辨率，但保持宽高比 + height, width = find_nearest_bucket(target_height, target_width, resolution=target_height) + print(f"Target resolution: {width}x{height}, Aspect ratio: {width/height:.3f}") + + input_image_np = resize_and_center_crop(input_image, target_width=width, target_height=height) + input_image_pt = torch.from_numpy(input_image_np).float() / 127.5 - 1 + input_image_pt = input_image_pt.permute(2, 0, 1)[None, :, None] + + # Processing end image if provided + has_end_image = end_image is not None + end_image_np = None + end_image_pt = None + + if has_end_image: + print('Processing end frame ...') + end_image_np = resize_and_center_crop(end_image, target_width=width, target_height=height) + end_image_pt = torch.from_numpy(end_image_np).float() / 127.5 - 1 + end_image_pt = end_image_pt.permute(2, 0, 1)[None, :, None] + + # VAE encoding + print('VAE encoding ...') + + if not self.high_vram: + load_model_as_complete(self.vae, target_device=gpu) + + start_latent = vae_encode(input_image_pt, self.vae) + end_latent = None + if has_end_image: + end_latent = vae_encode(end_image_pt, self.vae) + + # CLIP Vision + print('CLIP Vision encoding ...') + + if not self.high_vram: + load_model_as_complete(self.image_encoder, target_device=gpu) + + # Start image encoding + image_encoder_output = hf_clip_vision_encode(input_image_np, self.feature_extractor, self.image_encoder) + image_encoder_last_hidden_state = image_encoder_output.last_hidden_state + + # End image encoding if available + if has_end_image: + end_image_encoder_output = hf_clip_vision_encode(end_image_np, self.feature_extractor, self.image_encoder) + end_image_encoder_last_hidden_state = end_image_encoder_output.last_hidden_state + + # 根据end_condition_strength和enable_feature_fusion参数处理 + if end_condition_strength <= 0.0: + # 不使用结束特征 + pass # 保持image_encoder_last_hidden_state不变 + elif end_condition_strength >= 1.0: + # 完全按照原始方式融合 + if enable_feature_fusion: + image_encoder_last_hidden_state = (image_encoder_last_hidden_state + end_image_encoder_last_hidden_state) / 2 + else: + # 按照强度比例混合 + if enable_feature_fusion: + image_encoder_last_hidden_state = (1.0 - end_condition_strength) * image_encoder_last_hidden_state + end_condition_strength * end_image_encoder_last_hidden_state + + # Dtype + llama_vec = llama_vec.to(self.transformer.dtype) + llama_vec_n = llama_vec_n.to(self.transformer.dtype) + clip_l_pooler = clip_l_pooler.to(self.transformer.dtype) + clip_l_pooler_n = clip_l_pooler_n.to(self.transformer.dtype) + image_encoder_last_hidden_state = image_encoder_last_hidden_state.to(self.transformer.dtype) + + print('Start Sample') + + rnd = torch.Generator("cpu").manual_seed(seed) + num_frames = latent_window_size * 4 - 3 + + history_latents = torch.zeros(size=(1, 16, 1 + 2 + 16, height // 8, width // 8), dtype=torch.float32).cpu() + history_pixels = None + total_generated_latent_frames = 0 + + # 根据padding_mode选择不同的padding策略 + if padding_mode == "default": + # 原始模式: 直接使用倒序数列 + latent_paddings = list(reversed(range(total_latent_sections))) + elif padding_mode == "constant": + # 常量模式: 保持固定值2，最后两块特殊处理 + if total_latent_sections <= 2: + # 对于短视频，使用[1, 0]或[0] + latent_paddings = list(reversed(range(total_latent_sections))) + else: + # 对于长视频，除了最后两块外都使用常量2 + latent_paddings = [2] * (total_latent_sections - 2) + [1, 0] + elif padding_mode == "optimized": + # 优化模式: 对于长视频使用自定义分布 + if total_latent_sections > 4: + latent_paddings = [3] + [2] * (total_latent_sections - 3) + [1, 0] + else: + latent_paddings = list(reversed(range(total_latent_sections))) + elif padding_mode == "waterfall": + # 瀑布模式 + if total_latent_sections <= 2: + latent_paddings = list(reversed(range(total_latent_sections))) + else: + # 最大padding值 + max_padding = 4 + # 头部增长阶段的长度 + head_size = min(4, total_latent_sections // 3) + # 尾部（值为1和0）的长度固定为2 + tail_size = 2 + # 中间部分长度 + middle_size = total_latent_sections - head_size - tail_size + + # 创建头部：线性增加到最大值 + head_paddings = [] + for i in range(head_size): + # 线性插值从3到max_padding + ratio = i / (head_size - 1) if head_size > 1 else 1 + padding = 3 + (max_padding - 3) * ratio + head_paddings.append(round(padding)) + + # 中间部分全部为2 + middle_paddings = [2] * middle_size + + # 尾部固定为[1, 0] + tail_paddings = [1, 0] + + # 组合所有部分 + latent_paddings = head_paddings + middle_paddings + tail_paddings + elif padding_mode == "center_focus": + # 中心聚焦模式 + if total_latent_sections <= 3: + latent_paddings = list(reversed(range(total_latent_sections))) + else: + # 更平滑的padding过渡 + max_padding = 3 + min_padding = 1 + + if total_latent_sections % 2 == 0: # 偶数个sections + half_size = total_latent_sections // 2 + # 创建半边，然后镜像 + first_half = [] + for i in range(half_size): + position_ratio = i / (half_size - 1) if half_size > 1 else 0 + padding = min_padding + (max_padding - min_padding) * position_ratio ** 2 + first_half.append(round(padding)) + # 镜像创建 + second_half = first_half.copy() + second_half.reverse() + latent_paddings = first_half + second_half + else: # 奇数个sections + mid_idx = total_latent_sections // 2 + latent_paddings = [] + for i in range(total_latent_sections): + distance_ratio = abs(i - mid_idx) / mid_idx if mid_idx > 0 else 0 + padding = max_padding - (max_padding - min_padding) * distance_ratio ** 2 + latent_paddings.append(round(padding)) + + # 确保第一个元素不要太大 + latent_paddings[0] = min(latent_paddings[0], 2) + + # 确保最后几个section的padding有平滑过渡 + if len(latent_paddings) >= 4: + latent_paddings[-4] = min(latent_paddings[-4], 3) + latent_paddings[-3] = 2 + latent_paddings[-2] = 1 + latent_paddings[-1] = 0 + elif len(latent_paddings) == 3: + latent_paddings[0] = 2 + latent_paddings[1] = 1 + latent_paddings[2] = 0 + else: + # 默认回退到优化模式 + if total_latent_sections > 4: + latent_paddings = [3] + [2] * (total_latent_sections - 3) + [1, 0] + else: + latent_paddings = list(reversed(range(total_latent_sections))) + + # 确保最后一个padding值为0 + if len(latent_paddings) > 0: + latent_paddings[-1] = 0 + + print(f"Selected padding mode: {padding_mode}") + print(f"Total sections: {total_latent_sections}") + print(f"Latent paddings: {latent_paddings}") + + # 计算每个section的history weight，实现衰减效果 + section_weights = [] + current_weight = history_weight + + # 找出最大padding值，用于归一化padding影响 + max_padding = max(latent_paddings) if latent_paddings else 3 + if max_padding == 0: # 防止除以零错误 + max_padding = 1 + + for i, padding in enumerate(latent_paddings): + section_weights.append(current_weight) + + # 根据当前section的padding值调整衰减率 + padding_factor = padding / max_padding # 归一化到0-1范围 + + # 调整衰减率：padding_factor越大，实际衰减越小 + adjusted_decay = history_decay * (1.0 - padding_factor * 0.8) # 0.8是补偿系数 + + # 应用调整后的衰减率，并确保不低于最小值 + current_weight = max(history_weight_min, current_weight * (1.0 - adjusted_decay)) + + print(f"Section {i+1} padding: {padding}, padding_factor: {padding_factor:.2f}, " + + f"adjusted_decay: {adjusted_decay:.4f}, weight: {current_weight:.4f}") + + print(f"Section history weights: {section_weights}") + + for i, latent_padding in enumerate(latent_paddings): + # 获取当前section的history weight + current_history_weight = section_weights[i] + print(f"Section {i+1} history weight: {current_history_weight:.4f}") + + is_last_section = latent_padding == 0 + is_first_section = latent_padding == latent_paddings[0] + latent_padding_size = latent_padding * latent_window_size + + print(f'latent_padding_size = {latent_padding_size}, is_last_section = {is_last_section}, is_first_section = {is_first_section}') + + indices = torch.arange(0, sum([1, latent_padding_size, latent_window_size, 1, 2, 16])).unsqueeze(0) + clean_latent_indices_pre, blank_indices, latent_indices, clean_latent_indices_post, clean_latent_2x_indices, clean_latent_4x_indices = indices.split([1, latent_padding_size, latent_window_size, 1, 2, 16], dim=1) + clean_latent_indices = torch.cat([clean_latent_indices_pre, clean_latent_indices_post], dim=1) + + # 确保所有latent都是5维的 [1, 16, T, h, w] + if len(start_latent.shape) == 4: # [1, 16, h, w] + clean_latents_pre = start_latent.unsqueeze(2).to(history_latents) # 变为 [1, 16, 1, h, w] + else: + clean_latents_pre = start_latent.to(history_latents) # 已经是5维 + + # 从history_latents中提取 + clean_latents_post, clean_latents_2x, clean_latents_4x = history_latents[:, :, :1 + 2 + 16, :, :].split([1, 2, 16], dim=2) + + # 应用当前section的history_weight，调整历史潜变量的影响力 + if current_history_weight < 1.0: + # 如果启用渐进式衰减，对latent_indices涉及到的帧应用不同的权重 + if progressive_decay and latent_window_size > 1 and not is_first_section: + # 只对clean_latents_post进行操作，它是历史信息 + weighted_clean_latents_post = clean_latents_post.clone() + + # 获取下一个section的权重 + next_section_idx = min(i + 1, len(section_weights) - 1) + next_weight = section_weights[next_section_idx] if next_section_idx < len(section_weights) else max(history_weight_min, current_history_weight * 0.5) + + # 在前一个section的后部分应用过渡衰减 + # 这创建了一个从当前权重到下一个section权重的平滑过渡 + middle_weight = (current_history_weight + next_weight) / 2 + weighted_clean_latents_post = weighted_clean_latents_post * middle_weight + + print(f"Progressive decay from {current_history_weight:.4f} to {middle_weight:.4f} for section {i+1}") + else: + weighted_clean_latents_post = clean_latents_post * current_history_weight + + # 默认合并clean_latents + clean_latents = torch.cat([clean_latents_pre, weighted_clean_latents_post], dim=2) + else: + # 默认合并clean_latents + clean_latents = torch.cat([clean_latents_pre, clean_latents_post], dim=2) + + # 使用结束帧latent + if has_end_image and is_first_section: + # 确保end_latent也是5维的 + if len(end_latent.shape) == 4: # [1, 16, h, w] + end_latent_to_use = end_latent.unsqueeze(2) # 变为 [1, 16, 1, h, w] + else: + end_latent_to_use = end_latent + + # 根据强度参数决定如何应用end_latent + if end_condition_strength <= 0.0: + # 强度为0，不使用end_latent，保持clean_latents不变 + print(f"End condition strength is 0, not using end reference") + elif end_condition_strength >= 1.0: + # 强度为1，完全使用end_latent + print(f"End condition strength is 1, fully using end reference") + clean_latents = torch.cat([clean_latents_pre, end_latent_to_use.to(history_latents)], dim=2) + else: + # 混合start_latent和end_latent + print(f"Mixing start and end with strength {end_condition_strength}") + # 进行插值混合 + mixed_latent = (1.0 - end_condition_strength) * clean_latents_post + end_condition_strength * end_latent_to_use.to(history_latents) + clean_latents = torch.cat([clean_latents_pre, mixed_latent], dim=2) + + if not self.high_vram: + unload_complete_models() + move_model_to_device_with_memory_preservation(self.transformer, target_device=gpu, preserved_memory_gb=gpu_memory_preservation) + + if use_teacache: + self.transformer.initialize_teacache(enable_teacache=True, num_steps=steps) + else: + self.transformer.initialize_teacache(enable_teacache=False) + + def callback(d): + self.update(1) + return + + generated_latents = sample_hunyuan( + transformer=self.transformer, + sampler='unipc', + width=width, + height=height, + frames=num_frames, + real_guidance_scale=cfg, + distilled_guidance_scale=gs, + guidance_rescale=rs, + num_inference_steps=steps, + generator=rnd, + prompt_embeds=llama_vec, + prompt_embeds_mask=llama_attention_mask, + prompt_poolers=clip_l_pooler, + negative_prompt_embeds=llama_vec_n, + negative_prompt_embeds_mask=llama_attention_mask_n, + negative_prompt_poolers=clip_l_pooler_n, + device=gpu, + dtype=torch.bfloat16, + image_embeddings=image_encoder_last_hidden_state, + latent_indices=latent_indices, + clean_latents=clean_latents, + clean_latent_indices=clean_latent_indices, + clean_latents_2x=clean_latents_2x, + clean_latent_2x_indices=clean_latent_2x_indices, + clean_latents_4x=clean_latents_4x, + clean_latent_4x_indices=clean_latent_4x_indices, + callback=callback, + ) + + if is_last_section: + # 确保start_latent是5维的，与generated_latents维度匹配 + if len(start_latent.shape) == 4: # [1, 16, h, w] + start_latent_5d = start_latent.unsqueeze(2) # 变为 [1, 16, 1, h, w] + else: + start_latent_5d = start_latent + generated_latents = torch.cat([start_latent_5d.to(generated_latents), generated_latents], dim=2) + + total_generated_latent_frames += int(generated_latents.shape[2]) + history_latents = torch.cat([generated_latents.to(history_latents), history_latents], dim=2) + + if not self.high_vram: + offload_model_from_device_for_memory_preservation(self.transformer, target_device=gpu, preserved_memory_gb=8) + load_model_as_complete(self.vae, target_device=gpu) + + real_history_latents = history_latents[:, :, :total_generated_latent_frames, :, :] + + if history_pixels is None: + history_pixels = vae_decode(real_history_latents, self.vae).cpu() + else: + section_latent_frames = (latent_window_size * 2 + 1) if is_last_section else (latent_window_size * 2) + # 使用用户指定的overlap_frames替代固定值 + overlapped_frames = overlap_frames + + current_pixels = vae_decode(real_history_latents[:, :, :section_latent_frames], self.vae).cpu() + history_pixels = soft_append_bcthw(current_pixels, history_pixels, overlapped_frames, blend_mode) + + if not self.high_vram: + unload_complete_models() + + # 如果是最后一个section，保存视频 + if is_last_section: + save_bcthw_as_mp4(history_pixels, video_path, fps=30) + break + + except Exception as e: + print(f"Error in exec: {str(e)}") + traceback.print_exc() + finally: + unload_complete_models() + + def update(self, in_progress): + self.pbar.update(in_progress) + + def extract_frames_as_pil(self, video_path): + video, _, _ = torchvision.io.read_video(video_path, pts_unit='sec') # (T, H, W, C) + frames = [to_pil_image(frame.permute(2, 0, 1)) for frame in video] + frames = [torch.from_numpy(np.array(frame).astype(np.float32) / 255.0) for frame in frames] + return frames + + def get_fps_with_torchvision(self, video_path): + _, _, info = torchvision.io.read_video(video_path, pts_unit='sec') + return info['video_fps'] + +NODE_CLASS_MAPPINGS = { + "TTPlanet_FramePack": Kiki_FramePack, +} + +def soft_append_bcthw(history, current, overlap=0, blend_mode='linear'): + if overlap <= 0: + return torch.cat([history, current], dim=2) + + assert history.shape[2] >= overlap, f"History length ({history.shape[2]}) must be >= overlap ({overlap})" + assert current.shape[2] >= overlap, f"Current length ({current.shape[2]}) must be >= overlap ({overlap})" + + # 根据不同的混合模式选择权重曲线 + if blend_mode == 'cosine': + # 使用余弦曲线，提供更平滑的过渡 + x = torch.linspace(0, math.pi, overlap, dtype=history.dtype, device=history.device) + weights = (torch.cos(x) + 1) / 2 + elif blend_mode == 'sigmoid': + # 使用sigmoid函数，在中间区域提供更平缓的过渡 + x = torch.linspace(-6, 6, overlap, dtype=history.dtype, device=history.device) + weights = torch.sigmoid(x) + else: # 'linear' + # 默认线性 + weights = torch.linspace(1, 0, overlap, dtype=history.dtype, device=history.device) + + weights = weights.view(1, 1, -1, 1, 1) + blended = weights * history[:, :, -overlap:] + (1 - weights) * current[:, :, :overlap] + output = torch.cat([history[:, :, :-overlap], blended, current[:, :, overlap:]], dim=2) + + return output.to(history) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..a695c9e --- /dev/null +++ b/requirements.txt @@ -0,0 +1,7 @@ +accelerate>=1.6.0 +diffusers>=0.33.1 +transformers>=4.46.2 +scipy>=1.12.0 +torchsde>=0.2.6 +einops +safetensors