Setup the project

__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']

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
+

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

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

diffusers_helper/gradio/progress_bar.py

@@ -0,0 +1,86 @@
+progress_html = '''
+<div class="loader-container">
+  <div class="loader"></div>
+  <div class="progress-container">
+    <progress value="*number*" max="100"></progress>
+  </div>
+  <span>*text*</span>
+</div>
+'''
+
+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

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)

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

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)

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

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

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

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()

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)

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)

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