nodes.py

import nodes
import torch
import comfy.model_management as mm
from PIL import ImageDraw, Image, ImageOps, ImageChops, ImageColor
import numpy as np
from typing import List, Set, Union, Optional, Dict, Optional, Tuple
import struct
import folder_paths
from pathlib import Path
import yaml

from comfy import utils as utils
from comfy import samplers as samplers
from comfy import sampler_helpers as sampler_helpers
from comfy import patcher_extension as patcher_extension
from comfy import model_patcher as comfy_model_patcher
from comfy import hooks as hooks
from comfy import model_sampling as model_sampling

from comfy.samplers import preprocess_conds_hooks as preprocess_conds_hooks
from comfy.samplers import get_total_hook_groups_in_conds as get_total_hook_groups_in_conds
from comfy.samplers import filter_registered_hooks_on_conds as filter_registered_hooks_on_conds
from comfy.samplers import sampling_function as sampling_function
from comfy.samplers import cast_to_load_options as cast_to_load_options
from comfy.samplers import process_conds as process_conds

from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
from enum import Enum

import os
import gc
import time
import json
import latent_preview
from decimal import *
import traceback
import sys
import cv2
import glob
import comfy
from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel
from comfy.cli_args import args
from PIL.PngImagePlugin import PngInfo
import random
from tqdm import tqdm
import io

import node_helpers

from torch import Tensor
from einops import repeat
from unittest.mock import patch

from comfy.ldm.flux.layers import timestep_embedding, apply_mod
from comfy.ldm.lightricks.model import precompute_freqs_cis
from comfy.ldm.lightricks.symmetric_patchifier import latent_to_pixel_coords
from comfy.ldm.common_dit import rms_norm
from comfy.ldm.wan.model import sinusoidal_embedding_1d

from .diffusers_helper.memory import move_model_to_device_with_memory_preservation, offload_model_from_device_for_memory_preservation
from .diffusers_helper.k_diffusion_hunyuan import sample_hunyuan, sample_hunyuan2
from .diffusers_helper.utils import crop_or_pad_yield_mask, soft_append_bcthw
import math

import threading
import copy

import safetensors
import einops
import socket
import re

import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

ALLOWED_EXT = ('.jpeg', '.jpg', '.png', '.tiff', '.gif', '.bmp', '.webp', '.avif', '.vfif')
ALLOWED_VIDEO_EXT = ('mp4', 'flv', 'mov', 'avi', 'mpg', 'webm', 'mkv')
ALLOWED_CAPTION_EXT = ('.txt')

vae_scaling_factor = 0.476986

class VideoGenerationType(Enum):
    T2V = "t2v"
    I2V = "i2v"
    V2V = "v2v"

def get_offload_device():
    return torch.device("cpu")

def tensor2pil(image):
    # _tensor_check_image(image)
    return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(0), 0, 255).astype(np.uint8))

def pil2tensor(image, device=None):
    if device is None:
        return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)

    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).to(device).unsqueeze(0)

def pil2tensorSwap(image: Union[Image.Image, List[Image.Image]], device=None) -> torch.Tensor:
    if isinstance(image, list):
        if device is None:
            return torch.cat([pil2tensor(img) for img in image], dim=0)

        new_tensor = torch.cat([pil2tensor(img, device=device) for img in image], dim=0)

        return new_tensor

    if device is None:
        return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)

    new_tensor = torch.from_numpy(np.array(image).astype(np.float32) / 255.0).to(device).unsqueeze(0)

    return new_tensor

def tensor2pilSwap(image: torch.Tensor) -> List[Image.Image]:
    batch_count = image.size(0) if len(image.shape) > 3 else 1
    if batch_count > 1:
        out = []
        for i in range(batch_count):
            out.extend(tensor2pil(image[i]))
        return out

    return [
        Image.fromarray(
            np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
        )
    ]

def get_lora_list():
    templist = folder_paths.get_filename_list("loras")

    i = 0
    while i < len(templist):
        templist[i] = templist[i].lower()
        i += 1

    return templist

def get_default_gen_output_path():
    default_path = "{}\\MyLoras\\new_lora_hy.safetensors".format(folder_paths.output_directory)
    return default_path

def get_default_pt_st_output_folder():
    default_path = "{}\\pt_conversions".format(folder_paths.output_directory)
    return default_path

def get_default_output_path():
    default_path = "{}\\MergedHunyuanLoras\\new_lora_hy.safetensors".format(folder_paths.output_directory)
    return default_path

def get_default_output_folder():
    default_folder = "{}\\MergedHunyuanLoras".format(folder_paths.output_directory)
    return default_folder

def get_default_xl_output_folder():
    default_folder = "{}\\MergedXLLoras".format(folder_paths.output_directory)
    return default_folder

def get_default_wan_output_path():
    default_path = "{}\\MergedWanLoras\\new_lora_wan.safetensors".format(folder_paths.output_directory)
    return default_path

def get_default_wan_output_folder():
    default_folder = "{}\\MergedWanLoras".format(folder_paths.output_directory)
    return default_folder

SUPPORTED_MODELS_COEFFICIENTS = {
    "hunyuan_video": [7.33226126e+02, -4.01131952e+02, 6.75869174e+01, -3.14987800e+00, 9.61237896e-02],
    "wan2.1_t2v_1.3B": [2.39676752e+03, -1.31110545e+03, 2.01331979e+02, -8.29855975e+00, 1.37887774e-01],
    "wan2.1_t2v_14B": [-5784.54975374, 5449.50911966, -1811.16591783, 256.27178429, -13.02252404],
    "wan2.1_i2v_480p_14B": [-3.02331670e+02, 2.23948934e+02, -5.25463970e+01, 5.87348440e+00, -2.01973289e-01],
    "wan2.1_i2v_720p_14B": [-114.36346466, 65.26524496, -18.82220707, 4.91518089, -0.23412683]
}

INITIAL_COEFFICIENTS = [1.0, 1.0, 1.0, 1.0, 1.0]

def poly1d(coefficients, x):
    result = torch.zeros_like(x)
    for i, coeff in enumerate(coefficients):
        result += coeff * (x ** (len(coefficients) - 1 - i))
    return result

def load_upscale_model(model_name):
    model_path = folder_paths.get_full_path_or_raise("upscale_models", model_name)
    print("Upscale Model Path for {}: {}".format(model_name, model_path))
    sd = utils.load_torch_file(model_path, safe_load=True)
    if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd:
        sd = utils.state_dict_prefix_replace(sd, {"module.":""})
    out = ModelLoader().load_from_state_dict(sd).eval()

    if not isinstance(out, ImageModelDescriptor):
        raise Exception("Upscale model must be a single-image model.")

    return out

def upscale_with_model(upscale_model, images, upscale_by, rescale_method):
    __imageScaler = ImageUpscaleWithModel()
    upscaled_images = []

    samples = images.movedim(-1,1)

    width = round(samples.shape[3])
    height = round(samples.shape[2])

    target_width = round(samples.shape[3] * upscale_by)
    target_height = round(samples.shape[2] * upscale_by)

    samples = __imageScaler.upscale(upscale_model, images)[0].movedim(-1,1)

    upscaled_width = round(samples.shape[3])
    upscaled_height = round(samples.shape[2])

    if upscaled_width > target_width or upscaled_height > target_height:
        samples = utils.common_upscale(samples, target_width, target_height, rescale_method, "disabled")

    samples = samples.movedim(1,-1)

    return samples

def upscale_latents_by(latent, upscale_method, scale_by):
    s = latent.clone()
    width = round(latent.shape[-1] * scale_by)
    height = round(latent.shape[-2] * scale_by)

    s = utils.common_upscale(latent, width, height, upscale_method, "disabled")
    return s

def partial_encode_basic(vae, image):
    latents = vae.encode(image[:,:,:,:3])
    return latents

def partial_encode_tiled(vae, image, tile_size=256, overlap=64, temporal_size=64, temporal_overlap=8, use_type=torch.float32, unload_after=True):
    latents = vae.encode_tiled(image[:,:,:,:3], tile_x=tile_size, tile_y=tile_size, overlap=overlap, tile_t=temporal_size, overlap_t=temporal_overlap)
    latents = latents.to(dtype=use_type, device=get_offload_device())

    if unload_after:
        mm.unload_all_models()
        mm.soft_empty_cache()
        gc.collect()
        time.sleep(1)

    return latents

def partial_decode_basic(vae, latents):
    images = vae.decode(latents)
    if len(images.shape) == 5: #Combine batches
        images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1])
    return images

def partial_decode_tiled(vae, latents, tile_size, overlap=64, temporal_size=64, temporal_overlap=8, unload_after=True):
    if tile_size < overlap * 4:
        overlap = tile_size // 4
    if temporal_size < temporal_overlap * 2:
        temporal_overlap = temporal_overlap // 2
    temporal_compression = vae.temporal_compression_decode()
    if temporal_compression is not None:
        temporal_size = max(2, temporal_size // temporal_compression)
        temporal_overlap = max(1, min(temporal_size // 2, temporal_overlap // temporal_compression))
    else:
        temporal_size = None
        temporal_overlap = None

    compression = vae.spacial_compression_decode()
    images = vae.decode_tiled(latents, tile_x=tile_size // compression, tile_y=tile_size // compression, overlap=overlap // compression, tile_t=temporal_size, overlap_t=temporal_overlap)
    if len(images.shape) == 5: #Combine batches
        images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1])

    if unload_after:
        mm.unload_all_models()
        mm.soft_empty_cache()
        gc.collect()
        time.sleep(1)

    return images

def convert_key_format(key: str) -> str:
    """Standardize LoRA key format by removing prefixes."""
    prefixes = ["diffusion_model.", "transformer."]
    for prefix in prefixes:
        if key.startswith(prefix):
            key = key[len(prefix):]
            break
    return key

def filter_lora_keys(lora: Dict[str, torch.Tensor], blocks_type: str) -> Dict[str, torch.Tensor]:
    """Filter LoRA weights based on block type."""
    if blocks_type == "all":
        return lora

    filtered_lora = {}

    for key, value in lora.items():
        base_key = convert_key_format(key)
        if blocks_type in base_key:
            filtered_lora[key] = value
    return filtered_lora

def convert_to_musubi(lora: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    """Checks for and converts to Musubi Tuner format which supports Network Alpha and uses different naming."""
    prefix = "lora_unet_"
    musubi = False
    lora_alphas = {}
    for key, value in lora.items():
        if key.startswith(prefix):
            lora_name = key.split(".", 1)[0]  # before first dot

            if lora_name not in lora_alphas and "alpha" in key:
                lora_alphas[lora_name] = value
                musubi = True

    if musubi:
        print("Lora is already in musubi format. Nothing to convert.")
        return lora

    print("Converting lora to musubi format.")

    temp_lora = convert_lora(lora, do_check_for_musuibi=False)

    for key in temp_lora.keys():
        print("Intermediate Lora Key: {}".format(key))

    print("\n")

    converted_lora = {}
    prefix = "diffusion_model."
    has_double_blocks = False
    has_single_blocks = False

    for key, weight in temp_lora.items():
        print("Converting Key: {}".format(key))

        if "linear" in key:
            print("Skipping Key.")
            continue

        if ("double.blocks" in key) or ("double_blocks" in key):
            has_double_blocks = True

        elif ("single.blocks" in key) or ("single_blocks" in key):
            has_single_blocks = True

        # lora_name = key.split("_", 1)[0]  # before first dot

        # HunyuanVideo lora name to module name: ugly but works
        # module_name = lora_name[len(prefix) :]  # remove "diffusion_model."

        module_name = key.replace(prefix, "")

        print("module_name: {}".format(module_name))

        module_name = module_name.replace(".", "_")  # replace "_" with "."
        module_name = module_name.replace("double.blocks.", "double_blocks_")  # fix double blocks
        module_name = module_name.replace("single.blocks.", "single_blocks_")  # fix single blocks
        module_name = module_name.replace("img.", "img_")  # fix img
        module_name = module_name.replace("txt.", "txt_")  # fix txt
        module_name = module_name.replace("attn.", "attn_")  # fix attn

        musubi_prefix = "lora_unet"

        if "lora_A" in key:
            new_key = f"{musubi_prefix}_{module_name.replace('_lora_A_weight', '.lora_down.weight')}"
        elif "lora_B" in key:
            new_key = f"{musubi_prefix}_{module_name.replace('_lora_B_weight', '.lora_up.weight')}"
        else:
            print(f"unexpected key: {key} in diffusion_model LoRA format")
            continue

        print("New Key: {}".format(new_key))

        converted_lora[new_key] = weight.to(dtype=torch.bfloat16)

    if has_double_blocks:
        print("Lora Has Double Blocks.")

        for block_num in range(20):
            converted_lora[f"lora_unet_double_blocks_{block_num}_img_attn_proj.alpha"] = torch.empty([]).to(dtype=torch.bfloat16)
            converted_lora[f"lora_unet_double_blocks_{block_num}_img_attn_qkv.alpha"] = torch.empty([]).to(dtype=torch.bfloat16)
            converted_lora[f"lora_unet_double_blocks_{block_num}_img_mlp_fc1.alpha"] = torch.empty([]).to(dtype=torch.bfloat16)
            converted_lora[f"lora_unet_double_blocks_{block_num}_img_mlp_fc2.alpha"] = torch.empty([]).to(dtype=torch.bfloat16)
            converted_lora[f"lora_unet_double_blocks_{block_num}_txt_attn_proj.alpha"] = torch.empty([]).to(dtype=torch.bfloat16)
            converted_lora[f"lora_unet_double_blocks_{block_num}_txt_attn_qkv.alpha"] = torch.empty([]).to(dtype=torch.bfloat16)
            converted_lora[f"lora_unet_double_blocks_{block_num}_txt_mlp_fc1.alpha"] = torch.empty([]).to(dtype=torch.bfloat16)
            converted_lora[f"lora_unet_double_blocks_{block_num}_txt_mlp_fc2.alpha"] = torch.empty([]).to(dtype=torch.bfloat16)

    if has_single_blocks:
        print("Lora Has Single Blocks.")

        for block_num in range(40):
            converted_lora[f"lora_unet_single_blocks_{block_num}_linear1.alpha"] = torch.Tensor([]).to(dtype=torch.bfloat16)
            converted_lora[f"lora_unet_single_blocks_{block_num}_linear2.alpha"] = torch.Tensor([]).to(dtype=torch.bfloat16)

    return converted_lora

def check_for_musubi(lora: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    """Checks for and converts from Musubi Tuner format which supports Network Alpha and uses different naming. Largely copied from that project"""
    prefix = "lora_unet_"
    musubi = False
    lora_alphas = {}
    for key, value in lora.items():
        if key.startswith(prefix):
            lora_name = key.split(".", 1)[0]  # before first dot

            if lora_name not in lora_alphas and "alpha" in key:
                lora_alphas[lora_name] = value
                musubi = True
    if musubi:
        print("Loading Musubi Tuner format LoRA...")

        converted_lora = {}

        for key, weight in lora.items():
            if key.startswith(prefix):
                if "alpha" in key:
                    continue

            lora_name = key.split(".", 1)[0]  # before first dot

            # HunyuanVideo lora name to module name: ugly but works
            module_name = lora_name[len(prefix) :]  # remove "lora_unet_"
            module_name = module_name.replace("_", ".")  # replace "_" with "."
            module_name = module_name.replace("double.blocks.", "double_blocks.")  # fix double blocks
            module_name = module_name.replace("single.blocks.", "single_blocks.")  # fix single blocks
            module_name = module_name.replace("img.", "img_")  # fix img
            module_name = module_name.replace("txt.", "txt_")  # fix txt
            module_name = module_name.replace("attn.", "attn_")  # fix attn

            diffusers_prefix = "diffusion_model"

            if "lora_down" in key:
                new_key = f"{diffusers_prefix}.{module_name}.lora_A.weight"
                dim = weight.shape[0]
            elif "lora_up" in key:
                new_key = f"{diffusers_prefix}.{module_name}.lora_B.weight"
                dim = weight.shape[1]
            else:
                print(f"unexpected key: {key} in Musubi LoRA format")
                continue
            # scale weight by alpha
            if lora_name in lora_alphas:
                # we scale both down and up, so scale is sqrt
                scale = lora_alphas[lora_name] / dim
                scale = scale.sqrt()
                weight = weight * scale
            else:
                print(f"missing alpha for {lora_name}")

            converted_lora[new_key] = weight

        return converted_lora
    else:
        print("Loading Diffusers format LoRA...")

        return lora

def convert_lora(lora, convert_to="diffusion_model", do_check_for_musuibi=True):
    if do_check_for_musuibi:
        temp_lora = check_for_musubi(lora)
    else:
        temp_lora = lora

    new_lora = {}

    for key in temp_lora.keys():
        if convert_to in key:
            new_lora = temp_lora
            break

        if key.startswith("transformer.") and (convert_to == "diffusion_model"):
            new_key = key.replace("transformer.", "diffusion_model.")
            new_lora[new_key] = temp_lora[key]
            continue

        if key.startswith("diffusion_model.") and  (convert_to == "transformer"):
            new_key = key.replace("diffusion_model.", "transformer.")
            new_lora[new_key] = temp_lora[key]
            continue

        if key.startswith("lora_unet_"):
            new_key = key.replace("lora_unet_", "{}.".format(convert_to))

            if "double" in new_key:
                new_key = new_key.replace("double_blocks_", "double_blocks.")
                new_key = new_key.replace("_img_attn", ".img_attn")
                new_key = new_key.replace("_img_mlp", ".img_mlp")
                new_key = new_key.replace("_txt_attn", ".txt_attn")
                new_key = new_key.replace("_txt_mlp", ".txt_mlp")
                new_key = new_key.replace(".lora_up.", ".lora_A.")
                new_key = new_key.replace(".lora_down.", ".lora_B.")

                continue

            if "single" in new_key:
                new_key = new_key.replace("single_blocks_", "single_blocks.")
                new_key = new_key.replace("_linear", ".linear")
                new_key = new_key.replace(".lora_up.", ".lora_A.")
                new_key = new_key.replace(".lora_down.", ".lora_B.")

                continue

    return new_lora

class HyVideoModel(comfy.model_base.BaseModel):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.pipeline = {}
        self.load_device = mm.get_torch_device()

    def __getitem__(self, k):
        return self.pipeline[k]

    def __setitem__(self, k, v):
        self.pipeline[k] = v

class HyVideoModelConfig:
    def __init__(self, dtype):
        self.unet_config = {}
        self.unet_extra_config = {}
        self.latent_format = comfy.latent_formats.HunyuanVideo
        self.latent_format.latent_channels = 16
        self.manual_cast_dtype = dtype
        self.sampling_settings = {"multiplier": 1.0}
        self.memory_usage_factor = 2.0
        self.unet_config["disable_unet_model_creation"] = True

def warped_prepare_noise(latent_image, seed, generator=None):
    """
    creates random noise given a latent image and a seed.
    optional arg skip can be used to skip and discard x number of noise generations for a given seed
    """
    if generator is None:
        generator = torch.manual_seed(seed)

    output_result = torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu")

    return output_result

def warped_prepare_noise_images(images, seed, generator=None):
    """
    creates random noise given a latent image and a seed.
    optional arg skip can be used to skip and discard x number of noise generations for a given seed
    """
    if generator is None:
        generator = torch.manual_seed(seed)

    output_result = torch.randn(images.size(), dtype=images.dtype, layout=images.layout, generator=generator, device="cpu")

    return output_result

def get_upscale_methods():
    return ["nearest-exact", "bilinear", "area", "bicubic", "bislerp"]

def get_rescale_methods():
    return ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]

def convert_lora_dimensions(max_dimension, lora):
    new_lora = {}

    for key in lora.keys():
        temp_weights = lora[key]

        # if (temp_weights.shape[0] == max_dimension) or (temp_weights.shape[1] == max_dimension):
        #     return lora

        if temp_weights.shape[0] < temp_weights.shape[1]:
            padding = torch.zeros([max_dimension, temp_weights.shape[1]])

            if temp_weights.shape[0] <= max_dimension:
                padding[:temp_weights.shape[0],:] = temp_weights
                new_lora[key] = padding
            else:
                padding[:max_dimension,:] = temp_weights[:max_dimension,:]
                new_lora[key] = padding
        else:
            padding = torch.zeros([temp_weights.shape[0], max_dimension])

            if temp_weights.shape[1] <= max_dimension:
                padding[:,:temp_weights.shape[1]] = temp_weights
                new_lora[key] = padding
            else:
                padding[:,:max_dimension] = temp_weights[:,:max_dimension]
                new_lora[key] = padding
    lora = None

    return new_lora

class WarpedHunyuanLoraMerge:
    def __init__(self):
        self.base_output_dir = get_default_output_folder()
        os.makedirs(self.base_output_dir, exist_ok = True)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "save_path": ("STRING", {"default": get_default_output_path()}),
                "lora_1": (['None'] + get_lora_list(),),
                "strength_1": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}),
                "blocks_type_1": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "lora_2": (['None'] + get_lora_list(),),
                "strength_2": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}),
                "blocks_type_2": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "save_metadata": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    OUTPUT_IS_LIST = (True,)
    FUNCTION = "merge_multiple_loras"
    CATEGORY = "Warped/Hunyuan/Merge"
    DESCRIPTION = "Load and apply multiple LoRA models with different strengths and block types. Model input is required."

    def load_lora(self, lora_name: str, strength: float, blocks_type: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
        """Load and filter a single LoRA model."""
        if not lora_name or strength == 0:
            return {}, {}

        # Get the full path to the LoRA file
        lora_path = folder_paths.get_full_path("loras", lora_name)
        if not os.path.exists(lora_path):
            raise ValueError(f"LoRA file not found: {lora_path}")

        # Load the LoRA weights
        lora_weights = utils.load_torch_file(lora_path)

        return lora_weights

    def merge_multiple_loras(self, save_path, lora_1, strength_1, blocks_type_1, lora_2, strength_2, blocks_type_2, save_metadata=True):
        """Load and apply multiple LoRA models."""
        temp_loras = {}
        metadata = {"loras": "{} and {}".format(lora_1, lora_2)}
        metadata["strengths"] = "{} and {}".format(strength_1, strength_2)
        metadata["block_types"] = "{} and {}".format(blocks_type_1, blocks_type_2)

        if lora_1 != "None" and strength_1 != 0:
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_1, 1.0, "all")
            lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
            filtered_lora = filter_lora_keys(lora_weights, blocks_type_1)
            temp_loras["1"] = {"lora_weights": lora_weights, "strength": strength_1, "filtered_lora": filtered_lora}

        if lora_2 != "None" and strength_2 != 0:
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_2, 1.0, "all")
            lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
            filtered_lora = filter_lora_keys(lora_weights, blocks_type_2)
            temp_loras["2"] = {"lora_weights": lora_weights, "strength": strength_2, "filtered_lora": filtered_lora}

        loras = {}

        for lora_key in temp_loras.keys():
            loras[lora_key] = {"lora_weights": {}, "strength": temp_loras[lora_key]["strength"], "filtered_lora": temp_loras[lora_key]["filtered_lora"]}

            for key in temp_loras[lora_key]["lora_weights"].keys():
                new_key = key.replace("transformer.", "diffusion_model.")
                loras[lora_key]["lora_weights"][new_key] = temp_loras[lora_key]["lora_weights"][key]

        new_lora = {}

        for lora_key in loras.keys():
            for key in loras[lora_key]["lora_weights"].keys():
                if not key in new_lora.keys():
                    new_lora[key] = None
                print("Lora: {}  | Key: {}  |  Shape: {}".format(lora_key, key, loras[lora_key]["lora_weights"][key].shape))

        # Merge The Weighted Key Weights
        for key in new_lora.keys():
            for lora_key in loras.keys():
                if key in loras[lora_key]["lora_weights"].keys():
                    if not new_lora[key] is None:
                        temp_weights = torch.mul(loras[lora_key]["lora_weights"][key], loras[lora_key]["strength"])

                        if new_lora[key].shape[0] < new_lora[key].shape[1]:
                            if temp_weights.shape[0] < new_lora[key].shape[0]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([new_lora[key].shape[0], new_lora[key].shape[1]])
                                padding[:temp_weights.shape[0],:] = temp_weights
                                temp_weights = padding
                            elif temp_weights.shape[0] > new_lora[key].shape[0]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([temp_weights.shape[0], temp_weights.shape[1]])
                                padding[:new_lora[key].shape[0],:] = new_lora[key]
                                new_lora[key] = padding
                        else:
                            if temp_weights.shape[1] < new_lora[key].shape[1]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([new_lora[key].shape[0], new_lora[key].shape[1]])
                                padding[:,:temp_weights.shape[1]] = temp_weights
                                temp_weights = padding
                            elif temp_weights.shape[1] > new_lora[key].shape[1]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([temp_weights.shape[0], temp_weights.shape[1]])
                                padding[:,:new_lora[key].shape[1]] = new_lora[key]
                                new_lora[key] = padding

                        try:
                            new_lora[key] = torch.add(new_lora[key], temp_weights)
                        except Exception as e:
                            raise(e)
                    else:
                        new_lora[key] = torch.mul(loras[lora_key]["lora_weights"][key], loras[lora_key]["strength"])

        if not save_metadata:
            metadata = None

        utils.save_torch_file(new_lora, save_path, metadata=metadata)

        save_message = "Weights Saved To: {}".format(save_path)

        return {"ui": {"tags": [save_message]}}

class WarpedHunyuanMultiLoraMerge:
    def __init__(self):
        self.base_output_dir = get_default_output_folder()
        os.makedirs(self.base_output_dir, exist_ok = True)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "save_path": ("STRING", {"default": get_default_output_path()}),
                "lora_1": (['None'] + get_lora_list(),),
                "strength_1": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
                "blocks_type_1": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "lora_2": (['None'] + get_lora_list(),),
                "strength_2": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
                "blocks_type_2": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "lora_3": (['None'] + get_lora_list(),),
                "strength_3": ("FLOAT", {"default": 0.00, "min": 0.00, "max": 2.00, "step": 0.01}),
                "blocks_type_3": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "lora_4": (['None'] + get_lora_list(),),
                "strength_4": ("FLOAT", {"default": 0.00, "min": 0.00, "max": 2.00, "step": 0.01}),
                "blocks_type_4": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "save_metadata": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    OUTPUT_IS_LIST = (True,)
    FUNCTION = "merge_multiple_loras"
    CATEGORY = "Warped/Hunyuan/Merge"
    DESCRIPTION = "Load and apply multiple LoRA models with different strengths and block types. Model input is required."

    def load_lora(self, lora_name: str, strength: float, blocks_type: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
        """Load and filter a single LoRA model."""
        if not lora_name or strength == 0:
            return {}, {}

        # Get the full path to the LoRA file
        lora_path = folder_paths.get_full_path("loras", lora_name)
        if not os.path.exists(lora_path):
            raise ValueError(f"LoRA file not found: {lora_path}")

        # Load the LoRA weights
        lora_weights = utils.load_torch_file(lora_path)

        return lora_weights

    def merge_multiple_loras(self, save_path, lora_1, strength_1, blocks_type_1, lora_2, strength_2, blocks_type_2, lora_3, strength_3, blocks_type_3, lora_4, strength_4, blocks_type_4, save_metadata=True):
        temp_loras = {}
        metadata = {"loras": "{} and {} and {} and {}".format(lora_1, lora_2, lora_3, lora_4)}
        metadata["strengths"] = "{} and {} and {} and {}".format(strength_1, strength_2, strength_3, strength_4)
        metadata["block_types"] = "{} and {} and {} and {}".format(blocks_type_1, blocks_type_2, blocks_type_3, blocks_type_4)

        if lora_1 != "None" and strength_1 != 0:
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_1, 1.0, "all")
            lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
            filtered_lora = filter_lora_keys(lora_weights, blocks_type_1)
            temp_loras["1"] = {"lora_weights": lora_weights, "strength": strength_1, "filtered_lora": filtered_lora}

        if lora_2 != "None" and strength_2 != 0:
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_2, 1.0, "all")
            lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
            filtered_lora = filter_lora_keys(lora_weights, blocks_type_2)
            temp_loras["2"] = {"lora_weights": lora_weights, "strength": strength_2, "filtered_lora": filtered_lora}

        if lora_3 != "None" and strength_3 != 0:
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_3, 1.0, "all")
            lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
            filtered_lora = filter_lora_keys(lora_weights, blocks_type_3)
            temp_loras["3"] = {"lora_weights": lora_weights, "strength": strength_3, "filtered_lora": filtered_lora}

        if lora_4 != "None" and strength_4 != 0:
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_4, 1.0, "all")
            lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
            filtered_lora = filter_lora_keys(lora_weights, blocks_type_4)
            temp_loras["4"] = {"lora_weights": lora_weights, "strength": strength_4, "filtered_lora": filtered_lora}

        loras = {}

        for lora_key in temp_loras.keys():
            loras[lora_key] = {"lora_weights": {}, "strength": temp_loras[lora_key]["strength"], "filtered_lora": temp_loras[lora_key]["filtered_lora"]}

            for key in temp_loras[lora_key]["lora_weights"].keys():
                new_key = key.replace("transformer.", "diffusion_model.")
                loras[lora_key]["lora_weights"][new_key] = temp_loras[lora_key]["lora_weights"][key]

        new_lora = {}

        for lora_key in loras.keys():
            for key in loras[lora_key]["lora_weights"].keys():
                if not key in new_lora.keys():
                    new_lora[key] = None
                print("Lora: {}  | Key: {}".format(lora_key, key))

        # Merge The Weighted Key Weights
        for key in new_lora.keys():
            for lora_key in loras.keys():
                if key in loras[lora_key]["lora_weights"].keys():
                    if not new_lora[key] is None:
                        temp_weights = torch.mul(loras[lora_key]["lora_weights"][key], loras[lora_key]["strength"])

                        if new_lora[key].shape[0] < new_lora[key].shape[1]:
                            if temp_weights.shape[0] < new_lora[key].shape[0]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([new_lora[key].shape[0], new_lora[key].shape[1]])
                                padding[:temp_weights.shape[0],:] = temp_weights
                                temp_weights = padding.to(dtype=torch.bfloat16)
                            elif temp_weights.shape[0] > new_lora[key].shape[0]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([temp_weights.shape[0], temp_weights.shape[1]])
                                padding[:new_lora[key].shape[0],:] = new_lora[key]
                                new_lora[key] = padding.to(dtype=torch.bfloat16)
                        else:
                            if temp_weights.shape[1] < new_lora[key].shape[1]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([new_lora[key].shape[0], new_lora[key].shape[1]])
                                padding[:,:temp_weights.shape[1]] = temp_weights
                                temp_weights = padding.to(dtype=torch.bfloat16)
                            elif temp_weights.shape[1] > new_lora[key].shape[1]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([temp_weights.shape[0], temp_weights.shape[1]])
                                padding[:,:new_lora[key].shape[1]] = new_lora[key]
                                new_lora[key] = padding.to(dtype=torch.bfloat16)

                        try:
                            new_lora[key] = torch.add(new_lora[key], temp_weights).to(dtype=torch.bfloat16)
                        except Exception as e:
                            raise(e)
                    else:
                        new_lora[key] = torch.mul(loras[lora_key]["lora_weights"][key], loras[lora_key]["strength"]).to(dtype=torch.bfloat16)

        if not save_metadata:
            metadata = None

        utils.save_torch_file(new_lora, save_path, metadata=metadata)

        save_message = "Weights Saved To: {}".format(save_path)

        return {"ui": {"tags": [save_message]}}

class WarpedHunyuanLoraAvgMerge:
    def __init__(self):
        self.base_output_dir = get_default_output_folder()
        os.makedirs(self.base_output_dir, exist_ok = True)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "save_path": ("STRING", {"default": get_default_output_path()}),
                "lora_1": (['None'] + get_lora_list(),),
                "blocks_type_1": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "lora_2": (['None'] + get_lora_list(),),
                "blocks_type_2": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "save_metadata": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    OUTPUT_IS_LIST = (True,)
    FUNCTION = "merge_multiple_loras"
    CATEGORY = "Warped/Hunyuan/Merge"
    DESCRIPTION = "Load and apply multiple LoRA models with different strengths and block types. Model input is required."

    def load_lora(self, lora_name: str, strength: float, blocks_type: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
        """Load and filter a single LoRA model."""
        if not lora_name or strength == 0:
            return {}, {}

        # Get the full path to the LoRA file
        lora_path = folder_paths.get_full_path("loras", lora_name)
        if not os.path.exists(lora_path):
            raise ValueError(f"LoRA file not found: {lora_path}")

        # Load the LoRA weights
        lora_weights = utils.load_torch_file(lora_path)

        return lora_weights

    def merge_multiple_loras(self, save_path, lora_1, blocks_type_1, lora_2, blocks_type_2, save_metadata=True):
        """Load and apply multiple LoRA models."""
        strength = 1.0000
        temp_loras = {}
        metadata = {"loras": "{} and {}".format(lora_1, lora_2)}
        metadata["strengths"] = "{} and {}".format(strength, strength)
        metadata["block_types"] = "{} and {}".format(blocks_type_1, blocks_type_2)

        if lora_1 != "None":
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_1, 1.0, "all")
            lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
            filtered_lora = filter_lora_keys(lora_weights, blocks_type_1)
            temp_loras["1"] = {"lora_weights": lora_weights, "strength": strength, "filtered_lora": filtered_lora}

        if lora_2 != "None":
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_2, 1.0, "all")
            lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
            filtered_lora = filter_lora_keys(lora_weights, blocks_type_2)
            temp_loras["2"] = {"lora_weights": lora_weights, "strength": strength, "filtered_lora": filtered_lora}

        loras = {}

        for lora_key in temp_loras.keys():
            loras[lora_key] = {"lora_weights": {}, "strength": temp_loras[lora_key]["strength"], "filtered_lora": temp_loras[lora_key]["filtered_lora"]}

            for key in temp_loras[lora_key]["lora_weights"].keys():
                new_key = key.replace("transformer.", "diffusion_model.")
                loras[lora_key]["lora_weights"][new_key] = temp_loras[lora_key]["lora_weights"][key]

        new_lora = {}

        for lora_key in loras.keys():
            for key in loras[lora_key]["lora_weights"].keys():
                if not key in new_lora.keys():
                    new_lora[key] = None
                print("Lora: {}  | Key: {}".format(lora_key, key))

        # Merge The Weighted Key Weights
        for key in new_lora.keys():
            for lora_key in loras.keys():
                if key in loras[lora_key]["lora_weights"].keys():
                    if not new_lora[key] is None:
                        temp_weights = torch.mul(loras[lora_key]["lora_weights"][key], loras[lora_key]["strength"])

                        if new_lora[key].shape[0] < new_lora[key].shape[1]:
                            if temp_weights.shape[0] < new_lora[key].shape[0]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([new_lora[key].shape[0], new_lora[key].shape[1]])
                                padding[:temp_weights.shape[0],:] = temp_weights
                                temp_weights = padding
                            elif temp_weights.shape[0] > new_lora[key].shape[0]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([temp_weights.shape[0], temp_weights.shape[1]])
                                padding[:new_lora[key].shape[0],:] = new_lora[key]
                                new_lora[key] = padding
                        else:
                            if temp_weights.shape[1] < new_lora[key].shape[1]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([new_lora[key].shape[0], new_lora[key].shape[1]])
                                padding[:,:temp_weights.shape[1]] = temp_weights
                                temp_weights = padding
                            elif temp_weights.shape[1] > new_lora[key].shape[1]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([temp_weights.shape[0], temp_weights.shape[1]])
                                padding[:,:new_lora[key].shape[1]] = new_lora[key]
                                new_lora[key] = padding

                        try:
                            new_lora[key] = torch.add(new_lora[key], temp_weights)
                            new_lora[key] = torch.div(new_lora[key], 2.0000)
                        except Exception as e:
                            raise(e)
                    else:
                        new_lora[key] = torch.mul(loras[lora_key]["lora_weights"][key], loras[lora_key]["strength"])

        if not save_metadata:
            metadata = None

        utils.save_torch_file(new_lora, save_path, metadata=metadata)

        save_message = "Weights Saved To: {}".format(save_path)

        return {"ui": {"tags": [save_message]}}

class WarpedHunyuanMultiLoraAvgMerge:
    def __init__(self):
        self.base_output_dir = get_default_output_folder()
        os.makedirs(self.base_output_dir, exist_ok = True)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "save_path": ("STRING", {"default": get_default_output_path()}),
                "lora_1": (['None'] + get_lora_list(),),
                "blocks_type_1": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "lora_2": (['None'] + get_lora_list(),),
                "blocks_type_2": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "lora_3": (['None'] + get_lora_list(),),
                "blocks_type_3": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "lora_4": (['None'] + get_lora_list(),),
                "blocks_type_4": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "save_metadata": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    OUTPUT_IS_LIST = (True,)
    FUNCTION = "merge_multiple_loras"
    CATEGORY = "Warped/Hunyuan/Merge"
    DESCRIPTION = "Load and apply multiple LoRA models with different strengths and block types. Model input is required."

    def load_lora(self, lora_name: str, strength: float, blocks_type: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
        """Load and filter a single LoRA model."""
        if not lora_name or strength == 0:
            return {}, {}

        # Get the full path to the LoRA file
        lora_path = folder_paths.get_full_path("loras", lora_name)
        if not os.path.exists(lora_path):
            raise ValueError(f"LoRA file not found: {lora_path}")

        # Load the LoRA weights
        lora_weights = utils.load_torch_file(lora_path)

        return lora_weights

    def merge_multiple_loras(self, save_path, lora_1, blocks_type_1, lora_2, blocks_type_2, lora_3, blocks_type_3, lora_4, blocks_type_4, save_metadata=True):
        strength = 1.0000
        temp_loras = {}
        metadata = {"loras": "{} and {} and {} and {}".format(lora_1, lora_2, lora_3, lora_4)}
        metadata["strengths"] = "{} and {} and {} and {}".format(strength, strength, strength, strength)
        metadata["block_types"] = "{} and {} and {} and {}".format(blocks_type_1, blocks_type_2, blocks_type_3, blocks_type_4)

        if lora_1 != "None":
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_1, 1.0, "all")
            lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
            filtered_lora = filter_lora_keys(lora_weights, blocks_type_1)
            temp_loras["1"] = {"lora_weights": lora_weights, "strength": strength, "filtered_lora": filtered_lora}

        if lora_2 != "None":
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_2, 1.0, "all")
            lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
            filtered_lora = filter_lora_keys(lora_weights, blocks_type_2)
            temp_loras["2"] = {"lora_weights": lora_weights, "strength": strength, "filtered_lora": filtered_lora}

        if lora_3 != "None":
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_3, 1.0, "all")
            lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
            filtered_lora = filter_lora_keys(lora_weights, blocks_type_3)
            temp_loras["3"] = {"lora_weights": lora_weights, "strength": strength, "filtered_lora": filtered_lora}

        if lora_4 != "None":
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_4, 1.0, "all")
            lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
            filtered_lora = filter_lora_keys(lora_weights, blocks_type_4)
            temp_loras["4"] = {"lora_weights": lora_weights, "strength": strength, "filtered_lora": filtered_lora}

        loras = {}

        for lora_key in temp_loras.keys():
            loras[lora_key] = {"lora_weights": {}, "strength": temp_loras[lora_key]["strength"], "filtered_lora": temp_loras[lora_key]["filtered_lora"]}

            for key in temp_loras[lora_key]["lora_weights"].keys():
                new_key = key.replace("transformer.", "diffusion_model.")
                loras[lora_key]["lora_weights"][new_key] = temp_loras[lora_key]["lora_weights"][key]

        new_lora = {}
        num_loras = int(len(loras))

        for lora_key in loras.keys():
            for key in loras[lora_key]["lora_weights"].keys():
                if not key in new_lora.keys():
                    new_lora[key] = {"weights": None, "count": 0}
                print("Lora: {}  | Key: {}".format(lora_key, key))

        # Merge The Weighted Key Weights
        for key in new_lora.keys():
            for lora_key in loras.keys():
                if key in loras[lora_key]["lora_weights"].keys():
                    if not new_lora[key]["weights"] is None:
                        temp_weights = torch.mul(loras[lora_key]["lora_weights"][key], loras[lora_key]["strength"])

                        if new_lora[key]["weights"].shape[0] < new_lora[key]["weights"].shape[1]:
                            if temp_weights.shape[0] < new_lora[key]["weights"].shape[0]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key]["weights"] = new_lora[key]["weights"].clone().detach()

                                padding = torch.zeros([new_lora[key]["weights"].shape[0], new_lora[key]["weights"].shape[1]])
                                padding[:temp_weights.shape[0],:] = temp_weights
                                temp_weights = padding
                            elif temp_weights.shape[0] > new_lora[key]["weights"].shape[0]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key]["weights"] = new_lora[key]["weights"].clone().detach()

                                padding = torch.zeros([temp_weights.shape[0], temp_weights.shape[1]])
                                padding[:new_lora[key]["weights"].shape[0],:] = new_lora[key]["weights"]
                                new_lora[key]["weights"] = padding
                        else:
                            if temp_weights.shape[1] < new_lora[key]["weights"].shape[1]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key]["weights"] = new_lora[key]["weights"].clone().detach()

                                padding = torch.zeros([new_lora[key]["weights"].shape[0], new_lora[key]["weights"].shape[1]])
                                padding[:,:temp_weights.shape[1]] = temp_weights
                                temp_weights = padding
                            elif temp_weights.shape[1] > new_lora[key]["weights"].shape[1]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key]["weights"] = new_lora[key]["weights"].clone().detach()

                                padding = torch.zeros([temp_weights.shape[0], temp_weights.shape[1]])
                                padding[:,:new_lora[key]["weights"].shape[1]] = new_lora[key]["weights"]
                                new_lora[key]["weights"] = padding

                        try:
                            new_lora[key]["weights"] = torch.add(new_lora[key]["weights"], temp_weights)
                            new_lora[key]["count"] += 1
                        except Exception as e:
                            raise(e)
                    else:
                        new_lora[key]["weights"] = torch.mul(loras[lora_key]["lora_weights"][key], loras[lora_key]["strength"])
                        new_lora[key]["count"] = 1

        final_lora = {}
        for key in new_lora.keys():
            final_lora[key] = torch.div(new_lora[key]["weights"], new_lora[key]["count"])

        if not save_metadata:
            metadata = None

        utils.save_torch_file(final_lora, save_path, metadata=metadata)

        save_message = "Weights Saved To: {}".format(save_path)

        return {"ui": {"tags": [save_message]}}

# class WarpedHunyuanMultiLoraMixer:
#     @classmethod
#     def INPUT_TYPES(s):
#         return {
#             "required": {
#                 "save_folder": ("STRING", {"default": get_default_output_folder()}),
#                 "model_prefix": ("STRING", {"default": "new_model_hy"}),
#                 "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
#                 "num_output": ("INT", {"default": 1, "min": 1, "max": 100}),
#                 "lora_1": (['None'] + get_lora_list(),),
#                 "lora_2": (['None'] + get_lora_list(),),
#                 "lora_3": (['None'] + get_lora_list(),),
#                 "lora_4": (['None'] + get_lora_list(),),
#                 "lora_5": (['None'] + get_lora_list(),),
#                 "lora_6": (['None'] + get_lora_list(),),
#                 "lora_7": (['None'] + get_lora_list(),),
#                 "lora_8": (['None'] + get_lora_list(),),
#                 "save_metadata": ("BOOLEAN", {"default": True}),
#             },
#         }
#
#     RETURN_TYPES = ()
#     OUTPUT_NODE = True
#     OUTPUT_IS_LIST = (True,)
#     FUNCTION = "merge_multiple_loras"
#     CATEGORY = "Warped/Hunyuan/Mixers"
#     DESCRIPTION = "Load and apply multiple LoRA models with different strengths and block types. Model input is required."
#
#     def convert_key_format(self, key: str) -> str:
#         """Standardize LoRA key format by removing prefixes."""
#         prefixes = ["diffusion_model.", "transformer."]
#         for prefix in prefixes:
#             if key.startswith(prefix):
#                 key = key[len(prefix):]
#                 break
#         return key
#
#     def load_lora(self, lora_name: str, strength: float, blocks_type: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
#         """Load and filter a single LoRA model."""
#         if not lora_name or strength == 0:
#             return {}, {}
#
#         # Get the full path to the LoRA file
#         lora_path = folder_paths.get_full_path("loras", lora_name)
#         if not os.path.exists(lora_path):
#             raise ValueError(f"LoRA file not found: {lora_path}")
#
#         # Load the LoRA weights
#         lora_weights = utils.load_torch_file(lora_path)
#
#         return lora_weights
#
#     def get_mixtures(self, seed, num_output, lora_keys):
#         random.seed(seed)
#         mixtures = {}
#
#         for i in range(num_output):
#             mixtures["{}".format(i + 1)] = {}
#
#         for key in lora_keys:
#             for mixture_key in mixtures.keys():
#                 mixtures[mixture_key][key] = {"single": [], "double": []}
#
#         for mixture_key in mixtures.keys():
#             for j in range(40):
#                 temp_key = "{}".format(random.randint(1, len(lora_keys)))
#                 mixtures[mixture_key][temp_key]["single"].append(j)
#
#             for j in range(20):
#                 temp_key = "{}".format(random.randint(1, len(lora_keys)))
#                 mixtures[mixture_key][temp_key]["double"].append(j)
#
#             i += 1
#
#         print("\nMixtures\n")
#
#         block_metadata = ""
#
#         for mixture_key in mixtures.keys():
#             for key in mixtures[mixture_key]:
#                 print("{} | {}: {}".format(mixture_key, key, mixtures[mixture_key][key]))
#
#                 if len(block_metadata) > 0:
#                     block_metadata = "{}  |  {}".format(block_metadata, "{}: {}: {}".format(mixture_key, key, mixtures[mixture_key][key]))
#                 else:
#                     block_metadata = "{}: {}: {}".format(mixture_key, key, mixtures[mixture_key][key])
#
#             print("\n")
#
#         return mixtures, block_metadata
#
#     def merge_multiple_loras(self, save_folder, model_prefix, seed, num_output, lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8, save_metadata=True):
#         print("Save_folder: {}".format(save_folder))
#         os.makedirs(save_folder, exist_ok = True)
#
#         temp_loras = {}
#         metadata = {"loras": "{} and {} and {} and {} and {} and {} and {} and {}".format(lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8)}
#         metadata["seed"] = "{}".format(seed)
#         metadata["num_output"] = "{}".format(num_output)
#
#         if lora_1 != "None":
#             print(lora_1)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_1, 1.0, "all")
#             lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
#             temp_loras["1"] = {"lora_weights": lora_weights}
#
#         if lora_2 != "None":
#             print(lora_2)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_2, 1.0, "all")
#             lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
#             temp_loras["2"] = {"lora_weights": lora_weights}
#
#         if lora_3 != "None":
#             print(lora_3)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_3, 1.0, "all")
#             lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
#             temp_loras["3"] = {"lora_weights": lora_weights}
#
#         if lora_4 != "None":
#             print(lora_4)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_4, 1.0, "all")
#             lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
#             temp_loras["4"] = {"lora_weights": lora_weights}
#
#         if lora_5 != "None":
#             print(lora_5)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_5, 1.0, "all")
#             lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
#             temp_loras["5"] = {"lora_weights": lora_weights}
#
#         if lora_6 != "None":
#             print(lora_6)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_6, 1.0, "all")
#             lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
#             temp_loras["6"] = {"lora_weights": lora_weights}
#
#         if lora_7 != "None":
#             print(lora_7)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_7, 1.0, "all")
#             lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
#             temp_loras["7"] = {"lora_weights": lora_weights}
#
#         if lora_8 != "None":
#             print(lora_8)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_8, 1.0, "all")
#             lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
#             temp_loras["8"] = {"lora_weights": lora_weights}
#
#         loras = {}
#         max_dimension = 0
#
#         for lora_key in temp_loras.keys():
#             # print(lora_key)
#             loras[lora_key] = {"lora_weights": {}}
#
#             for key in temp_loras[lora_key]["lora_weights"].keys():
#                 new_key = key.replace("transformer.", "diffusion_model.")
#                 loras[lora_key]["lora_weights"][new_key] = temp_loras[lora_key]["lora_weights"][key]
#
#                 temp_dimension = min(loras[lora_key]["lora_weights"][new_key].shape[0], loras[lora_key]["lora_weights"][new_key].shape[1])
#
#                 if temp_dimension > max_dimension:
#                     max_dimension = temp_dimension
#
#         merge_mixtures, block_metadata = self.get_mixtures(seed, num_output, loras.keys())
#         metadata["max_dimension"] = "{}".format(max_dimension)
#
#         print("Max Dimension: {}".format(max_dimension))
#
#         # convert the rank/dims for each lora to be mixed
#         for lora_key in temp_loras.keys():
#             temp_loras[lora_key]["lora_weights"] = convert_lora_dimensions(max_dimension, temp_loras[lora_key]["lora_weights"])
#
#         save_message = ""
#
#         for mixture_key in merge_mixtures:
#             new_lora = {}
#             output_filename = os.path.join(save_folder, "{}_{:05}.safetensors".format(model_prefix, int(mixture_key)))
#
#             metadata["merge_mixture"] = "{}".format(merge_mixtures[mixture_key])
#             # metadata["block_metadata"] = "{}".format(block_metadata[int(mixture_key)])
#
#             for lora_key in loras.keys():
#                 mixture_single_blocks = merge_mixtures[mixture_key][lora_key]["single"]
#                 mixture_double_blocks = merge_mixtures[mixture_key][lora_key]["double"]
#
#                 for key in loras[lora_key]["lora_weights"].keys():
#                     temp_strings = str(key).split('.')
#                     temp_block_num = int(temp_strings[2])
#
#                     if temp_strings[1] == "single_blocks":
#                         if temp_block_num in mixture_single_blocks:
#                             new_lora[key] = loras[lora_key]["lora_weights"][key].to(dtype=torch.bfloat16)
#                         continue
#
#                     if temp_strings[1] == "double_blocks":
#                         if temp_block_num in mixture_double_blocks:
#                             new_lora[key] = loras[lora_key]["lora_weights"][key].to(dtype=torch.bfloat16)
#
#             if not save_metadata:
#                 metadata = None
#
#             print("Saving Model To: {}...".format(output_filename))
#             utils.save_torch_file(new_lora, output_filename, metadata=metadata)
#             print("Saving Model To: {}...Done.".format(output_filename))
#
#             save_message = "{}\n{}".format(save_message, "Weights Saved To: {}".format(output_filename))
#
#             new_lora = None
#             mm.soft_empty_cache()
#             gc.collect()
#             time.sleep(1)
#
#         return {"ui": {"tags": ["save_message"]}}
#
# class WarpedHunyuanMultiLoraMixerExt:
#     @classmethod
#     def INPUT_TYPES(s):
#         return {
#             "required": {
#                 "save_folder": ("STRING", {"default": get_default_output_folder()}),
#                 "model_prefix": ("STRING", {"default": "new_model_hy"}),
#                 "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
#                 "num_output": ("INT", {"default": 1, "min": 1, "max": 100}),
#                 "lora_1": (['None'] + get_lora_list(),),
#                 "strength_1": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
#                 "lora_2": (['None'] + get_lora_list(),),
#                 "strength_2": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
#                 "lora_3": (['None'] + get_lora_list(),),
#                 "strength_3": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
#                 "lora_4": (['None'] + get_lora_list(),),
#                 "strength_4": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
#                 "lora_5": (['None'] + get_lora_list(),),
#                 "strength_5": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
#                 "lora_6": (['None'] + get_lora_list(),),
#                 "strength_6": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
#                 "lora_7": (['None'] + get_lora_list(),),
#                 "strength_7": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
#                 "lora_8": (['None'] + get_lora_list(),),
#                 "strength_8": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
#                 "save_metadata": ("BOOLEAN", {"default": True}),
#                 "convert_to": (["diffusion_model", "transformer"], {"default": "diffusion_model"}),
#                 "max_dimension": ([32, 64, 128], {"default": 64}),
#             },
#         }
#
#     RETURN_TYPES = ()
#     OUTPUT_NODE = True
#     OUTPUT_IS_LIST = (True,)
#     FUNCTION = "merge_multiple_loras"
#     CATEGORY = "Warped/Hunyuan/Mixers"
#     DESCRIPTION = "Load and apply multiple LoRA models with different strengths and block types. Model input is required."
#
#     def load_lora(self, lora_name: str, strength: float, blocks_type: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
#         """Load and filter a single LoRA model."""
#         if not lora_name or strength == 0:
#             return {}, {}
#
#         # Get the full path to the LoRA file
#         lora_path = folder_paths.get_full_path("loras", lora_name)
#         if not os.path.exists(lora_path):
#             raise ValueError(f"LoRA file not found: {lora_path}")
#
#         # Load the LoRA weights
#         lora_weights = utils.load_torch_file(lora_path)
#
#         lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
#
#         return lora_weights
#
#     def get_random_key(self, keys):
#         if len(keys) > 0:
#             random_key = random.randint(0, len(keys) - 1)
#             # print("Keys: {}  |  Random Key: {}".format(random_key))
#             return keys[random_key]
#
#         return -1
#
#     def get_mixtures(self, seed, num_output, lora_keys, block_types):
#         random.seed(seed)
#         mixtures = {}
#
#         for i in range(num_output):
#             mixtures["{}".format(i + 1)] = {}
#
#         single_block_loras = []
#         double_block_loras = []
#
#         for key in lora_keys:
#             for mixture_key in mixtures.keys():
#                 mixtures[mixture_key][key] = {"single": [], "double": []}
#
#                 if block_types[key]["has_single_blocks"]:
#                     single_block_loras.append(int(key))
#
#                 if block_types[key]["has_double_blocks"]:
#                     double_block_loras.append(int(key))
#
#         for mixture_key in mixtures.keys():
#             if len(single_block_loras) > 0:
#                 for j in range(40):
#                     random_key = self.get_random_key(single_block_loras)
#                     temp_key = "{}".format(random_key)
#                     mixtures[mixture_key][temp_key]["single"].append(j)
#
#             if len(double_block_loras) > 0:
#                 for j in range(20):
#                     random_key = self.get_random_key(double_block_loras)
#                     temp_key = "{}".format(random_key)
#                     mixtures[mixture_key][temp_key]["double"].append(j)
#
#             i += 1
#
#         print("\nMixtures\n")
#
#         block_metadata = ""
#
#         for mixture_key in mixtures.keys():
#             for key in mixtures[mixture_key]:
#                 print("{} | {}: {}".format(mixture_key, key, mixtures[mixture_key][key]))
#
#                 if len(block_metadata) > 0:
#                     block_metadata = "{}  |  {}".format(block_metadata, "{}: {}: {}".format(mixture_key, key, mixtures[mixture_key][key]))
#                 else:
#                     block_metadata = "{}: {}: {}".format(mixture_key, key, mixtures[mixture_key][key])
#
#             print("\n")
#
#         return mixtures, block_metadata
#
#     def determine_lora_block_types(self, loras):
#         block_types = {}
#
#         for lora_key in loras.keys():
#             block_types[lora_key] = { "has_single_blocks": False, "has_double_blocks": False }
#
#             for key in loras[lora_key]["lora_weights"].keys():
#                 if "single_blocks" in key:
#                     block_types[lora_key]["has_single_blocks"] = True
#                 elif "double_blocks" in key:
#                     block_types[lora_key]["has_double_blocks"] = True
#
#                 if block_types[lora_key]["has_single_blocks"] and block_types[lora_key]["has_double_blocks"]:
#                     break
#
#         return block_types
#
#     def merge_multiple_loras(self, save_folder, model_prefix, seed, num_output, lora_1, strength_1, lora_2, strength_2, lora_3, strength_3, lora_4, strength_4,
#                             lora_5, strength_5, lora_6, strength_6, lora_7, strength_7, lora_8, strength_8, save_metadata=True, convert_to="diffusion_model", max_dimension=64):
#         print("Save_folder: {}".format(save_folder))
#         os.makedirs(save_folder, exist_ok = True)
#
#         temp_loras = {}
#         metadata = {"loras": "{} and {} and {} and {} and {} and {} and {} and {}".format(lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8)}
#         metadata["strengths"] = "{} and {} and {} and {} and {} and {} and {} and {}".format(strength_1, strength_2, strength_3, strength_4, strength_5, strength_6, strength_7, strength_8)
#         metadata["seed"] = "{}".format(seed)
#         metadata["num_output"] = "{}".format(num_output)
#
#         if (lora_1 != "None") and (strength_1 > 0.0):
#             print(lora_1)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_1, 1.0, "all")
#             lora_weights = convert_lora_dimensions(max_dimension, lora_weights)
#
#             temp_loras["1"] = {"lora_weights": lora_weights, "strength": strength_1}
#
#             for key in lora_weights.keys():
#                 print("LORA 1: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
#                 break
#
#             lora_weights = None
#
#         if (lora_2 != "None") and (strength_2 > 0.0):
#             print(lora_2)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_2, 1.0, "all")
#             lora_weights = convert_lora_dimensions(max_dimension, lora_weights)
#
#             temp_loras["2"] = {"lora_weights": lora_weights, "strength": strength_2}
#
#             for key in lora_weights.keys():
#                 print("LORA 2: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
#                 break
#
#             lora_weights = None
#
#         if (lora_3 != "None") and (strength_3 > 0.0):
#             print(lora_3)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_3, 1.0, "all")
#             lora_weights = convert_lora_dimensions(max_dimension, lora_weights)
#
#             temp_loras["3"] = {"lora_weights": lora_weights, "strength": strength_3}
#
#             for key in lora_weights.keys():
#                 print("LORA 3: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
#                 break
#
#             lora_weights = None
#
#         if (lora_4 != "None") and (strength_4 > 0.0):
#             print(lora_4)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_4, 1.0, "all")
#             lora_weights = convert_lora_dimensions(max_dimension, lora_weights)
#
#             temp_loras["4"] = {"lora_weights": lora_weights, "strength": strength_4}
#
#             for key in lora_weights.keys():
#                 print("LORA 4: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
#                 break
#
#             lora_weights = None
#
#         if (lora_5 != "None") and (strength_5 > 0.0):
#             print(lora_5)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_5, 1.0, "all")
#             lora_weights = convert_lora_dimensions(max_dimension, lora_weights)
#
#             temp_loras["5"] = {"lora_weights": lora_weights, "strength": strength_5}
#
#             for key in lora_weights.keys():
#                 print("LORA 5: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
#                 break
#
#             lora_weights = None
#
#         if (lora_6 != "None") and (strength_6 > 0.0):
#             print(lora_6)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_6, 1.0, "all")
#             lora_weights = convert_lora_dimensions(max_dimension, lora_weights)
#
#             temp_loras["6"] = {"lora_weights": lora_weights, "strength": strength_6}
#
#             for key in lora_weights.keys():
#                 print("LORA 6: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
#                 break
#
#             lora_weights = None
#
#         if (lora_7 != "None") and (strength_7 > 0.0):
#             print(lora_7)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_7, 1.0, "all")
#             lora_weights = convert_lora_dimensions(max_dimension, lora_weights)
#
#             temp_loras["7"] = {"lora_weights": lora_weights, "strength": strength_7}
#
#             for key in lora_weights.keys():
#                 print("LORA 7: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
#                 break
#
#             lora_weights = None
#
#         if (lora_8 != "None") and (strength_8 > 0.0):
#             print(lora_8)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_8, 1.0, "all")
#             lora_weights = convert_lora_dimensions(max_dimension, lora_weights)
#
#             temp_loras["8"] = {"lora_weights": lora_weights, "strength": strength_8}
#
#             for key in lora_weights.keys():
#                 print("LORA 8: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
#                 break
#
#             lora_weights = None
#
#         loras = {}
#
#         for lora_key in temp_loras.keys():
#             loras[lora_key] = {"lora_weights": {}, "strength": temp_loras[lora_key]["strength"]}
#
#             for key in temp_loras[lora_key]["lora_weights"].keys():
#                 new_key = key.replace("transformer.", "diffusion_model.")
#                 loras[lora_key]["lora_weights"][new_key] = temp_loras[lora_key]["lora_weights"][key].to(dtype=torch.bfloat16)
#
#         block_types = self.determine_lora_block_types(loras)
#         merge_mixtures, block_metadata = self.get_mixtures(seed, num_output, loras.keys(), block_types)
#
#         metadata["max_dimension"] = "{}".format(max_dimension)
#         metadata["block_types"] = "{}".format(block_types)
#
#         print("Max Dimension: {}".format(max_dimension))
#
#         save_message = ""
#
#         for mixture_key in merge_mixtures:
#             new_lora = {}
#             output_filename = os.path.join(save_folder, "{}_{:05}.safetensors".format(model_prefix, int(mixture_key)))
#
#             metadata["merge_mixture"] = "{}".format(merge_mixtures[mixture_key])
#             # metadata["block_metadata"] = "{}".format(block_metadata[int(mixture_key)])
#
#             for lora_key in loras.keys():
#                 mixture_single_blocks = merge_mixtures[mixture_key][lora_key]["single"]
#                 mixture_double_blocks = merge_mixtures[mixture_key][lora_key]["double"]
#
#                 for key in loras[lora_key]["lora_weights"].keys():
#                     temp_strings = str(key).split('.')
#                     temp_block_num = int(temp_strings[2])
#
#                     if temp_strings[1] == "single_blocks":
#                         if temp_block_num in mixture_single_blocks:
#                             new_lora[key] = torch.mul(loras[lora_key]["lora_weights"][key], loras[lora_key]["strength"]).to(dtype=torch.bfloat16)
#                         continue
#
#                     if temp_strings[1] == "double_blocks":
#                         if temp_block_num in mixture_double_blocks:
#                             new_lora[key] = torch.mul(loras[lora_key]["lora_weights"][key], loras[lora_key]["strength"]).to(dtype=torch.bfloat16)
#
#             if not save_metadata:
#                 metadata = None
#
#             print("Saving Model To: {}...".format(output_filename))
#             utils.save_torch_file(new_lora, output_filename, metadata=metadata)
#             print("Saving Model To: {}...Done.".format(output_filename))
#
#             save_message = "{}\n{}".format(save_message, "Weights Saved To: {}".format(output_filename))
#
#             new_lora = None
#             mm.soft_empty_cache()
#             gc.collect()
#             time.sleep(1)
#
#         return {"ui": {"tags": ["save_message"]}}

class WarpedHunyuanMultiLoraMixer:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "save_folder": ("STRING", {"default": get_default_output_folder()}),
                "model_prefix": ("STRING", {"default": "new_model_hy"}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                "num_output": ("INT", {"default": 1, "min": 1, "max": 100}),
                "lora_1": (['None'] + get_lora_list(),),
                "strength_1": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
                "blocks_1": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
                "lora_2": (['None'] + get_lora_list(),),
                "strength_2": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
                "blocks_2": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
                "lora_3": (['None'] + get_lora_list(),),
                "strength_3": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
                "blocks_3": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
                "lora_4": (['None'] + get_lora_list(),),
                "strength_4": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
                "blocks_4": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
                "lora_5": (['None'] + get_lora_list(),),
                "strength_5": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
                "blocks_5": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
                "lora_6": (['None'] + get_lora_list(),),
                "strength_6": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
                "blocks_6": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
                "lora_7": (['None'] + get_lora_list(),),
                "strength_7": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
                "blocks_7": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
                "lora_8": (['None'] + get_lora_list(),),
                "strength_8": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
                "blocks_8": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
                "save_metadata": ("BOOLEAN", {"default": True}),
                "convert_to": (["diffusion_model", "transformer"], {"default": "diffusion_model"}),
                "max_dimension": ([32, 64, 128], {"default": 128}),
                "discard_linear": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    OUTPUT_IS_LIST = (True,)
    FUNCTION = "merge_multiple_loras"
    CATEGORY = "Warped/Hunyuan/Mixers"
    DESCRIPTION = "Load and apply multiple LoRA models with different strengths and block types. Model input is required."

    def load_lora(self, lora_name: str, strength: float, blocks_type: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
        """Load and filter a single LoRA model."""
        if not lora_name or strength == 0:
            return {}, {}

        # # Load the LoRA weights
        lora_weights = warped_load_lora_weights(lora_name)
        lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")

        return lora_weights

    def get_random_key(self, keys):
        if len(keys) > 0:
            random_key = random.randint(0, len(keys) - 1)
            # print("Keys: {}  |  Random Key: {}".format(random_key))
            return keys[random_key]

        return -1

    def get_mixtures(self, seed, num_output, lora_keys, block_types):
        random.seed(seed)
        mixtures = {}

        for i in range(num_output):
            mixtures["{}".format(i + 1)] = {}

        single_block_loras = []
        double_block_loras = []

        for key in lora_keys:
            for mixture_key in mixtures.keys():
                mixtures[mixture_key][key] = {"single": [], "double": []}

                if block_types[key]["has_single_blocks"]:
                    single_block_loras.append(int(key))

                if block_types[key]["has_double_blocks"]:
                    double_block_loras.append(int(key))

        for mixture_key in mixtures.keys():
            if len(single_block_loras) > 0:
                for j in range(40):
                    random_key = self.get_random_key(single_block_loras)
                    temp_key = "{}".format(random_key)
                    mixtures[mixture_key][temp_key]["single"].append(j)

            if len(double_block_loras) > 0:
                for j in range(20):
                    random_key = self.get_random_key(double_block_loras)
                    temp_key = "{}".format(random_key)
                    mixtures[mixture_key][temp_key]["double"].append(j)

            i += 1

        at_least_one = False

        while not at_least_one:
            if len(double_block_loras) > 0:
                for key in lora_keys:
                    lowest = 999
                    highest = -1
                    lowest_mixture_key = None
                    highest_mixture_key = None

                    for mixture_key in mixtures:
                        if len(mixtures[mixture_key][key]["double"]) < lowest:
                            lowest = len(mixtures[mixture_key][key]["double"])
                            lowest_mixture_key = mixture_key

                        if len(mixtures[mixture_key][key]["double"]) > highest:
                            highest = len(mixtures[mixture_key][key]["double"])
                            highest_mixture_key = mixture_key

                    if lowest < 1:
                        print("low key: {}".format(lowest_mixture_key))
                        print("high key: {}".format(highest_mixture_key))
                        mixtures[lowest_mixture_key][key]["double"].append(mixtures[highest_mixture_key][key]["double"][0])
                        print("low mix: {}".format(mixtures[lowest_mixture_key][key]["double"]))
                        print("high mixes before: {}".format(mixtures[highest_mixture_key][key]["double"]))
                        mixtures[highest_mixture_key][key]["double"].pop(0)
                        print("high mixes after: {}".format(mixtures[highest_mixture_key][key]["double"]))
                    else:
                        at_least_one = True
            else:
                at_least_one = True

        print("\nMixtures\n")

        block_metadata = ""

        for mixture_key in mixtures.keys():
            for key in mixtures[mixture_key]:
                print("{} | {}: {}".format(mixture_key, key, mixtures[mixture_key][key]))

                if len(block_metadata) > 0:
                    block_metadata = "{}  |  {}".format(block_metadata, "{}: {}: {}".format(mixture_key, key, mixtures[mixture_key][key]))
                else:
                    block_metadata = "{}: {}: {}".format(mixture_key, key, mixtures[mixture_key][key])

            print("\n")

        return mixtures, block_metadata

    def determine_lora_block_types(self, loras):
        block_types = {}

        for lora_key in loras.keys():
            block_types[lora_key] = { "has_single_blocks": False, "has_double_blocks": False }

            for key in loras[lora_key]["lora_weights"].keys():
                if "single_blocks" in key:
                    block_types[lora_key]["has_single_blocks"] = True
                elif "double_blocks" in key:
                    block_types[lora_key]["has_double_blocks"] = True

                if block_types[lora_key]["has_single_blocks"] and block_types[lora_key]["has_double_blocks"]:
                    break

        return block_types

    def merge_multiple_loras(self, save_folder, model_prefix, seed, num_output, lora_1, strength_1, blocks_1, lora_2, strength_2, blocks_2, lora_3, strength_3, blocks_3, lora_4, strength_4, blocks_4,
                            lora_5, strength_5, blocks_5, lora_6, strength_6, blocks_6, lora_7, strength_7, blocks_7, lora_8, strength_8, blocks_8, save_metadata=True, convert_to="diffusion_model",
                            max_dimension=128, discard_linear=True):
        print("Save_folder: {}".format(save_folder))
        os.makedirs(save_folder, exist_ok = True)

        temp_loras = {}
        metadata = {"loras": "{} and {} and {} and {} and {} and {} and {} and {}".format(lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8)}
        metadata["strengths"] = "{} and {} and {} and {} and {} and {} and {} and {}".format(strength_1, strength_2, strength_3, strength_4, strength_5, strength_6, strength_7, strength_8)
        metadata["seed"] = "{}".format(seed)
        metadata["num_output"] = "{}".format(num_output)

        if (lora_1 != "None") and (strength_1 > 0.0):
            print(lora_1)
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_1, 1.0, "all")

            if (blocks_1 == "single_blocks") or (blocks_1 == "double_blocks"):
                lora_weights = filter_lora_keys(lora_weights, blocks_1)

            lora_weights = convert_lora_dimensions(max_dimension, lora_weights)

            print_it = True
            for key in lora_weights.keys():
                if print_it:
                    print("LORA 1: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
                    print_it = False

                lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)

            temp_loras["1"] = {"lora_weights": lora_weights, "strength": strength_1}

            lora_weights = None

        if (lora_2 != "None") and (strength_2 > 0.0):
            print(lora_2)
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_2, 1.0, "all")

            if (blocks_2 == "single_blocks") or (blocks_2 == "double_blocks"):
                lora_weights = filter_lora_keys(lora_weights, blocks_2)

            lora_weights = convert_lora_dimensions(max_dimension, lora_weights)

            print_it = True
            for key in lora_weights.keys():
                if print_it:
                    print("LORA 2: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
                    print_it = False

                lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)

            temp_loras["2"] = {"lora_weights": lora_weights, "strength": strength_2}

            lora_weights = None

        if (lora_3 != "None") and (strength_3 > 0.0):
            print(lora_3)
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_3, 1.0, "all")

            if (blocks_3 == "single_blocks") or (blocks_3 == "double_blocks"):
                lora_weights = filter_lora_keys(lora_weights, blocks_3)

            lora_weights = convert_lora_dimensions(max_dimension, lora_weights)

            print_it = True
            for key in lora_weights.keys():
                if print_it:
                    print("LORA 3: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
                    print_it = False

                lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)

            temp_loras["3"] = {"lora_weights": lora_weights, "strength": strength_3}

            lora_weights = None

        if (lora_4 != "None") and (strength_4 > 0.0):
            print(lora_4)
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_4, 1.0, "all")

            if (blocks_4 == "single_blocks") or (blocks_4 == "double_blocks"):
                lora_weights = filter_lora_keys(lora_weights, blocks_4)

            lora_weights = convert_lora_dimensions(max_dimension, lora_weights)

            print_it = True
            for key in lora_weights.keys():
                if print_it:
                    print("LORA 4: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
                    print_it = False

                lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)

            temp_loras["4"] = {"lora_weights": lora_weights, "strength": strength_4}

            lora_weights = None

        if (lora_5 != "None") and (strength_5 > 0.0):
            print(lora_5)
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_5, 1.0, "all")

            if (blocks_5 == "single_blocks") or (blocks_5 == "double_blocks"):
                lora_weights = filter_lora_keys(lora_weights, blocks_5)

            lora_weights = convert_lora_dimensions(max_dimension, lora_weights)

            print_it = True
            for key in lora_weights.keys():
                if print_it:
                    print("LORA 5: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
                    print_it = False

                lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)

            temp_loras["5"] = {"lora_weights": lora_weights, "strength": strength_5}

            lora_weights = None

        if (lora_6 != "None") and (strength_6 > 0.0):
            print(lora_6)
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_6, 1.0, "all")

            if (blocks_6 == "single_blocks") or (blocks_6 == "double_blocks"):
                lora_weights = filter_lora_keys(lora_weights, blocks_6)

            lora_weights = convert_lora_dimensions(max_dimension, lora_weights)

            print_it = True
            for key in lora_weights.keys():
                if print_it:
                    print("LORA 6: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
                    print_it = False

                lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)

            temp_loras["6"] = {"lora_weights": lora_weights, "strength": strength_6}

            lora_weights = None

        if (lora_7 != "None") and (strength_7 > 0.0):
            print(lora_7)
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_7, 1.0, "all")

            if (blocks_7 == "single_blocks") or (blocks_7 == "double_blocks"):
                lora_weights = filter_lora_keys(lora_weights, blocks_7)

            lora_weights = convert_lora_dimensions(max_dimension, lora_weights)

            print_it = True
            for key in lora_weights.keys():
                if print_it:
                    print("LORA 7: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
                    print_it = False

                lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)

            temp_loras["7"] = {"lora_weights": lora_weights, "strength": strength_7}

            lora_weights = None

        if (lora_8 != "None") and (strength_8 > 0.0):
            print(lora_8)
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_8, 1.0, "all")

            if (blocks_8 == "single_blocks") or (blocks_8 == "double_blocks"):
                lora_weights = filter_lora_keys(lora_weights, blocks_8)

            lora_weights = convert_lora_dimensions(max_dimension, lora_weights)

            print_it = True
            for key in lora_weights.keys():
                if print_it:
                    print("LORA 8: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
                    print_it = False

                lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)

            temp_loras["8"] = {"lora_weights": lora_weights, "strength": strength_8}

            lora_weights = None

        loras = {}

        for lora_key in temp_loras.keys():
            loras[lora_key] = {"lora_weights": {}, "strength": temp_loras[lora_key]["strength"]}

            for key in temp_loras[lora_key]["lora_weights"].keys():
                new_key = key.replace("transformer.", "diffusion_model.")
                loras[lora_key]["lora_weights"][new_key] = temp_loras[lora_key]["lora_weights"][key].to(dtype=torch.bfloat16)

        block_types = self.determine_lora_block_types(loras)
        merge_mixtures, block_metadata = self.get_mixtures(seed, num_output, loras.keys(), block_types)

        metadata["max_dimension"] = "{}".format(max_dimension)
        metadata["block_types"] = "{}".format(block_types)

        print("Max Dimension: {}".format(max_dimension))

        save_message = ""

        for mixture_key in merge_mixtures:
            time.sleep(1)
            new_lora = {}
            output_filename = os.path.join(save_folder, "{}_{:05}.safetensors".format(model_prefix, int(mixture_key)))

            metadata["merge_mixture"] = "{}".format(merge_mixtures[mixture_key])
            # metadata["block_metadata"] = "{}".format(block_metadata[int(mixture_key)])

            for lora_key in loras.keys():
                mixture_single_blocks = merge_mixtures[mixture_key][lora_key]["single"]
                mixture_double_blocks = merge_mixtures[mixture_key][lora_key]["double"]

                for key in loras[lora_key]["lora_weights"].keys():
                    temp_strings = str(key).split('.')
                    temp_block_num = int(temp_strings[2])

                    if temp_strings[1] == "single_blocks":
                        if temp_block_num in mixture_single_blocks:
                            new_lora[key] = torch.mul(loras[lora_key]["lora_weights"][key], loras[lora_key]["strength"]).to(dtype=torch.bfloat16)
                        continue

                    if temp_strings[1] == "double_blocks":
                        if temp_block_num in mixture_double_blocks:
                            new_lora[key] = torch.mul(loras[lora_key]["lora_weights"][key], loras[lora_key]["strength"]).to(dtype=torch.bfloat16)

            if not save_metadata:
                metadata = None

            if discard_linear:
                temp_lora = new_lora.copy()

                for temp_key in temp_lora:
                    if "linear" in temp_key:
                        del new_lora[temp_key]

            print("Saving Model To: {}...".format(output_filename))
            utils.save_torch_file(new_lora, output_filename, metadata=metadata)
            print("Saving Model To: {}...Done.".format(output_filename))

            save_message = "{}\n{}".format(save_message, "Weights Saved To: {}".format(output_filename))

            new_lora = None
            mm.soft_empty_cache()
            gc.collect()
            time.sleep(1)

        return {"ui": {"tags": ["save_message"]}}

# class WarpedHunyuanMultiLoraMixerPlatinum:
#     @classmethod
#     def INPUT_TYPES(s):
#         return {
#             "required": {
#                 "save_folder": ("STRING", {"default": get_default_output_folder()}),
#                 "model_prefix": ("STRING", {"default": "new_model_hy"}),
#                 "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
#                 "num_output": ("INT", {"default": 1, "min": 1, "max": 100}),
#                 "lora_1": (['None'] + get_lora_list(),),
#                 "strength_1": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
#                 "blocks_1": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
#                 "lora_2": (['None'] + get_lora_list(),),
#                 "strength_2": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
#                 "blocks_2": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
#                 "lora_3": (['None'] + get_lora_list(),),
#                 "strength_3": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
#                 "blocks_3": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
#                 "lora_4": (['None'] + get_lora_list(),),
#                 "strength_4": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
#                 "blocks_4": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
#                 "lora_5": (['None'] + get_lora_list(),),
#                 "strength_5": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
#                 "blocks_5": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
#                 "lora_6": (['None'] + get_lora_list(),),
#                 "strength_6": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
#                 "blocks_6": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
#                 "lora_7": (['None'] + get_lora_list(),),
#                 "strength_7": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
#                 "blocks_7": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
#                 "lora_8": (['None'] + get_lora_list(),),
#                 "strength_8": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
#                 "blocks_8": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
#                 "save_metadata": ("BOOLEAN", {"default": True}),
#                 "convert_to": (["diffusion_model", "transformer"], {"default": "diffusion_model"}),
#                 "max_dimension": ([32, 64, 128], {"default": 128}),
#                 "remove_linear": ("BOOLEAN", {"default": True}),
#                 "remove_single_blocks": ("BOOLEAN", {"default": True}),
#             },
#         }
#
#     RETURN_TYPES = ()
#     OUTPUT_NODE = True
#     OUTPUT_IS_LIST = (True,)
#     FUNCTION = "merge_multiple_loras"
#     CATEGORY = "Warped/Hunyuan/Mixers"
#     DESCRIPTION = "Load and apply multiple LoRA models with different strengths and block types. Model input is required."
#
#     def load_lora(self, lora_name: str, strength: float, blocks_type: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
#         """Load and filter a single LoRA model."""
#         if not lora_name or strength == 0:
#             return {}, {}
#
#         # Get the full path to the LoRA file
#         lora_path = folder_paths.get_full_path("loras", lora_name)
#         if not os.path.exists(lora_path):
#             raise ValueError(f"LoRA file not found: {lora_path}")
#
#         # Load the LoRA weights
#         lora_weights = utils.load_torch_file(lora_path)
#
#         lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
#
#         return lora_weights
#
#     def get_random_key(self, keys):
#         if len(keys) > 0:
#             random_key = random.randint(0, len(keys) - 1)
#             # print("Keys: {}  |  Random Key: {}".format(random_key))
#             return keys[random_key]
#
#         return -1
#
#     def get_mixtures(self, seed, num_output, lora_keys, block_types):
#         random.seed(seed)
#         mixtures = {}
#
#         for i in range(num_output):
#             mixtures["{}".format(i + 1)] = {}
#
#         single_block_loras = []
#         double_block_loras = []
#
#         for key in lora_keys:
#             for mixture_key in mixtures.keys():
#                 mixtures[mixture_key][key] = {"single": [], "double": []}
#
#                 if block_types[key]["has_single_blocks"]:
#                     single_block_loras.append(int(key))
#
#                 if block_types[key]["has_double_blocks"]:
#                     double_block_loras.append(int(key))
#
#         for mixture_key in mixtures.keys():
#             if len(single_block_loras) > 0:
#                 for j in range(40):
#                     random_key = self.get_random_key(single_block_loras)
#                     temp_key = "{}".format(random_key)
#                     mixtures[mixture_key][temp_key]["single"].append(j)
#
#             if len(double_block_loras) > 0:
#                 for j in range(20):
#                     random_key = self.get_random_key(double_block_loras)
#                     temp_key = "{}".format(random_key)
#                     mixtures[mixture_key][temp_key]["double"].append(j)
#
#             i += 1
#
#         at_least_one = False
#
#         while not at_least_one:
#             if len(double_block_loras) > 0:
#                 for key in lora_keys:
#                     lowest = 999
#                     highest = -1
#                     lowest_mixture_key = None
#                     highest_mixture_key = None
#
#                     for mixture_key in mixtures:
#                         if len(mixtures[mixture_key][key]["double"]) < lowest:
#                             lowest = len(mixtures[mixture_key][key]["double"])
#                             lowest_mixture_key = mixture_key
#
#                         if len(mixtures[mixture_key][key]["double"]) > highest:
#                             highest = len(mixtures[mixture_key][key]["double"])
#                             highest_mixture_key = mixture_key
#
#                     if lowest < 1:
#                         print("low key: {}".format(lowest_mixture_key))
#                         print("high key: {}".format(highest_mixture_key))
#                         mixtures[lowest_mixture_key][key]["double"].append(mixtures[highest_mixture_key][key]["double"][0])
#                         print("low mix: {}".format(mixtures[lowest_mixture_key][key]["double"]))
#                         print("high mixes before: {}".format(mixtures[highest_mixture_key][key]["double"]))
#                         mixtures[highest_mixture_key][key]["double"].pop(0)
#                         print("high mixes after: {}".format(mixtures[highest_mixture_key][key]["double"]))
#                     else:
#                         at_least_one = True
#             else:
#                 at_least_one = True
#
#         print("\nMixtures\n")
#
#         block_metadata = ""
#
#         for mixture_key in mixtures.keys():
#             for key in mixtures[mixture_key]:
#                 print("{} | {}: {}".format(mixture_key, key, mixtures[mixture_key][key]))
#
#                 if len(block_metadata) > 0:
#                     block_metadata = "{}  |  {}".format(block_metadata, "{}: {}: {}".format(mixture_key, key, mixtures[mixture_key][key]))
#                 else:
#                     block_metadata = "{}: {}: {}".format(mixture_key, key, mixtures[mixture_key][key])
#
#             print("\n")
#
#         return mixtures, block_metadata
#
#     def determine_lora_block_types(self, loras):
#         block_types = {}
#
#         for lora_key in loras.keys():
#             block_types[lora_key] = { "has_single_blocks": False, "has_double_blocks": False }
#
#             for key in loras[lora_key]["lora_weights"].keys():
#                 if "single_blocks" in key:
#                     block_types[lora_key]["has_single_blocks"] = True
#                 elif "double_blocks" in key:
#                     block_types[lora_key]["has_double_blocks"] = True
#
#                 if block_types[lora_key]["has_single_blocks"] and block_types[lora_key]["has_double_blocks"]:
#                     break
#
#         return block_types
#
#     def merge_multiple_loras(self, save_folder, model_prefix, seed, num_output, lora_1, strength_1, blocks_1, lora_2, strength_2, blocks_2, lora_3, strength_3, blocks_3, lora_4, strength_4, blocks_4,
#                             lora_5, strength_5, blocks_5, lora_6, strength_6, blocks_6, lora_7, strength_7, blocks_7, lora_8, strength_8, blocks_8, save_metadata=True, convert_to="diffusion_model",
#                             max_dimension=64, remove_linear=True, remove_single_blocks=True):
#         print("Save_folder: {}".format(save_folder))
#         os.makedirs(save_folder, exist_ok = True)
#
#         temp_loras = {}
#         metadata = {"loras": "{} and {} and {} and {} and {} and {} and {} and {}".format(lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8)}
#         metadata["strengths"] = "{} and {} and {} and {} and {} and {} and {} and {}".format(strength_1, strength_2, strength_3, strength_4, strength_5, strength_6, strength_7, strength_8)
#         metadata["seed"] = "{}".format(seed)
#         metadata["num_output"] = "{}".format(num_output)
#
#         if (lora_1 != "None") and (strength_1 > 0.0):
#             print(lora_1)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_1, 1.0, "all")
#
#             if (blocks_1 == "single_blocks") or (blocks_1 == "double_blocks"):
#                 lora_weights = filter_lora_keys(lora_weights, blocks_1)
#
#             lora_weights = convert_lora_dimensions(max_dimension, lora_weights)
#
#             print_it = True
#             for key in lora_weights.keys():
#                 if print_it:
#                     print("LORA 1: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
#                     print_it = False
#
#                 lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)
#
#             temp_loras["1"] = {"lora_weights": lora_weights, "strength": strength_1}
#
#             lora_weights = None
#
#         if (lora_2 != "None") and (strength_2 > 0.0):
#             print(lora_2)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_2, 1.0, "all")
#
#             if (blocks_2 == "single_blocks") or (blocks_2 == "double_blocks"):
#                 lora_weights = filter_lora_keys(lora_weights, blocks_2)
#
#             lora_weights = convert_lora_dimensions(max_dimension, lora_weights)
#
#             print_it = True
#             for key in lora_weights.keys():
#                 if print_it:
#                     print("LORA 2: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
#                     print_it = False
#
#                 lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)
#
#             temp_loras["2"] = {"lora_weights": lora_weights, "strength": strength_2}
#
#             lora_weights = None
#
#         if (lora_3 != "None") and (strength_3 > 0.0):
#             print(lora_3)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_3, 1.0, "all")
#
#             if (blocks_3 == "single_blocks") or (blocks_3 == "double_blocks"):
#                 lora_weights = filter_lora_keys(lora_weights, blocks_3)
#
#             lora_weights = convert_lora_dimensions(max_dimension, lora_weights)
#
#             print_it = True
#             for key in lora_weights.keys():
#                 if print_it:
#                     print("LORA 3: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
#                     print_it = False
#
#                 lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)
#
#             temp_loras["3"] = {"lora_weights": lora_weights, "strength": strength_3}
#
#             lora_weights = None
#
#         if (lora_4 != "None") and (strength_4 > 0.0):
#             print(lora_4)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_4, 1.0, "all")
#
#             if (blocks_4 == "single_blocks") or (blocks_4 == "double_blocks"):
#                 lora_weights = filter_lora_keys(lora_weights, blocks_4)
#
#             lora_weights = convert_lora_dimensions(max_dimension, lora_weights)
#
#             print_it = True
#             for key in lora_weights.keys():
#                 if print_it:
#                     print("LORA 4: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
#                     print_it = False
#
#                 lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)
#
#             temp_loras["4"] = {"lora_weights": lora_weights, "strength": strength_4}
#
#             lora_weights = None
#
#         if (lora_5 != "None") and (strength_5 > 0.0):
#             print(lora_5)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_5, 1.0, "all")
#
#             if (blocks_5 == "single_blocks") or (blocks_5 == "double_blocks"):
#                 lora_weights = filter_lora_keys(lora_weights, blocks_5)
#
#             lora_weights = convert_lora_dimensions(max_dimension, lora_weights)
#
#             print_it = True
#             for key in lora_weights.keys():
#                 if print_it:
#                     print("LORA 5: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
#                     print_it = False
#
#                 lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)
#
#             temp_loras["5"] = {"lora_weights": lora_weights, "strength": strength_5}
#
#             lora_weights = None
#
#         if (lora_6 != "None") and (strength_6 > 0.0):
#             print(lora_6)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_6, 1.0, "all")
#
#             if (blocks_6 == "single_blocks") or (blocks_6 == "double_blocks"):
#                 lora_weights = filter_lora_keys(lora_weights, blocks_6)
#
#             lora_weights = convert_lora_dimensions(max_dimension, lora_weights)
#
#             print_it = True
#             for key in lora_weights.keys():
#                 if print_it:
#                     print("LORA 6: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
#                     print_it = False
#
#                 lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)
#
#             temp_loras["6"] = {"lora_weights": lora_weights, "strength": strength_6}
#
#             lora_weights = None
#
#         if (lora_7 != "None") and (strength_7 > 0.0):
#             print(lora_7)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_7, 1.0, "all")
#
#             if (blocks_7 == "single_blocks") or (blocks_7 == "double_blocks"):
#                 lora_weights = filter_lora_keys(lora_weights, blocks_7)
#
#             lora_weights = convert_lora_dimensions(max_dimension, lora_weights)
#
#             print_it = True
#             for key in lora_weights.keys():
#                 if print_it:
#                     print("LORA 7: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
#                     print_it = False
#
#                 lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)
#
#             temp_loras["7"] = {"lora_weights": lora_weights, "strength": strength_7}
#
#             lora_weights = None
#
#         if (lora_8 != "None") and (strength_8 > 0.0):
#             print(lora_8)
#             # Load and filter the LoRA weights
#             lora_weights = self.load_lora(lora_8, 1.0, "all")
#
#             if (blocks_8 == "single_blocks") or (blocks_8 == "double_blocks"):
#                 lora_weights = filter_lora_keys(lora_weights, blocks_8)
#
#             lora_weights = convert_lora_dimensions(max_dimension, lora_weights)
#
#             print_it = True
#             for key in lora_weights.keys():
#                 if print_it:
#                     print("LORA 8: {}  |  Sample Shape: {}".format(key, lora_weights[key].shape))
#                     print_it = False
#
#                 lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)
#
#             temp_loras["8"] = {"lora_weights": lora_weights, "strength": strength_8}
#
#             lora_weights = None
#
#         loras = {}
#
#         for lora_key in temp_loras.keys():
#             loras[lora_key] = {"lora_weights": {}, "strength": temp_loras[lora_key]["strength"]}
#
#             for key in temp_loras[lora_key]["lora_weights"].keys():
#                 new_key = key.replace("transformer.", "diffusion_model.")
#                 loras[lora_key]["lora_weights"][new_key] = temp_loras[lora_key]["lora_weights"][key].to(dtype=torch.bfloat16)
#
#         block_types = self.determine_lora_block_types(loras)
#         merge_mixtures, block_metadata = self.get_mixtures(seed, num_output, loras.keys(), block_types)
#
#         metadata["max_dimension"] = "{}".format(max_dimension)
#         metadata["block_types"] = "{}".format(block_types)
#
#         print("Max Dimension: {}".format(max_dimension))
#
#         save_message = ""
#
#         for mixture_key in merge_mixtures:
#             time.sleep(1)
#             new_lora = {}
#             output_filename = os.path.join(save_folder, "{}_{:05}.safetensors".format(model_prefix, int(mixture_key)))
#
#             metadata["merge_mixture"] = "{}".format(merge_mixtures[mixture_key])
#             # metadata["block_metadata"] = "{}".format(block_metadata[int(mixture_key)])
#
#             for lora_key in loras.keys():
#                 mixture_single_blocks = merge_mixtures[mixture_key][lora_key]["single"]
#                 mixture_double_blocks = merge_mixtures[mixture_key][lora_key]["double"]
#
#                 for key in loras[lora_key]["lora_weights"].keys():
#                     temp_strings = str(key).split('.')
#                     temp_block_num = int(temp_strings[2])
#
#                     if temp_strings[1] == "single_blocks":
#                         if temp_block_num in mixture_single_blocks:
#                             new_lora[key] = torch.mul(loras[lora_key]["lora_weights"][key], loras[lora_key]["strength"]).to(dtype=torch.bfloat16)
#                         continue
#
#                     if temp_strings[1] == "double_blocks":
#                         if temp_block_num in mixture_double_blocks:
#                             new_lora[key] = torch.mul(loras[lora_key]["lora_weights"][key], loras[lora_key]["strength"]).to(dtype=torch.bfloat16)
#
#             dummy_lora = new_lora,copy()
#
#             for key in dummy_lora:
#                 if "single_blocks" in key:
#                     if remove_single_blocks:
#                         del new_lora[key]
#                     elif remove_linear:
#                         if "linear" in key:
#                             del new_lora[key]
#
#                     continue
#
#                 if remove_linear:
#                     if "linear" in key:
#                         del new_lora[key]
#
#             if not save_metadata:
#                 metadata = None
#
#             print("Saving Model To: {}...".format(output_filename))
#             utils.save_torch_file(new_lora, output_filename, metadata=metadata)
#             print("Saving Model To: {}...Done.".format(output_filename))
#
#             save_message = "{}\n{}".format(save_message, "Weights Saved To: {}".format(output_filename))
#
#             new_lora = None
#             mm.soft_empty_cache()
#             gc.collect()
#             time.sleep(1)
#
#         return {"ui": {"tags": ["save_message"]}}

class WarpedWanLoraMerge:
    def __init__(self):
        self.base_output_dir = get_default_wan_output_folder()
        os.makedirs(self.base_output_dir, exist_ok = True)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "save_path": ("STRING", {"default": get_default_wan_output_path()}),
                "lora_1": (['None'] + get_lora_list(),),
                "strength_1": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}),
                "lora_2": (['None'] + get_lora_list(),),
                "strength_2": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}),
                "save_metadata": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    OUTPUT_IS_LIST = (True,)
    FUNCTION = "merge_multiple_loras"
    CATEGORY = "Warped/Wan/Merge"
    DESCRIPTION = "Load and apply multiple LoRA models with different strengths and block types. Model input is required."

    def load_lora(self, lora_name: str, strength: float, blocks_type: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
        """Load and filter a single LoRA model."""
        if not lora_name or strength == 0:
            return {}, {}

        # Get the full path to the LoRA file
        lora_path = folder_paths.get_full_path("loras", lora_name)
        if not os.path.exists(lora_path):
            raise ValueError(f"LoRA file not found: {lora_path}")

        # Load the LoRA weights
        lora_weights = utils.load_torch_file(lora_path)

        return lora_weights

    def merge_multiple_loras(self, save_path, lora_1, strength_1, lora_2, strength_2, save_metadata=True):
        """Load and apply multiple LoRA models."""
        temp_loras = {}
        metadata = {"loras": "{} and {}".format(lora_1, lora_2)}
        metadata["strengths"] = "{} and {}".format(strength_1, strength_2)

        if lora_1 != "None" and strength_1 != 0:
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_1, 1.0, "all")
            temp_loras["1"] = {"lora_weights": lora_weights, "strength": strength_1}

        if lora_2 != "None" and strength_2 != 0:
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora_2, 1.0, "all")
            temp_loras["2"] = {"lora_weights": lora_weights, "strength": strength_2}

        loras = {}

        for lora_key in temp_loras.keys():
            loras[lora_key] = {"lora_weights": {}, "strength": temp_loras[lora_key]["strength"]}

            for key in temp_loras[lora_key]["lora_weights"].keys():
                # new_key = key.replace("transformer.", "diffusion_model.")
                loras[lora_key]["lora_weights"][key] = temp_loras[lora_key]["lora_weights"][key]

        new_lora = {}

        for lora_key in loras.keys():
            for key in loras[lora_key]["lora_weights"].keys():
                if not key in new_lora.keys():
                    new_lora[key] = None
                print("Lora: {}  | Key: {}  |  Shape: {}".format(lora_key, key, loras[lora_key]["lora_weights"][key].shape))

        # Merge The Weighted Key Weights
        for key in new_lora.keys():
            for lora_key in loras.keys():
                if key in loras[lora_key]["lora_weights"].keys():
                    if not new_lora[key] is None:
                        temp_weights = torch.mul(loras[lora_key]["lora_weights"][key], loras[lora_key]["strength"])

                        if new_lora[key].shape[0] < new_lora[key].shape[1]:
                            if temp_weights.shape[0] < new_lora[key].shape[0]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([new_lora[key].shape[0], new_lora[key].shape[1]])
                                padding[:temp_weights.shape[0],:] = temp_weights
                                temp_weights = padding
                            elif temp_weights.shape[0] > new_lora[key].shape[0]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([temp_weights.shape[0], temp_weights.shape[1]])
                                padding[:new_lora[key].shape[0],:] = new_lora[key]
                                new_lora[key] = padding
                        else:
                            if temp_weights.shape[1] < new_lora[key].shape[1]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([new_lora[key].shape[0], new_lora[key].shape[1]])
                                padding[:,:temp_weights.shape[1]] = temp_weights
                                temp_weights = padding
                            elif temp_weights.shape[1] > new_lora[key].shape[1]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([temp_weights.shape[0], temp_weights.shape[1]])
                                padding[:,:new_lora[key].shape[1]] = new_lora[key]
                                new_lora[key] = padding

                        try:
                            new_lora[key] = torch.add(new_lora[key], temp_weights)
                        except Exception as e:
                            raise(e)
                    else:
                        new_lora[key] = torch.mul(loras[lora_key]["lora_weights"][key], loras[lora_key]["strength"])

        if not save_metadata:
            metadata = None

        utils.save_torch_file(new_lora, save_path, metadata=metadata)

        save_message = "Weights Saved To: {}".format(save_path)

        return {"ui": {"tags": [save_message]}}

class WarpedCreateSpecialImageBatch:
    def __init__(self, device="cpu"):
        self.device = device

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    { "image": ("IMAGE", ),
                      "color": ("STRING", {"default": "#000000"}),
                      "batch_size": ("INT", {"default": 1, "min": 1, "max": 1001, "step": 4}),
                      "all_same_image": ("BOOLEAN", {"default": False}),
                    }
                }
    RETURN_TYPES = ("IMAGE", "INT", )
    RETURN_NAMES = ("image", "num_images", )
    FUNCTION = "generate"

    CATEGORY = "Warped/General/Image"

    def generate(self, image, color="#000000", batch_size=1, all_same_image=False):
        image_color = ImageColor.getrgb(color)

        temp_image = tensor2pilSwap(image)
        temp_image = pil2tensorSwap(temp_image[0])
        temp_image = tensor2pilSwap(temp_image)
        temp_image = temp_image[0]

        if not all_same_image:
            dummy_image = temp_image.copy()
            dummy_image.paste((0,0,0), (0, 0, dummy_image.size[0], dummy_image.size[1]))

            image_batch = [temp_image]

            i = 0
            while i < (batch_size - 1):
                image_batch.append(dummy_image.copy())

                i += 1

            intermediate_images = pil2tensorSwap(image_batch)

            final_images = None

            print_it = True
            for image in intermediate_images:
                image = image.unsqueeze(0)
                if not final_images is None:
                    final_images = torch.cat((final_images, image), dim=0)

                    # if print_it:
                    #     print_it = False
                    #     print(final_images)
                else:
                    final_images = image

            return (final_images, batch_size, )

        image_batch = [temp_image]

        i = 0
        while i < (batch_size - 1):
            image_batch.append(temp_image.copy())

            i += 1

        intermediate_images = pil2tensorSwap(image_batch)

        final_images = None

        print_it = True
        for image in intermediate_images:
            image = image.unsqueeze(0)
            if not final_images is None:
                final_images = torch.cat((final_images, image), dim=0)
            else:
                final_images = image

        return (final_images, batch_size, )

class WarpedCreateSpecialImageBatchExt:
    def __init__(self, device="cpu"):
        self.device = device

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    { "image": ("IMAGE", ),
                      "color": ("STRING", {"default": "#000000"}),
                      "batch_size": ("INT", {"default": 1, "min": 1, "max": 1001, "step": 4}),
                      "all_same_image": ("BOOLEAN", {"default": False}),
                    },
                "optional":
                    {
                      "start_image": ("IMAGE", ),
                    }
                }
    RETURN_TYPES = ("IMAGE", "INT", )
    RETURN_NAMES = ("image", "num_images", )
    FUNCTION = "generate"

    CATEGORY = "Warped/General/Image"

    def generate(self, image, color="#000000", batch_size=1, all_same_image=False, start_image=None):
        image_color = ImageColor.getrgb(color)

        if len(image.shape) < 4:
            image = image.unsqueeze(0)

        height = image.shape[1]
        width  = image.shape[2]

        temp_image = image.clone().detach()

        if not all_same_image:
            dummy_image = torch.zeros([batch_size - 1, self.height, self.width, 3], dtype=torch.float32, device=self.offload_device)
            final_images = torch.cat((image, dummy_image), 0)

            return (final_images, batch_size, )

        if start_image is None:
            final_images = temp_image
        else:
            final_images = start_image

        i = 0
        while i < (batch_size - 1):
            final_images = torch.cat((final_images, temp_image.clone().detach()), 0)
            i += 1

        return (final_images, batch_size, )

class WarpedCreateSpecialImageBatchExp:
    def __init__(self, device="cpu"):
        self.device = device

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    { "image": ("IMAGE", ),
                      "batch_size": ("INT", {"default": 1, "min": 1, "max": 1001, "step": 4}),
                      "all_same_image": ("BOOLEAN", {"default": False}),
                    }
                }
    RETURN_TYPES = ("IMAGE", "INT", )
    RETURN_NAMES = ("image", "num_images", )
    FUNCTION = "generate"

    CATEGORY = "Warped/General/Image"

    def generate(self, image, batch_size=1, all_same_image=False):
        if len(image.shape) < 4:
            image = image.unsqueeze(0)

        height = image.shape[1]
        width = image.shape[2]

        if not all_same_image:
            temp_image = torch.ones((batch_size - 1, height, width, image.shape[-1]), device=image.device, dtype=image.dtype) * 0.5
            final_images = torch.cat((image, temp_image), 0)

            return (final_images, batch_size, )

        final_images = image.clone().detach()

        i = 0
        while i < (batch_size - 1):
            final_images = torch.cat((final_images, image.clone().detach()), 0)
            i += 1

        return (final_images, batch_size, )

class WarpedCreateEmptyImageBatch:
    def __init__(self, device="cpu"):
        self.device = device

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    { "batch_size": ("INT", {"default": 1, "min": 1, "max": 1001, "step": 4}),
                      "width": ("INT", {"default": 320, "min": 256, "max": 4096, "step": 16}),
                      "height": ("INT", {"default": 512, "min": 256, "max": 4096, "step": 16}),
                    }
                }
    RETURN_TYPES = ("IMAGE", "INT", )
    RETURN_NAMES = ("image", "num_images", )
    FUNCTION = "generate"

    CATEGORY = "Warped/General/Image"

    def generate(self, batch_size, width, height):
        color="#000000"
        image_color = ImageColor.getrgb(color)

        temp_image = Image.new(mode="RGB", size=(width, height))
        temp_image = pil2tensorSwap(temp_image)

        if len(temp_image.shape) < 4:
            temp_image = temp_image.unsqueeze(0)

        final_images = None
        count = 0

        while count < batch_size:
            if not final_images is None:
                final_images = torch.cat((final_images, temp_image.clone().detach()), dim=0)
            else:
                final_images = temp_image.clone().detach()

            count += 1

        return (final_images, batch_size, )

class WarpedCreateEmptyLatentBatch:
    def __init__(self, device="cpu"):
        self.device = device
        self.offload_device = get_offload_device()

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    { "batch_size": ("INT", {"default": 1, "min": 1, "max": 1001, "step": 4}),
                      "width": ("INT", {"default": 320, "min": 256, "max": 4096, "step": 16}),
                      "height": ("INT", {"default": 512, "min": 256, "max": 4096, "step": 16}),
                    }
                }
    RETURN_TYPES = ("LATENT", "INT", )
    RETURN_NAMES = ("latents", "num_images", )
    FUNCTION = "generate"

    CATEGORY = "Warped/General/Latent"

    def generate(self, batch_size, width, height):
        temp_latent = torch.zeros([1, 16, int(((batch_size - 1) / 4) + 1), int(height // 8), int(width // 8)], dtype=torch.float32, device=self.offload_device)

        print("Empty Latent Batch Shape: {}".format(temp_latent.shape))

        if len(temp_latent.shape) < 5:
            temp_latent = temp_latent.unsqueeze(0)

        return ({"samples": temp_latent}, batch_size, )

class WarpedHunyuanVideoToVideo:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"video_batch": ("IMAGE", ),
                             "positive": ("CONDITIONING", ),
                             "vae": ("VAE", ),
                             "guidance_type": (["v1 (concat)", "v2 (replace)", "custom"], )
                            },
                }

    RETURN_TYPES = ("CONDITIONING", "LATENT")
    RETURN_NAMES = ("positive", "latent")
    FUNCTION = "encode"

    CATEGORY = "Warped/Hunyuan/Conditioning"

    def encode(self, video_batch, positive, vae, guidance_type):
        mm.unload_all_models()

        out_latent = {}

        batch_size = 1
        width = video_batch.shape[2]
        height = video_batch.shape[1]
        length = video_batch.shape[0]

        print("width: {}  |  height: {}  |  length: {}".format(width, height, length))

        video_batch = comfy.utils.common_upscale(video_batch[:length, :, :, :3].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
        latent = partial_encode_tiled(vae, video_batch) #vae.encode(video_batch)

        print("latents Shape: {}".format(latent.shape))

        video_tuple = torch.split(video_batch, 1, dim=0)
        video_split = [item for item in video_tuple]

        start_image = video_split[0]

        video_tuple = None
        video_split = None

        concat_latent_image = vae.encode(start_image)
        mask = torch.ones((1, 1, latent.shape[2], concat_latent_image.shape[-2], concat_latent_image.shape[-1]), device=start_image.device, dtype=start_image.dtype)
        mask[:, :, :((start_image.shape[0] - 1) // 4) + 1] = 0.0

        if guidance_type == "v1 (concat)":
            cond = {"concat_latent_image": concat_latent_image, "concat_mask": mask}
        elif guidance_type == "v2 (replace)":
            cond = {'guiding_frame_index': 0}
            latent[:, :, :concat_latent_image.shape[2]] = concat_latent_image
            out_latent["noise_mask"] = mask
        elif guidance_type == "custom":
            cond = {"ref_latent": concat_latent_image}

        positive = node_helpers.conditioning_set_values(positive, cond)

        out_latent["samples"] = latent
        return (positive, out_latent)

class WarpedSamplerCustomAdv:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"image": ("IMAGE", ),
                    "vae": ("VAE", ),
                    "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                    "guider": ("GUIDER", ),
                    "sampler": ("SAMPLER", ),
                    "sigmas": ("SIGMAS", ),
                    "enc_tile_size": ("INT", {"default": 128, "min": 64, "max": 4096, "step": 64}),
                    "enc_overlap": ("INT", {"default": 64, "min": 0, "max": 4096, "step": 32}),
                    "enc_temporal_size": ("INT", {"default": 64, "min": 8, "max": 4096, "step": 4, "tooltip": "Only used for video VAEs: Amount of frames to encode at a time."}),
                    "enc_temporal_overlap": ("INT", {"default": 8, "min": 4, "max": 4096, "step": 4, "tooltip": "Only used for video VAEs: Amount of frames to overlap."}),
                    "dec_tile_size": ("INT", {"default": 256, "min": 64, "max": 4096, "step": 32}),
                    "dec_overlap": ("INT", {"default": 64, "min": 0, "max": 4096, "step": 32}),
                    "dec_temporal_size": ("INT", {"default": 64, "min": 8, "max": 4096, "step": 4, "tooltip": "Only used for video VAEs: Amount of frames to decode at a time."}),
                    "dec_temporal_overlap": ("INT", {"default": 8, "min": 4, "max": 4096, "step": 4, "tooltip": "Only used for video VAEs: Amount of frames to overlap."}),
                    "skip_frames": ("INT", {"default": 0, "min": 0, "max": 32, "step": 4}),
                    "noise_scale": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 1.00, "step": 0.01}),
                    },
                "optional":
                    {"scaling_strength": ("FLOAT", {"default": 1.0}),
                    "output_latents": ("BOOLEAN", {"default": False}),
                    }
                }

    RETURN_TYPES = ("IMAGE", "LATENT", "INT", "STRING", "BOOLEAN", )
    RETURN_NAMES = ("images", "latents", "seed", "generation_status", "valid_output", )

    FUNCTION = "sample"

    CATEGORY = "Warped/General/Sampling"

    def sample(self, image, vae, seed, guider, sampler, sigmas, enc_tile_size, enc_overlap, enc_temporal_size, enc_temporal_overlap,
                    dec_tile_size, dec_overlap, dec_temporal_size, dec_temporal_overlap, skip_frames, noise_scale, scaling_strength=1.0, output_latents=False):
        self.device = mm.get_torch_device()
        self.offload_device = get_offload_device()
        mm.unload_all_models()
        gc.collect()
        time.sleep(1)

        self.vae = vae
        self.seed = seed
        self.guider = guider
        self.sampler = sampler
        self.sigmas = sigmas
        self.noise_scale = noise_scale
        self.enc_tile_size = enc_tile_size
        self.enc_overlap = enc_overlap
        self.enc_temporal_size = enc_temporal_size
        self.enc_temporal_overlap = enc_temporal_overlap
        self.dec_tile_size = dec_tile_size
        self.dec_overlap = dec_overlap
        self.dec_temporal_size = dec_temporal_size
        self.dec_temporal_overlap = dec_temporal_overlap
        self.g_output = {}

        callback = self.setup_callbacks()
        disable_pbar = not utils.PROGRESS_BAR_ENABLED

        if len(image.shape) < 4:
            image = image.unsqueeze(0)

        num_frames = image.shape[0]

        self.width = image.shape[2]
        self.height = image.shape[1]
        print("\nWidth is {}  |  Height is {}".format(self.width, self.height))

        generation_status = ""

        latents, noise_latents = self.initialize_frames(image)

        print("-------------------------------------------------------------------------------------------")
        print("WarpedSamplerCustomAdv: Latents Shape: {}  |  Noise Latents Shape: {}".format(latents.shape, noise_latents.shape))
        print("-------------------------------------------------------------------------------------------")

        output_images = None
        output_images_latents = None
        interrupted = False
        valid_output = False

        try:
            mm.soft_empty_cache()
            gc.collect()
            time.sleep(0.1)

            latent = {"samples": latents}

            noise_mask = None
            if "noise_mask" in latent:
                noise_mask = latent["noise_mask"]

            print("Noise Shape: {}  |  Latents Shape: {}".format(noise_latents.shape, latents.shape))
            print("WarpedSamplerCustomAdv: Generating {} Frames in {} Latents....".format(num_frames, latents.shape[2]))

            samples = guider.sample(noise_latents, latents, sampler, sigmas, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=self.seed)
            samples = samples.to(mm.intermediate_device())

            mm.unload_all_models()
            mm.soft_empty_cache()
            gc.collect()
            time.sleep(1)

            if len(samples.shape) < 5:
                samples = samples.unsqueeze(0)

            samples = samples.clone().detach() / scaling_strength

            print("WarpedSamplerCustomAdv: Decoding Latents...")
            output_images = self.decode_tiled(samples)
            samples = None

            output_images = self.process_skip_images(output_images, skip_frames)
            valid_output = True
            print("WarpedSamplerCustomAdv: Decoded Images Shape: {}".format(output_images.shape[0]))

            if output_latents:
                output_images_latents = self.encode_tiled(output_images)

            print("WarpedSamplerCustomAdv: Generating {} frames in {} latents...Done.".format(num_frames, latents.shape[2]))

            samples = None
            latents = None

            mm.soft_empty_cache()
            gc.collect()
            time.sleep(0.1)

            print("*******************************************")
            print("****** WarpedSamplerCustomAdv: Total Images Generated {}".format(output_images.shape[0]))
            print("*******************************************\n")

            if len(output_images.shape) < 4:
                output_images = output_images.unsqueeze(0)

            generation_status = "****** WarpedSamplerCustomAdv: Total Images Generated {} ******".format(output_images.shape[0])

            interrupted = False
        except mm.InterruptProcessingException as ie:
            interrupted = True
            print(f"\nWarpedSamplerCustomAdv: Processing Interrupted.")
            print("WarpedSamplerCustomAdv: Returning only partial results (if any).\n If zero images generated, a blank yellow image will be returned, and valid_output will indicate False.\n")
            mm.unload_all_models()
            mm.soft_empty_cache()
            gc.collect()
            time.sleep(1)

            generation_status = f"\nWarpedSamplerCustomAdv: Processing Interrupted."

            traceback.print_tb(ie.__traceback__, limit=99, file=sys.stdout)

            raise mm.InterruptProcessingException(f"WarpedSamplerCustomAdv: Processing Interrupted.")

            pass

        except BaseException as e:
            print(f"\nWarpedSamplerCustomAdv: Exception During Processing: {str(e)}")
            print("WarpedSamplerCustomAdv: Returning only partial results (if any).\n If zero images generated, a blank red image will be returned, and valid_output will indicate False.\n")
            mm.unload_all_models()
            mm.soft_empty_cache()
            gc.collect()
            time.sleep(1)

            generation_status = f"WarpedSamplerCustomAdv: Exception During Processing: {str(e)}"
            generation_status = "{}{}".format(generation_status, "WarpedSamplerCustomAdv: Returning only partial results (if any).\nIf zero images generated, a blank red image will be returned, and valid_output will indicate False.")

            traceback.print_tb(e.__traceback__, limit=99, file=sys.stdout)

            pass

        callback = None
        guider.model_patcher.model.to(get_offload_device())

        latent = None
        latent_image = None
        noise_mask = None
        samples = None

        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
        gc.collect()
        time.sleep(1)

        mm.unload_all_models()
        mm.soft_empty_cache()
        gc.collect()
        time.sleep(1)

        image = None

        if interrupted:
            temp_image = Image.new('RGB', (self.width, self.height), color = 'yellow')
            image = pil2tensorSwap(temp_image)
            output_images = image
        elif output_images is None:
            temp_image = Image.new('RGB', (self.width, self.height), color = 'red')
            image = pil2tensorSwap(temp_image)
            output_images = image

        if output_images_latents is None:
            output_images_latents = torch.zeros([1, 16, 1, self.height // 8, self.width // 8], dtype=torch.float32, device=self.offload_device)

        return (output_images, {"samples": output_images_latents}, self.seed, generation_status, valid_output, )

    def generate_noise(self, input_latent, generator=None):
        latent_image = input_latent["samples"]
        return warped_prepare_noise(latent_image, self.seed, generator=generator)

    def process_skip_images(self, frames, skip_count):
        if len(frames.shape) < 4:
            frames = frames.unsqueeze(0)

        num_frames = frames.shape[0]

        image_batches_tuple = torch.split(frames, 1, dim=0)
        image_batches_split = [item for item in image_batches_tuple]

        new_video = None
        i = 0

        while i < len(image_batches_split):
            if i < skip_count:
                i += 1
                continue

            if not new_video is None:
                new_video = torch.cat((new_video, image_batches_split[i]), 0)
            else:
                new_video = image_batches_split[i]

            i += 1

        return new_video

    def get_blank_image(self, length=1):
        new_image = torch.zeros([length, self.height, self.width, 3], dtype=torch.float32, device=self.offload_device)
        return new_image

    def get_new_noise(self, length):
        new_noise = torch.zeros([length, self.height, self.width, 3], dtype=torch.float32, device=self.offload_device)

        new_noise = self.encode_tiled(new_noise)

        new_noise = comfy.sample.fix_empty_latent_channels(self.guider.model_patcher, new_noise)

        if len(new_noise) < 5:
            new_latent = new_noise.unsqueeze(0)

        new_noise = self.generate_noise({"samples": new_noise})

        return new_noise

    def pad_noise(self, latent, num_frames=1):
        pad_frames = torch.zeros([1, 16, num_frames, self.height, self.width], dtype=torch.float32, device=self.offload_device)
        pad_frames = torch.cat((latent, pad_frames), 2)

        return pad_frames

    def setup_callbacks(self):
        callback = latent_preview.prepare_callback(self.guider.model_patcher, self.sigmas.shape[-1] - 1, self.g_output)

        return callback

    def encode(self, images):
        if len(images.shape) < 4:
            images = images.unsqueze(0)

        encoded_data = partial_encode_basic(self.vae, images)

        if len(encoded_data.shape) < 5:
            encoded_data.unsqueeze(0)

        return encoded_data

    def encode_as_batched(self, images):
        if len(images.shape) < 4:
            images = images.unsqueze(0)

        print("WarpedSamplerCustomAdv: latents Shape Before Split: {}".format(latents.shape))
        image_batches_tuple = torch.split(images, self.enc_split_sizes, dim=0)
        image_batches_split = [item for item in image_batches_tuple]

        if not (image_batches_split is None) and (len(image_batches_split) > 0):
            print("WarpedSamplerCustomAdv: Splits Length: {}".format(len(image_batches_split)))
            for entry in image_batches_split:
                print(entry.shape)
            print("WarpedSamplerCustomAdv: Splits End")
        else:
            print("++++++ WarpedSamplerCustomAdv: Empty Results Splits ++++++")
            return None

        result_images = None

        print("WarpedSamplerCustomAdv: Encoding Batch Latents...")
        for entry in image_batches_split:
            encoded_data = self.encode(entry)

            if len(encoded_data.shape) < 4:
                encoded_data = encoded_data.unsqueeze(0)

            print("WarpedSamplerCustomAdv: Entry Shape: {}  |  Encoded Entry Shape: {}".format(entry.shape, encoded_data.shape))

            if not result_images is None:
                result_images = torch.cat((result_images, encoded_data), 0)
            else:
                result_images = decoded_data

        print("WarpedSamplerCustomAdv: Encoded Images Shape: {}".format(result_images.shape))
        print("WarpedSamplerCustomAdv: Encoding Batch Latents...Done.")

        return result_images

    def encode_tiled(self, images):
        if len(images.shape) < 4:
            images = images.unsqueze(0)

        encoded_data = partial_encode_tiled(self.vae, images, self.enc_tile_size, self.enc_overlap, self.enc_temporal_size, self.enc_temporal_overlap)

        if len(encoded_data.shape) < 5:
            encoded_data.unsqueeze(0)

        return encoded_data

    def decode_tiled(self, latents):
        decoded_data = partial_decode_tiled(self.vae, latents, self.dec_tile_size, self.dec_overlap, self.dec_temporal_size, self.dec_temporal_overlap)

        if len(decoded_data.shape) < 4:
            decoded_data.unsqueeze(0)

        return decoded_data

    def initialize_noise(self, frame_count, clear_cache=True):
        noise_latents_full = torch.zeros([1, 16, int(frame_count), self.height // 8, self.width // 8], dtype=torch.float32, device=self.offload_device)
        print("WarpedSamplerCustomAdv: Encoded noise_latents_full Shape: {}".format(noise_latents_full.shape))

        if Decimal(self.noise_scale).compare(Decimal(0.00)) != 0:
            noise_latents_full = warped_prepare_noise(noise_latents_full, self.seed)
            print("WarpedSamplerCustomAdv: noise_latents_full Shape: {}".format(noise_latents_full.shape))

            noise_latents_full = torch.mul(noise_latents_full, self.noise_scale)

        if len(noise_latents_full.shape) < 5:
            noise_latents_full.unsqueeze(0)

        noise_latents_full = noise_latents_full.to(dtype=torch.float32, device=self.offload_device)

        if clear_cache:
            mm.soft_empty_cache()
            gc.collect()
            time.sleep(0.1)

        return noise_latents_full

    def initialize_frames(self, image):
        if len(image.shape) < 4:
            image = image.unsqueeze(0)

        print("WarpedSamplerCustomAdv: Decoded latents_full Shape: {}".format(image.shape))
        latents_full = self.encode_tiled(image)

        if len(latents_full.shape) < 5:
            latents_full.unsqueeze(0)

        print("WarpedSamplerCustomAdv: Encoded latents_full Shape: {}".format(latents_full.shape))

        latents_full = latents_full.to(dtype=torch.float32, device=self.offload_device)

        noise_latents_full = self.initialize_noise(latents_full.shape[2])

        mm.soft_empty_cache()
        gc.collect()
        time.sleep(0.1)

        return latents_full, noise_latents_full

class WarpedSamplerCustomAdvLatent:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"latent": ("LATENT", ),
                    "vae": ("VAE", ),
                    "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                    "guider": ("GUIDER", ),
                    "sampler": ("SAMPLER", ),
                    "sigmas": ("SIGMAS", ),
                    "dec_tile_size": ("INT", {"default": 256, "min": 64, "max": 4096, "step": 32}),
                    "dec_overlap": ("INT", {"default": 64, "min": 0, "max": 4096, "step": 32}),
                    "dec_temporal_size": ("INT", {"default": 64, "min": 8, "max": 4096, "step": 4, "tooltip": "Only used for video VAEs: Amount of frames to decode at a time."}),
                    "dec_temporal_overlap": ("INT", {"default": 8, "min": 4, "max": 4096, "step": 4, "tooltip": "Only used for video VAEs: Amount of frames to overlap."}),
                    "skip_frames": ("INT", {"default": 0, "min": 0, "max": 32, "step": 4}),
                    "noise_scale": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 1.00, "step": 0.01}),
                    },
                "optional":
                    {"scaling_strength": ("FLOAT", {"default": 1.0}),
                    "is_v2v": ("BOOLEAN", {"default": False}),
                    "output_latents": ("BOOLEAN", {"default": False}),
                    }
                }

    RETURN_TYPES = ("IMAGE", "LATENT", "INT", "STRING", "BOOLEAN", )
    RETURN_NAMES = ("images", "latents", "seed", "generation_status", "valid_output", )

    FUNCTION = "sample"

    CATEGORY = "Warped/General/Sampling"

    def sample(self, latent, vae, seed, guider, sampler, sigmas, dec_tile_size, dec_overlap, dec_temporal_size, dec_temporal_overlap,
                    skip_frames, noise_scale, scaling_strength=1.0, is_v2v=False, output_latents=False):
        self.device = mm.get_torch_device()
        self.offload_device = get_offload_device()
        mm.unload_all_models()
        gc.collect()
        time.sleep(1)

        self.vae = vae
        self.seed = seed
        self.guider = guider
        self.sampler = sampler
        self.sigmas = sigmas
        self.noise_scale = noise_scale
        self.dec_tile_size = dec_tile_size
        self.dec_overlap = dec_overlap
        self.dec_temporal_size = dec_temporal_size
        self.dec_temporal_overlap = dec_temporal_overlap
        self.g_output = {}
        self.init_mode = "image"

        if is_v2v:
            self.init_mode = "video"

        callback = self.setup_callbacks()
        disable_pbar = not utils.PROGRESS_BAR_ENABLED

        print("\nSigmas: {}".format(self.sigmas))

        latents = latent["samples"]

        if len(latents.shape) < 5:
            latents = latents.unsqueeze(0)

        num_frames = int(((latents.shape[2] - 1) * 4) + 1)
        self.latents_depth = latents.shape[1]

        print("latents_depth: {}".format(self.latents_depth))

        self.width = latents.shape[4]
        self.height = latents.shape[3]
        print("\nDecoded Width is {}  |  Decoded Height is {}".format(int(self.width * 8), int(self.height * 8)))

        generation_status = ""

        noise_latents = self.initialize_frames(latents, init_mode=self.init_mode)

        print("-------------------------------------------------------------------------------------------")
        print("WarpedSamplerCustomAdvLatent: Latents Shape: {}  |  Noise Latents Shape: {}".format(latents.shape, noise_latents.shape))
        print("-------------------------------------------------------------------------------------------")

        output_images = None
        output_images_latents = None
        interrupted = False
        valid_output = False

        try:
            mm.soft_empty_cache()
            gc.collect()
            time.sleep(0.1)

            noise_mask = None
            if "noise_mask" in latent:
                noise_mask = latent["noise_mask"]

            print("Noise Shape: {}  |  Latents Shape: {}".format(noise_latents.shape, latents.shape))
            print("WarpedSamplerCustomAdvLatent: Generating {} Frames in {} Latents....".format(num_frames, latents.shape[2]))

            samples = guider.sample(noise_latents, latents, sampler, sigmas, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=self.seed)
            samples = samples.to(mm.intermediate_device())

            mm.unload_all_models()
            mm.soft_empty_cache()
            gc.collect()
            time.sleep(1)

            if len(samples.shape) < 5:
                samples = samples.unsqueeze(0)

            samples = samples.clone().detach() / scaling_strength

            print("WarpedSamplerCustomAdvLatent: Decoding Latents...")
            output_images = self.decode_tiled(samples)
            samples = None
            valid_output = True

            output_images = self.process_skip_images(output_images, skip_frames)
            print("WarpedSamplerCustomAdvLatent: Decoded Images Shape: {}".format(output_images.shape[0]))

            if output_latents:
                output_images_latents = self.encode_tiled(output_images)

            print("WarpedSamplerCustomAdvLatent: Generating {} frames in {} latents...Done.".format(num_frames, latents.shape[2]))

            samples = None
            latents = None

            mm.soft_empty_cache()
            gc.collect()
            time.sleep(0.1)

            print("*******************************************")
            print("****** WarpedSamplerCustomAdvLatent: Total Images Generated {}".format(output_images.shape[0]))
            print("*******************************************\n")

            if len(output_images.shape) < 4:
                output_images = output_images.unsqueeze(0)

            generation_status = "****** WarpedSamplerCustomAdvLatent: Total Images Generated {} ******".format(output_images.shape[0])

            interrupted = False
        except mm.InterruptProcessingException as ie:
            interrupted = True
            print(f"\nWarpedSamplerCustomAdvLatent: Processing Interrupted.")
            print("WarpedSamplerCustomAdvLatent: Returning only partial results (if any).\n If zero images generated, a blank yellow image will be returned, and valid_output will indicate False.\n")
            mm.unload_all_models()
            mm.soft_empty_cache()
            gc.collect()
            time.sleep(1)

            generation_status = f"\nWarpedSamplerCustomAdvLatent: Processing Interrupted."

            traceback.print_tb(ie.__traceback__, limit=99, file=sys.stdout)

            raise mm.InterruptProcessingException(f"WarpedSamplerCustomAdvLatent: Processing Interrupted.")
            pass

        except BaseException as e:
            print(f"\nWarpedSamplerCustomAdvLatent: Exception During Processing: {str(e)}")
            print("WarpedSamplerCustomAdvLatent: Returning only partial results (if any).\n If zero images generated, a blank red image will be returned, and valid_output will indicate False.\n")
            mm.unload_all_models()
            mm.soft_empty_cache()
            gc.collect()
            time.sleep(1)

            generation_status = f"WarpedSamplerCustomAdvLatent: Exception During Processing: {str(e)}"
            generation_status = "{}{}".format(generation_status, "WarpedSamplerCustomAdvLatent: Returning only partial results (if any).\nIf zero images generated, a blank red image will be returned, and valid_output will indicate False.")

            traceback.print_tb(e.__traceback__, limit=99, file=sys.stdout)

            pass

        callback = None
        guider.model_patcher.model.to(get_offload_device())

        latent = None
        latent_image = None
        noise_mask = None
        samples = None

        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
        gc.collect()
        time.sleep(1)

        mm.unload_all_models()
        mm.soft_empty_cache()
        gc.collect()
        time.sleep(1)

        image = None

        if interrupted and (output_images is None):
            temp_image = Image.new('RGB', (self.width * 8, self.height * 8), color = 'yellow')
            image = pil2tensorSwap(temp_image)
            output_images = image
        elif output_images is None:
            temp_image = Image.new('RGB', (self.width * 8, self.height * 8), color = 'red')
            image = pil2tensorSwap(temp_image)
            output_images = image

        if output_images_latents is None:
            output_images_latents = torch.zeros([1, self.latents_depth, 1, self.height, self.width], dtype=torch.float32, device=self.offload_device)

        return (output_images, {"samples": output_images_latents}, self.seed, generation_status, valid_output, )

    def generate_noise(self, input_latent, generator=None):
        latent_image = input_latent["samples"]
        return warped_prepare_noise(latent_image, self.seed, generator=generator)

    def process_skip_images(self, frames, skip_count):
        if len(frames.shape) < 4:
            frames = frames.unsqueeze(0)

        num_frames = frames.shape[0]

        image_batches_tuple = torch.split(frames, 1, dim=0)
        image_batches_split = [item for item in image_batches_tuple]

        new_video = None
        i = 0

        while i < len(image_batches_split):
            if i < skip_count:
                i += 1
                continue

            if not new_video is None:
                new_video = torch.cat((new_video, image_batches_split[i]), 0)
            else:
                new_video = image_batches_split[i]

            i += 1

        return new_video

    def get_blank_image(self, length=1):
        new_image = torch.zeros([length, self.height, self.width, 3], dtype=torch.float32, device=self.offload_device)
        return new_image

    def get_new_noise(self, length):
        new_noise = torch.zeros([length, self.height, self.width, 3], dtype=torch.float32, device=self.offload_device)

        new_noise = self.encode_tiled(new_noise)

        new_noise = comfy.sample.fix_empty_latent_channels(self.guider.model_patcher, new_noise)

        if len(new_noise) < 5:
            new_latent = new_noise.unsqueeze(0)

        new_noise = self.generate_noise({"samples": new_noise})

        return new_noise

    def pad_noise(self, latent, num_frames=1):
        pad_frames = torch.zeros([1, self.latents_depth, num_frames, self.height, self.width], dtype=torch.float32, device=self.offload_device)
        pad_frames = torch.cat((latent, pad_frames), 2)

        return pad_frames

    def setup_callbacks(self):
        callback = latent_preview.prepare_callback(self.guider.model_patcher, self.sigmas.shape[-1] - 1, self.g_output)

        return callback

    def decode_tiled(self, latents):
        decoded_data = partial_decode_tiled(self.vae, latents, self.dec_tile_size, self.dec_overlap, self.dec_temporal_size, self.dec_temporal_overlap)

        if len(decoded_data.shape) < 4:
            decoded_data.unsqueeze(0)

        return decoded_data

    def initialize_noise(self, frame_count, clear_cache=True):
        noise_latents_full = torch.zeros([1, self.latents_depth, int(frame_count), self.height, self.width], dtype=torch.float32, device=self.offload_device)

        print("WarpedSamplerCustomAdvLatent: Encoded noise_latents_full Shape: {}".format(noise_latents_full.shape))

        if Decimal(self.noise_scale).compare(Decimal(0.00)) != 0:
            noise_latents_full = warped_prepare_noise(noise_latents_full, self.seed)
            print("WarpedSamplerCustomAdvLatent: noise_latents_full Shape: {}".format(noise_latents_full.shape))

            noise_latents_full = torch.mul(noise_latents_full, self.noise_scale)

        if len(noise_latents_full.shape) < 5:
            noise_latents_full.unsqueeze(0)

        noise_latents_full = noise_latents_full.to(dtype=torch.float32, device=self.offload_device)

        if clear_cache:
            mm.soft_empty_cache()
            gc.collect()
            time.sleep(0.1)

        return noise_latents_full

    def initialize_frames(self, latents, init_mode="image"):
        if len(latents.shape) < 5:
            latents = latents.unsqueeze(0)

        print("WarpedSamplerCustomAdvLatent: Encoded latents_full Shape: {}".format(latents.shape))
        latents_full = latents.clone().detach()
        latents_full = latents_full.to(dtype=torch.float32, device=self.offload_device)

        if init_mode == "image":
            noise_latents_full = self.initialize_noise(latents_full.shape[2])
        else:
            if Decimal(self.noise_scale).compare(Decimal(0.00)) != 0:
                noise_latents_full = warped_prepare_noise(latents_full, self.seed)
                print("WarpedSamplerCustomAdvLatent: noise_latents_full Shape: {}".format(noise_latents_full.shape))

                noise_latents_full = torch.mul(noise_latents_full, self.noise_scale)

            if len(noise_latents_full.shape) < 5:
                noise_latents_full.unsqueeze(0)

            noise_latents_full = noise_latents_full.to(dtype=torch.float32, device=self.offload_device)

        mm.soft_empty_cache()
        gc.collect()
        time.sleep(0.1)

        return noise_latents_full

class WarpedWanImageToVideo:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"positive": ("CONDITIONING", ),
                             "negative": ("CONDITIONING", ),
                             "vae": ("VAE", ),
                             "width": ("INT", {"default": 832, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}),
                             "height": ("INT", {"default": 480, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}),
                             "length": ("INT", {"default": 81, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 4}),
                             "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
                },
                "optional": {"clip_vision_output": ("CLIP_VISION_OUTPUT", ),
                             "start_image": ("IMAGE", ),
                }}

    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT", "INT", )
    RETURN_NAMES = ("positive", "negative", "latent", "num_frames", )
    FUNCTION = "encode"

    CATEGORY = "Warped/Wan/Conditioning"

    def encode(self, positive, negative, vae, width, height, length, batch_size, start_image=None, clip_vision_output=None):
        latent = torch.zeros([batch_size, 16, ((length - 1) // 4) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device())
        if start_image is not None:
            start_image = comfy.utils.common_upscale(start_image[:length].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
            image = torch.mul(torch.ones((length, height, width, start_image.shape[-1]), device=start_image.device, dtype=start_image.dtype), 0.5)
            image[:start_image.shape[0]] = start_image

            concat_latent_image = vae.encode(image[:, :, :, :3])

            mask = torch.ones((1, 1, latent.shape[2], concat_latent_image.shape[-2], concat_latent_image.shape[-1]), device=start_image.device, dtype=start_image.dtype)
            mask[:, :, :((start_image.shape[0] - 1) // 4) + 1] = 0.0

            positive = node_helpers.conditioning_set_values(positive, {"concat_latent_image": concat_latent_image, "concat_mask": mask})
            negative = node_helpers.conditioning_set_values(negative, {"concat_latent_image": concat_latent_image, "concat_mask": mask})

        if clip_vision_output is not None:
            positive = node_helpers.conditioning_set_values(positive, {"clip_vision_output": clip_vision_output})
            negative = node_helpers.conditioning_set_values(negative, {"clip_vision_output": clip_vision_output})

        out_latent = {}
        out_latent["samples"] = latent
        return (positive, negative, out_latent, length, )

PROMPT_TEMPLATE_ENCODE_VIDEO_I2V = (
    "<|start_header_id|>system<|end_header_id|>\n\n<image>\nDescribe the video by detailing the following aspects according to the reference image: "
    "1. The main content and theme of the video."
    "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects."
    "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects."
    "4. background environment, light, style and atmosphere."
    "5. camera angles, movements, and transitions used in the video:<|eot_id|>\n\n"
    "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>"
    "<|start_header_id|>assistant<|end_header_id|>\n\n"
)

class WarpedSamplerCustomScripted:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {
                    "positive": ("STRING", {"default": None} ),
                    "negative": ("STRING", {"default": None}),
                    "clip": ("CLIP", ),
                    "preferred_frame_count": ("INT", {"default": 81, "min": 17, "max": nodes.MAX_RESOLUTION, "step": 4}),
                    "preferred_batch_size": ("INT", {"default": 17, "min": 17, "max": 301, "step": 4}),
                    "use_batch_size": (["next_lowest", "next_highest", "closest", "exact"], {"default": "next_lowest", "tooltip": "Number of frames generated may be impacted by choice."}),
                    "vae": ("VAE", ),
                    "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                    "model": ("MODEL", ),
                    "sampler": ("SAMPLER", ),
                    "sigmas": ("SIGMAS", ),
                    "dec_tile_size": ("INT", {"default": 256, "min": 64, "max": 4096, "step": 32}),
                    "dec_overlap": ("INT", {"default": 64, "min": 0, "max": 4096, "step": 32}),
                    "dec_temporal_size": ("INT", {"default": 64, "min": 8, "max": 4096, "step": 4, "tooltip": "Only used for video VAEs: Amount of frames to decode at a time."}),
                    "dec_temporal_overlap": ("INT", {"default": 8, "min": 4, "max": 4096, "step": 4, "tooltip": "Only used for video VAEs: Amount of frames to overlap."}),
                    "skip_frames": ("INT", {"default": 0, "min": 0, "max": 32, "step": 4}),
                    "noise_scale": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 1.00, "step": 0.01}),
                    "mode": (["Hunyuan", "Wan 2.1", "Wan 2.2 TI2V", "Wan 2.2 Standard"], {"default": "Hunyuan"}),
                    "hunyuan_guidance_type": (["v1 (concat)", "v2 (replace)", "custom"], {"default": "v2 (replace)", "tooltip": "Only used if mode is \"Hunyuan\"."}),
                    "hunyuan_image_interleave": ("INT", {"default": 1, "min": 1, "max": 512, "tooltip": "How much the image influences things vs the text prompt. Higher number means more influence from the text prompt."}),
                    },
                "optional":
                    {"noise_strength": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}),
                    "blend_frames": ("INT", {"default":0, "min":0, "max": 16, "step": 1}),
                    "t2v_width": ("INT", {"default":480, "min":256, "max": 1280, "step": 16}),
                    "t2v_height": ("INT", {"default":720, "min":256, "max": 1280, "step": 16}),
                    "scaling_strength": ("FLOAT", {"default": 1.0}),
                    "clip_vision_model": ("CLIP_VISION", ),
                    "start_image": ("IMAGE", ),
                    "flux_guidance": ("FLOAT", {"default": 7, "min": 0.0, "max": 100.0, "step": 0.1}),
                    "use_flux_guidance": ("BOOLEAN", {"default": False}),
                    "dummy_frames": ("INT", {"default":17, "min":17, "max": 301, "step": 4, "tooltip": "Number of frames to generate in dummy batch."}),
                    "gen_dummy": ("BOOLEAN", {"default": False, "tooltip": "For t2v or i2v only. Will generate a dummy batch to obtain a starting image for main generation."}),
                    "gen_dummy_only": ("BOOLEAN", {"default": False, "tooltip": "Will generate dummy batch only."}),
                    "use_dummy_image": (["first", "middle", "last", "random", "all"], {"default": "last", "tooltip": "Which dummy batch image to start main generation."}),
                    "output_latents": ("BOOLEAN", {"default": False}),
                    "verbose_messaging": ("BOOLEAN", {"default": False}),
                    "secondary_sigmas": ("SIGMAS", {"default": None}),
                    "secondary_model": ("MODEL", {"default": None}),
                    "batches_script": ("WARPEDSCRIPTS", {"default": None}),
                    }
                }

    RETURN_TYPES = ("IMAGE", "LATENT", "INT", "STRING", "IMAGE", "BOOLEAN",)
    RETURN_NAMES = ("images", "latents", "seed", "generation_status", "image_used", "valid_output",)

    FUNCTION = "sample"

    CATEGORY = "Warped/General/Sampling"

    def sample(self, positive, negative, clip, preferred_frame_count, preferred_batch_size, use_batch_size, vae, seed, model, sampler, sigmas, dec_tile_size, dec_overlap, dec_temporal_size, dec_temporal_overlap,
                    skip_frames, noise_scale, mode="Hunyuan", hunyuan_guidance_type="v2 (replace)", hunyuan_image_interleave=2, noise_strength=1.0, blend_frames=5, t2v_width=480, t2v_height=720, scaling_strength=1.0,
                    clip_vision_model=None, start_image=None, flux_guidance=3.5, dummy_frames=17, gen_dummy=False, gen_dummy_only=False, use_dummy_image="last", use_flux_guidance=False, output_latents=False, verbose_messaging=False,
                    secondary_sigmas=None, secondary_model=None, batches_script=None):

        batch_scripts = batches_script

        if (positive is None) and (batches_script is None):
            raise ValueError("positive and batches_script cannot both be None.")
        # elif not batches_script is None:
        #     if verbose_messaging:
        #         print("\n{}\n".format(batches_script))
        #
        #     batch_scripts = {}
        #     temp_strings = batches_script.split("||")
        #
        #     for entry in temp_strings:
        #         temp_script = entry.strip().split("|")
        #         batch_scripts[temp_script[0].strip()] = temp_script[1].strip()
        #
        #     if verbose_messaging:
        #         print("\n{}\n".format(batch_scripts))

        self.device = mm.get_torch_device()
        self.offload_device = get_offload_device()
        mm.unload_all_models()
        gc.collect()
        time.sleep(1)

        self.positive = positive
        self.negative = negative
        self.clip = clip
        self.use_clip = clip.clone()
        self.preferred_frame_count = preferred_frame_count
        self.preferred_batch_size = preferred_batch_size
        self.vae = vae
        self.seed = seed
        self.sampler = sampler
        self.sigmas = sigmas
        self.secondary_sigmas = secondary_sigmas
        self.noise_scale = noise_scale
        self.mode = mode
        self.dec_tile_size = dec_tile_size
        self.dec_overlap = dec_overlap
        self.dec_temporal_size = dec_temporal_size
        self.dec_temporal_overlap = dec_temporal_overlap
        self.clip_vision = clip_vision_model
        self.noise_strength = noise_strength
        self.blend_frames = blend_frames
        self.t2v_width = t2v_width
        self.t2v_height = t2v_height
        self.dummy_frames = dummy_frames
        self.dummy_latents = int((dummy_frames - 1) // 4) + 1
        self.gen_dummy = gen_dummy
        self.gen_dummy_only = gen_dummy_only
        self.use_dummy_image = use_dummy_image
        self.skip_frames = skip_frames
        self.verbose_messaging = verbose_messaging
        self.model = model
        self.secondary_model = secondary_model
        self.batches_script = batches_script
        self.g_output = {}

        disable_pbar = not utils.PROGRESS_BAR_ENABLED

        self.use_sigmas = self.sigmas

        self.guider = WarpedGuider_Basic()
        self.guider.set_model(self.model)
        callback = self.setup_callbacks()

        self.latent_window_size, self.batch_count, self.truncated_frame_count = self.get_latent_window_size(preferred_batch_size, self.preferred_frame_count, use_batch_size=use_batch_size)
        self.batch_frame_count = int(((self.truncated_frame_count - 1) // self.batch_count) + 1)
        print("latent_window_size: {}  |  batch_count: {}  |  truncated_frame_count: {}  | batch_frame_count: {}".format(self.latent_window_size, self.batch_count, self.truncated_frame_count, self.batch_frame_count))

        is_i2v = True
        is_v2v = False

        if start_image is None:
            start_image = torch.zeros([1, self.t2v_height, self.t2v_width, 3], dtype=torch.float32, device=self.offload_device)
            is_i2v = False
        else:
            if start_image.shape[0] > 1:
                is_v2v = True
                is_i2v = False

        if self.mode == "Wan 2.2 TI2V":
            self.latent_depth = 48
        else:
            self.latent_depth = 16

        print("start_image Shape: {}".format(start_image.shape))

        if self.mode != "Wan 2.2 TI2V":
            self.width = int(start_image.shape[2] // 8)
            self.height = int(start_image.shape[1] // 8)
        else:
            self.width = int(start_image.shape[2] // 16)
            self.height = int(start_image.shape[1] // 16)

        noise, dummy_noise = self.setup_latent_noise()

        output_images = []
        output_images_latents = None
        interrupted = False
        generation_status = ""
        total_images_generated = 0
        is_dummy_section = False
        is_first_section = False
        is_last_section = False
        switch_guider = True
        last_image = None
        positive_prompt = positive
        valid_output = False

        if not self.gen_dummy:
            print("Generating {} Batches".format(self.batch_count))
        else:
            if not self.gen_dummy_only:
                print("Generating Dummy Batch Plus {} Batches".format(self.batch_count))
            else:
                print("Generating Dummy Batch Only")

        generation_batches = list(range(self.batch_count))

        if self.gen_dummy and (not self.gen_dummy_only):
            generation_batches = [0] + generation_batches
            is_dummy_section = True
        else:
            generation_batches = [0]
            is_dummy_section = True

        print("generation_batches: {}".format(generation_batches))

        try:
            for i, generation_batch in enumerate(generation_batches):
                batch_number = generation_batches[i]
                batch_key = "{}".format(batch_number)

                if is_dummy_section:
                    print("Generating Dummy Batch.")
                    is_first_section = False

                    if not self.gen_dummy_only:
                        is_last_section = False
                    else:
                        is_last_section = True

                    if not batch_scripts is None:
                        positive_prompt = batch_scripts["dummy"]
                        print("\nDummy Section Positive Prompt: {}\n".format(positive_prompt))
                        last_prompt = positive_prompt
                else:
                    print("Generating Batch {}.".format(batch_number))

                    is_first_section = batch_number == min(generation_batches)
                    is_last_section = batch_number == max(generation_batches)

                    if not batch_scripts is None:
                        if batch_key in batch_scripts.keys():
                            positive_prompt = batch_scripts[batch_key]
                            last_prompt = positive_prompt
                        else:
                            positive_prompt = last_prompt

                        print("\nBatch: {} Positive Prompt: {}\n".format(batch_number, positive_prompt))

                print("is_first_section: {}  |  is_last_section: {}  |  is_dummy_section: {}".format(is_first_section, is_last_section, is_dummy_section))

                mm.soft_empty_cache()
                gc.collect()
                time.sleep(0.1)

                image_embeds = None

                print("start_image Shape: {}".format(start_image.shape))

                if is_i2v or is_v2v:
                    # print("GENERAL HERE 1")

                    if len(start_image.shape) < 4:
                        start_image = start_image.unsqueeze(0)

                    if mode != "Wan 2.2 TI2V":
                        print("Performing clip_vision_encode.")
                        image_embeds = self.clip_vision_encode(start_image.clone().detach())
                        print("Performing clip_vision_encode...Done!")

                if not is_dummy_section:
                    noise_latents = noise[batch_number]
                else:
                    noise_latents = dummy_noise

                print("noise_latents Shape: {}".format(noise_latents.shape))

                noise_mask = None

                if is_i2v:
                    if (not self.secondary_model is None) and switch_guider:
                        switch_guider = False
                        callback = None
                        self.guider = WarpedGuider_Basic()
                        self.guider.set_model(self.secondary_model)
                        callback = self.setup_callbacks()
                        self.use_clip = self.clip.clone()

                if mode == "Hunyuan":
                    if is_i2v or is_v2v:
                        temp_cond = self.hunyuan_text_encode(self.use_clip, image_embeds, positive_prompt, hunyuan_image_interleave)
                        positive_cond, latents = self.hunyuan_encode(temp_cond, self.vae, start_image.shape[2], start_image.shape[1], ((noise_latents.shape[2] - 1) * 4) + 1, 1, hunyuan_guidance_type, start_image=start_image)

                        if "noise_mask" in latents:
                            noise_mask = latents["noise_mask"] #* scaling_strength

                        if is_v2v:
                            scaling_strength = vae_scaling_factor

                        latents = latents["samples"] * scaling_strength
                    else:
                        positive_cond = self.do_text_encode(self.use_clip, positive_prompt)
                        latents = torch.zeros([1, 16, noise_latents.shape[2], self.height, self.width], dtype=torch.float32, device=self.offload_device)

                        latents = latents * scaling_strength

                    if use_flux_guidance:
                        positive_cond = self.apply_flux_guidance(positive_cond, flux_guidance)

                    self.guider.set_conds_single(positive_cond)
                else:
                    if not is_i2v:
                        # print("HERE WAN 1")
                        positive_cond, negative_cond = self.wan_text_encode(self.use_clip, positive_text=positive_prompt, negative_text=negative)
                        positive_cond, negative_cond, latents = self.wan_encode(positive_cond, negative_cond, self.vae, start_image.shape[2], start_image.shape[1], ((noise_latents.shape[2] - 1) * 4) + 1, 1, start_image=None, clip_vision_output=None)

                        if "noise_mask" in latents:
                            noise_mask = latents["noise_mask"] * scaling_strength

                        latents = latents["samples"] * scaling_strength
                    else:
                        # print("HERE WAN 2")

                        temp_cond, negative_cond = self.wan_text_encode(self.use_clip, positive_text=positive_prompt, negative_text=negative)
                        positive_cond, negative_cond, latents = self.wan_encode(temp_cond, negative_cond, self.vae, start_image.shape[2], start_image.shape[1], ((noise_latents.shape[2] - 1) * 4) + 1, 1, start_image=start_image, clip_vision_output=image_embeds)

                        if "noise_mask" in latents:
                            noise_mask = latents["noise_mask"] * scaling_strength

                        latents = latents["samples"] * scaling_strength

                    self.guider.set_conds_both(positive_cond, negative_cond)


                print("latents Shape: {}".format(latents.shape))

                num_frames = int(((latents.shape[2] - 1) * 4) + 1)

                print("\nDecoded Width is {}  |  Decoded Height is {}".format(int(self.width * 8), int(self.height * 8)))

                print("-------------------------------------------------------------------------------------------")
                print("WarpedSamplerCustomScripted: Latents Shape: {}  |  Noise Latents Shape: {}".format(latents.shape, noise_latents.shape))
                print("-------------------------------------------------------------------------------------------")

                print("Noise Shape: {}  |  Latents Shape: {}".format(noise_latents.shape, latents.shape))
                print("WarpedSamplerCustomScripted: Generating {} Frames in {} Latents....".format(num_frames, latents.shape[2]))

                samples = self.guider.sample(noise_latents, latents, self.sampler, self.use_sigmas, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=self.seed)
                samples = samples.to(mm.intermediate_device())

                mm.unload_all_models()
                mm.soft_empty_cache()
                gc.collect()
                time.sleep(1)

                if len(samples.shape) < 5:
                    samples = samples.unsqueeze(0)

                samples = samples.clone().detach() / scaling_strength # (scaling_strength * vae_scaling_factor)

                print("WarpedSamplerCustomScripted: Decoding Latents...")
                decoded_images = self.decode_tiled(samples)

                if not is_dummy_section or (is_dummy_section and ((self.use_dummy_image == "all") or self.gen_dummy_only)):
                    output_images.append(decoded_images)
                    total_images_generated += decoded_images.shape[0]
                    valid_output = True

                samples = None

                # output_images = self.process_skip_images(output_images, skip_frames)
                print("WarpedSamplerCustomScripted: Decoded Images Shape: {}".format(decoded_images.shape[0]))

                # if output_latents:
                #     output_images_latents = self.encode_tiled(output_images)

                print("WarpedSamplerCustomScripted: Generating {} frames in {} latents...Done.".format(num_frames, latents.shape[2]))

                if not is_i2v and not is_v2v:
                    if not self.secondary_model is None:
                        self.guider.model_patcher.model.to(get_offload_device())

                        if not self.secondary_sigmas is None:
                            self.use_sigmas = self.secondary_sigmas

                if not is_last_section:
                    decoded_tuple = torch.split(decoded_images, 1, dim=0)
                    decoded_split = [item for item in decoded_tuple]
                    decoded_images = decoded_images.to(get_offload_device())
                    decoded_images = None

                    if (not is_dummy_section) or (is_dummy_section and ((self.use_dummy_image == "last") or (self.use_dummy_image == "all"))):
                        start_image = decoded_split[len(decoded_split) - 1].clone().detach()
                    elif is_dummy_section:
                        if self.use_dummy_image == "first":
                            start_image = decoded_split[0].clone().detach()
                        elif self.use_dummy_image == "middle":
                            start_image = decoded_split[int(len(decoded_split) // 2)].clone().detach()
                        else:
                            random.seed(self.seed)
                            start_image = dummy_split[random.randrange(0, len(decoded_split) - 1, 1)].clone().detach()

                    if not is_v2v:
                        is_i2v = True

                if last_image is None:
                    last_image = start_image.clone().detach()

                samples = None
                latents = None

                mm.soft_empty_cache()
                gc.collect()
                time.sleep(0.1)

                print("*******************************************")
                print("****** WarpedSamplerCustomScripted: Total Images Generated {}".format(total_images_generated))
                print("*******************************************\n")

                generation_status = "****** WarpedSamplerCustomScripted: Total Images Generated {} ******".format(total_images_generated)

                interrupted = False

                if is_dummy_section:
                    if self.gen_dummy_only:
                        break

                    is_dummy_section = False
        except mm.InterruptProcessingException as ie:
            interrupted = True
            print(f"\nWarpedSamplerCustomScripted: Processing Interrupted.")
            print("WarpedSamplerCustomScripted: Returning only partial results (if any).\n If zero images generated, a blank yellow image will be returned, and valid_output will return False.\n")
            mm.unload_all_models()
            mm.soft_empty_cache()
            gc.collect()
            time.sleep(1)

            generation_status = f"\nWarpedSamplerCustomScripted: Processing Interrupted."

            traceback.print_tb(ie.__traceback__, limit=99, file=sys.stdout)

            # raise mm.InterruptProcessingException(f"WarpedSamplerCustomScripted: Processing Interrupted.")
            pass

        except BaseException as e:
            print(f"\nWarpedSamplerCustomScripted: Exception During Processing: {str(e)}")
            print("WarpedSamplerCustomScripted: Returning only partial results (if any).\n If zero images generated, a blank red image will be returned, and valid_output will return False.\n")
            mm.unload_all_models()
            mm.soft_empty_cache()
            gc.collect()
            time.sleep(1)

            generation_status = f"WarpedSamplerCustomScripted: Exception During Processing: {str(e)}"
            generation_status = "{}{}".format(generation_status, "WarpedSamplerCustomScripted: Returning only partial results (if any).\nIf zero images generated, a blank red image will be returned, and valid_output will return False.")

            traceback.print_tb(e.__traceback__, limit=99, file=sys.stdout)

            pass

        callback = None
        self.guider.model_patcher.model.to(get_offload_device())

        latent = None
        latent_image = None
        noise_mask = None
        samples = None

        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
        gc.collect()
        time.sleep(1)

        mm.unload_all_models()
        mm.soft_empty_cache()
        gc.collect()
        time.sleep(1)

        image = None

        if interrupted and not valid_output:
            temp_image = Image.new('RGB', (self.width * 8, self.height * 8), color = 'yellow')
            image = pil2tensorSwap(temp_image)
            final_images = image
        elif not valid_output:
            temp_image = Image.new('RGB', (self.width * 8, self.height * 8), color = 'red')
            image = pil2tensorSwap(temp_image)
            final_images = image
        else:
            final_images = self.assemble_final_result(output_images)

        if output_images_latents is None:
            output_images_latents = torch.zeros([1, 16, 1, self.height, self.width], dtype=torch.float32, device=self.offload_device)

        return (final_images, {"samples": output_images_latents}, self.seed, generation_status, last_image, valid_output,)

    def assemble_final_result(self, image_batches):
        if self.blend_frames < 1:
            resulting_images = None
            for entry in image_batches:
                if not resulting_images is None:
                    resulting_images = torch.cat((resulting_images, entry), 0)
                else:
                    resulting_images = entry

                entry.to(device=self.offload_device)
                entry = None
        else:
            blend_value = 1.0 / self.blend_frames
            i = 0
            while i < (len(image_batches) - 1):
                alpha_blend_val = blend_value
                blend_count = self.blend_frames

                image_batches_tuple = torch.split(image_batches[i], 1, dim=0)
                image_batches_split = [item for item in image_batches_tuple]
                image1 = image_batches_split[len(image_batches_split) - 1]
                # image1 = image_batches_split[len(image_batches_split) - 4]
                image_batches_tuple = None
                image_batches_split = None

                image_batches_tuple = torch.split(image_batches[i + 1], 1, dim=0)
                image_batches_split = [item for item in image_batches_tuple]
                image2 = image_batches_split[0]
                image_batches_tuple = None
                image_batches_split = None

                image1 = tensor2pilSwap(image1)[0]
                image2 = tensor2pilSwap(image2)[0]

                blend_latents = None

                while blend_count > 0:
                    blended_image = Image.blend(image1, image2, alpha_blend_val)
                    temp_latent = pil2tensorSwap(blended_image)
                    blended_image = None

                    if len(temp_latent.shape) < 4:
                        temp_latent = temp_latent.unsqueeze(0)

                    if not blend_latents is None:
                        blend_latents = torch.cat((blend_latents, temp_latent), 0)
                    else:
                        blend_latents = temp_latent

                    alpha_blend_val += blend_value
                    blend_count -= 1

                image_batches_tuple = torch.split(image_batches[i], image_batches[i].shape[0] - 3, dim=0)
                image_batches_split = [item for item in image_batches_tuple]
                image_batches_tuple = None

                image_batches[i] = torch.cat((image_batches_split[0], blend_latents), 0)
                blend_latents = None
                image_batches_split = None

                self.cleanup(unload_models=False, cleanup_models=False, cleanup_cuda=True)

                i += 1

            resulting_images = None
            for entry in image_batches:
                if not resulting_images is None:
                    resulting_images = torch.cat((resulting_images, entry), 0)
                else:
                    resulting_images = entry

                entry.to(device=self.offload_device)

            image_batches = None

            self.cleanup(unload_models=False, cleanup_models=False, cleanup_cuda=True)


        print("assemble_final_result: Full decoded images count: {}".format(resulting_images.shape[0]))

        if self.skip_frames < 1:
            return resulting_images

        skipped_frames = 1

        image_batches_tuple = torch.split(resulting_images, 1, dim=0)
        image_batches_split = [item for item in image_batches_tuple]

        resulting_images = None

        for image in image_batches_split:
            if skipped_frames <= self.skip_frames:
                skipped_frames += 1
                continue

            if not resulting_images is None:
                resulting_images = torch.cat((resulting_images, image), 0)
            else:
                resulting_images = image

        print("assemble_final_result: Final decoded images count: {}".format(resulting_images.shape[0]))

        return resulting_images

    def cleanup(self, unload_models=False, cleanup_models=False, cleanup_cuda=False):
        if unload_models:
            mm.unload_all_models()

        if cleanup_models:
            mm.cleanup_models()

        if cleanup_cuda:
            mm.soft_empty_cache()

        gc.collect()
        time.sleep(1)

        return

    def generate_noise(self, input_latent, generator=None):
        latent_image = input_latent["samples"]
        return warped_prepare_noise(latent_image, self.seed, generator=generator)

    def setup_latent_noise(self):
        dummy_noise = None

        total_noise_frames = 0
        rnd = torch.Generator("cpu").manual_seed(self.seed)

        if self.gen_dummy:
            total_noise_frames += self.dummy_latents

        total_noise_frames += (self.latent_window_size * self.batch_count) - self.batch_count + 1

        temp_noise = torch.randn((1, self.latent_depth, total_noise_frames, self.height, self.width), generator=rnd, device=rnd.device).to(device=rnd.device, dtype=torch.float32)

        noise = []

        latent_batches_tuple = torch.split(temp_noise, self.dummy_latents, dim=2)
        latent_batches_split = [item for item in latent_batches_tuple]

        dummy_noise = latent_batches_split[0].clone().detach()

        temp_noise = None

        i = 0

        for entry in latent_batches_split:
            if i != 0:
                if not temp_noise is None:
                    temp_noise = torch.cat((temp_noise, entry.clone().detach()), 2)
                else:
                    temp_noise = entry.clone().detach()

            i += 1

        latent_batches_tuple = None
        latent_batches_tuple = None

        latent_batches_tuple = torch.split(temp_noise, self.latent_window_size - 1, dim=2)
        latent_batches_split = [item for item in latent_batches_tuple]
        latent_batches_tuple = None

        for entry in latent_batches_split:
            noise.append(entry.clone().detach())

        noise[len(noise) - 2] = torch.cat((noise[len(noise) - 2], noise[len(noise) - 1]), 2)
        del noise[-1]

        latent_batches_split = None

        if self.verbose_messaging:
            if not dummy_noise is None:
                print("setup_latent_noise: Dummy Noise Shape: {}".format(dummy_noise.shape))

            i = 0
            for entry in noise:
                print("setup_latent_noise: Batch: {}  |  noise Shape: {}".format(i, entry.shape))
                i += 1

            print("\n")

        return noise, dummy_noise

    def get_latent_window_size(self, preferred_batch_size, frame_count, use_batch_size="next_lowest"):
        latent_size_factor = 4

        if self.verbose_messaging:
            print("get_latent_window_size: preferred_batch_size: {}".format(preferred_batch_size))
            print("get_latent_window_size: frame_count: {}".format(frame_count))
            print("get_latent_window_size: use_batch_size: {}".format(use_batch_size))

        num_frames = int(((frame_count - 1) // 4) * 4) + 1

        if num_frames != frame_count:
            print(f"Truncating video from {frame_count} to {num_frames} an because odd number of frames is not allowed.")

        if ((num_frames - 1) % (preferred_batch_size - 1)) == 0:
            print("(1) latent_window_size set to: {}".format(self.decoded_to_encoded_length(preferred_batch_size)))
            print("(1) batch_count set to: {}".format(int((num_frames - 1) / (preferred_batch_size - 1))))
            return self.decoded_to_encoded_length(preferred_batch_size), int((num_frames - 1) / (preferred_batch_size - 1)), num_frames

        if use_batch_size == "exact":
            num_frames_final = int(((num_frames - 1) // (preferred_batch_size - 1)) + 1) * (preferred_batch_size - 1)

            if num_frames_final != frame_count:
                print(f"Truncating video from {num_frames} to {num_frames_final} frames for preferred_batch_size compatibility.")

            print("(2) latent_window_size set to: {}".format(self.decoded_to_encoded_length(preferred_batch_size + 1)))
            print("(2) batch_count set to: {}".format(int((num_frames_final - 1) / (preferred_batch_size - 1))))
            return self.decoded_to_encoded_length(preferred_batch_size), int((num_frames_final - 1) / (preferred_batch_size - 1)), num_frames_final

        next_lowest_found = False
        next_highest_found = False
        next_lowest = preferred_batch_size - 1
        next_highest = preferred_batch_size - 1

        if self.verbose_messaging:
            print("get_latent_window_size: Next Lowest Initialized To: {}".format(next_lowest))
            print("get_latent_window_size: Next Highest Initialized To: {}".format(next_highest))

        num_frames_final = int(((num_frames - 1) // 4) * 4) + 1

        if num_frames != num_frames_final:
            print(f"Truncating video from {num_frames} to {num_frames_final} frames for latent_window_size compatibility.")

        if (use_batch_size == "closest") or (use_batch_size == "next_lowest"):
            while next_lowest >= 12:
                next_lowest -= 4

                if (int((num_frames_final - 1) // 4) % next_lowest) == 0:
                    next_lowest_found = True
                    break

            next_lowest += 1

            if next_lowest_found and (use_batch_size == "next_lowest"):
                print("(3) latent_window_size set to: {}".format(self.decoded_to_encoded_length(next_lowest + 1)))
                print("(3) batch_count set to: {}".format(int((num_frames_final - 1) / next_lowest)))
                return self.decoded_to_encoded_length(next_lowest + 1), int((num_frames_final - 1) / next_lowest), num_frames_final

        while next_highest <= 156:
            next_highest += 4

            if (int((num_frames_final - 1) // 4) % next_highest) == 0:
                next_highest_found = True
                break

        if next_highest_found and (use_batch_size == "next_highest"):
            print("(4) latent_window_size set to: {}".format(self.decoded_to_encoded_length(next_highest + 1)))
            print("(4) batch_count set to: {}".format(int((num_frames_final - 1) / next_highest)))
            return self.decoded_to_encoded_length(next_highest + 1), int((num_frames_final - 1) / next_highest), num_frames_final

        if next_highest_found and next_lowest_found:
            if (preferred_batch_size - next_lowest) <= (next_highest - preferred_batch_size):
                print("(5) latent_window_size set to: {}".format(self.decoded_to_encoded_length(next_lowest + 1)))
                print("(5) batch_count set to: {}".format(int((num_frames_final - 1) / next_lowest)))
                return self.decoded_to_encoded_length(next_lowest + 1), int((num_frames_final - 1) / next_lowest), num_frames_final
            elif (next_highest - preferred_batch_size) < (preferred_batch_size - next_lowest):
                print("(6) latent_window_size set to: {}".format(self.decoded_to_encoded_length(next_highest + 1)))
                print("(6) batch_count set to: {}".format(int((num_frames_final - 1) / next_highest)))
                return self.decoded_to_encoded_length(next_highest + 1), int((num_frames_final - 1) / next_highest), num_frames_final

        print("Unable to find a compatible latent_window_size for number of frames = {} and preferred_batch_size = {}.".format(frame_count, preferred_batch_size))
        print("Recalculating Number Of Frames Based On preferred_batch_size of: {}".format(preferred_batch_size))

        return self.calculate_new_number_of_frames(preferred_batch_size, (((frame_count - 1) // 4) * 4) + 1, use_batch_size)

    def calculate_new_number_of_frames(self, preferred_batch_size, frame_count, use_batch_size):
        working_batch_size = preferred_batch_size - 1
        working_frame_count = frame_count - 1

        next_lowest = next_highest = working_frame_count
        next_lowest_found = False
        next_highest_found = False

        while next_lowest > 37:
            next_lowest -= 4

            if int(next_lowest % working_batch_size) == 0:
                next_lowest_found = True
                break

        if next_lowest_found and (use_batch_size == "next_lowest"):
            return self.decoded_to_encoded_length(preferred_batch_size), int(next_lowest // working_batch_size), next_lowest + 1

        while next_highest < 999997:
            next_highest += 4

            if int(next_highest % working_batch_size) == 0:
                next_highest_found = True
                break

        if next_highest_found and (use_batch_size == "next_highest"):
            return self.decoded_to_encoded_length(preferred_batch_size), int(next_highest // working_batch_size), next_highest + 1

        if next_lowest_found and next_highest_found:
            if (working_frame_count - next_lowest) <= (next_highest - working_frame_count):
                return self.decoded_to_encoded_length(preferred_batch_size), int(next_lowest // working_batch_size), next_lowest + 1

            return self.decoded_to_encoded_length(preferred_batch_size), int(next_highest // working_batch_size), next_highest + 1

        if next_lowest_found:
            return self.decoded_to_encoded_length(preferred_batch_size), int(next_lowest // working_batch_size), next_lowest + 1

        if next_highest_found:
            return self.decoded_to_encoded_length(preferred_batch_size), int(next_highest // working_batch_size), next_highest + 1

        raise ValueError("Unable to find a compatible latent_window_size for number of frames = {} and preferred_batch_size = {}.".format(frame_count, preferred_batch_size))

    def clip_vision_encode(self, image, crop="center"):
        # image_device = image.get_device()
        # image = image.to(device=self.device)

        crop_image = False
        if crop != "center":
            crop_image = False
        output = self.clip_vision.encode_image(image, crop=crop_image)

        # image = image.to(device=image_device)

        mm.soft_empty_cache()
        gc.collect()
        time.sleep(0.25)

        return output

    def encoded_to_decoded_length(self, latent_length):
        if latent_length <= 0:
            return 0

        result_length = ((latent_length - 1) * 4) + 1

        return result_length

    def decoded_to_encoded_length(self, image_length):
        if image_length <= 0:
            return 0

        result_length = int(((image_length - 1) / 4) + 1)

        return result_length

    def apply_flux_guidance(self, conditioning, guidance):
        c = node_helpers.conditioning_set_values(conditioning, {"guidance": guidance})
        return c

    def hunyuan_text_encode(self, clip, clip_vision_output, prompt, image_interleave):
        tokens = clip.tokenize(prompt, llama_template=PROMPT_TEMPLATE_ENCODE_VIDEO_I2V, image_embeds=clip_vision_output.mm_projected, image_interleave=image_interleave)
        return clip.encode_from_tokens_scheduled(tokens)

    def hunyuan_encode(self, positive, vae, width, height, length, batch_size, guidance_type, start_image=None):
        latent = torch.zeros([batch_size, 16, ((length - 1) // 4) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device())
        out_latent = {}

        if start_image is not None:
            start_image = comfy.utils.common_upscale(start_image[:length, :, :, :3].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)

            concat_latent_image = partial_encode_tiled(vae, start_image) # vae.encode(start_image)
            mask = torch.ones((1, 1, latent.shape[2], concat_latent_image.shape[-2], concat_latent_image.shape[-1]), device=start_image.device, dtype=start_image.dtype)
            mask[:, :, :((start_image.shape[0] - 1) // 4) + 1] = 0.0

            if guidance_type == "v1 (concat)":
                cond = {"concat_latent_image": concat_latent_image, "concat_mask": mask}
            elif guidance_type == "v2 (replace)":
                cond = {'guiding_frame_index': 0}
                latent[:, :, :concat_latent_image.shape[2]] = concat_latent_image
                out_latent["noise_mask"] = mask
            elif guidance_type == "custom":
                cond = {"ref_latent": concat_latent_image}

            positive = node_helpers.conditioning_set_values(positive, cond)

        out_latent["samples"] = latent
        return positive, out_latent

    def wan_text_encode(self, clip, positive_text="", negative_text=""):
        print("wan_text_encode: Loading clip model to device: {}".format(clip.patcher.load_device))
        clip.patcher.model.to(device=clip.patcher.load_device)

        print("wan_text_encode: Encoding Prompts...")
        positive_conditioning = self.do_text_encode(clip, positive_text)
        negative_conditioning = self.do_text_encode(clip, negative_text)
        print("wan_text_encode: Encoding Prompts...Done.")

        print("wan_text_encode: Unloading clip model to device: {}".format(get_offload_device()))
        clip.patcher.model.to(device=get_offload_device())

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        gc.collect()
        time.sleep(1)

        return positive_conditioning, negative_conditioning

    def wan_encode_image(self, vae, width, height, length, batch_size, start_image=None):
        latent = torch.zeros([1, self.latent_depth, ((length - 1) // 4) + 1, height // 16, width // 16], device=comfy.model_management.intermediate_device())

        if start_image is None:
            out_latent = {}
            out_latent["samples"] = latent
            return out_latent

        mask = torch.ones([latent.shape[0], 1, ((length - 1) // 4) + 1, latent.shape[-2], latent.shape[-1]], device=comfy.model_management.intermediate_device())

        if start_image is not None:
            start_image = comfy.utils.common_upscale(start_image[:length].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
            latent_temp = vae.encode(start_image)
            latent[:, :, :latent_temp.shape[-3]] = latent_temp
            mask[:, :, :latent_temp.shape[-3]] *= 0.0

        out_latent = {}
        latent_format = comfy.latent_formats.Wan22()
        latent = latent_format.process_out(latent) * mask + latent * (1.0 - mask)
        out_latent["samples"] = latent.repeat((batch_size, ) + (1,) * (latent.ndim - 1))
        out_latent["noise_mask"] = mask.repeat((batch_size, ) + (1,) * (mask.ndim - 1))

        return out_latent


    def encode(self, clip, text):
        if clip is None:
            raise RuntimeError("ERROR: clip input is invalid: None\n\nIf the clip is from a checkpoint loader node your checkpoint does not contain a valid clip or text encoder model.")
        tokens = clip.tokenize(text)
        return clip.encode_from_tokens_scheduled(tokens)

    def do_text_encode(self, clip, text):

        if clip is None:
            raise RuntimeError("ERROR: clip input is invalid: None\n\nIf the clip is from a checkpoint loader node your checkpoint does not contain a valid clip or text encoder model.")

        tokens = clip.tokenize(text)
        return_encoding = clip.encode_from_tokens_scheduled(tokens)

        return return_encoding

    def wan_encode(self, positive, negative, vae, width, height, length, batch_size, start_image=None, clip_vision_output=None):
        if self.mode == "Wan 2.1" or self.mode == "Wan 2.2 Standard":
            latent = torch.zeros([batch_size, 16, ((length - 1) // 4) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device())
            if start_image is not None:
                start_image = comfy.utils.common_upscale(start_image[:length].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
                image = torch.mul(torch.ones((length, height, width, start_image.shape[-1]), device=start_image.device, dtype=start_image.dtype), 0.5)
                image[:start_image.shape[0]] = start_image

                concat_latent_image = vae.encode(image[:, :, :, :3])

                mask = torch.ones((1, 1, latent.shape[2], concat_latent_image.shape[-2], concat_latent_image.shape[-1]), device=start_image.device, dtype=start_image.dtype)
                mask[:, :, :((start_image.shape[0] - 1) // 4) + 1] = 0.0

                positive = node_helpers.conditioning_set_values(positive, {"concat_latent_image": concat_latent_image, "concat_mask": mask})
                negative = node_helpers.conditioning_set_values(negative, {"concat_latent_image": concat_latent_image, "concat_mask": mask})

            if clip_vision_output is not None:
                positive = node_helpers.conditioning_set_values(positive, {"clip_vision_output": clip_vision_output})
                negative = node_helpers.conditioning_set_values(negative, {"clip_vision_output": clip_vision_output})

            out_latent = {}
            out_latent["samples"] = latent
        else:
            out_latent = self.wan_encode_image(vae, width, height, length, batch_size, start_image)

        return positive, negative, out_latent

    def process_skip_images(self, frames, skip_count):
        if len(frames.shape) < 4:
            frames = frames.unsqueeze(0)

        num_frames = frames.shape[0]

        image_batches_tuple = torch.split(frames, 1, dim=0)
        image_batches_split = [item for item in image_batches_tuple]

        new_video = None
        i = 0

        while i < len(image_batches_split):
            if i < skip_count:
                i += 1
                continue

            if not new_video is None:
                new_video = torch.cat((new_video, image_batches_split[i]), 0)
            else:
                new_video = image_batches_split[i]

            i += 1

        return new_video

    def get_blank_image(self, length=1):
        new_image = torch.zeros([length, self.height, self.width, 3], dtype=torch.float32, device=self.offload_device)
        return new_image

    def get_new_noise(self, length):
        new_noise = torch.zeros([length, self.height, self.width, 3], dtype=torch.float32, device=self.offload_device)

        new_noise = self.encode_tiled(new_noise)

        new_noise = comfy.sample.fix_empty_latent_channels(self.guider.model_patcher, new_noise)

        if len(new_noise) < 5:
            new_latent = new_noise.unsqueeze(0)

        new_noise = self.generate_noise({"samples": new_noise})

        return new_noise

    def pad_noise(self, latent, num_frames=1):
        pad_frames = torch.zeros([1, 16, num_frames, self.height, self.width], dtype=torch.float32, device=self.offload_device)
        pad_frames = torch.cat((latent, pad_frames), 2)

        return pad_frames

    def setup_callbacks(self):
        callback = latent_preview.prepare_callback(self.guider.model_patcher, self.sigmas.shape[-1] - 1, self.g_output)

        return callback

    def decode_tiled(self, latents):
        decoded_data = partial_decode_tiled(self.vae, latents, self.dec_tile_size, self.dec_overlap, self.dec_temporal_size, self.dec_temporal_overlap)

        if len(decoded_data.shape) < 4:
            decoded_data.unsqueeze(0)

        return decoded_data

    def initialize_noise(self, frame_count, clear_cache=True):
        noise_latents_full = torch.zeros([1, 16, int(frame_count), self.height, self.width], dtype=torch.float32, device=self.offload_device)
        print("WarpedSamplerCustomScripted: Encoded noise_latents_full Shape: {}".format(noise_latents_full.shape))

        if Decimal(self.noise_scale).compare(Decimal(0.00)) != 0:
            noise_latents_full = warped_prepare_noise(noise_latents_full, self.seed)
            print("WarpedSamplerCustomScripted: noise_latents_full Shape: {}".format(noise_latents_full.shape))

            noise_latents_full = torch.mul(noise_latents_full, self.noise_scale)

        if len(noise_latents_full.shape) < 5:
            noise_latents_full.unsqueeze(0)

        noise_latents_full = noise_latents_full.to(dtype=torch.float32, device=self.offload_device)

        if clear_cache:
            mm.soft_empty_cache()
            gc.collect()
            time.sleep(0.1)

        return noise_latents_full

    def initialize_frames(self, latents):
        if len(latents.shape) < 5:
            latents = latents.unsqueeze(0)

        print("WarpedSamplerCustomScripted: Encoded latents_full Shape: {}".format(latents.shape))
        latents_full = latents.clone().detach()
        latents_full = latents_full.to(dtype=torch.float32, device=self.offload_device)

        noise_latents_full = self.initialize_noise(latents_full.shape[2])

        mm.soft_empty_cache()
        gc.collect()
        time.sleep(0.1)

        return noise_latents_full

def sd3_patch(model, shift, multiplier=1000):
    m = model.clone()

    sampling_base = model_sampling.ModelSamplingDiscreteFlow
    sampling_type = model_sampling.CONST

    class ModelSamplingAdvanced(sampling_base, sampling_type):
        pass

    result_model_sampling = ModelSamplingAdvanced(model.model.model_config)
    result_model_sampling.set_parameters(shift=shift, multiplier=multiplier)
    m.add_object_patch("model_sampling", result_model_sampling)
    return m

class WarpedCreateSpecialImageBatchFromVideo:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"video_path": ("STRING", {"default": ""}),
                    "last_n_frames": ("INT", {"default": 16}),
                    "num_frames": ("INT", {"default": 61, "min": 5, "max": 1000001, "step": 4}),
                    },
                }
    RETURN_TYPES = ("IMAGE", "IMAGE", "INT", )
    RETURN_NAMES = ("image_batch", "first_image", "num_frames",)
    FUNCTION = "generate"

    CATEGORY = "Warped/General/Video"

    def generate(self, video_path, last_n_frames, num_frames):
        cap = cv2.VideoCapture(video_path)

        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        print('WarpedCreateSpecialImageBatchFromVideo: width = %d' % width)
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        print('WarpedCreateSpecialImageBatchFromVideo: height = %d' % height)
        length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        print('WarpedCreateSpecialImageBatchFromVideo: length = %d' % length)

        batched_images = None
        first_image = None
        last_image = None
        frame_count = 0
        starting_index = int(length - last_n_frames - 1)

        if starting_index < 0:
            starting_index = 0

        if ((starting_index + num_frames) > length) or (num_frames == 0):
            temp_num_frames = length - starting_index

        try:
            skip = 0

            if starting_index > 0:
                skip = length - starting_index

            if starting_index > 0:
                while(cap.isOpened()) and (skip < starting_index):
                    _, frameorig = cap.read()
                    skip += 1

            while (cap.isOpened()) and (frame_count < temp_num_frames):
                frame_count += 1

                # Take each frame
                _, frameorig = cap.read()

                color_coverted = cv2.cvtColor(frameorig, cv2.COLOR_BGR2RGB)
                pil_image = Image.fromarray(color_coverted)

                if first_image is None:
                    first_image_pil = pil_image.copy()
                    first_image = pil2tensorSwap(first_image_pil)

                last_image_pil = pil_image.copy()
                temp_image = pil2tensorSwap(pil_image)
                last_image = pil2tensorSwap(last_image_pil)

                if len(temp_image.shape) < 4:
                    temp_image = temp_image.unsqueeze(0)

                if not batched_images is None:
                    batched_images = torch.cat((batched_images, temp_image), 0)
                else:
                    batched_images = temp_image
        except:
            print("WarpedCreateSpecialImageBatchFromVideo: Exception During Video File Read.")
        finally:
            cap.release()

        if len(last_image.shape) < 4:
            last_image = last_image.unsqueeze(0)

        while frame_count < num_frames:
            batched_images = torch.cat((batched_images, last_image.clone().detach()), 0)
            frame_count += 1

        print("WarpedCreateSpecialImageBatchFromVideo: Batched Images Shape: {}".format(batched_images.shape))

        return (batched_images, first_image, int(batched_images.shape[0]))

class WarpedBundleAllVideoImages:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"video_path": ("STRING", {"default": ""}),
                    "use_gpu": ("BOOLEAN", {"default": True}),
                    },
                }
    RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "INT", "INT", "INT", "FLOAT", )
    RETURN_NAMES = ("image_batch", "first_image", "last_image", "num_frames", "width", "height", "fps", )
    FUNCTION = "generate"

    CATEGORY = "Warped/General/Video"

    def generate(self, video_path, use_gpu=False):
        if use_gpu:
            device = mm.get_torch_device()
        else:
            device = get_offload_device()

        print("WarpedBundleAllVideoImages: device: {}".format(device))

        cap = cv2.VideoCapture(video_path)

        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        print('WarpedBundleAllVideoImages: width = %d' % width)
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        print('WarpedBundleAllVideoImages: height = %d' % height)
        length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        print('WarpedBundleAllVideoImages: length = %d' % length)
        fps = float(cap.get(cv2.CAP_PROP_FPS))
        fps = round(fps)
        print('WarpedBundleAllVideoImages: fps = %d' % fps)

        print("WarpedBundleAllVideoImages: Width: {}  |  Height: {}".format(width, height))

        batched_images = None
        num_frames = 0
        last_image = None
        first_image = None

        try:
            while(cap.isOpened()):
                # Take each frame
                _, frameorig = cap.read()

                color_coverted = cv2.cvtColor(frameorig, cv2.COLOR_BGR2RGB)
                pil_image = Image.fromarray(color_coverted)
                temp_image = pil2tensorSwap(pil_image, device=device)

                if len(temp_image.shape) < 4:
                    temp_image = temp_image.unsqueeze(0)

                if first_image is None:
                    first_image = temp_image.clone().detach()

                last_image = temp_image.clone().detach()

                if not batched_images is None:
                    batched_images = torch.cat((batched_images, temp_image), 0)
                else:
                    batched_images = temp_image

                temp_image.to(get_offload_device())
                temp_image = None

                num_frames += 1

                if num_frames % 20 == 0:
                    print("WarpedBundleAllVideoImages: Frames Read: {}".format(num_frames))
        except:
            print("WarpedBundleAllVideoImages: Exception During Video File Read.")
        finally:
            cap.release()

        batched_images = batched_images.to(get_offload_device())

        mm.soft_empty_cache()
        gc.collect()
        time.sleep(1)

        return (batched_images, first_image, last_image, num_frames, width, height, float(fps), )

def augmentation_add_noise(image, noise_aug_strength, seed=None):
    if not seed is None:
        torch.manual_seed(seed)

    sigma = torch.ones((image.shape[0],)).to(image.device, image.dtype) * noise_aug_strength
    image_noise = torch.randn_like(image) * sigma[:, None, None, None]
    image_noise = torch.where(image==-1, torch.zeros_like(image), image_noise)
    image_out = image + image_noise
    return image_out

class WarpedImageNoiseAugmentation:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
                "noise_aug_strength": ("FLOAT", {"default": None, "min": 0.0, "max": 100.0, "step": 0.001}),
                "seed": ("INT", {"default": 123,"min": 0, "max": 0xffffffffffffffff, "step": 1}),
            }
        }

    RETURN_TYPES = ("IMAGE", "INT", )
    RETURN_NAMES = ("image", "seed", )
    FUNCTION = "add_noise"

    CATEGORY = "Warped/General/Image"

    def add_noise(self, image, noise_aug_strength, seed):
        image_out = augmentation_add_noise(image, noise_aug_strength, seed)
        return image_out, seed,

def augmentation_patch_model(model, latent, index, strength, start_percent, end_percent):
    def outer_wrapper(samples, index, start_percent, end_percent):
        def unet_wrapper(apply_model, args):
            steps = args["c"]["transformer_options"]["sample_sigmas"]
            inp, timestep, c = args["input"], args["timestep"], args["c"]
            matched_step_index = (steps == timestep).nonzero()
            if len(matched_step_index) > 0:
                current_step_index = matched_step_index.item()
            else:
                for i in range(len(steps) - 1):
                    # walk from beginning of steps until crossing the timestep
                    if (steps[i] - timestep) * (steps[i + 1] - timestep) <= 0:
                        current_step_index = i
                        break
                else:
                    current_step_index = 0
            current_percent = current_step_index / (len(steps) - 1)
            if samples is not None:
                if start_percent <= current_percent <= end_percent:
                    inp[:, :, [index], :, :] = samples[:, :, [0], :, :].to(inp)
                else:
                    inp[:, :, [index], :, :] = torch.zeros(1)
            return apply_model(inp, timestep, **c)
        return unet_wrapper

    samples = latent["samples"] * 0.476986 * strength
    m = model.clone()
    m.set_model_unet_function_wrapper(outer_wrapper(samples, index, start_percent, end_percent))

    return m, 0.476986 * strength

class WarpedLeapfusionHunyuanI2V:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL",),
                "latent": ("LATENT",),
                "index": ("INT", {"default": 0, "min": -1, "max": 1000, "step": 1,"tooltip": "The index of the latent to be replaced. 0 for first frame and -1 for last"}),
                "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The start percentage of steps to apply"}),
                "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The end percentage of steps to apply"}),
                "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.001}),
            }
        }

    RETURN_TYPES = ("MODEL", "FLOAT", )
    RETURN_NAMES = ("model", "scale_factor", )
    FUNCTION = "patch"

    CATEGORY = "Warped/Hunyuan/LeapFusion"

    def patch(self, model, latent, index, strength, start_percent, end_percent):
        new_model, scaling_strength = augmentation_patch_model(model, latent, index, strength, start_percent, end_percent)
        return (new_model, scaling_strength,)

class WarpedSaveAnimatedPng:
    def __init__(self):
        self.output_dir = folder_paths.get_output_directory()
        self.temp_dir = folder_paths.get_temp_directory()
        self.type = "output"
        self.prefix_append = ""

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"images": ("IMAGE", ),
                     "png_filename": ("STRING", {"default": ""}),
                     "fps": ("FLOAT", {"default": 6.0, "min": 0.01, "max": 1000.0, "step": 0.01}),
                     "compress_level": ("INT", {"default": 4, "min": 0, "max": 9})
                     },
                }

    RETURN_TYPES = ()
    FUNCTION = "save_images"
    OUTPUT_NODE = True
    CATEGORY = "Warped/General/Image/Animation"

    def save_images(self, images, fps, compress_level, png_filename=""):
        temp_filename = png_filename.split('.')

        i = 1
        file = temp_filename[0]
        while i < len(temp_filename) - 2:
            file = "{}.{}".format(file, temp_filename[i])
            i += 1

        file = "{}_{:05}_.webp".format(file, 1)

        filename_path = os.path.join(self.output_dir, png_filename)
        temp_path = os.path.join(self.temp_dir, file)
        filename = os.path.abspath(filename_path)

        print("Output Filename: {}".format(filename))

        preview_result = None

        results = list()
        pil_images = []
        for image in images:
            i = 255. * image.cpu().numpy()
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))

            if (img.width > 768) or (img.height > 768):
                img = self.scale_image(img.copy(), 768)

            pil_images.append(img)

        metadata = None
        num_frames = len(pil_images)
        print("Number Of Frames: {}".format(num_frames))

        pil_images[0].save(filename, pnginfo=metadata, compress_level=compress_level, save_all=True, duration=int(1000.0/fps), append_images=pil_images[1:num_frames])
        pil_images[0].save(temp_path, pnginfo=metadata, compress_level=compress_level, save_all=True, duration=int(1000.0/fps), append_images=pil_images[1:num_frames], lossless=False, quality=80, method=4)

        results.append({
            "filename": file,
            "subfolder": "",
            "type": "temp",
        })

        return { "ui": { "images": results, "animated": (True,)} }

    def scale_image(self, image, length=1024):
        img = image

        if img.height >= img.width:
            newHeight = length
            newWidth = int(float(length / img.height) * img.width)
        else:
            newWidth = length
            newHeight = int(float(length / img.width) * img.height)

        newImage = img.resize((newWidth, newHeight), resample=Image.BILINEAR)

        return newImage

class WarpedCFGGuider:
    def __init__(self):
        self.original_conds = {}
        self.cfg = 1.0
        self.inner_executor = None
        self.i2v_model = None

    def set_conds(self, positive, negative):
        self.inner_set_conds({"positive": positive, "negative": negative})

    def set_cfg(self, cfg):
        self.cfg = cfg

    def inner_set_conds(self, conds):
        for k in conds:
            self.original_conds[k] = sampler_helpers.convert_cond(conds[k])

    def __call__(self, *args, **kwargs):
        return self.predict_noise(*args, **kwargs)

    def predict_noise(self, x, timestep, model_options={}, seed=None):
        return sampling_function(self.inner_model, x, timestep, self.conds.get("negative", None), self.conds.get("positive", None), self.cfg, model_options=model_options, seed=seed)

    def inner_sample(self, noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed):
        if latent_image is not None and torch.count_nonzero(latent_image) > 0: #Don't shift the empty latent image.
            latent_image = self.inner_model.process_latent_in(latent_image)

        self.conds = process_conds(self.inner_model, noise, self.conds, device, latent_image, denoise_mask, seed)

        extra_model_options = comfy_model_patcher.create_model_options_clone(self.model_options)
        extra_model_options.setdefault("transformer_options", {})["sample_sigmas"] = sigmas
        extra_args = {"model_options": extra_model_options, "seed": seed}

        executor = patcher_extension.WrapperExecutor.new_class_executor(
            sampler.sample,
            sampler,
            patcher_extension.get_all_wrappers(patcher_extension.WrappersMP.SAMPLER_SAMPLE, extra_args["model_options"], is_model_options=True)
        )
        samples = executor.execute(self, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
        return self.inner_model.process_latent_out(samples.to(torch.float32))

    def outer_sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
        self.inner_model, self.conds, self.loaded_models = sampler_helpers.prepare_sampling(self.model_patcher, noise.shape, self.conds, self.model_options)
        device = self.model_patcher.load_device

        if denoise_mask is not None:
            denoise_mask = sampler_helpers.prepare_mask(denoise_mask, noise.shape, device)

        noise = noise.to(device)
        latent_image = latent_image.to(device)
        sigmas = sigmas.to(device)
        cast_to_load_options(self.model_options, device=device, dtype=self.model_patcher.model_dtype())

        try:
            self.model_patcher.pre_run()
            output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
        finally:
            self.model_patcher.cleanup()

        sampler_helpers.cleanup_models(self.conds, self.loaded_models)
        del self.inner_model
        del self.loaded_models
        return output

    def set_model(self, model):
        self.model_patcher = model
        self.model_options = model.model_options
        self.inner_executor = None

        return

    def set_models(self, t2v_model, i2v_model):
        self.model_patcher = t2v_model
        self.model_options = t2v_model.model_options
        self.inner_executor = None
        self.i2v_model = i2v_model

        return

    def set_i2v_model(self):
        if not self.i2v_model is None:
            self.model_patcher = self.i2v_model
            self.model_options = self.i2v_model.model_options
            self.inner_executor = None

        return

    def sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
        if sigmas.shape[-1] == 0:
            return latent_image

        self.conds = {}
        for k in self.original_conds:
            self.conds[k] = list(map(lambda a: a.copy(), self.original_conds[k]))
        preprocess_conds_hooks(self.conds)

        try:
            orig_model_options = self.model_options
            self.model_options = comfy_model_patcher.create_model_options_clone(self.model_options)
            # if one hook type (or just None), then don't bother caching weights for hooks (will never change after first step)
            orig_hook_mode = self.model_patcher.hook_mode
            if get_total_hook_groups_in_conds(self.conds) <= 1:
                self.model_patcher.hook_mode = hooks.EnumHookMode.MinVram
            sampler_helpers.prepare_model_patcher(self.model_patcher, self.conds, self.model_options)
            filter_registered_hooks_on_conds(self.conds, self.model_options)
            executor = patcher_extension.WrapperExecutor.new_class_executor(
                self.outer_sample,
                self,
                patcher_extension.get_all_wrappers(patcher_extension.WrappersMP.OUTER_SAMPLE, self.model_options, is_model_options=True)
            )
            output = executor.execute(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
        finally:
            cast_to_load_options(self.model_options, device=self.model_patcher.offload_device)
            self.model_options = orig_model_options
            self.model_patcher.hook_mode = orig_hook_mode
            self.model_patcher.restore_hook_patches()

        del self.conds
        return output

class WarpedGuider_Basic(WarpedCFGGuider):
    def set_conds_single(self, positive):
        self.inner_set_conds({"positive": positive})

    def set_conds_both(self, positive, negative):
        self.inner_set_conds({"positive": positive, "negative": negative})

class WarpedBasicGuider:
    def __init__(self):
        self.guider = None

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"model": ("MODEL", ),
                    }
            }

    RETURN_TYPES = ("WARPED_GUIDER",)
    RETURN_NAMES = ("guider",)

    FUNCTION = "get_guider"
    CATEGORY = "Warped/General/Sampling/Guiders"

    def get_guider(self, model):
        self.guider = WarpedGuider_Basic()
        self.guider.set_model(model)
        return (self.guider,)

class WarpedDualGuider:
    def __init__(self):
        self.guider = None

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"t2v_model": ("MODEL", ),
                    "i2v_model": ("MODEL", ),
                    }
               }

    RETURN_TYPES = ("WARPED_GUIDER",)
    RETURN_NAMES = ("guider",)

    FUNCTION = "get_guider"
    CATEGORY = "Warped/General/Sampling/Guiders"

    def get_guider(self, t2v_model, i2v_model):
        self.guider = WarpedGuider_Basic()
        self.guider.set_models(t2v_model, i2v_model)

        return (self.guider,)

class WarpedUpscaleWithModel:
    rescale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]

    RETURN_TYPES = ("IMAGE",)

    FUNCTION = "upscale"
    CATEGORY = "image/upscaling"

    def __init__(self):
        self.__imageScaler = ImageUpscaleWithModel()

    @classmethod
    def INPUT_TYPES(self):
        return {
            "required": {
                "upscale_model": ("UPSCALE_MODEL",),
                "images": ("IMAGE",),
                "upscale_by": ("FLOAT", {
                    "default": 2.0,
                    "min": 1.0,
                    "max": 8.0,
                    "step": 0.05,
                }),
                "rescale_method": (self.rescale_methods,),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("images",)
    FUNCTION = "upscale"

    CATEGORY = "Warped/General/Image"


    def upscale(self, upscale_model, images, upscale_by, rescale_method):
        upscaled_images = []

        samples = images.movedim(-1,1)

        width = round(samples.shape[3])
        height = round(samples.shape[2])

        target_width = round(samples.shape[3] * upscale_by)
        target_height = round(samples.shape[2] * upscale_by)


        samples = self.__imageScaler.upscale(upscale_model, images)[0].movedim(-1,1)

        upscaled_width = round(samples.shape[3])
        upscaled_height = round(samples.shape[2])

        if upscaled_width > target_width or upscaled_height > target_height:
            samples = comfy.utils.common_upscale(samples, target_width, target_height, rescale_method, "disabled")

        samples = samples.movedim(1,-1)

        return (samples,)

class WarpedLoadVideosBatch:
    def __init__(self):
        self.index = 0
        self.output_dir = folder_paths.get_output_directory()

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": ("STRING", {"default": '', "multiline": False}),
            },
        }

    RETURN_TYPES = ("STRING", "STRING", "STRING", )
    RETURN_NAMES = ("mp4_filename", "png_filename", "webp_filename", )

    FUNCTION = "load_batch_videos"

    CATEGORY = "Warped/General/Video"

    def load_batch_videos(self, path):
        if not os.path.exists(path):
            return ("", "", )

        index=0
        mode="incremental_image"
        label='Batch 001'
        suffix=""

        retry = False

        try:
            return self.do_the_load(path, index, mode, label, suffix)
        except:
            self.index = 0
            retry = True

        if retry:
            return self.do_the_load(path, index, mode, label, suffix)

        return ("", "", )


    def do_the_load(self, path, index, mode, label, suffix):
        fl = self.BatchVideoLoader(path, label, '*', index)
        new_paths = fl.video_paths

        filename = fl.video_paths[self.index]

        tempStrings = filename.split('\\')
        png_filename = tempStrings[len(tempStrings) - 1]
        png_filename = png_filename.replace(".mp4", ".png")
        webp_filename = tempStrings[len(tempStrings) - 1]
        webp_filename = webp_filename.replace(".mp4", ".webp")

        print("Filename: {}".format(filename))
        print("Png Filename: {}".format(png_filename))
        print("Webp Filename: {}".format(webp_filename))

        self.index += 1

        if self.index >= len(fl.video_paths):
            self.index = 0

        return (filename, png_filename, webp_filename, )


    class BatchVideoLoader:
        def __init__(self, directory_path, label, pattern, index):
            self.video_paths = []
            self.load_videos(directory_path, pattern)
            self.video_paths.sort()

            self.index = index
            self.label = label

        def load_videos(self, directory_path, pattern):
            for file_name in glob.glob(os.path.join(directory_path, pattern), recursive=True):
                if file_name.lower().endswith(ALLOWED_VIDEO_EXT):
                    abs_file_path = os.path.abspath(file_name)
                    self.video_paths.append(abs_file_path)

        def get_video_by_id(self, video_id):
            if video_id < 0 or video_id >= len(self.video_paths):
                cstr(f"Invalid video index `{video_id}`").error.print()
                return

            return self.video_paths[video_id]

        def get_next_video(self):
            if self.index >= len(self.video_paths):
                self.index = 0

            video_path = self.video_paths[self.index]
            self.index += 1

            if self.index == len(self.video_paths):
                self.index = 0

            cstr(f'{cstr.color.YELLOW}{self.label}{cstr.color.END} Index: {self.index}').msg.print()

            return video_path

        def get_current_video(self):
            if self.index >= len(self.video_paths):
                self.index = 0
            video_path = self.video_paths[self.index]

            return video_path

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return float("NaN")

class WarpedBundleVideoImages:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"video_path": ("STRING", {"default": ""}),
                    "starting_index": ("INT", {"default": 0, "min": 0, "max": 25000000, "step": 1}),
                    "num_frames": ("INT", {"default": 61, "min": 5, "max": 1000001, "step": 4}),
                    "use_gpu": ("BOOLEAN", {"default": True}),
                    },
                }
    RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "INT", "INT", "INT", "FLOAT", )
    RETURN_NAMES = ("image_batch", "first_image", "last_image", "num_frames", "width", "height", "fps", )
    FUNCTION = "generate"

    CATEGORY = "Warped/General/Video"

    def generate(self, video_path, starting_index, num_frames, use_gpu=False):
        if use_gpu:
            device = mm.get_torch_device()
        else:
            device = get_offload_device()

        print("WarpedBundleVideoImages: device: {}".format(device))

        cap = cv2.VideoCapture(video_path)

        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        print('WarpedBundleVideoImages: width = %d' % width)
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        print('WarpedBundleVideoImages: height = %d' % height)
        length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        print('WarpedBundleVideoImages: length = %d' % length)
        fps = float(cap.get(cv2.CAP_PROP_FPS))
        fps = round(fps)
        print('WarpedBundleVideoImages: fps = %d' % fps)

        batched_images = None
        # print_it = True
        last_image = None
        first_image = None

        if starting_index > length:
            starting_index = 0

        if ((starting_index + num_frames) > length) or (num_frames == 0):
            num_frames = length - starting_index

        try:
            if starting_index > 0:
                cap.set(cv2.CAP_PROP_POS_FRAMES, starting_index)

            video_frames = []
            frame_count = 0
            for i in tqdm(range(num_frames), desc="Reading Video Frames: "):
                frame_count += 1

                # Take each frame
                _, frameorig = cap.read()

                color_coverted = cv2.cvtColor(frameorig, cv2.COLOR_BGR2RGB)
                video_frames.append(color_coverted)
        except:
            print("WarpedBundleVideoImages: Exception During Video File Read.")
        finally:
            cap.release()

        try:
            pil_image = Image.fromarray(video_frames[0])
            first_image = pil2tensorSwap(pil_image, device=device)
            first_image = first_image.to(get_offload_device())

            pil_image = Image.fromarray(video_frames[len(video_frames) - 1])
            last_image = pil2tensorSwap(pil_image, device=device)
            last_image = last_image.to(get_offload_device())

            for i in tqdm(range(len(video_frames)), desc="Preprocessing Video Frames: "):
                pil_image = Image.fromarray(video_frames[i])
                temp_image = pil2tensorSwap(pil_image, device=device)

                if len(temp_image.shape) < 4:
                    temp_image = temp_image.unsqueeze(0)

                if not batched_images is None:
                    batched_images = torch.cat((batched_images, temp_image), 0)
                else:
                    batched_images = temp_image

                temp_image.to(get_offload_device())
                temp_image = None
        except:
            print("WarpedBundleVideoImages: Exception During Video File Processing.")

        batched_images = batched_images.to(get_offload_device())

        mm.soft_empty_cache()
        gc.collect()
        time.sleep(1)

        return (batched_images, first_image, last_image, int(batched_images.shape[0]), width, height, float(fps), )

class WarpedGetImageFromVideo:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"video_path": ("STRING", {"default": ""}),
                    "image_index": ("INT", {"default": 0}),
                    },
                }
    RETURN_TYPES = ("IMAGE", "INT", "INT", )
    RETURN_NAMES = ("image", "width", "height", )
    FUNCTION = "generate"

    CATEGORY = "Warped/General/Video"

    def generate(self, video_path, image_index):
        cap = cv2.VideoCapture(video_path)

        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        print('WarpedGetImageFromVideo: width = %d' % width)
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        print('WarpedGetImageFromVideo: height = %d' % height)
        length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        print('WarpedGetImageFromVideo: length = %d' % length)

        return_image = None

        if image_index > length:
            image_index = length - 1

        try:
            cap.set(cv2.CAP_PROP_POS_FRAMES, image_index)

            # Take each frame
            _, frameorig = cap.read()

            color_coverted = cv2.cvtColor(frameorig, cv2.COLOR_BGR2RGB)

            pil_image = Image.fromarray(color_coverted)
            return_image = pil2tensorSwap(pil_image, device=get_offload_device())

            if len(return_image.shape) < 4:
                return_image = return_image.unsqueeze(0)

            print("return_image: Shape: {}".format(return_image.shape))

            frameorig = None
            color_coverted = None
            pil_image = None
        except:
            print("WarpedGetImageFromVideo: Exception During Video File Read.")
        finally:
            cap.release()

        mm.soft_empty_cache()
        gc.collect()
        time.sleep(1)

        return (return_image, width, height, )

class WarpedGetTwoImagesFromVideo:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"video_path": ("STRING", {"default": ""}),
                    "first_image_index": ("INT", {"default": 0}),
                    "second_image_index": ("INT", {"default": 1}),
                    },
                }
    RETURN_TYPES = ("IMAGE", "INT", "INT", )
    RETURN_NAMES = ("first_image", "second_image", "width", "height", )
    FUNCTION = "generate"

    CATEGORY = "Warped/General/Video"

    def generate(self, video_path, first_image_index, second_image_index):
        cap = cv2.VideoCapture(video_path)

        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        print('WarpedGetImageFromVideo: width = %d' % width)
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        print('WarpedGetImageFromVideo: height = %d' % height)
        length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        print('WarpedGetImageFromVideo: length = %d' % length)

        return_image_one = None
        return_image_two = None

        if first_image_index > length:
            first_image_index = 0

        if second_image_index > length:
            second_image_index = length - 1

        try:
            cap.set(cv2.CAP_PROP_POS_FRAMES, first_image_index)

            # Take each frame
            _, frameorig = cap.read()

            color_coverted = cv2.cvtColor(frameorig, cv2.COLOR_BGR2RGB)

            pil_image = Image.fromarray(color_coverted)
            return_image_one = pil2tensorSwap(pil_image, device=get_offload_device())

            if len(return_image_one.shape) < 4:
                return_image_one = return_image_one.unsqueeze(0)

            print("return_image_one: Shape: {}".format(return_image_one.shape))

            cap.set(cv2.CAP_PROP_POS_FRAMES, second_image_index)

            # Take each frame
            _, frameorig = cap.read()

            color_coverted = cv2.cvtColor(frameorig, cv2.COLOR_BGR2RGB)

            pil_image = Image.fromarray(color_coverted)
            return_image_two = pil2tensorSwap(pil_image, device=get_offload_device())

            if len(return_image_two.shape) < 4:
                return_image_two = return_image_two.unsqueeze(0)

            print("return_image_two: Shape: {}".format(return_image_two.shape))

            frameorig = None
            color_coverted = None
            pil_image = None
        except:
            print("WarpedGetImageFromVideo: Exception During Video File Read.")
        finally:
            cap.release()

        mm.soft_empty_cache()
        gc.collect()
        time.sleep(1)

        return (return_image_one, return_image_two, width, height, )

class WarpedWanImageToVideo:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"positive": ("CONDITIONING", ),
                             "negative": ("CONDITIONING", ),
                             "vae": ("VAE", ),
                             "width": ("INT", {"default": 832, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}),
                             "height": ("INT", {"default": 480, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}),
                             "length": ("INT", {"default": 81, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 4}),
                             "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
                },
                "optional": {"clip_vision_output": ("CLIP_VISION_OUTPUT", ),
                             "start_image": ("IMAGE", ),
                }}

    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT", "INT", )
    RETURN_NAMES = ("positive", "negative", "latent", "num_frames", )
    FUNCTION = "encode"

    CATEGORY = "Warped/Wan/Conditioning"

    def encode(self, positive, negative, vae, width, height, length, batch_size, start_image=None, clip_vision_output=None):
        latent = torch.zeros([batch_size, 16, ((length - 1) // 4) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device())
        if start_image is not None:
            start_image = comfy.utils.common_upscale(start_image[:length].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
            image = torch.mul(torch.ones((length, height, width, start_image.shape[-1]), device=start_image.device, dtype=start_image.dtype), 0.5)
            image[:start_image.shape[0]] = start_image

            concat_latent_image = vae.encode(image[:, :, :, :3])

            mask = torch.ones((1, 1, latent.shape[2], concat_latent_image.shape[-2], concat_latent_image.shape[-1]), device=start_image.device, dtype=start_image.dtype)
            mask[:, :, :((start_image.shape[0] - 1) // 4) + 1] = 0.0

            positive = node_helpers.conditioning_set_values(positive, {"concat_latent_image": concat_latent_image, "concat_mask": mask})
            negative = node_helpers.conditioning_set_values(negative, {"concat_latent_image": concat_latent_image, "concat_mask": mask})

        if clip_vision_output is not None:
            positive = node_helpers.conditioning_set_values(positive, {"clip_vision_output": clip_vision_output})
            negative = node_helpers.conditioning_set_values(negative, {"clip_vision_output": clip_vision_output})

        out_latent = {}
        out_latent["samples"] = latent
        return (positive, negative, out_latent, length, )

class WarpedHunyuanImageToVideo:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
                    "images": ("IMAGE", ),
                    "positive": ("CONDITIONING", ),
                    "vae": ("VAE", ),
                    "num_frames": ("INT", {"default": 53, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 4}),
                    "guidance_type": (["v1 (concat)", "v2 (replace)"], )
                    },
                }

    RETURN_TYPES = ("CONDITIONING", "LATENT")
    RETURN_NAMES = ("positive", "latent")
    FUNCTION = "encode"

    CATEGORY = "Warped/Hunyuan/Conditioning"

    def encode(self, images, positive, vae, num_frames, guidance_type):
        images.to(dtype=torch.float32, device=get_offload_device())

        if len(images.shape) < 4:
            images = images.unsqueeze()

        height = images.shape[1]
        width  = images.shape[2]

        out_latent = {}

        if images.shape[0] == 1:
            print("Single Image")
            latent = torch.zeros([1, 16, ((num_frames - 1) // 4) + 1, height // 8, width // 8], dtype=torch.float32, device=get_offload_device())

            images = comfy.utils.common_upscale(images[:num_frames, :, :, :3].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)

            concat_latent_image = vae.encode(images)
            mask = torch.ones((1, 1, latent.shape[2], concat_latent_image.shape[-2], concat_latent_image.shape[-1]), device=images.device, dtype=images.dtype)
            mask[:, :, :((images.shape[0] - 1) // 4) + 1] = 0.0

            if guidance_type == "v1 (concat)":
                cond = {"concat_latent_image": concat_latent_image, "concat_mask": mask}
            else:
                cond = {'guiding_frame_index': 0}
                latent[:, :, :concat_latent_image.shape[2]] = concat_latent_image
                out_latent["noise_mask"] = mask

            positive = node_helpers.conditioning_set_values(positive, cond)

            out_latent["samples"] = latent

            return (positive, out_latent)

        print("Shape Before: {}".format(images.shape))

        images = comfy.utils.common_upscale(images.movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
        # images = torch.mul(torch.ones((num_frames, height, width, images.shape[-1]), device=images.device, dtype=images.dtype), 0.5)

        print("Shape After: {}".format(images.shape))

        concat_latent_image = partial_encode_tiled(vae, images)

        if len(concat_latent_image.shape) < 5:
            concat_latent_image = concat_latent_image.unsqueeze(0)

        mask = torch.ones((1, 1, concat_latent_image.shape[2], concat_latent_image.shape[-2], concat_latent_image.shape[-1]), device=images.device, dtype=images.dtype)
        # mask[:, :, :((images.shape[0] - 1) // 4) + 1] = 0.0

        out_latent["samples"] = concat_latent_image

        if guidance_type == "v1 (concat)":
            cond = {"concat_latent_image": concat_latent_image, "concat_mask": mask}
        else:
            cond = {'guiding_frame_index': 0}
            # out_latent["noise_mask"] = mask

        positive = node_helpers.conditioning_set_values(positive, cond)

        return (positive, out_latent)

class WarpedImageFromBatch:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "image": ("IMAGE",),
                              "batch_index": ("INT", {"default": 0, "min": 0, "max": 4095}),
                              "length": ("INT", {"default": 1, "min": 1, "max": 4096}),
                              }}
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "frombatch"

    CATEGORY = "Warped/General/Image"

    def frombatch(self, image, batch_index, length):
        # print("\n-------------------------------------------------------------------------------")
        # print(image)
        # print("-------------------------------------------------------------------------------\n")

        s_in = image

        if isinstance(s_in, list):
            s_in = s_in[0]

        batch_index = min(s_in.shape[0] - 1, batch_index)
        length = min(s_in.shape[0] - batch_index, length)
        s = s_in[batch_index:batch_index + length].clone()
        return (s,)

class WarpedBaseWanLoraLoader:
    """Base class for Wan LoRA loading functionality"""

    def __init__(self):
        self.loaded_lora: Optional[Tuple[str, Dict[str, torch.Tensor]]] = None

    @classmethod
    def get_cache_dir(cls) -> str:
        """Get or create the cache directory for block settings"""
        try:
            from folder_paths import base_path, folder_names_and_paths
            cache_dir = Path(folder_names_and_paths["custom_nodes"][0][0]) / "ComfyUI-WarpedToolset" / "cache"
            if not os.path.exists(cache_dir):
                os.makedirs(cache_dir)
            return cache_dir
        except Exception as e:
            logger.error(f"Failed to create or access cache directory: {str(e)}")
            raise

    def get_settings_filename(self, lora_name: str) -> str:
        """Generate the settings filename for a given LoRA"""
        base_name = os.path.splitext(lora_name)[0]
        return os.path.join(self.get_cache_dir(), f"{base_name}_blocks.yaml")

    def get_block_settings(self, lora_name: str, use_block_cache: bool = True, include_single_blocks: bool = False) -> dict:
        """Load block settings from cache or return defaults"""
        # Initialize with all double blocks enabled and single blocks based on parameter
        default_settings = {
            **{f"blocks.{i}.": True for i in range(40)},
        }

        if not use_block_cache:
            return default_settings

        try:
            settings_file = self.get_settings_filename(lora_name)
            if os.path.exists(settings_file):
                cached_settings = yaml.safe_load(open(settings_file, 'r'))
                # Merge cached settings with default single block settings
                return {
                    **default_settings,
                    **cached_settings,
                }
            return default_settings
        except Exception as e:
            logger.error(f"Failed to load block settings for {lora_name}: {str(e)}")
            return default_settings

    def save_block_settings(self, lora_name: str, block_settings: dict):
        """Save block settings to cache"""
        try:
            settings_file = self.get_settings_filename(lora_name)
            # Ensure directory exists
            os.makedirs(os.path.dirname(settings_file), exist_ok=True)
            save_settings = {k: v for k, v in block_settings.items() if k.startswith('blocks.')}
            with open(settings_file, 'w') as f:
                yaml.safe_dump(save_settings, f)
        except Exception as e:
            logger.error(f"Failed to save block settings for {lora_name}: {str(e)}")

    def filter_lora_keys(self, lora: Dict[str, torch.Tensor], block_settings: dict) -> Dict[str, torch.Tensor]:
        """Filter LoRA keys based on block settings"""
        filtered_blocks = {k: v for k, v in block_settings.items() if v is True}
        return {key: value for key, value in lora.items()
                if any(block in key for block in filtered_blocks)}

    def load_lora_file(self, lora_name: str) -> Dict[str, torch.Tensor]:
        """Load LoRA file and cache it"""
        from comfy.utils import load_torch_file

        lora_path = folder_paths.get_full_path("loras", lora_name)
        if not os.path.exists(lora_path):
            raise Exception(f"LoRA {lora_name} not found at {lora_path}")

        if self.loaded_lora is not None and self.loaded_lora[0] == lora_path:
            return self.loaded_lora[1]

        lora = load_torch_file(lora_path)
        self.loaded_lora = (lora_path, lora)
        return lora

    def get_file_mtime(self, filepath: str) -> str:
        """Get modification time of file as string"""
        try:
            return str(os.path.getmtime(filepath))
        except:
            return "0"

    def get_lora_mtime(self, lora_name: str) -> str:
        """Get modification time of LoRA file"""
        try:
            lora_path = folder_paths.get_full_path("loras", lora_name)
            return self.get_file_mtime(lora_path)
        except:
            return "0"

    def get_cache_mtime(self, lora_name: str) -> str:
        """Get modification time of cache file"""
        try:
            cache_file = self.get_settings_filename(lora_name)
            return self.get_file_mtime(cache_file)
        except:
            return "0"

    def apply_lora(self, model, clip, lora_name: str, strength_model: float, strength_clip: float, block_settings: Optional[dict] = None) -> torch.nn.Module:
        """Apply LoRA to model with given settings"""
        from comfy.sd import load_lora_for_models

        if not lora_name:
            return model

        try:
            lora = self.load_lora_file(lora_name)
            if block_settings is None:
                block_settings = self.get_block_settings(lora_name, True)  # Always use cache for direct loading

            filtered_lora = self.filter_lora_keys(lora, block_settings)
            new_model, new_clip = load_lora_for_models(model, clip, filtered_lora, strength_model, strength_clip)

            return (new_model, new_clip,)

        except Exception as e:
            logger.error(f"Error applying LoRA {lora_name}: {str(e)}")

            return (model, clip, )

class WarpedWanLoadAndEditLoraBlocks(WarpedBaseWanLoraLoader):
    """Interactive LoRA block editor"""

    @classmethod
    def INPUT_TYPES(s):
        arg_dict = {
            "required": {
                "model": ("MODEL",),
                "clip": ("CLIP", {"tooltip": "The CLIP model the LoRA will be applied to."}),
                "lora_name": (folder_paths.get_filename_list("loras"),),
                "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the diffusion model. This value can be negative."}),
                "strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the CLIP model. This value can be negative."}),
                "save_settings": ("BOOLEAN", {"default": True}),
            }
        }

        for i in range(40):
            arg_dict["required"][f"blocks.{i}."] = ("BOOLEAN", {"default": True})

        return arg_dict

    RETURN_TYPES = ("MODEL", "CLIP", )
    RETURN_NAMES = ("model", "clip", )
    FUNCTION = "load_lora"
    CATEGORY = "Warped/Wan/Lora"

    @classmethod
    def IS_CHANGED(s, model, clip, lora_name: str, strength_model: float, strength_clip: float, save_settings: bool, **kwargs):
        instance = s()
        lora_mtime = instance.get_lora_mtime(lora_name)
        return f"{lora_name}_{strength_model}_{strength_clip}_{lora_mtime}"

    def load_lora(self, model, clip, lora_name: str, strength_model: float, strength_clip: float, save_settings: bool, **kwargs):
        if not lora_name:
            return (model,)

        # Add single blocks settings based on the parameter
        block_settings = {
            **kwargs,
        }

        if save_settings:
            self.save_block_settings(lora_name, block_settings)

        return_model, return_clip = self.apply_lora(model, clip, lora_name, strength_model, strength_clip, block_settings)

        return (return_model, return_clip,)

class WarpedImageScaleToSide:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"image": ("IMAGE", ),
                     "length": ("INT", {"default": 1024}),
                     "scale_to": (["long_side", "short_side"], {"default": "long_side"}),
                     "use_gpu": ("BOOLEAN", {"default": False}),
                    },
                }

    CATEGORY = "Warped/General/Image"

    RETURN_TYPES = ("IMAGE", "INT", "INT", )
    RETURN_NAMES = ("image", "width", "height")

    FUNCTION = "scale_image"

    def scale_image(self, image, length=1024, scale_to="long_side", use_gpu=False):
        if use_gpu:
            device = mm.get_torch_device()
        else:
            device = get_offload_device()

        print("WarpedImageScaleToSide: device: {}".format(device))

        if len(image.shape) < 4:
            image = image.unsqueeze(0)

        image_batches_tuple = torch.split(image, 1, dim=0)
        image_batches_split = [item for item in image_batches_tuple]

        final_image = None

        # for i in tqdm(range(num_frames), desc="Preprocessing Video Frames: "):
        for single_image in image_batches_split:
            img = tensor2pilSwap(single_image)
            img = img[0]

            if scale_to == "long_side":
                if img.height >= img.width:
                    newHeight = length
                    newWidth = int(float(length / img.height) * img.width)
                else:
                    newWidth = length
                    newHeight = int(float(length / img.width) * img.height)
            else:
                if img.height <= img.width:
                    newHeight = length
                    newWidth = int(float(length / img.height) * img.width)
                else:
                    newWidth = length
                    newHeight = int(float(length / img.width) * img.height)

            tempImage = img.resize((newWidth, newHeight), resample=Image.BILINEAR)

            newImage = pil2tensorSwap([tempImage], device=device)

            if len(newImage.shape) < 4:
                newImage = newImage.unsqueeze(0)

            if not final_image is None:
                final_image = torch.cat((final_image, newImage), 0)
            else:
                final_image = newImage

        newImage.to(device=get_offload_device())
        newImage = None
        final_image.to(device=get_offload_device())

        mm.soft_empty_cache()
        gc.collect()
        time.sleep(1)

        return (final_image, final_image.shape[2], final_image.shape[1],)

def get_base_lora_dirs():
    return folder_paths.get_folder_paths("loras")

def get_lora_directories():
    lora_dirs = get_base_lora_dirs()

    result_lora_dirs = []

    for lora_dir in lora_dirs:
        temp_dirs = [x[0] for x in os.walk(lora_dir)]
        result_lora_dirs = result_lora_dirs + temp_dirs

    result_lora_dirs.sort()

    return(result_lora_dirs)

def get_lora_path_parts(path):
    temp_base_dirs = get_base_lora_dirs()
    base_dir = ""
    lora_name = ""

    for temp_dir in temp_base_dirs:
        if path.startswith(temp_dir):
            base_dir = temp_dir
            lora_name = path[len(base_dir) + 1:]
            print("get_lora_path_parts: base_dir: {}  |  lora_name: {}".format(base_dir, lora_name))
            break

    return base_dir, lora_name

class WarpedLoadLorasBatchByPrefix:
    def __init__(self):
        self.index = 0
        self.lora_dir = ""
        self.last_prefix = ""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "lora_dir": (get_lora_directories(), ),
                "lora_prefix": ("STRING", {"default": '', "multiline": False}),
            },
        }

    RETURN_TYPES = ("STRING", "STRING", )
    RETURN_NAMES = ("lora_name", "full_lora_path", )
    FUNCTION = "load_batch_loras"

    CATEGORY = "Warped/General/Lora"

    def load_batch_loras(self, lora_dir, lora_prefix):
        self.lora_dir = lora_dir
        path = lora_dir
        print(path)

        if not os.path.exists(path):
            return ("", "", )

        if not (self.last_prefix == lora_prefix):
            self.last_prefix = lora_prefix
            self.index = 0

        retry = False
        index = 0

        try:
            filename, full_filename = self.do_the_load(path, lora_prefix, index)
            print("WarpedLoadLorasBatchByPrefix: Filename: {}  |  Full File Path: {}".format(filename, full_filename))
            return (filename, full_filename, )
        except:
            self.index = 0
            retry = True

        if retry:
            retry = False
            filename, full_filename = self.do_the_load(path, lora_prefix, index)
            print("WarpedLoadLorasBatchByPrefix: Retrying: Filename: {}  |  Full File Path: {}".format(filename, full_filename))
            return (filename, full_filename, )

        return ("", "", )


    def do_the_load(self, path, prefix, index):
        prefix = prefix.strip(' ')

        if (len(prefix) == 1) and (prefix == '*'):
            fl = self.BatchLoraLoader(path, '*', index)
        else:
            prefix = prefix.strip('*')
            fl = self.BatchLoraLoader(path, "{}*".format(prefix), index)

        new_paths = fl.lora_paths

        filename = fl.lora_paths[self.index]

        # filename = os.path.join(self.sub_folder, filename)
        full_filename = os.path.join(path, filename)
        base_dir, lora_name = get_lora_path_parts(full_filename)

        self.index += 1

        if self.index >= len(fl.lora_paths):
            self.index = 0

        return lora_name, full_filename


    class BatchLoraLoader:
        def __init__(self, directory_path, pattern, index):
            self.lora_paths = []
            self.load_loras(directory_path, pattern)
            self.lora_paths.sort()

            self.index = index

        def load_loras(self, directory_path, pattern):
            for file_name in glob.glob(os.path.join(directory_path, pattern), recursive=True):
                temp_strings = file_name.split('\\')
                file_name = temp_strings[len(temp_strings) - 1]

                if file_name.lower().endswith("safetensors"):
                    self.lora_paths.append(file_name)

        def get_lora_by_id(self, lora_id):
            if lora_id < 0 or lora_id >= len(self.lora_paths):
                cstr(f"WarpedLoadLorasBatchByPrefix: Invalid lora index `{lora_id}`").error.print()
                return

            return self.lora_paths[lora_id]

        def get_next_lora(self):
            if self.index >= len(self.lora_paths):
                self.index = 0

            lora_path = self.lora_paths[self.index]
            self.index += 1

            if self.index == len(self.lora_paths):
                self.index = 0

            cstr(f'{cstr.color.YELLOW}{cstr.color.END} Index: {self.index}').msg.print()

            return lora_path

        def get_current_lora(self):
            if self.index >= len(self.lora_paths):
                self.index = 0
            lora_path = self.lora_paths[self.index]

            return lora_path

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return float("NaN")

class WarpedHunyuanVideoLoraLoader:
    def __init__(self):
        self.blocks_type = ["all", "single_blocks", "double_blocks"]
        self.loaded_lora = None

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL",),
                "lora_name": ("STRING", {"forceInput": True}),
                "strength": ("FLOAT", {
                    "default": 1.0,
                    "min": -10.0,
                    "max": 10.0,
                    "step": 0.01,
                    "display": "number"
                }),
                "blocks_type": (["all", "single_blocks", "double_blocks"],),
            }
        }

    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("model",)
    FUNCTION = "load_lora"
    CATEGORY = "Warped/Hunyuan/Lora"
    OUTPUT_NODE = False
    DESCRIPTION = "LoRA, single blocks double blocks"

    def load_lora(self, model, lora_name: str, strength: float, blocks_type: str):
        """
        Parameters
        ----------
        model : ModelPatcher
        lora_name : str
        strength : float
        blocks_type : str
            blocks: "all", "single_blocks" "double_blocks"

        Returns
        -------
        tuple
            LoRA
        """
        if not lora_name:
            return (model,)

        from comfy.utils import load_torch_file
        from comfy.sd import load_lora_for_models
        from comfy.lora import load_lora

        lora_path = folder_paths.get_full_path("loras", lora_name)
        if not os.path.exists(lora_path):
            raise Exception(f"Lora {lora_name} not found at {lora_path}")

        if self.loaded_lora is not None:
            if self.loaded_lora[0] == lora_path:
                lora = self.loaded_lora[1]
            else:
                self.loaded_lora = None

        if self.loaded_lora is None:
            lora = load_torch_file(lora_path)
            self.loaded_lora = (lora_path, lora)

        diffusers_lora = convert_lora(lora, convert_to="diffusion_model")
        filtered_lora = filter_lora_keys(diffusers_lora, blocks_type)

        new_model, _ = load_lora_for_models(model, None, filtered_lora, strength, 0)
        if new_model is not None:
            return (new_model,)

        return (model,)

    @classmethod
    def IS_CHANGED(s, model, lora_name, strength, blocks_type):
        return f"{lora_name}_{strength}_{blocks_type}"

class WarpedFramepackMultiLoraSelect:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
               "lora_01": (['None'] + folder_paths.get_filename_list("loras"), {"tooltip": "LORA models are expected to have .safetensors extension"}),
                "strength_01": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.0001, "tooltip": "LORA strength, set to 0.0 to unmerge the LORA"}),
                "fuse_lora_01": ("BOOLEAN", {"default": False, "tooltip": "Fuse the LORA model with the base model. This is recommended for better performance."}),
               "lora_02": (['None'] + folder_paths.get_filename_list("loras"), {"tooltip": "LORA models are expected to have .safetensors extension"}),
                "strength_02": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.0001, "tooltip": "LORA strength, set to 0.0 to unmerge the LORA"}),
                "fuse_lora_02": ("BOOLEAN", {"default": False, "tooltip": "Fuse the LORA model with the base model. This is recommended for better performance."}),
               "lora_03": (['None'] + folder_paths.get_filename_list("loras"), {"tooltip": "LORA models are expected to have .safetensors extension"}),
                "strength_03": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.0001, "tooltip": "LORA strength, set to 0.0 to unmerge the LORA"}),
                "fuse_lora_03": ("BOOLEAN", {"default": False, "tooltip": "Fuse the LORA model with the base model. This is recommended for better performance."}),
               "lora_04": (['None'] + folder_paths.get_filename_list("loras"), {"tooltip": "LORA models are expected to have .safetensors extension"}),
                "strength_04": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.0001, "tooltip": "LORA strength, set to 0.0 to unmerge the LORA"}),
                "fuse_lora_04": ("BOOLEAN", {"default": False, "tooltip": "Fuse the LORA model with the base model. This is recommended for better performance."}),
            },
            "optional": {
                "prev_lora":("FPLORA", {"default": None, "tooltip": "For loading multiple LoRAs"}),
            }
        }

    RETURN_TYPES = ("FPLORA",)
    RETURN_NAMES = ("lora", )
    FUNCTION = "select_multiple_loras"
    CATEGORY = "Warped/Framepack/Lora"
    DESCRIPTION = "Select a Hunyuan LoRA models"

    def select_multiple_loras(self, **kwargs):
        loras_list = []

        prev_lora = kwargs.get(f"prev_lora")

        if prev_lora is not None:
            loras_list.extend(prev_lora)

        for i in range(1, 5):
            temp_lora_name = kwargs.get(f"lora_0{i}")
            temp_strength = kwargs.get(f"strength_0{i}")
            temp_fuse_lora = kwargs.get(f"fuse_lora_0{i}")

            if (temp_lora_name != "None") and (Decimal(temp_strength).compare(Decimal(0.0)) != 0):
                lora = {
                    "path": folder_paths.get_full_path("loras", temp_lora_name),
                    "strength": temp_strength,
                    "name": temp_lora_name.split(".")[0],
                    "fuse_lora": temp_fuse_lora,
                }

                loras_list.append(lora)

        if len(loras_list) > 0:
            return (loras_list,)
        else:
            return (None,)

class WarpedFramepackMultiLoraSelectExt:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
               "lora_01": (['None'] + folder_paths.get_filename_list("loras"), {"tooltip": "LORA models are expected to have .safetensors extension"}),
                "strength_01": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.0001, "tooltip": "LORA strength, set to 0.0 to unmerge the LORA"}),
                "blocks_01": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
                "fuse_lora_01": ("BOOLEAN", {"default": False, "tooltip": "Fuse the LORA model with the base model. This is recommended for better performance."}),
               "lora_02": (['None'] + folder_paths.get_filename_list("loras"), {"tooltip": "LORA models are expected to have .safetensors extension"}),
                "strength_02": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.0001, "tooltip": "LORA strength, set to 0.0 to unmerge the LORA"}),
                "blocks_02": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
                "fuse_lora_02": ("BOOLEAN", {"default": False, "tooltip": "Fuse the LORA model with the base model. This is recommended for better performance."}),
               "lora_03": (['None'] + folder_paths.get_filename_list("loras"), {"tooltip": "LORA models are expected to have .safetensors extension"}),
                "strength_03": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.0001, "tooltip": "LORA strength, set to 0.0 to unmerge the LORA"}),
                "blocks_03": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
                "fuse_lora_03": ("BOOLEAN", {"default": False, "tooltip": "Fuse the LORA model with the base model. This is recommended for better performance."}),
               "lora_04": (['None'] + folder_paths.get_filename_list("loras"), {"tooltip": "LORA models are expected to have .safetensors extension"}),
                "strength_04": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.0001, "tooltip": "LORA strength, set to 0.0 to unmerge the LORA"}),
                "blocks_04": (["all", "double_blocks", "single_blocks"], {"default": "all", "tooltip": "all, single only, or double only block."}),
                "fuse_lora_04": ("BOOLEAN", {"default": False, "tooltip": "Fuse the LORA model with the base model. This is recommended for better performance."}),
            },
            "optional": {
                "prev_lora":("FPLORA", {"default": None, "tooltip": "For loading multiple LoRAs"}),
            }
        }

    RETURN_TYPES = ("FPLORA",)
    RETURN_NAMES = ("lora", )
    FUNCTION = "select_multiple_loras"
    CATEGORY = "Warped/Framepack/Lora"
    DESCRIPTION = "Select a Hunyuan LoRA models"

    def select_multiple_loras(self, **kwargs):
        loras_list = []

        prev_lora = kwargs.get(f"prev_lora")

        if prev_lora is not None:
            loras_list.extend(prev_lora)

        for i in range(1, 5):
            temp_lora_name = kwargs.get(f"lora_0{i}")
            temp_strength = kwargs.get(f"strength_0{i}")
            temp_blocks =  kwargs.get(f"blocks_0{i}")
            temp_fuse_lora = kwargs.get(f"fuse_lora_0{i}")

            if (temp_lora_name != "None") and (Decimal(temp_strength).compare(Decimal(0.0)) != 0):
                lora = {
                    "path": folder_paths.get_full_path("loras", temp_lora_name),
                    "strength": temp_strength,
                    "blocks": temp_blocks,
                    "name": temp_lora_name.split(".")[0],
                    "fuse_lora": temp_fuse_lora,
                }

                loras_list.append(lora)

        if len(loras_list) > 0:
            return (loras_list,)
        else:
            return (None,)

class WarpedHunyuanMultiLoraLoader:
    """
    Hunyuan Multi-Lora Loader
    This node works like the original lora_loader.py, with a required model input and output.
    It does not output LoRA information in HYVIDLORA format.
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL",),
                "lora_01": (['None'] + folder_paths.get_filename_list("loras"),),
                "strength_01": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                "blocks_type_01": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "lora_02": (['None'] + folder_paths.get_filename_list("loras"),),
                "strength_02": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                "blocks_type_02": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "lora_03": (['None'] + folder_paths.get_filename_list("loras"),),
                "strength_03": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                "blocks_type_03": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "lora_04": (['None'] + folder_paths.get_filename_list("loras"),),
                "strength_04": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                "blocks_type_04": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
            },
          "optional": {
                "lora_name": ("STRING", {"default": None, "forceInput": True}),
                "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                "blocks_type": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
            },
        }

    RETURN_TYPES = ("MODEL", "STRING", )
    RETURN_NAMES = ("model", "lora_metadata")
    OUTPUT_IS_LIST = (False, True)
    FUNCTION = "load_multiple_loras"
    CATEGORY = "Warped/Hunyuan/Lora"
    DESCRIPTION = "Load and apply multiple LoRA models with different strengths and block types. Model input is required."

    def convert_key_format(self, key: str) -> str:
        """Standardize LoRA key format by removing prefixes."""
        prefixes = ["diffusion_model.", "transformer."]
        for prefix in prefixes:
            if key.startswith(prefix):
                key = key[len(prefix):]
                break
        return key

    def load_lora(self, lora_name: str, strength: float, blocks_type: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
        """Load and filter a single LoRA model."""
        if not lora_name or strength == 0:
            return {}, {}

        # # Get the full path to the LoRA file
        # lora_path = folder_paths.get_full_path("loras", lora_name)
        # if not os.path.exists(lora_path):
        #     raise ValueError(f"LoRA file not found: {lora_path}")

        # Load the LoRA weights
        # lora_weights = utils.load_torch_file(lora_path)
        lora_weights = warped_load_lora_weights(lora_name)
        lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
        new_weights = {}

        for key in lora_weights:
            new_weights[key] = torch.mul(lora_weights[key], strength).to(dtype=torch.bfloat16)

        # Filter the LoRA weights based on the block type
        filtered_lora = filter_lora_keys(new_weights, blocks_type)

        return new_weights, filtered_lora

    def load_multiple_loras(self, model, **kwargs):
        """Load and apply multiple LoRA models."""
        from comfy.sd import load_lora_for_models

        temp_lora_name = kwargs.get(f"lora_name")
        temp_strength = kwargs.get(f"strength")
        temp_blocks_type = kwargs.get(f"blocks_type")

        lora_metadata = []

        if not temp_lora_name is None and temp_strength != 0:
            print("\n**** Lora Name: {}  |  Strength: {}  |  Block Types: {}\n".format(temp_lora_name, temp_strength, temp_blocks_type))
            lora_metadata.append("{}".format("Lora: {} | Strength: {} | Block Types: {}".format(temp_lora_name, temp_strength, temp_blocks_type)))

            lora_weights, filtered_lora = self.load_lora(temp_lora_name, temp_strength, temp_blocks_type)

            # Apply the LoRA weights to the model
            if filtered_lora:
                model, _ = load_lora_for_models(model, None, filtered_lora, 1.0, 0)
            else:
                model, _ = load_lora_for_models(model, None, lora_weights, 1.0, 0)

        for i in range(1, 5):
            temp_lora_name = kwargs.get(f"lora_0{i}")
            temp_strength = kwargs.get(f"strength_0{i}")
            temp_blocks_type = kwargs.get(f"blocks_type_0{i}")

            if temp_lora_name != "None" and temp_strength != 0:
                temp_message = "{}".format("Lora: {} | Strength: {} | Block Types: {}".format(temp_lora_name, temp_strength, temp_blocks_type))
                print(temp_message)

                lora_metadata.append(temp_message)

                # Load and filter the LoRA weights
                lora_weights, filtered_lora = self.load_lora(temp_lora_name, temp_strength, temp_blocks_type)

                # Apply the LoRA weights to the model
                if filtered_lora:
                    model, _ = load_lora_for_models(model, None, filtered_lora, 1.0, 0)
                else:
                    model, _ = load_lora_for_models(model, None, lora_weights, 1.0, 0)

        return (model, lora_metadata, )

    @classmethod
    def IS_CHANGED(s, **kwargs):
        return f"{kwargs.get('lora_name')}_{kwargs.get('strength')}_{kwargs.get('blocks_type')}_" \
               f"{kwargs.get('lora_01')}_{kwargs.get('strength_01')}_{kwargs.get('blocks_type_01')}_" \
               f"{kwargs.get('lora_02')}_{kwargs.get('strength_02')}_{kwargs.get('blocks_type_02')}_" \
               f"{kwargs.get('lora_03')}_{kwargs.get('strength_03')}_{kwargs.get('blocks_type_03')}_" \
               f"{kwargs.get('lora_04')}_{kwargs.get('strength_04')}_{kwargs.get('blocks_type_04')}"

class WarpedMultiLoraLoader:
    def __init__(self):
        self.loaded_lora = None

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL", {"tooltip": "The diffusion model the LoRA will be applied to."}),
                "clip": ("CLIP", {"tooltip": "The CLIP model the LoRA will be applied to."}),
                "lora_01": (['None'] + folder_paths.get_filename_list("loras"),),
                "strength_model_01": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                "strength_clip_01": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the CLIP model. This value can be negative."}),
                "lora_02": (['None'] + folder_paths.get_filename_list("loras"),),
                "strength_model_02": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                "strength_clip_02": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the CLIP model. This value can be negative."}),
                "lora_03": (['None'] + folder_paths.get_filename_list("loras"),),
                "strength_model_03": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                "strength_clip_03": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the CLIP model. This value can be negative."}),
                "lora_04": (['None'] + folder_paths.get_filename_list("loras"),),
                "strength_model_04": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                "strength_clip_04": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the CLIP model. This value can be negative."}),
            },
          "optional": {
                "lora_name": ("STRING", {"default": None, "forceInput": True}),
                "strength_model": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                "strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the CLIP model. This value can be negative."}),
            },
        }

    RETURN_TYPES = ("MODEL", "CLIP", "STRING", )
    RETURN_NAMES = ("model", "clip", "lora_metadata",)
    OUTPUT_IS_LIST = (False, False, True)
    FUNCTION = "load_multiple_loras"
    CATEGORY = "Warped/General/Loaders"
    DESCRIPTION = "Load and apply multiple LoRA models with different strengths. Model input is required."

    def load_lora(self, model, clip, lora_name, strength_model, strength_clip):
        if strength_model == 0 and strength_clip == 0:
            return model, clip

        lora_path = folder_paths.get_full_path("loras", lora_name)
        lora_path = os.path.abspath(lora_path)

        if not os.path.exists(lora_path):
            raise ValueError(f"LoRA file not found: {lora_path}")

        lora = None
        if self.loaded_lora is not None:
            if self.loaded_lora[0] == lora_path:
                lora = self.loaded_lora[1]
            else:
                self.loaded_lora = None

        if lora is None:
            lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
            self.loaded_lora = (lora_path, lora)

        model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora, strength_model, strength_clip)

        return model_lora, clip_lora

    def load_multiple_loras(self, model, clip, **kwargs):
        """Load and apply multiple LoRA models."""
        from comfy.sd import load_lora_for_models

        temp_lora_name = kwargs.get(f"lora_name")
        temp_strength_model = kwargs.get(f"strength_model")
        temp_strength_clip = kwargs.get(f"strength_clip")

        temp_model = model
        temp_clip = clip

        lora_metadata = []

        if not temp_lora_name is None and ((temp_strength_model != 0) or (temp_strength_clip != 0)):
            print("Lora Name: {}  |  Strength Model: {}  |  Strength Clip: {}".format(temp_lora_name, temp_strength_model, temp_strength_clip))
            lora_metadata.append("Lora: {} | Strength Model: {}  |  Strength clip: {}".format(temp_lora_name, temp_strength_model, temp_strength_clip))

            temp_model, temp_clip = self.load_lora(temp_model, temp_clip, temp_lora_name, temp_strength_model, temp_strength_clip)

        for i in range(1, 5):
            temp_lora_name = kwargs.get(f"lora_0{i}")

            if temp_lora_name == "None":
                continue

            temp_strength_model = kwargs.get(f"strength_model_0{i}")
            temp_strength_clip = kwargs.get(f"strength_clip_0{i}")

            print("Lora Name: {}  |  Strength Model: {}  |  Strength Clip: {}".format(temp_lora_name, temp_strength_model, temp_strength_clip))

            if not temp_lora_name is None and ((temp_strength_model != 0) or (temp_strength_clip != 0)):
                lora_metadata.append("Lora: {} | Strength Model: {}  |  Strength clip: {}".format(temp_lora_name, temp_strength_model, temp_strength_clip))
                # Load LoRA weights
                temp_model, temp_clip = self.load_lora(temp_model, temp_clip, temp_lora_name, temp_strength_model, temp_strength_clip)

        return (temp_model, temp_clip, lora_metadata, )

    @classmethod
    def IS_CHANGED(s, **kwargs):
        return f"{kwargs.get('lora_name')}_{kwargs.get('strength')}_{kwargs.get('blocks_type')}_" \
               f"{kwargs.get('lora_01')}_{kwargs.get('strength_01')}_{kwargs.get('blocks_type_01')}_" \
               f"{kwargs.get('lora_02')}_{kwargs.get('strength_02')}_{kwargs.get('blocks_type_02')}_" \
               f"{kwargs.get('lora_03')}_{kwargs.get('strength_03')}_{kwargs.get('blocks_type_03')}_" \
               f"{kwargs.get('lora_04')}_{kwargs.get('strength_04')}_{kwargs.get('blocks_type_04')}"

def get_save_lora_path(filename_prefix, output_dir):
    def map_filename(filename):
        prefix_len = len(os.path.basename(filename_prefix))
        prefix = filename[:prefix_len + 1]

        try:
            temp_strings = filename.split("_")
            temp_strings2 = temp_strings[len(temp_strings) - 1].split('.')
            digits = int(temp_strings2[0])
        except:
            digits = 0

        return (digits, prefix)

    subfolder = os.path.dirname(os.path.normpath(filename_prefix))
    filename = os.path.basename(os.path.normpath(filename_prefix))

    full_output_folder = os.path.join(output_dir, subfolder)

    full_folder_contents = os.listdir(full_output_folder)
    relevant_folder_contents = []

    for temp in full_folder_contents:
        if temp.startswith(filename_prefix):
            relevant_folder_contents.append(temp)

    if len(relevant_folder_contents) > 0:
        try:
            counter = max(filter(lambda a: a[1][:-1] == filename and a[1][-1] == "_", map(map_filename, relevant_folder_contents)))[0] + 1
        except ValueError:
            counter = 1
        except FileNotFoundError:
            os.makedirs(full_output_folder, exist_ok=True)
            counter = 1
    else:
        counter = 1

    return full_output_folder, filename, counter, subfolder, filename_prefix

class WarpedHunyuanLoraBatchMerge:
    def __init__(self):
        self.base_output_dir = get_default_output_folder()
        os.makedirs(self.base_output_dir, exist_ok = True)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "save_folder": ("STRING", {"default": get_default_output_folder()}),
                "model_prefix": ("STRING", {"default": "new_model_hy"}),
                "lora_1": ("STRING", {"default": None, "forceInput": True}),
                "strength_1": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}),
                "blocks_type_1": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "lora_2": (['None'] + get_lora_list(),),
                "strength_2": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}),
                "blocks_type_2": (["all", "single_blocks", "double_blocks"], {"default": "all"}),
                "save_metadata": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    OUTPUT_IS_LIST = (True,)
    FUNCTION = "merge_multiple_loras"
    CATEGORY = "Warped/Hunyuan/Merge"
    DESCRIPTION = "Load and apply multiple LoRA models with different strengths and block types. Model input is required."

    def load_lora(self, lora_name: str, strength: float, blocks_type: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
        """Load and filter a single LoRA model."""
        if not lora_name or strength == 0:
            return {}, {}

        # Get the full path to the LoRA file
        lora_path = folder_paths.get_full_path("loras", lora_name)
        if not os.path.exists(lora_path):
            raise ValueError(f"LoRA file not found: {lora_path}")

        # Load the LoRA weights
        lora_weights = utils.load_torch_file(lora_path)
        lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")
        filtered_lora = filter_lora_keys(lora_weights, blocks_type)

        return lora_weights, filtered_lora

    def merge_multiple_loras(self, save_folder, model_prefix, lora_1, strength_1, blocks_type_1, lora_2, strength_2, blocks_type_2, save_metadata=True):
        """Load and apply multiple LoRA models."""
        temp_loras = {}
        metadata = {"loras": "{} and {}".format(lora_1, lora_2)}
        metadata["strengths"] = "{} and {}".format(strength_1, strength_2)
        metadata["block_types"] = "{} and {}".format(blocks_type_1, blocks_type_2)

        print("Processing Lora: {}".format(lora_1))

        if lora_1 != "None" and strength_1 != 0:
            # Load and filter the LoRA weights
            lora_weights, filtered_lora = self.load_lora(lora_1, 1.0, blocks_type_1)
            temp_loras["1"] = {"lora_weights": lora_weights, "strength": strength_1, "filtered_lora": filtered_lora}

        if lora_2 != "None" and strength_2 != 0:
            # Load and filter the LoRA weights
            lora_weights, filtered_lora = self.load_lora(lora_2, 1.0, blocks_type_2)
            temp_loras["2"] = {"lora_weights": lora_weights, "strength": strength_2, "filtered_lora": filtered_lora}

        new_lora = {}

        for lora_key in temp_loras.keys():
            for key in temp_loras[lora_key]["filtered_lora"].keys():
                if not key in new_lora.keys():
                    new_lora[key] = None

        # Merge The Weighted Key Weights
        for key in new_lora.keys():
            for lora_key in temp_loras.keys():
                if key in temp_loras[lora_key]["filtered_lora"].keys():
                    if not new_lora[key] is None:
                        temp_weights = torch.mul(temp_loras[lora_key]["filtered_lora"][key], temp_loras[lora_key]["strength"])

                        if new_lora[key].shape[0] < new_lora[key].shape[1]:
                            if temp_weights.shape[0] < new_lora[key].shape[0]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([new_lora[key].shape[0], new_lora[key].shape[1]])
                                padding[:temp_weights.shape[0],:] = temp_weights
                                temp_weights = padding
                            elif temp_weights.shape[0] > new_lora[key].shape[0]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([temp_weights.shape[0], temp_weights.shape[1]])
                                padding[:new_lora[key].shape[0],:] = new_lora[key]
                                new_lora[key] = padding
                        else:
                            if temp_weights.shape[1] < new_lora[key].shape[1]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([new_lora[key].shape[0], new_lora[key].shape[1]])
                                padding[:,:temp_weights.shape[1]] = temp_weights
                                temp_weights = padding
                            elif temp_weights.shape[1] > new_lora[key].shape[1]:
                                temp_weights = temp_weights.clone().detach()
                                new_lora[key] = new_lora[key].clone().detach()

                                padding = torch.zeros([temp_weights.shape[0], temp_weights.shape[1]])
                                padding[:,:new_lora[key].shape[1]] = new_lora[key]
                                new_lora[key] = padding

                        try:
                            new_lora[key] = torch.add(new_lora[key], temp_weights)
                        except Exception as e:
                            raise(e)
                    else:
                        new_lora[key] = torch.mul(temp_loras[lora_key]["filtered_lora"][key], temp_loras[lora_key]["strength"])

        if not save_metadata:
            metadata = None

        full_output_path, filename, counter, subfolder, filename_prefix = get_save_lora_path(model_prefix, self.base_output_dir)

        output_filename = os.path.join(full_output_path, "{}_{:05}.safetensors".format(model_prefix, counter))
        utils.save_torch_file(new_lora, output_filename, metadata=metadata)

        save_message = "Weights Saved To: {}".format(output_filename)
        print(save_message)

        return {"ui": {"tags": [save_message]}}

class WarpedHunyuanLoraConvertToMusubi:
    def __init__(self):
        self.base_output_dir = get_default_output_folder()
        os.makedirs(self.base_output_dir, exist_ok = True)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "save_path": ("STRING", {"default": get_default_output_path()}),
                "lora": (get_lora_list(),),
                "save_metadata": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    OUTPUT_IS_LIST = (True,)
    FUNCTION = "convert_lora"
    CATEGORY = "Warped/Hunyuan/Lora/Experimental"
    DESCRIPTION = "Convert Keys For Hunyuan LORA and Save Modified LORA."

    def load_lora(self, lora_name: str) -> Tuple[Dict[str, torch.Tensor],]:
        """Load and filter a single LoRA model."""
        if not lora_name:
            return {}

        # Get the full path to the LoRA file
        lora_path = folder_paths.get_full_path("loras", lora_name)
        if not os.path.exists(lora_path):
            raise ValueError(f"LoRA file not found: {lora_path}")

        # Load the LoRA weights
        lora_weights = utils.load_torch_file(lora_path)

        return lora_weights

    def convert_lora(self, save_path, lora, save_metadata=True):
        metadata = {"original_lora": "{}".format(lora)}

        # Load the LoRA weights
        temp_lora = self.load_lora(lora)
        new_lora = convert_to_musubi(temp_lora)

        if not save_metadata:
            metadata = None

        if len(new_lora) < 1:
            utils.save_torch_file(temp_lora, save_path, metadata=metadata)
        else:
            utils.save_torch_file(new_lora, save_path, metadata=metadata)

        save_message = "Weights Saved To: {}".format(save_path)

        return {"ui": {"tags": [save_message]}}

class WarpedHunyuanLoraConvertKeys:
    def __init__(self):
        self.base_output_dir = get_default_output_folder()
        os.makedirs(self.base_output_dir, exist_ok = True)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "save_path": ("STRING", {"default": get_default_output_path()}),
                "lora": (get_lora_list(),),
                "convert_to": (["diffusion_model", "transformer"], {"default": "diffusion_model"}),
                "save_metadata": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    OUTPUT_IS_LIST = (True,)
    FUNCTION = "convert_lora"
    CATEGORY = "Warped/Hunyuan/Lora"
    DESCRIPTION = "Convert Keys For Hunyuan LORA and Save Modified LORA."

    def load_lora(self, lora_name: str) -> Tuple[Dict[str, torch.Tensor],]:
        """Load and filter a single LoRA model."""
        if not lora_name:
            return {}

        # Get the full path to the LoRA file
        lora_path = folder_paths.get_full_path("loras", lora_name)
        if not os.path.exists(lora_path):
            raise ValueError(f"LoRA file not found: {lora_path}")

        # Load the LoRA weights
        lora_weights = utils.load_torch_file(lora_path)

        return lora_weights

    def convert_lora(self, save_path, lora, convert_to, save_metadata=True):
        metadata = {"original_lora": "{}".format(lora)}

        # Load the LoRA weights
        temp_lora = self.load_lora(lora)
        new_lora = convert_lora(temp_lora, convert_to=convert_to)

        if not save_metadata:
            metadata = None

        if len(new_lora) < 1:
            utils.save_torch_file(temp_lora, save_path, metadata=metadata)
        else:
            utils.save_torch_file(new_lora, save_path, metadata=metadata)

        save_message = "Weights Saved To: {}".format(save_path)

        return {"ui": {"tags": [save_message]}}

# class WarpedHunyuanLoraConvertKeys2:
#     def __init__(self):
#         self.base_output_dir = get_default_output_folder()
#         os.makedirs(self.base_output_dir, exist_ok = True)
#
#     @classmethod
#     def INPUT_TYPES(s):
#         return {
#             "required": {
#                 "save_path": ("STRING", {"default": get_default_output_path()}),
#                 "lora": (get_lora_list(),),
#                 "convert_to": (["diffusion_model", "transformer", "framepack"], {"default": "diffusion_model"}),
#                 "save_metadata": ("BOOLEAN", {"default": True}),
#             },
#         }
#
#     RETURN_TYPES = ()
#     OUTPUT_NODE = True
#     OUTPUT_IS_LIST = (True,)
#     FUNCTION = "convert_lora"
#     CATEGORY = "Warped/Hunyuan/Lora"
#     DESCRIPTION = "Convert Keys For Hunyuan LORA and Save Modified LORA."
#
#     def load_lora(self, lora_name: str) -> Tuple[Dict[str, torch.Tensor],]:
#         """Load and filter a single LoRA model."""
#         if not lora_name:
#             return {}
#
#         # Get the full path to the LoRA file
#         lora_path = folder_paths.get_full_path("loras", lora_name)
#         if not os.path.exists(lora_path):
#             raise ValueError(f"LoRA file not found: {lora_path}")
#
#         # Load the LoRA weights
#         lora_weights = utils.load_torch_file(lora_path)
#
#         return lora_weights
#
#     def convert_lora(self, save_path, lora, convert_to, save_metadata=True):
#         metadata = {"original_lora": "{}".format(lora)}
#
#         # Load the LoRA weights
#         temp_lora = self.load_lora(lora)
#
#         if convert_to != "framepack":
#             new_lora = convert_lora(temp_lora, convert_to=convert_to)
#         else:
#             new_lora = convert_lora(temp_lora, convert_to="diffusion_model")
#
#         for key in new_lora.keys():
#             print("LORA Key: {}".format(key))
#
#         if not save_metadata:
#             metadata = None
#
#         if len(new_lora) < 1:
#             utils.save_torch_file(temp_lora, save_path, metadata=metadata)
#         else:
#             utils.save_torch_file(new_lora, save_path, metadata=metadata)
#
#         save_message = "Weights Saved To: {}".format(save_path)
#
#         return {"ui": {"tags": [save_message]}}

class WarpedLoraKeysAndMetadataReader:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "lora_name": (folder_paths.get_filename_list("loras"),),
            },
        }

    RETURN_TYPES = ("STRING", "STRING", )
    RETURN_NAMES = ("keys", "metadata", )
    OUTPUT_IS_LIST = (True, True, )
    FUNCTION = "read_data"
    CATEGORY = "Warped/General/Lora"
    DESCRIPTION = "Read Metadata From Lora."

    def get_metadata(self, lora_path):
        # Open the file in binary mode
        with open(lora_path, 'rb') as file:
            length_of_header_bytes = file.read(8)
            # Interpret the bytes as a little-endian unsigned 64-bit integer
            length_of_header = struct.unpack('<Q', length_of_header_bytes)[0]
            header_bytes = file.read(length_of_header)
            #header = json.loads(header_bytes.decode('utf-8'))
            metadata = json.loads(header_bytes)

            try:
                return metadata["__metadata__"]
            except:
                pass

        return {"metadata": "No Metadata" }

    def load_lora(self, lora_name: str, strength: 1.0, blocks_type: "all") -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
        """Load and filter a single LoRA model."""
        if not lora_name:
            return "", ""

        # Get the full path to the LoRA file
        lora_path = folder_paths.get_full_path("loras", lora_name)
        if not os.path.exists(lora_path):
            raise ValueError(f"LoRA file not found: {lora_path}")

        lora_weights = comfy.utils.load_torch_file(lora_path)

        metadata = self.get_metadata(lora_path)

        return lora_weights, metadata

    def read_data(self, lora_name):
        metadata = {}
        keys = []

        if lora_name != "None":
            # Load and filter the LoRA weights
            lora_weights, metadata = self.load_lora(lora_name, 1.0, "all")

        lora_keys = []
        for key in lora_weights.keys():
            lora_keys.append("{}  |  Shape: {}\n".format(key, lora_weights[key].shape))
            # print(key)

        lora_metadata = []
        if len(metadata.keys()) > 0:
            for key in metadata.keys():
                lora_metadata.append("{}: {}".format(key, metadata[key]))
                # print("{}: {}".format(key, metadata[key]))

        lora_weights = None

        mm.soft_empty_cache()
        gc.collect()
        time.sleep(1)

        return { "ui": { "lora_keys": lora_keys, "lora_metadata": lora_metadata }, "result": (lora_keys, lora_metadata,), }

class WarpedHunyuanLoraConvert:
    def __init__(self):
        self.base_output_dir = get_default_output_folder()
        os.makedirs(self.base_output_dir, exist_ok = True)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "save_path": ("STRING", {"default": get_default_output_path()}),
                "lora": (get_lora_list(),),
                "convert_to": (['32', '64', '128'], {"default": "32"}),
                "save_metadata": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    OUTPUT_IS_LIST = (True,)
    FUNCTION = "convert_lora"
    CATEGORY = "Warped/Hunyuan/Lora"
    DESCRIPTION = "Convert Hunyuan Lora."

    def load_lora(self, lora_name: str) -> Tuple[Dict[str, torch.Tensor]]:
        if not lora_name:
            return {}

        # Get the full path to the LoRA file
        lora_path = folder_paths.get_full_path("loras", lora_name)
        if not os.path.exists(lora_path):
            raise ValueError(f"LoRA file not found: {lora_path}")

        # Load the LoRA weights
        lora_weights = utils.load_torch_file(lora_path)

        return lora_weights

    def convert_lora(self, save_path, lora, convert_to="32", save_metadata=True):
        metadata = {"original_lora": "{}".format(lora)}

        if lora != "None":
            # Load and filter the LoRA weights
            lora_weights = self.load_lora(lora)
            lora_weights = convert_lora(lora_weights, convert_to="diffusion_model")

        for key in lora_weights.keys():
            sample_shape = lora_weights[key].shape
            print("Lora Shape Before: {}".format(sample_shape))

            if sample_shape[0] < sample_shape[1]:
                sample_res = int(sample_shape[0])
            else:
                sample_res = int(sample_shape[1])

            break

        target_res = int(convert_to)

        if target_res == sample_res:
            save_message = "LORA Resolution and Target Resolution are the same. Nothing to convert."
            print(save_message)
            return {"ui": {"tags": [save_message]}}

        new_lora = {}

        # Convert The LORA Weights
        for key in lora_weights.keys():
            temp_weights = lora_weights[key].clone().detach().to(dtype=torch.bfloat16)

            if temp_weights.shape[0] < temp_weights.shape[1]:
                padding = torch.zeros([target_res, temp_weights.shape[1]]).to(dtype=torch.bfloat16)

                # if upscale
                if temp_weights.shape[0] < target_res:
                    padding[:temp_weights.shape[0],:] = temp_weights
                    new_lora[key] = padding
                # if downscale
                else:
                    padding[:target_res,:] = temp_weights[:target_res,:]
                    new_lora[key] = padding
            else:
                padding = torch.zeros([temp_weights.shape[0], target_res]).to(dtype=torch.bfloat16)

                # if upscale
                if temp_weights.shape[1] < target_res:
                    padding[:,:temp_weights.shape[1]] = temp_weights
                    new_lora[key] = padding
                # if downscale
                else:
                    padding[:,:target_res] = temp_weights[:,:target_res]
                    new_lora[key] = padding

        if not save_metadata:
            metadata = None

        utils.save_torch_file(new_lora, save_path, metadata=metadata)

        save_message = "Weights Saved To: {}".format(save_path)

        return {"ui": {"tags": [save_message]}}

class WarpedFramepackSampler:
    def __init__(self):
        self.clip_vision = None

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("FramePackMODEL",),
                "vae": ("VAE",),
                "positive": ("CONDITIONING",),
                "negative": ("CONDITIONING",),
                "steps": ("INT", {"default": 25, "min": 1}),
                "cache_mode": (["disabled", "use_teacache", "use_magcache"], {"default": "use_magcache", "tooltip": "Whether or not to use magcache or teacache for faster sampling."}),
                "teacache_rel_l1_thresh": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The threshold for the relative L1 loss."}),
                "magcache_thresh": ("FLOAT", {"default": 0.1250, "min": 0.0000, "max": 0.3000, "step": 0.0010, "tooltip": "How strongly to cache the output of diffusion model. This value must be non-negative."}),
                "magcache_retention_ratio": ("FLOAT", {"default": 0.225, "min": 0.100, "max": 0.300, "step": 0.001, "tooltip": "The start percentage of the steps that will apply MagCache."}),
                "magcache_K": ("INT", {"default": 3, "min": 0, "max": 6, "step": 1, "tooltip": "The maxium skip steps of MagCache."}),
                "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 30.0, "step": 0.01}),
                "guidance_scale": ("FLOAT", {"default": 12.0, "min": 0.0, "max": 32.0, "step": 0.01}),
                "shift": ("FLOAT", {"default": 24.0, "min": 0.0, "max": 1000.0, "step": 0.01}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                "preferred_frame_count": ("INT", {"default": 301, "min": 33, "max": 1000001, "step": 4, "tooltip": "For I2V and T2V, The total frames in the video. Disreguarded for V2V"}),
                "preferred_batch_size": ("INT", {"default": 61, "min": 17, "max": 161, "step": 4, "tooltip": "The preferred number of frames to use for sampling."}),
                "use_batch_size": (["next_lowest", "next_highest", "closest", "exact"], {"default": "next_lowest", "tooltip": "Number of frames generated may be impacted by choice."}),
                "gpu_memory_preservation": ("FLOAT", {"default": 12.0, "min": 0.0, "max": 128.0, "step": 0.1, "tooltip": "The amount of GPU memory to preserve."}),
                "sampler": (["unipc_bh1", "unipc_bh2"],
                    {
                        "default": 'unipc_bh1'
                    }),
                "dec_tile_size": ("INT", {"default": 256, "min": 64, "max": 4096, "step": 32}),
                "dec_overlap": ("INT", {"default": 64, "min": 0, "max": 4096, "step": 32}),
                "dec_temporal_size": ("INT", {"default": 64, "min": 8, "max": 4096, "step": 4, "tooltip": "Only used for video VAEs: Amount of frames to decode at a time."}),
                "dec_temporal_overlap": ("INT", {"default": 8, "min": 4, "max": 4096, "step": 4, "tooltip": "Only used for video VAEs: Amount of frames to overlap."}),
                "skip_frames": ("INT", {"default": 0, "min": 0, "max": 128, "step": 4, "tooltip": "Number of frames to skip from beginning for output."}),
                "clip_vision_model": ("CLIP_VISION", ),
            },
            "optional": {
                "start_image": ("IMAGE", {"tooltip": "init image to use for image2video or video2video"} ),
                "end_image": ("IMAGE", {"tooltip": "end image to use for image2video"} ),
                "embed_interpolation": (["disabled", "weighted_average", "linear"], {"default": 'disabled', "tooltip": "Image embedding interpolation type. If linear, will smoothly interpolate with time, else it'll be weighted average with the specified weight."}),
                "start_embed_strength": ("FLOAT", {"default": 1.20, "min": 0.00, "max": 2.00, "step": 0.01, "tooltip": "Weighted average constant for image embed interpolation. If end image is not set, the embed's strength won't be affected"}),
                "secondary_embed_strength": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01, "tooltip": "Weighted average constant for image embed interpolation. If end image is not set, the embed's strength won't be affected"}),
                "video_image_batch": ("IMAGE", {"tooltip": "init Latents to use for video2video"} ),
                "denoise_strength": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 1.00, "step": 0.01}),
                "noise_strength": ("FLOAT", {"default": 1.00, "min": 0.10, "max": 1.00, "step": 0.01}),
                "blend_frames": ("INT", {"default":5, "min":0, "max": 16, "step": 1}),
                "t2v_width": ("INT", {"default":480, "min":256, "max": 1280, "step": 16}),
                "t2v_height": ("INT", {"default":720, "min":256, "max": 1280, "step": 16}),
                "dummy_frames": ("INT", {"default":17, "min":17, "max": 161, "step": 4, "tooltip": "Number of frames to generate in dummy batch."}),
                "gen_dummy": ("BOOLEAN", {"default": False, "tooltip": "For t2v or i2v only. Will generate a dummy batch to obtain a starting image for main generation."}),
                "gen_dummy_only": ("BOOLEAN", {"default": False, "tooltip": "Will generate dummy batch only."}),
                "dummy_cache_mode": (["disabled", "use_teacache", "use_magcache"], {"default": "use_magcache", "tooltip": "Whether or not to use magcache or teacache on dummy generation for faster sampling."}),
                "use_dummy_image": (["first", "middle", "last", "random", "all"], {"default": "last", "tooltip": "Which dummy batch image to start main generation."}),
                "v2v_context_count": ("INT", {"default":5, "min":3, "max": 10, "step": 1}),
                "verbose_messaging": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ("IMAGE", "STRING", "INT", "BOOLEAN",)
    RETURN_NAMES = ("images", "generation_status", "seed", "valid_output",)
    FUNCTION = "process"
    CATEGORY = "Warped/Framepack/Sampling"

    def process(self, model, vae, shift, positive, negative, preferred_frame_count, preferred_batch_size, use_batch_size, cache_mode, teacache_rel_l1_thresh,
                magcache_thresh, magcache_retention_ratio, magcache_K, steps, cfg,
                guidance_scale, seed, sampler, dec_tile_size, dec_overlap, dec_temporal_size, dec_temporal_overlap, skip_frames, clip_vision_model,
                gpu_memory_preservation, start_image=None, end_image=None, embed_interpolation="linear", start_embed_strength=1.0, secondary_embed_strength=1.0, video_image_batch=None,
                denoise_strength=1.00, noise_strength=1.00, blend_frames=0, t2v_width=640, t2v_height=640, dummy_frames=5, gen_dummy=False, gen_dummy_only=False, dummy_cache_mode="disabled", use_dummy_image="last",
                v2v_context_count=5, verbose_messaging=False):
        self.dec_tile_size = dec_tile_size
        self.dec_overlap = dec_overlap
        self.dec_temporal_size = dec_temporal_size
        self.dec_temporal_overlap = dec_temporal_overlap
        self.skip_frames = skip_frames
        self.vae = vae
        self.preferred_frame_count = preferred_frame_count
        self.seed = seed
        self.steps = steps
        self.cfg = cfg
        self.use_teacache = False
        self.use_magcache = False
        self.teacache_rel_l1_thresh = teacache_rel_l1_thresh
        self.magcache_thresh = magcache_thresh
        self.magcache_retention_ratio = magcache_retention_ratio
        self.magcache_K = magcache_K
        self.guidance_scale = guidance_scale
        self.sampler = sampler
        self.skip_frames = skip_frames
        self.gpu_memory_preservation = gpu_memory_preservation
        self.embed_interpolation = embed_interpolation
        self.start_embed_strength = start_embed_strength
        self.secondary_embed_strength = secondary_embed_strength
        self.denoise_strength = denoise_strength
        self.noise_strength = noise_strength
        self.transformer = model["transformer"]
        self.base_dtype = model["dtype"]
        self.positive = positive
        self.negative = negative
        self.device = mm.get_torch_device()
        self.offload_device = get_offload_device()
        self.clip_vision = clip_vision_model
        self.blend_frames = blend_frames
        self.shift = shift
        self.t2v_width=t2v_width
        self.t2v_height=t2v_height
        self.dummy_latents = int((dummy_frames - 1) // 4) + 1
        self.gen_dummy = gen_dummy
        self.gen_dummy_only = gen_dummy_only
        self.dummy_teacache = False
        self.dummy_magcache = False
        self.use_dummy_image = use_dummy_image
        self.v2v_context_count = v2v_context_count
        self.verbose_messaging = verbose_messaging

        if cache_mode == "use_teacache":
            self.use_teacache = True

        if cache_mode == "use_magcache":
            self.use_magcache = True

        if dummy_cache_mode == "use_teacache":
            self.dummy_teacache = True

        if dummy_cache_mode == "use_magcache":
            self.dummy_magcache = True

        print("Device: {}  | Offload Device: {}".format(self.device, self.offload_device))

        self.llama_vec = self.positive[0][0].to(self.base_dtype).to(self.device)
        self.clip_l_pooler = self.positive[0][1]["pooled_output"].to(self.base_dtype).to(self.device)

        if not math.isclose(self.cfg, 1.0):
            self.llama_vec_n = self.negative[0][0].to(self.base_dtype)
            self.clip_l_pooler_n = self.negative[0][1]["pooled_output"].to(self.base_dtype).to(self.device)
        else:
            self.llama_vec_n = torch.zeros_like(self.llama_vec, device=self.device)
            self.clip_l_pooler_n = torch.zeros_like(self.clip_l_pooler, device=self.device)

        self.llama_vec, self.llama_attention_mask = crop_or_pad_yield_mask(self.llama_vec, length=512)
        self.llama_vec_n, self.llama_attention_mask_n = crop_or_pad_yield_mask(self.llama_vec_n, length=512)

        if not video_image_batch is None:
            if len(video_image_batch.shape) < 4:
                video_image_batch = video_image_batch.unsqueeze(0)

            self.preferred_frame_count = video_image_batch.shape[0]

            self.latent_window_size, self.batch_count, truncated_frame_count = self.get_latent_window_size(preferred_batch_size, self.preferred_frame_count, use_batch_size=use_batch_size)

            if video_image_batch.shape[0] != truncated_frame_count:
                image_batches_tuple = torch.split(video_image_batch, truncated_frame_count, dim=0)
                image_batches_split = [item for item in image_batches_tuple]
                video_image_batch = image_batches_split[0]
        else:
            print("preferred_frame_count: {}".format(self.preferred_frame_count))
            self.latent_window_size, self.batch_count, truncated_frame_count = self.get_latent_window_size(preferred_batch_size, self.preferred_frame_count, use_batch_size=use_batch_size)

        self.num_frames = self.encoded_to_decoded_length(self.latent_window_size)
        self.total_frames = truncated_frame_count

        if self.verbose_messaging:
            print("num_frames: {}".format(self.num_frames))

        self.cleanup(unload_models=True, cleanup_models=True, cleanup_cuda=True)

        if not video_image_batch is None:
            self.mode = VideoGenerationType.V2V

            return self.process_v2v(video_image_batch)

        if not start_image is None:
            self.mode = VideoGenerationType.I2V
            return self.process_i2v(start_image, end_image)

        self.mode = VideoGenerationType.T2V

        return self.process_t2v()

    def process_t2v(self):
        return self.process_i2v(None, None)

    def process_i2v(self, start_image, end_image):
        start_latent = None
        end_latent = None

        if self.mode == VideoGenerationType.I2V:
            if not start_image is None:
                image_embeds = self.clip_vision_encode(start_image)
                start_latent = self.encode_batched(start_image, per_batch=self.latent_window_size)

                start_latent = start_latent["samples"] * vae_scaling_factor
                print("start_latent Shape: {}".format(start_latent.shape))

            if not end_image is None:
                if (start_image.shape[1] != end_image.shape[1]) or (start_image.shape[2] != end_image.shape[2]):
                    raise ValueError("Unable to continue: end_image Height/Width Does Not Match start_image Height/Width.")

                end_latent = self.encode_batched(end_image, per_batch=self.latent_window_size)
                print("end_latent Shape: {}".format(end_latent["samples"].shape))
                end_image_embeds = self.clip_vision_encode(end_image)
        else:
            start_image = torch.zeros(size=(1, self.t2v_height, self.t2v_width, 3), dtype=torch.float32).cpu()
            start_latent = torch.zeros(size=(1, 16, 1, int(self.t2v_height // 8), int(self.t2v_width // 8)), dtype=torch.float32).cpu()
            image_embeds = self.clip_vision_encode(start_image)

        self.width = start_image.shape[2]
        self.height = start_image.shape[1]

        total_latent_sections = self.batch_count

        if self.verbose_messaging:
            print("total_latent_sections: ", total_latent_sections)

        if end_latent is not None:
            end_latent = end_latent["samples"] * vae_scaling_factor
        has_end_image = end_latent is not None

        if self.verbose_messaging:
            print("start_latent", start_latent.shape)

        B, C, T, H, W = start_latent.shape

        if self.verbose_messaging:
            print("process_i2v: Image Width: {}".format(self.width))
            print("process_i2v: Image Height: {}".format(self.height))
            print("process_i2v: total_frames: {}".format(self.total_frames))
            print("process_i2v: batch_count: {}".format(self.batch_count))
            print("process_i2v: latent_window_size: {}".format(self.latent_window_size))

        start_image_encoder_last_hidden_state = image_embeds["last_hidden_state"].to(self.base_dtype).to(self.device)

        if has_end_image:
            assert end_image_embeds is not None
            end_image_encoder_last_hidden_state = end_image_embeds["last_hidden_state"].to(self.base_dtype).to(self.device)
        else:
            end_image_encoder_last_hidden_state = torch.zeros_like(start_image_encoder_last_hidden_state)

        rnd = torch.Generator("cpu").manual_seed(self.seed)
        total_generated_latent_frames = 0

        if self.mode == VideoGenerationType.I2V:
            if start_latent is None:
                raise ValueError("A start_image value is required for I2V.")

            cat_list = []

            history_latents = torch.zeros(size=(1, 16, 19, H, W), dtype=torch.float32).cpu()
            original_history_latents = history_latents.clone().detach()

            cat_list.append(history_latents)
            cat_list.append(start_latent.to(history_latents))

            history_latents = torch.cat(cat_list, dim=2)
        else:
            history_latents = torch.zeros(size=(1, 16, 20, H, W), dtype=torch.float32).cpu()
            original_history_latents = torch.zeros(size=(1, 16, 19, H, W), dtype=torch.float32).cpu()

        latent_paddings_list = list(range(total_latent_sections))
        latent_paddings = latent_paddings_list.copy()  # Create a copy for iteration

        comfy_model, patcher, callback = self.initialize_comfy_model()

        latent_batches_gend = []

        is_dummy_section = False
        was_dummy_used = False

        noise, dummy_noise = self.setup_latent_noise(self.latent_window_size * self.batch_count, W, H)

        if self.gen_dummy:
            # dummy_noise = self.setup_dummy_noise(W, H)

            if not self.gen_dummy_only:
                latent_paddings.insert(0, 0)
            else:
                latent_paddings = [0]
                latent_paddings_list = latent_paddings.copy()
                total_latent_sections = 1

            is_dummy_section = True
            dummy_frame_count = ((self.dummy_latents - 1) * 4) + 1

        if self.verbose_messaging:
            print("latent_paddings: {}".format(latent_paddings))
            print("noise batches length: {}".format(len(noise)))

        generated_latents = None
        interrupted = False
        real_history_latents = None
        context_latents = None
        generation_status = ""
        temp_history_latents = None
        dummy_images = None
        dummy_gen_latents = None
        generated_image_batches = []
        valid_output = False

        self.cleanup(unload_models=True, cleanup_models=False, cleanup_cuda=True)

        try:
            for i, latent_padding in enumerate(latent_paddings):
                if not self.gen_dummy_only:
                    is_first_section = latent_padding == min(latent_paddings)
                    is_last_section = latent_padding == max(latent_paddings)
                else:
                    is_first_section = True
                    is_last_section = True

                if self.verbose_messaging:
                    print("history_latents Shape: {}".format(history_latents.shape))

                if not is_dummy_section:
                    noise_latent = noise[latent_padding]
                else:
                    noise_latent = dummy_noise

                latent_padding_size = latent_padding * self.latent_window_size

                if self.verbose_messaging:
                    print("latent_padding_size {}: {}  |  latent_padding: {} latent_window_size: {}".format(latent_padding, latent_padding_size, latent_padding, self.latent_window_size))

                if self.embed_interpolation != "disabled":
                    if self.embed_interpolation == "linear":
                        if total_latent_sections <= 1:
                            frac = 1.0  # Handle case with only one section
                        else:
                            frac = 1 - (latent_padding / (total_latent_sections - 1))  # going backwards
                    else:
                        frac = self.start_embed_strength if has_end_image else 1.0

                    image_encoder_last_hidden_state = ((start_image_encoder_last_hidden_state * frac) + ((1 - frac) * end_image_encoder_last_hidden_state)) * self.start_embed_strength
                else:
                    if is_dummy_section or (is_first_section and not self.gen_dummy):
                        image_encoder_last_hidden_state = start_image_encoder_last_hidden_state * self.start_embed_strength
                    else:
                        image_encoder_last_hidden_state = start_image_encoder_last_hidden_state * self.secondary_embed_strength

                start_latent_frames = T  # 0 or 1

                if self.verbose_messaging:
                    print(f'latent_padding_size = {latent_padding_size}, is_last_section = {is_last_section}, is_first_section = {is_first_section}, is_dummy_section = {is_dummy_section}')

                if not is_dummy_section:
                    indices = torch.arange(0, sum([1, 16, 2, 1, self.latent_window_size])).unsqueeze(0)
                    clean_latent_indices_start, clean_latent_4x_indices, clean_latent_2x_indices, clean_latent_1x_indices, latent_indices = indices.split([1, 16, 2, 1, self.latent_window_size], dim=1)
                    clean_latent_indices = torch.cat([clean_latent_indices_start, clean_latent_1x_indices], dim=1)
                    clean_latents_4x, clean_latents_2x, clean_latents_1x = history_latents[:, :, -sum([16, 2, 1]):, :, :].split([16, 2, 1], dim=2)
                else:
                    indices = torch.arange(0, sum([1, 16, 2, 1, self.dummy_latents])).unsqueeze(0)
                    clean_latent_indices_start, clean_latent_4x_indices, clean_latent_2x_indices, clean_latent_1x_indices, latent_indices = indices.split([1, 16, 2, 1, self.dummy_latents], dim=1)
                    clean_latent_indices = torch.cat([clean_latent_indices_start, clean_latent_1x_indices], dim=1)
                    clean_latents_4x, clean_latents_2x, clean_latents_1x = history_latents[:, :, -sum([16, 2, 1]):, :, :].split([16, 2, 1], dim=2)

                if is_last_section and (not end_latent is None):
                    clean_latents = torch.cat([start_latent.to(history_latents), end_latent.to(history_latents)], dim=2)
                else:
                    clean_latents = torch.cat([start_latent.to(history_latents), clean_latents_1x], dim=2)

                if self.verbose_messaging:
                    print("history_latents Shape: {}\n".format(history_latents.shape))

                    print("indices: {}".format(indices))
                    print("latent_indices: {}\n".format(latent_indices))
                    print("clean_latent_2x_indices: {}".format(clean_latent_2x_indices))
                    print("clean_latent_4x_indices: {}".format(clean_latent_4x_indices))
                    print("clean_latent_indices: {}\n".format(clean_latent_indices))

                    print("clean_latents_2x Shape: {}".format(clean_latents_2x.shape))
                    print("clean_latents_4x Shape: {}".format(clean_latents_4x.shape))
                    print("clean_latents Shape: {}\n".format(clean_latents.shape))
                    print("noise Shape: {}\n".format(noise_latent.shape))

                if (self.use_teacache and not is_dummy_section) or (is_dummy_section and self.dummy_teacache):
                    self.transformer.__class__.forward = self.transformer.orig_forward
                    self.transformer.initialize_teacache(enable_teacache=True, num_steps=self.steps, rel_l1_thresh=self.teacache_rel_l1_thresh, verbose_messaging=self.verbose_messaging)
                    self.transformer.initialize_magcache(enable_magcache=False)

                    if self.verbose_messaging:
                        print("TeaCache Enabled")
                elif (self.use_magcache and not is_dummy_section) or (is_dummy_section and self.dummy_magcache):
                    self.transformer.__class__.forward = self.transformer.__class__.magcache_framepack_forward
                    self.transformer.initialize_magcache(enable_magcache=True, num_steps=self.steps, magcache_thresh=self.magcache_thresh, K=self.magcache_K, retention_ratio=self.magcache_retention_ratio, verbose_messaging=self.verbose_messaging)
                    self.transformer.initialize_teacache(enable_teacache=False)

                    if self.verbose_messaging:
                        print("MagCache Enabled")
                else:
                    self.transformer.__class__.forward = self.transformer.orig_forward
                    self.transformer.initialize_teacache(enable_teacache=False)
                    self.transformer.initialize_magcache(enable_magcache=False)

                    if self.verbose_messaging:
                        print("Both TeaCache and MagCache Disabled")

                if not is_dummy_section:
                    print("Generating Batch {}".format(latent_padding))
                else:
                    print("Generating Dummy Batch.")

                generated_latents = self.generate_video(W * 8, H * 8, self.num_frames if not is_dummy_section else dummy_frame_count, rnd, noise_latent, self.llama_vec, self.llama_attention_mask, self.clip_l_pooler, self.llama_vec_n,
                                                        self.llama_attention_mask_n, self.clip_l_pooler_n, image_encoder_last_hidden_state, latent_indices, clean_latents, clean_latent_indices, clean_latents_2x, clean_latent_2x_indices,
                                                        clean_latents_4x, clean_latent_4x_indices, callback)
                if not is_dummy_section:
                    noise[latent_padding] = None
                else:
                    dummy_noise = None

                if self.verbose_messaging:
                    print("generated_latents {} | Shape: {}".format(latent_padding, generated_latents.shape))

                if not is_dummy_section:
                    offload_model_from_device_for_memory_preservation(self.transformer, self.device, preserved_memory_gb=0)
                    self.cleanup(unload_models=True, cleanup_models=False, cleanup_cuda=True)

                    decoded_images = partial_decode_tiled(self.vae, generated_latents.clone().detach()  / vae_scaling_factor, self.dec_tile_size, self.dec_overlap, self.dec_temporal_size, self.dec_temporal_overlap)
                    generated_image_batches.append(decoded_images.clone().detach())
                    valid_output = True

                    decoded_tuple = torch.split(decoded_images, 1, dim=0)
                    decoded_split = [item for item in decoded_tuple]
                    decoded_images = decoded_images.to(get_offload_device())
                    decoded_images = None

                    start_image = decoded_split[len(decoded_split) - 1]
                    image_embeds = self.clip_vision_encode(start_image)
                    start_latent = self.encode_batched(start_image, per_batch=self.latent_window_size)

                    start_latent = start_latent["samples"] * vae_scaling_factor

                    if self.verbose_messaging:
                        print("start_latent Shape: {}".format(start_latent.shape))

                    start_image_encoder_last_hidden_state = image_embeds["last_hidden_state"].to(self.base_dtype).to(self.device)
                    history_latents = torch.cat([start_latent.to(original_history_latents), original_history_latents], dim=2)

                    decoded_tuple = None
                    decoded_split = None

                    self.cleanup(unload_models=True, cleanup_models=False, cleanup_cuda=True)

                    if not is_last_section:
                        history_latents = torch.cat([original_history_latents, generated_latents.to(original_history_latents)], dim=2)
                        move_model_to_device_with_memory_preservation(self.transformer, target_device=self.device, preserved_memory_gb=self.gpu_memory_preservation)

                    # latent_batches_gend.append(generated_latents.clone().detach()  / vae_scaling_factor)

                    if self.verbose_messaging:
                        print("history_latents {} | Shape: {}".format(latent_padding, history_latents.shape))

                    total_generated_latent_frames += int(generated_latents.shape[2])
                else:
                    offload_model_from_device_for_memory_preservation(self.transformer, self.device, preserved_memory_gb=0)
                    self.cleanup(unload_models=True, cleanup_models=False, cleanup_cuda=True)

                    dummy_gen_latents = generated_latents.clone().detach()  / vae_scaling_factor
                    dummy_decoded = partial_decode_tiled(self.vae, dummy_gen_latents, self.dec_tile_size, self.dec_overlap, self.dec_temporal_size, self.dec_temporal_overlap)

                    if len(dummy_decoded) < 4:
                        dummy_decoded = dummy_decoded.unsqueeze(0)

                    dummy_images = dummy_decoded.clone().detach()

                    if (self.use_dummy_image == "all") or self.gen_dummy_only:
                        generated_image_batches.append(dummy_images)
                        valid_output = True

                    if not self.gen_dummy_only:
                        dummy_tuple = torch.split(dummy_decoded, 1, dim=0)
                        dummy_split = [item for item in dummy_tuple]
                        dummy_decoded = dummy_decoded.to(get_offload_device())
                        dummy_decoded = None

                        if (self.use_dummy_image == "last") or (self.use_dummy_image == "all"):
                            start_image = dummy_split[len(dummy_split) - 1]
                        elif self.use_dummy_image == "first":
                            start_image = dummy_split[0]
                        elif self.use_dummy_image == "middle":
                            start_image = dummy_split[int(math.ceil(len(dummy_split) / 2))]
                        else:
                            random.seed(self.seed)
                            start_image = dummy_split[random.randrange(0, len(dummy_split) - 1, 1)]

                        image_embeds = self.clip_vision_encode(start_image)
                        start_latent = self.encode_batched(start_image, per_batch=self.latent_window_size)

                        start_latent = start_latent["samples"] * vae_scaling_factor
                        print("start_latent Shape: {}".format(start_latent.shape))

                        start_image_encoder_last_hidden_state = image_embeds["last_hidden_state"].to(self.base_dtype).to(self.device)
                        history_latents = torch.cat([start_latent.to(original_history_latents), original_history_latents], dim=2)

                        is_dummy_section = False
                        dummy_tuple = None
                        dummpy_split = None
                        dummy_decoded = None

                    self.cleanup(unload_models=True, cleanup_models=False, cleanup_cuda=True)

                    if not self.gen_dummy_only:
                        move_model_to_device_with_memory_preservation(self.transformer, target_device=self.device, preserved_memory_gb=self.gpu_memory_preservation)

                generated_latents = None

                if self.verbose_messaging:
                    print("history_latents {} | Shape: {}".format(latent_padding, history_latents.shape))

                if is_last_section:
                    break
        except mm.InterruptProcessingException as ie:
            interrupted = True
            print(f"\nWarpedFramepackSampler: Processing Interrupted.")
            print("WarpedFramepackSampler: Returning only partial results (if any).\n If zero images generated, a blank yellow image will be returned.\n")

            generation_status = f"\nWarpedFramepackSampler: Processing Interrupted."

            pass
        except BaseException as e:
            self.cleanup(unload_models=True, cleanup_models=False, cleanup_cuda=True)

            print(f"\nWarpedFramepackSampler: Exception During Processing: {str(e)}")
            print("WarpedFramepackSampler: Returning only partial results (if any).\n If zero images generated, a blank red image will be returned.\n")
            generation_status = f"WarpedFramepackSampler: Exception During Processing: {str(e)}"
            generation_status = "{}{}".format(generation_status, "WarpedFramepackSampler: Returning only partial results (if any).\nIf zero images generated, a blank red image will be returned.")

            traceback.print_tb(e.__traceback__, limit=99, file=sys.stdout)

            pass

        history_latents= None
        latent_paddings = None
        noise = None
        indices = None
        clean_latent_indices_pre = None
        blank_indices = None
        latent_indices = None
        clean_latent_indices_post = None
        clean_latent_2x_indices = None
        clean_latent_4x_indices = None
        clean_latent_indices = None
        clean_latents_pre = None
        clean_latents_post = None
        clean_latents_2x = None
        clean_latents_4x = None
        clean_latents = None
        clean_latent_indices_start = None
        clean_latent_1x_indices = None
        clean_latents_1x = None
        video_latent_batches = None
        video_image_batch = None

        # if not self.gen_dummy_only:
        offload_model_from_device_for_memory_preservation(self.transformer, self.device, preserved_memory_gb=0)

        self.cleanup(unload_models=True, cleanup_models=False, cleanup_cuda=True)

        if (len(generated_image_batches) > 0):
            output_images = self.assemble_final_result(generated_image_batches)
        elif not dummy_images is None:
            output_images = dummy_images
        elif not dummy_gen_latents is None:
            output_images = partial_decode_tiled(self.vae, dummy_gen_latents, self.dec_tile_size, self.dec_overlap, self.dec_temporal_size, self.dec_temporal_overlap)
        elif interrupted:
            temp_image = Image.new('RGB', (self.width, self.height), color = 'yellow')
            output_images = pil2tensorSwap(temp_image)
        else:
            temp_image = Image.new('RGB', (self.width, self.height), color = 'red')
            output_images = pil2tensorSwap(temp_image)

        latent_batches_gend = None

        generation_status = "{}\nImages Generated: {}".format(generation_status, output_images.shape[0])

        if self.verbose_messaging:
            print("output_images Shape: {}".format(output_images.shape))

        return (output_images, generation_status, self.seed, valid_output,)

    def process_v2v(self, video_image_batch):
        if len(video_image_batch.shape) < 4:
            video_image_batch = video_image_batch.unsqueeze(0)

        self.width = video_image_batch.shape[2]
        self.height = video_image_batch.shape[1]
        self.buffer_length = 19

        image_batch_size = int((video_image_batch.shape[0] - 1) / self.batch_count) + 1
        latent_size_factor = 4
        latent_batch_size = self.decoded_to_encoded_length(image_batch_size)

        rnd = torch.Generator("cpu").manual_seed(self.seed)

        if self.verbose_messaging:
            print("process_v2v: Video Context Images Shape: {}".format(video_image_batch.shape))
            print("process_v2v: image_batch_size: {}".format(image_batch_size))
            print("process_v2v: latent_batch_size: {}".format(latent_batch_size))

        total_latent_sections = self.batch_count

        video_image_batch.to(dtype=torch.float32, device=self.device)
        video_latent_batches, video_encoding_batches = self.video_encode(video_image_batch, image_batch_size)
        video_image_batch.to(dtype=torch.float32, device=self.offload_device)

        B, C, _, H, W = video_latent_batches[0].shape
        T = 1

        output_images = None

        latent_paddings_list = list(range(total_latent_sections))

        total_generated_latent_frames = 0
        latent_paddings = latent_paddings_list.copy()  # Create a copy for iteration

        latent_embeds = self.get_video_latent_embeds(video_image_batch, image_batch_size, latent_paddings)

        comfy_model, patcher, callback = self.initialize_comfy_model()

        print("latent_paddings count: {}  |  latent_paddings: {}".format(len(latent_paddings), latent_paddings))

        # noise = self.setup_latent_noise(W, H)
        noise, _ = self.setup_latent_noise(self.latent_window_size * self.batch_count, W, H)

        if self.verbose_messaging:
            print("Noise Batches Length: {}".format(len(noise)))
            print("sample noise Shape: {}".format(noise[0].shape))

        has_end_image = True
        generated_latents = None
        interrupted = False
        history_latents = None
        real_history_latents = None
        context_latents = None
        generation_status = ""
        temp_history_latents = None
        original_history_latents = None
        valid_output = False
        is_dummy_section = False

        self.cleanup(unload_models=True, cleanup_models=False, cleanup_cuda=True)

        latent_batches_gend = []

        try:
            for padding_i, latent_padding in enumerate(latent_paddings):
                padding_key = "{}".format(latent_padding)
                next_padding_key = "{}".format(latent_padding + 1)

                is_first_section = latent_padding == min(latent_paddings)
                is_last_section = latent_padding == max(latent_paddings)

                start_latent = latent_embeds[padding_key]["start_latent"]
                start_image_encoder_last_hidden_state = latent_embeds[padding_key]["start_embedding"]["last_hidden_state"].to(self.base_dtype).to(self.device)
                image_encoder_last_hidden_state = torch.mul(start_image_encoder_last_hidden_state, self.start_embed_strength)

                if self.verbose_messaging:
                    print("\nlatent_padding: {}  |  latent_window_size: {}".format(latent_padding, self.latent_window_size))

                noise_latent = noise[latent_padding]

                if self.verbose_messaging:
                    print("noise_latent Shape: {}".format(noise_latent.shape))

                if is_first_section:
                    original_history_latents = torch.zeros(size=(1, 16, self.buffer_length, H, W), dtype=torch.float32).cpu()

                latent_padding_size = 0
                start_latent_frames = 1
                print(f'latent_padding_size = {latent_padding_size}, is_last_section = {is_last_section}, is_first_section = {is_first_section}')

                history_latents = torch.cat([original_history_latents, video_latent_batches[latent_padding].to(original_history_latents)], dim=2)

                is_plus_one = False
                indices = torch.arange(0, history_latents.shape[2]).unsqueeze(0)

                clean_latents, clean_latent_indices, latent_indices = self.get_video_clean_latents(history_latents, self.v2v_context_count, is_plus_one=is_plus_one)
                clean_latent_4x_indices, clean_latent_2x_indices, blank_indices = indices.split([history_latents.shape[2] - 2 - video_latent_batches[latent_padding].shape[2], 2, video_latent_batches[latent_padding].shape[2]], dim=1)
                clean_latents_4x, clean_latents_2x, blank_latents = history_latents[:, :, :history_latents.shape[2], :, :].split([history_latents.shape[2] - 2 - video_latent_batches[latent_padding].shape[2], 2, video_latent_batches[latent_padding].shape[2]], dim=2)

                if self.v2v_context_count > history_latents.shape[2]:
                    self.v2v_context_count = history_latents.shape[2]

                if self.verbose_messaging:
                    print("history_latents Shape: {}\n".format(history_latents.shape))
                    print("indices: {}".format(indices))
                    print("latent_indices: {}\n".format(latent_indices))
                    print("clean_latent_2x_indices: {}".format(clean_latent_2x_indices))
                    print("clean_latent_4x_indices: {}".format(clean_latent_4x_indices))
                    print("clean_latent_indices: {}\n".format(clean_latent_indices))
                    print("clean_latents_2x Shape: {}".format(clean_latents_2x.shape))
                    print("clean_latents_4x Shape: {}".format(clean_latents_4x.shape))
                    print("clean_latents Shape: {}\n".format(clean_latents.shape))
                    print("noise Shape: {}\n".format(noise[latent_padding].shape))

                if (self.use_teacache and not is_dummy_section) or (is_dummy_section and self.dummy_teacache):
                    self.transformer.__class__.forward = self.transformer.orig_forward
                    self.transformer.initialize_teacache(enable_teacache=True, num_steps=self.steps, rel_l1_thresh=self.teacache_rel_l1_thresh, verbose_messaging=self.verbose_messaging)
                    self.transformer.initialize_magcache(enable_magcache=False)

                    if self.verbose_messaging:
                        print("TeaCache Enabled")
                elif (self.use_magcache and not is_dummy_section) or (is_dummy_section and self.dummy_magcache):
                    self.transformer.__class__.forward = self.transformer.__class__.magcache_framepack_forward
                    self.transformer.initialize_magcache(enable_magcache=True, num_steps=self.steps, magcache_thresh=self.magcache_thresh, K=self.magcache_K, retention_ratio=self.magcache_retention_ratio, verbose_messaging=self.verbose_messaging)
                    self.transformer.initialize_teacache(enable_teacache=False)

                    if self.verbose_messaging:
                        print("MagCache Enabled")
                else:
                    self.transformer.__class__.forward = self.transformer.orig_forward
                    self.transformer.initialize_teacache(enable_teacache=False)
                    self.transformer.initialize_magcache(enable_magcache=False)

                    if self.verbose_messaging:
                        print("Both TeaCache and MagCache Disabled")

                clean_latents.to(dtype=torch.float32, device=self.device)
                clean_latents_2x.to(dtype=torch.float32, device=self.device)
                clean_latents_4x.to(dtype=torch.float32, device=self.device)
                noise_latent.to(dtype=torch.float32, device=self.device)

                print("Generating Batch {}".format(latent_padding))

                generated_latents = self.generate_video(W * 8, H * 8, self.num_frames, rnd, noise_latent, self.llama_vec, self.llama_attention_mask, self.clip_l_pooler, self.llama_vec_n, self.llama_attention_mask_n,
                                                        self.clip_l_pooler_n, image_encoder_last_hidden_state, latent_indices, clean_latents, clean_latent_indices, clean_latents_2x, clean_latent_2x_indices,
                                                        clean_latents_4x, clean_latent_4x_indices, callback)
                noise[latent_padding] = None

                if self.verbose_messaging:
                    print("generated_latents for batch {}: Shape: {}\n".format(latent_padding, generated_latents.shape))

                latent_batches_gend.append(generated_latents.clone().detach() / vae_scaling_factor)
                valid_output = True
                total_generated_latent_frames += int(generated_latents.shape[2])

                generated_latents.to(dtype=torch.float32, device=self.offload_device)
                generated_latents = None

                clean_latents = clean_latents.to(dtype=torch.float32, device=self.offload_device)
                clean_latents_2x = clean_latents_2x.to(dtype=torch.float32, device=self.offload_device)
                clean_latents_4x = clean_latents_4x.to(dtype=torch.float32, device=self.offload_device)
                noise_latent = noise_latent.to(dtype=torch.float32, device=self.offload_device)

                initial_latent = None
                clean_latents = None
                clean_latents_2x = None
                clean_latents_4x = None
                noise_latent = None
                history_latents = None
                history_latents2 = None

                video_latent_batches[latent_padding] = video_latent_batches[latent_padding].to(device=self.offload_device)
                video_latent_batches[latent_padding] = None

                image_encoder_last_hidden_state = None
                generated_latents = None

                mm.soft_empty_cache()
                gc.collect()
                time.sleep(1)

                if is_last_section:
                    break
        except mm.InterruptProcessingException as ie:
            interrupted = True
            print(f"\nWarpedFramepackSampler: Processing Interrupted.")
            print("WarpedFramepackSampler: Returning only partial results (if any).\n If zero images generated, a blank yellow image will be returned.\n")

            generation_status = f"\nWarpedFramepackSampler: Processing Interrupted."

            pass
        except BaseException as e:
            self.cleanup(unload_models=True, cleanup_models=False, cleanup_cuda=True)

            print(f"\nWarpedFramepackSampler: Exception During Processing: {str(e)}")
            print("WarpedFramepackSampler: Returning only partial results (if any).\n If zero images generated, a blank red image will be returned.\n")
            generation_status = f"WarpedFramepackSampler: Exception During Processing: {str(e)}"
            generation_status = "{}{}".format(generation_status, "WarpedFramepackSampler: Returning only partial results (if any).\nIf zero images generated, a blank red image will be returned.")

            traceback.print_tb(e.__traceback__, limit=99, file=sys.stdout)

            pass

        history_latents= None
        latent_paddings = None
        noise = None
        indices = None
        clean_latent_indices_pre = None
        blank_indices = None
        latent_indices = None
        clean_latent_indices_post = None
        clean_latent_2x_indices = None
        clean_latent_4x_indices = None
        clean_latent_indices = None
        clean_latents_pre = None
        clean_latents_post = None
        clean_latents_2x = None
        clean_latents_4x = None
        clean_latents = None
        clean_latent_indices_start = None
        clean_latent_1x_indices = None
        clean_latents_1x = None
        video_latent_batches = None
        video_image_batch = None
        latent_embeds = None

        # self.transformer.to(self.offload_device)
        offload_model_from_device_for_memory_preservation(self.transformer, self.device, preserved_memory_gb=0)

        self.cleanup(unload_models=True, cleanup_models=False, cleanup_cuda=True)

        if len(latent_batches_gend) > 0:
            output_images = self.decode_batches(latent_batches_gend, self.skip_frames)
        elif interrupted:
            temp_image = Image.new('RGB', (self.width, self.height), color = 'yellow')
            output_images = pil2tensorSwap(temp_image)
        else:
            temp_image = Image.new('RGB', (self.width, self.height), color = 'red')
            output_images = pil2tensorSwap(temp_image)

        output_images.to(device=self.offload_device)
        latent_batches_gend = None

        generation_status = "{}\nImages Generated: {}".format(generation_status, output_images.shape[0])

        if self.verbose_messaging:
            print("output_images Shape: {}".format(output_images.shape))

        return (output_images, generation_status, self.seed, valid_output,)

    def generate_video(self, width, height, num_frames, seed_generator, noise_latent, llama_vec, llama_attention_mask, clip_l_pooler, llama_vec_n, llama_attention_mask_n, clip_l_pooler_n, image_encoder_last_hidden_state,
                        latent_indices, clean_latents, clean_latent_indices, clean_latents_2x, clean_latent_2x_indices, clean_latents_4x, clean_latent_4x_indices, callback):
        with torch.autocast(device_type=mm.get_autocast_device(self.device), dtype=self.base_dtype, enabled=True):
            generated_latents = sample_hunyuan2(
                transformer=self.transformer,
                sampler=self.sampler,
                initial_latent=None,
                strength=self.denoise_strength,
                width=width,
                height=height,
                frames=num_frames,
                real_guidance_scale=self.cfg,
                distilled_guidance_scale=self.guidance_scale,
                guidance_rescale=0,
                shift=self.shift if Decimal(self.shift).compare(Decimal(0.0)) != 0 else None,
                num_inference_steps=self.steps,
                generator=seed_generator,
                noise=noise_latent,
                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=self.device,
                dtype=self.base_dtype,
                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,
            )

        return generated_latents

    def setup_latent_noise(self, noise_latent_length, width, height):
        rnd = torch.Generator("cpu").manual_seed(self.seed)
        dummy_noise = None

        if not self.gen_dummy or (self.mode == VideoGenerationType.V2V):
            temp_noise = torch.randn((1, 16, noise_latent_length, height, width), generator=rnd, device=rnd.device).to(device=rnd.device, dtype=torch.float32)
        else:
            dummy_noise = torch.randn((1, 16, self.dummy_latents, height, width), generator=rnd, device=rnd.device).to(device=rnd.device, dtype=torch.float32)

            rnd = torch.Generator("cpu").manual_seed(self.seed)
            temp_noise = torch.randn((1, 16, noise_latent_length + self.dummy_latents - 1, height, width), generator=rnd, device=rnd.device).to(device=rnd.device, dtype=torch.float32)

            temp_noise_tuple = torch.split(temp_noise, self.dummy_latents, dim=2)
            temp_noise_split = [item for item in temp_noise_tuple]
            temp_noise_tuple = None

            # dummy_noise = temp_noise_split[0]

            dummy_noise_tuple = torch.split(dummy_noise, 1, dim=2)
            dummy_noise_split = [item for item in dummy_noise_tuple]
            dummy_noise_tuple = None

            temp_noise = dummy_noise_split[len(dummy_noise_split) - 1]
            dummy_noise_split = None

            i = -1
            for entry in temp_noise_split:
                i += 1

                if i > 0:
                    if not temp_noise is None:
                        temp_noise = torch.cat((temp_noise, entry), 2)
                    else:
                        temp_noise = entry

            temp_noise_split = None

            gc.collect()
            time.sleep(1)

        if Decimal(self.noise_strength).compare(Decimal(1.00)) != 0:
            temp_noise = torch.mul(temp_noise, self.noise_strength)

        temp_noise_tuple = torch.split(temp_noise, self.latent_window_size, dim=2)
        noise = [item for item in temp_noise_tuple]
        temp_noise = None
        temp_noise_tuple = None

        if self.verbose_messaging:
            if not dummy_noise is None:
                print("\nsetup_latent_noise: dummy_noise Shape: {}".format(dummy_noise.shape))

            i = 0
            for entry in noise:
                print("setup_latent_noise: Batch: {}  |  noise Shape: {}".format(i, entry.shape))
                i += 1

            print("\n")

        return noise, dummy_noise

    def initialize_comfy_model(self):
        comfy_model = HyVideoModel(
                HyVideoModelConfig(self.base_dtype),
                model_type=comfy.model_base.ModelType.FLOW,
                device=self.device,
            )

        patcher = comfy.model_patcher.ModelPatcher(comfy_model, self.device, torch.device("cpu"))
        callback = latent_preview.prepare_callback(patcher, self.steps)

        move_model_to_device_with_memory_preservation(self.transformer, target_device=self.device, preserved_memory_gb=self.gpu_memory_preservation)

        return comfy_model, patcher, callback

    def get_video_clean_latents(self, video_latents, context_frames, is_plus_one=True):
        from torch import Tensor
        latents_tuple = torch.split(video_latents, 1, dim=2)
        latents_split = [item for item in latents_tuple]
        latents_tuple = None

        if self.verbose_messaging:
            print("get_video_clean_latents: context_frames: {}".format(context_frames))
            print("get_video_clean_latents: video_latents.shape[2]: {}".format(video_latents.shape[2]))

        if context_frames < 3:
            context_frames = 3
        elif context_frames > self.latent_window_size:
            context_frames = self.latent_window_size

        if context_frames >= video_latents.shape[2]:
            clean_latent_indices = []
            i = 0
            while i < video_latents.shape[2]:
                clean_latent_indices.append(i)

                i +=1
        elif context_frames > 3:
            if context_frames > (video_latents.shape[2] - 3):
                context_frames = video_latents.shape[2] - 3

            offset = math.ceil((video_latents.shape[2] - 2 - self.buffer_length) / (context_frames - 2))

            if offset < 1:
                offset = 1

            print("offset: {}".format(offset))
            index = 0
            i = 0
            clean_latent_indices = []
            temp_clean_latent_indices = {}

            while i < (context_frames - 2):
                if index < (video_latents.shape[2] - 1):
                    temp_clean_latent_indices["{}".format(int(index + self.buffer_length))] = int(index + self.buffer_length)

                index += offset
                i += 1

            temp_clean_latent_indices["{}".format(int(video_latents.shape[2] - 2))] = int(video_latents.shape[2] - 2)
            temp_clean_latent_indices["{}".format(int(video_latents.shape[2] - 1))] = int(video_latents.shape[2] - 1)

            if self.verbose_messaging:
                print("temp_clean_latent_indices.items(): {}".format(temp_clean_latent_indices.items()))

            for key in temp_clean_latent_indices:
                clean_latent_indices.append(temp_clean_latent_indices[key])

            clean_latent_indices.sort()

            clean_latents = None
            for index in clean_latent_indices:
                if not clean_latents is None:
                    clean_latents = torch.cat((clean_latents, latents_split[index]), 2)
                else:
                    clean_latents = latents_split[index]
        else:
            clean_latent_indices = [self.buffer_length, video_latents.shape[2] - 2, video_latents.shape[2] - 1]
            clean_latents = torch.cat([latents_split[self.buffer_length], latents_split[len(latents_split) - 2], latents_split[len(latents_split) - 1]], 2)

        return_indices = Tensor([clean_latent_indices])

        latent_indices = []

        i = self.buffer_length

        n = len(latents_split)

        while i < n:
            latent_indices.append(i)
            i += 1

        latent_indices = Tensor([latent_indices])
        latent_indices = latent_indices.to(dtype=torch.uint8, device=self.offload_device)

        latents_split = None

        return_indices = return_indices.to(latent_indices)

        return clean_latents, return_indices, latent_indices

    def setup_i2v_history_latents(self, history_latents, video_latents):
        cat_list = []

        cat_list.append(history_latents.to(dtype=torch.float32, device=self.offload_device))
        cat_list.append(video_latents.to(dtype=torch.float32, device=self.offload_device))

        history_latents = torch.cat(cat_list, dim=2)

        return history_latents

    def setup_dummy_noise(self, width, height):
        rnd = torch.Generator("cpu").manual_seed(self.seed)
        noise = torch.randn((1, 16, self.dummy_latents, height, width), generator=rnd, device=rnd.device).to(device=rnd.device, dtype=torch.float32)

        if Decimal(self.noise_strength).compare(Decimal(1.00)) != 0:
            noise = torch.mul(temp_noise, self.noise_strength)

        return noise

    def encoded_to_decoded_length(self, latent_length):
        if latent_length <= 0:
            return 0

        result_length = ((latent_length - 1) * 4) + 1

        return result_length

    def decoded_to_encoded_length(self, image_length):
        if image_length <= 0:
            return 0

        result_length = int(((image_length - 1) / 4) + 1)

        return result_length

    def get_latent_window_size(self, preferred_batch_size, frame_count, use_batch_size="next_lowest"):
        latent_size_factor = 4

        if self.verbose_messaging:
            print("get_latent_window_size: preferred_batch_size: {}".format(preferred_batch_size))
            print("get_latent_window_size: frame_count: {}".format(frame_count))
            print("get_latent_window_size: use_batch_size: {}".format(use_batch_size))

        num_frames = int(((frame_count - 1) // 4) * 4) + 1

        if num_frames != frame_count:
            print(f"Truncating video from {frame_count} to {num_frames} an because odd number of frames is not allowed.")

        if ((num_frames - 1) % (preferred_batch_size - 1)) == 0:
            print("(1) latent_window_size set to: {}".format(self.decoded_to_encoded_length(preferred_batch_size)))
            print("(1) batch_count set to: {}".format(int((num_frames - 1) / (preferred_batch_size - 1))))
            return self.decoded_to_encoded_length(preferred_batch_size), int((num_frames - 1) / (preferred_batch_size - 1)), num_frames

        if use_batch_size == "exact":
            num_frames_final = int(((num_frames - 1) // (preferred_batch_size - 1)) + 1) * (preferred_batch_size - 1)

            if num_frames_final != frame_count:
                print(f"Truncating video from {num_frames} to {num_frames_final} frames for preferred_batch_size compatibility.")

            print("(2) latent_window_size set to: {}".format(self.decoded_to_encoded_length(preferred_batch_size + 1)))
            print("(2) batch_count set to: {}".format(int((num_frames_final - 1) / (preferred_batch_size - 1))))
            return self.decoded_to_encoded_length(preferred_batch_size), int((num_frames_final - 1) / (preferred_batch_size - 1)), num_frames_final

        next_lowest_found = False
        next_highest_found = False
        next_lowest = preferred_batch_size - 1
        next_highest = preferred_batch_size - 1

        if self.verbose_messaging:
            print("get_latent_window_size: Next Lowest Initialized To: {}".format(next_lowest))
            print("get_latent_window_size: Next Highest Initialized To: {}".format(next_highest))

        num_frames_final = int(((num_frames - 1) // 4) * 4) + 1

        if num_frames != num_frames_final:
            print(f"Truncating video from {num_frames} to {num_frames_final} frames for latent_window_size compatibility.")

        if (use_batch_size == "closest") or (use_batch_size == "next_lowest"):
            while next_lowest >= 12:
                next_lowest -= 4

                if (int((num_frames_final - 1) // 4) % next_lowest) == 0:
                    next_lowest_found = True
                    break

            next_lowest += 1

            if next_lowest_found and (use_batch_size == "next_lowest"):
                print("(3) latent_window_size set to: {}".format(self.decoded_to_encoded_length(next_lowest + 1)))
                print("(3) batch_count set to: {}".format(int((num_frames_final - 1) / next_lowest)))
                return self.decoded_to_encoded_length(next_lowest + 1), int((num_frames_final - 1) / next_lowest), num_frames_final

        while next_highest <= 156:
            next_highest += 4

            if (int((num_frames_final - 1) // 4) % next_highest) == 0:
                next_highest_found = True
                break

        if next_highest_found and (use_batch_size == "next_highest"):
            print("(4) latent_window_size set to: {}".format(self.decoded_to_encoded_length(next_highest + 1)))
            print("(4) batch_count set to: {}".format(int((num_frames_final - 1) / next_highest)))
            return self.decoded_to_encoded_length(next_highest + 1), int((num_frames_final - 1) / next_highest), num_frames_final

        if next_highest_found and next_lowest_found:
            if (preferred_batch_size - next_lowest) <= (next_highest - preferred_batch_size):
                print("(5) latent_window_size set to: {}".format(self.decoded_to_encoded_length(next_lowest + 1)))
                print("(5) batch_count set to: {}".format(int((num_frames_final - 1) / next_lowest)))
                return self.decoded_to_encoded_length(next_lowest + 1), int((num_frames_final - 1) / next_lowest), num_frames_final
            elif (next_highest - preferred_batch_size) < (preferred_batch_size - next_lowest):
                print("(6) latent_window_size set to: {}".format(self.decoded_to_encoded_length(next_highest + 1)))
                print("(6) batch_count set to: {}".format(int((num_frames_final - 1) / next_highest)))
                return self.decoded_to_encoded_length(next_highest + 1), int((num_frames_final - 1) / next_highest), num_frames_final

        print("Unable to find a compatible latent_window_size for number of frames = {} and preferred_batch_size = {}.".format(frame_count, preferred_batch_size))
        print("Recalculating Number Of Frames Based On preferred_batch_size of: {}".format(preferred_batch_size))

        return self.calculate_new_number_of_frames(preferred_batch_size, (((frame_count - 1) // 4) * 4) + 1, use_batch_size)

    def calculate_new_number_of_frames(self, preferred_batch_size, frame_count, use_batch_size):
        working_batch_size = preferred_batch_size - 1
        working_frame_count = frame_count - 1

        next_lowest = next_highest = working_frame_count
        next_lowest_found = False
        next_highest_found = False

        while next_lowest > 37:
            next_lowest -= 4

            if int(next_lowest % working_batch_size) == 0:
                next_lowest_found = True
                break

        if next_lowest_found and (use_batch_size == "next_lowest"):
            return self.decoded_to_encoded_length(preferred_batch_size), int(next_lowest // working_batch_size), next_lowest + 1

        while next_highest < 999997:
            next_highest += 4

            if int(next_highest % working_batch_size) == 0:
                next_highest_found = True
                break

        if next_highest_found and (use_batch_size == "next_highest"):
            return self.decoded_to_encoded_length(preferred_batch_size), int(next_highest // working_batch_size), next_highest + 1

        if next_lowest_found and next_highest_found:
            if (working_frame_count - next_lowest) <= (next_highest - working_frame_count):
                return self.decoded_to_encoded_length(preferred_batch_size), int(next_lowest // working_batch_size), next_lowest + 1

            return self.decoded_to_encoded_length(preferred_batch_size), int(next_highest // working_batch_size), next_highest + 1

        if next_lowest_found:
            return self.decoded_to_encoded_length(preferred_batch_size), int(next_lowest // working_batch_size), next_lowest + 1

        if next_highest_found:
            return self.decoded_to_encoded_length(preferred_batch_size), int(next_highest // working_batch_size), next_highest + 1

        raise ValueError("Unable to find a compatible latent_window_size for number of frames = {} and preferred_batch_size = {}.".format(frame_count, preferred_batch_size))

    def clip_vision_encode(self, image, crop="center"):
            crop_image = True
            if crop != "center":
                crop_image = False
            output = self.clip_vision.encode_image(image, crop=crop_image)
            return output

    def get_video_latent_embeds(self, video_frames, batch_frame_count, paddings):
        if self.verbose_messaging:
            print("get_video_latent_embeds: video_frames Shape: {}".format(video_frames.shape))
            print("get_video_latent_embeds: batch_frame_count: {}".format(batch_frame_count))
            print("get_video_latent_embeds: paddings: {}".format(paddings))

        if len(video_frames.shape) < 4:
            video_frames = video_frames.unsqueeze(0)

        first_batch = video_frames[:batch_frame_count, :, :, :]
        remaining_batches = video_frames[batch_frame_count:, :, :, :]

        start_image = first_batch[1:2, :, :, :]
        end_image = first_batch[(batch_frame_count - 1):, :, :, :]
        next_start_image = first_batch[(first_batch.shape[0] - 4):(first_batch.shape[0] - 3), :, :, :]

        start_latent = self.encode_batched(start_image, self.latent_window_size)
        start_latent_embedding = self.clip_vision_encode(start_image)
        start_image = start_image.to(self.offload_device)
        end_latent = self.encode_batched(end_image, self.latent_window_size)
        end_latent_embedding = self.clip_vision_encode(end_image)
        end_image = end_image.to(self.offload_device)
        next_first_latent = end_latent["samples"].clone().detach().to(end_latent["samples"])
        next_first_embedding = copy.deepcopy(end_latent_embedding)

        if self.verbose_messaging:
            print("get_video_latent_embeds: first_batch Shape: {}".format(first_batch.shape))
            print("get_video_latent_embeds: remaining_batches Shape: {}".format(remaining_batches.shape))
            print("get_video_latent_embeds: batch 0: start_image Shape: {}".format(start_image.shape))
            print("get_video_latent_embeds: batch 0: start_latent Shape: {}".format(start_latent["samples"].shape))
            print("get_video_latent_embeds: batch 0: end_image Shape: {}".format(end_image.shape))
            print("get_video_latent_embeds: batch 0: end_latent Shape: {}".format(end_latent["samples"].shape))
            print("get_video_latent_embeds: batch 0: next_start_image Shape: {}".format(next_start_image.shape))
            print("get_video_latent_embeds: batch 0: next_first_latent Shape: {}".format(next_first_latent.shape))

        latent_embeds = {"0": {}}
        latent_embeds["0"]["start_latent"] = (start_latent["samples"] * vae_scaling_factor).to(self.offload_device)
        latent_embeds["0"]["start_embedding"] = start_latent_embedding
        latent_embeds["0"]["end_latent"] = (end_latent["samples"] * vae_scaling_factor).to(self.offload_device)
        latent_embeds["0"]["end_embedding"] = end_latent_embedding

        start_image = None
        start_latent = None
        start_latent_embedding = None
        end_image = None
        end_latent = None
        end_latent_embedding = None
        next_start_four = None
        next_start_image = None

        if remaining_batches.shape[0] > 0:
            image_batches_tuple = torch.split(remaining_batches, batch_frame_count - 1, dim=0)
            image_batches_split = [item for item in image_batches_tuple]
            image_batches_tuple = None

            if self.verbose_messaging:
                i = 0
                for entry in image_batches_split:
                    print("get_video_latent_embeds: batch: {}  |  split Shape: {}".format(i + 1, image_batches_split[i].shape))
                    i += 1

            print("get_video_latent_embeds Batch: Encoding Start/End Images...")

            batch_number = 1
            for batch in image_batches_split:
                batch_key = "{}".format(batch_number)

                if self.verbose_messaging:
                    print("get_video_latent_embeds: Processing Batch: {}".format(batch_number))

                latent_embeds[batch_key] = {"start_latent": None, "end_latent": None, "start_embedding": None, "end_embedding": None}

                start_latent = next_first_latent
                start_latent_embedding = next_first_embedding

                end_image = batch[(batch.shape[0] - 1):, :, :, :]
                end_image = end_image.to(self.device)
                end_latent = self.encode_batched(end_image, self.latent_window_size)
                # end_latent = end_latent["samples"]
                end_latent_embedding = self.clip_vision_encode(end_image)
                end_image = end_image.to(self.device)

                latent_embeds[batch_key]["start_latent"] = (start_latent * vae_scaling_factor).to(device=self.offload_device)
                latent_embeds[batch_key]["start_embedding"] = start_latent_embedding
                latent_embeds[batch_key]["end_latent"] = (end_latent["samples"] * vae_scaling_factor).to(device=self.offload_device)
                latent_embeds[batch_key]["end_embedding"] = end_latent_embedding

                if batch_number < max(paddings):
                    next_start_image = batch[(batch.shape[0] - 1):, :, :, :]
                    next_first_latent = self.encode_batched(next_start_image, self.latent_window_size)
                    next_first_latent = next_first_latent["samples"]
                    next_first_embedding = copy.deepcopy(end_latent_embedding)
                    next_start_image = None
                else:
                    next_first_embedding = None
                    next_start_image = None
                    next_first_latent = None

                start_image = None
                start_latent = None
                start_latent_embedding = None
                end_image = None
                end_latent = None
                end_latent_embedding = None
                batch_image_batches_split = None

                batch_number += 1

            print("get_video_latent_embeds Batch: Encoding Start/End Images...Done")

            if self.verbose_messaging:
                print("get_video_latent_embeds: latent_embeds length: {}".format(len(latent_embeds)))

                for batch_key in latent_embeds.keys():
                    print("get_video_latent_embeds: Batch {}: start_latent Shape: {}".format(batch_key, latent_embeds[batch_key]["start_latent"].shape))
                    print("get_video_latent_embeds: Batch {}: end_latent Shape: {}".format(batch_key, latent_embeds[batch_key]["end_latent"].shape))
                    print("get_video_latent_embeds: Batch {}: start_embedding Shape: {}".format(batch_key, latent_embeds[batch_key]["start_embedding"]["last_hidden_state"].shape))
                    print("get_video_latent_embeds: Batch {}: end_embedding Shape: {}".format(batch_key, latent_embeds[batch_key]["end_embedding"]["last_hidden_state"].shape))

        self.cleanup(unload_models=False, cleanup_models=False, cleanup_cuda=True)

        return latent_embeds

    def encode_tiled(self, images):
        if len(images.shape) < 4:
            images = images.unsqueze(0)

        encoded_data = partial_encode_tiled(self.vae, images)

        if len(encoded_data.shape) < 5:
            encoded_data.unsqueeze(0)

        return encoded_data
    def decode_tiled(self, latents, unload_after=True):
        decoded_data = partial_decode_tiled(self.vae, latents, self.dec_tile_size, self.dec_overlap, self.dec_temporal_size, self.dec_temporal_overlap, unload_after)

        if len(decoded_data.shape) < 4:
            decoded_data.unsqueeze(0)

        return decoded_data

    def decode_batches(self, latent_batches, skip_frames):
        if (latent_batches is None) or (len(latent_batches) < 1):
            print("decode_batches: Warning...nothing to decode.")
            return None

        self.cleanup(unload_models=True, cleanup_models=False, cleanup_cuda=False)

        resulting_images = None

        if self.blend_frames < 1:
            i = 0
            for entry in latent_batches:
                entry.to(dtype=torch.bfloat16, device=self.device)

                if i < (len(latent_batches) - 1):
                    temp_decoded = self.decode_tiled(entry, unload_after=False)
                else:
                    temp_decoded = self.decode_tiled(entry, unload_after=True)

                if len(temp_decoded.shape) < 4:
                    temp_decoded = temp_decoded.unsqueeze(0)

                # if i > 0:
                #     temp_decoded = temp_decoded[3:, :, :, :]

                if not resulting_images is None:
                    resulting_images = torch.cat((resulting_images, temp_decoded), 0)
                else:
                    resulting_images = temp_decoded

                entry.to(device=self.offload_device)

                self.cleanup(unload_models=False, cleanup_models=False, cleanup_cuda=True)

                i += 1
        else:
            temp_decoded_batches = []

            i = 0
            for entry in latent_batches:
                entry.to(dtype=torch.bfloat16, device=self.device)

                if i < (len(latent_batches) - 1):
                    temp_decoded = self.decode_tiled(entry, unload_after=False)
                else:
                    temp_decoded = self.decode_tiled(entry, unload_after=True)

                if len(temp_decoded.shape) < 4:
                    temp_decoded = temp_decoded.unsqueeze(0)

                temp_decoded_batches.append(temp_decoded)

                entry.to(device=self.offload_device)

                self.cleanup(unload_models=False, cleanup_models=False, cleanup_cuda=True)

                i += 1

            resulting_images = self.assemble_final_result(temp_decoded_batches)
            temp_decoded_batches = None
            self.cleanup(unload_models=False, cleanup_models=False, cleanup_cuda=True)

        print("decode_batches: Full decoded images count: {}".format(resulting_images.shape[0]))

        if skip_frames < 1:
            return resulting_images

        skipped_frames = 1

        image_batches_tuple = torch.split(resulting_images, 1, dim=0)
        image_batches_split = [item for item in image_batches_tuple]

        resulting_images = None

        for image in image_batches_split:
            if skipped_frames <= skip_frames:
                skipped_frames += 1
                continue

            if not resulting_images is None:
                resulting_images = torch.cat((resulting_images, image), 0)
            else:
                resulting_images = image

        print("decode_batches: Final decoded images count: {}".format(resulting_images.shape[0]))

        return resulting_images

    def assemble_final_result(self, image_batches):
        if self.blend_frames < 1:
            resulting_images = None
            for entry in image_batches:
                if not resulting_images is None:
                    resulting_images = torch.cat((resulting_images, entry), 0)
                else:
                    resulting_images = entry

                entry.to(device=self.offload_device)
                entry = None
        else:
            blend_value = 1.0 / self.blend_frames
            i = 0
            while i < (len(image_batches) - 1):
                alpha_blend_val = blend_value
                blend_count = self.blend_frames

                image_batches_tuple = torch.split(image_batches[i], 1, dim=0)
                image_batches_split = [item for item in image_batches_tuple]
                image1 = image_batches_split[len(image_batches_split) - 1]
                # image1 = image_batches_split[len(image_batches_split) - 4]
                image_batches_tuple = None
                image_batches_split = None

                image_batches_tuple = torch.split(image_batches[i + 1], 1, dim=0)
                image_batches_split = [item for item in image_batches_tuple]
                image2 = image_batches_split[0]
                image_batches_tuple = None
                image_batches_split = None

                image1 = tensor2pilSwap(image1)[0]
                image2 = tensor2pilSwap(image2)[0]

                blend_latents = None

                while blend_count > 0:
                    blended_image = Image.blend(image1, image2, alpha_blend_val)
                    temp_latent = pil2tensorSwap(blended_image)
                    blended_image = None

                    if len(temp_latent.shape) < 4:
                        temp_latent = temp_latent.unsqueeze(0)

                    if not blend_latents is None:
                        blend_latents = torch.cat((blend_latents, temp_latent), 0)
                    else:
                        blend_latents = temp_latent

                    alpha_blend_val += blend_value
                    blend_count -= 1

                image_batches_tuple = torch.split(image_batches[i], image_batches[i].shape[0] - 3, dim=0)
                image_batches_split = [item for item in image_batches_tuple]
                image_batches_tuple = None

                image_batches[i] = torch.cat((image_batches_split[0], blend_latents), 0)
                blend_latents = None
                image_batches_split = None

                self.cleanup(unload_models=False, cleanup_models=False, cleanup_cuda=True)

                i += 1

            resulting_images = None
            for entry in image_batches:
                if not resulting_images is None:
                    resulting_images = torch.cat((resulting_images, entry), 0)
                else:
                    resulting_images = entry

                entry.to(device=self.offload_device)

            image_batches = None

            self.cleanup(unload_models=False, cleanup_models=False, cleanup_cuda=True)


        print("assemble_final_result: Full decoded images count: {}".format(resulting_images.shape[0]))

        if self.skip_frames < 1:
            return resulting_images

        skipped_frames = 1

        image_batches_tuple = torch.split(resulting_images, 1, dim=0)
        image_batches_split = [item for item in image_batches_tuple]

        resulting_images = None

        for image in image_batches_split:
            if skipped_frames <= self.skip_frames:
                skipped_frames += 1
                continue

            if not resulting_images is None:
                resulting_images = torch.cat((resulting_images, image), 0)
            else:
                resulting_images = image

        print("assemble_final_result: Final decoded images count: {}".format(resulting_images.shape[0]))

        return resulting_images

    def encode_batched(self, image, per_batch=16):
        from nodes import VAEEncode
        from comfy.utils import ProgressBar

        t = []
        pbar = ProgressBar(image.shape[0])
        for start_idx in range(0, image.shape[0], per_batch):
            try:
                sub_pixels = self.vae.vae_encode_crop_pixels(image[start_idx:start_idx+per_batch])
            except:
                sub_pixels = VAEEncode.vae_encode_crop_pixels(image[start_idx:start_idx+per_batch])
            t.append(self.vae.encode(sub_pixels[:,:,:,:3]))
            pbar.update(per_batch)

        return {"samples": torch.cat(t, dim=0)}

    def cleanup(self, unload_models=False, cleanup_models=False, cleanup_cuda=False):
        if unload_models:
            mm.unload_all_models()

        if cleanup_models:
            mm.cleanup_models()

        if cleanup_cuda:
            mm.soft_empty_cache()

        gc.collect()
        time.sleep(1)

        return

    def video_encode(self, video_frames, batch_size):
        if len(video_frames.shape) < 4:
            video_frames = video_frames.unsqueeze(0)

        print(f"video_encode: Encoding input video frames in batch size {batch_size} (reduce preferred_batch_size if memory issues here or if forcing video resolution)")

        if self.verbose_messaging:
            print("video_encode: video_frames Shape: {}".format(video_frames.shape))
            print("video_encode: batch_size: {}".format(batch_size))

        first_batch = video_frames[:batch_size, :, :, :]
        remaining_batches = video_frames[batch_size:, :, :, :]
        # next_start_image = first_batch[(first_batch.shape[0] - 4):(first_batch.shape[0] - 3), :, :, :]
        next_start_four = first_batch[(first_batch.shape[0] - 4):, :, :, :]

        first_batch_latents = partial_encode_tiled(self.vae, first_batch, unload_after=False)

        if self.verbose_messaging:
            print("video_encode: first_batch Shape: {}".format(first_batch.shape))
            print("video_encode: remaining_batches Shape: {}".format(remaining_batches.shape))
            # print("video_encode: next_start_image Shape: {}".format(next_start_image.shape))
            print("video_encode: next_start_four Shape: {}".format(next_start_four.shape))

        if len(first_batch_latents.shape) < 5:
            first_batch_latents = first_batch_latents.unsqueeze(0)

        first_batch_latents = first_batch_latents * vae_scaling_factor
        final_latent_batches = [first_batch_latents.to(self.offload_device)]

        image_batches_tuple = torch.split(remaining_batches, batch_size - 1, dim=0)
        image_batches_split = [item for item in image_batches_tuple]

        final_image_batches = []
        last_frame = None

        latents = None
        batch_encodings = []

        with torch.no_grad():
            i = 0
            for entry in image_batches_split:
                entry = entry.to(dtype=torch.bfloat16, device=self.device)
                entry = torch.cat((next_start_four.to(entry), entry), 0)
                # next_start_image = entry[(entry.shape[0] - 1):, :, :, :]
                next_start_four = entry[(entry.shape[0] - 4):, :, :, :]

                if i < (len(image_batches_split) - 1):
                    batch_latent = partial_encode_tiled(self.vae, entry, unload_after=False)
                else:
                    batch_latent = partial_encode_tiled(self.vae, entry, unload_after=True)

                if len(batch_latent.shape) < 5:
                    batch_latent = batch_latent.unsqueeze(0)

                batch_latent = batch_latent * vae_scaling_factor
                batch_latent = batch_latent.to(self.offload_device)
                final_latent_batches.append(batch_latent.clone().detach())
                entry = entry.to(device=self.offload_device)
                entry = None
                batch_latent = None

        print(f"Encoding input video frames in batch size {batch_size} Done.")

        if self.verbose_messaging:
            i = 0
            for entry in final_latent_batches:
                print("video_encode: Batch: {}  |  batch Shape: {}".format(i, entry.shape))
                i += 1

        return final_latent_batches, batch_encodings

def cleanup(device=None, unload_models=False, cleanup_cuda=False):
    if unload_models:
        mm.free_memory(1e30, device)

    if cleanup_cuda:
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()

    gc.collect()
    time.sleep(1)

    return

class WarpedFramepackLoraSelectBatch:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
               "lora": ("STRING", {"default": "", "forceInput": True}),
                "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.0001, "tooltip": "LORA strength, set to 0.0 to unmerge the LORA"}),
                "fuse_lora": ("BOOLEAN", {"default": False, "tooltip": "Fuse the LORA model with the base model. This is recommended for better performance."}),
            },
            "optional": {
                "prev_lora":("FPLORA", {"default": None, "tooltip": "For loading multiple LoRAs"}),
            }
        }

    RETURN_TYPES = ("FPLORA",)
    RETURN_NAMES = ("lora", )
    FUNCTION = "getlorapath"
    CATEGORY = "Warped/Framepack/Lora"
    DESCRIPTION = "Select a LoRA model from ComfyUI/models/loras"

    def getlorapath(self, lora, strength, prev_lora=None, fuse_lora=True):
        loras_list = []

        lora = {
            "path": folder_paths.get_full_path("loras", lora),
            "strength": strength,
            "name": lora.split(".")[0],
            "fuse_lora": fuse_lora,
        }

        if prev_lora is not None:
            loras_list.extend(prev_lora)

        loras_list.append(lora)

        return (loras_list,)

def warped_load_torch_file(ckpt, return_metadata=False):
    from safetensors.torch import load as safeload

    metadata = {}

    try:
        f = safeload(ckpt)

        sd = {}
        for k in f.keys():
            if k != "__metadata__":
                sd[k] = f[k]
            elif return_metadata:
                metadata = f[k]
    except Exception as e:
        if len(e.args) > 0:
            message = e.args[0]
            if "HeaderTooLarge" in message:
                raise ValueError("{}\n\nFile path: {}\n\nThe safetensors file is corrupt or invalid. Make sure this is actually a safetensors file and not a ckpt or pt or other filetype.".format(message, ckpt))
            if "MetadataIncompleteBuffer" in message:
                raise ValueError("{}\n\nFile path: {}\n\nThe safetensors file is corrupt/incomplete. Check the file size and make sure you have copied/downloaded it correctly.".format(message, ckpt))
        raise e

    return (sd, metadata) if return_metadata else sd

def warped_load_lora_weights(lora_name, return_metadata=False):
    if (lora_name is None) or (len(lora_name) < 1):
        raise ValueError("lora_name cannot be None or Empty.")

    # Get the full path to the LoRA file
    lora_path = folder_paths.get_full_path("loras", lora_name)

    if not os.path.exists(lora_path):
        raise ValueError(f"LoRA file not found: {lora_path}")

    print("Reading LORA At Path: {}".format(lora_path))
    with open(lora_path, "rb") as file:
        lora_weights = file.read()

    print("Loading Lora: {}...".format(lora_name))
    lora_sd, metadata = warped_load_torch_file(lora_weights, return_metadata=True)
    print("Loading Lora: {}...Done".format(lora_name))

    if return_metadata:
        return lora_sd, metadata

    return lora_sd

def get_available_devices():
    available_devices = ["cpu"]

    if torch.cuda.is_available():
        available_devices.append("cuda")

        if torch.cuda.device_count() > 1:
            for i in range(torch.cuda.device_count()):
                temp_device = "cuda:{}".format(i)
                available_devices.append(temp_device)

    return available_devices

class WarpedDualEncoder(ComfyNodeABC):
    @classmethod
    def INPUT_TYPES(self):
        return {"required": { "clip": ("CLIP", ),
                              "positive_text": (IO.STRING, {"multiline": True, "dynamicPrompts": True, "tooltip": "The positive prompt to be encoded."}),
                              "negative_text": (IO.STRING, {"multiline": True, "dynamicPrompts": True, "tooltip": "The negative prompt to be encoded."}),
                            },
               }
    RETURN_TYPES = (IO.CONDITIONING, IO.CONDITIONING, "STRING", "STRING", )
    RETURN_NAMES = ("pos_conditioning", "neg_conditioning", "pos_prompt", "neg_prompt", )
    FUNCTION = "process"
    CATEGORY = "Warped/General/Conditioning"
    DESCRIPTION = "Encodes both positive and negative prompts."

    def process(self, clip, positive_text="", negative_text=""):
        print("WarpedDualEncoder: Loading clip model to device: {}".format(clip.patcher.load_device))
        clip.patcher.model.to(device=clip.patcher.load_device)

        print("WarpedDualEncoder: Encoding Prompts...")
        positive_conditioning = self.encode(clip, positive_text)
        negative_conditioning = self.encode(clip, negative_text)
        print("WarpedDualEncoder: Encoding Prompts...Done.")

        print("WarpedDualEncoder: Unloading clip model to device: {}".format(get_offload_device()))
        clip.patcher.model.to(device=get_offload_device())

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        gc.collect()
        time.sleep(1)

        return (positive_conditioning, negative_conditioning, positive_text, negative_text, )

    def encode(self, clip, text):

        if clip is None:
            raise RuntimeError("ERROR: clip input is invalid: None\n\nIf the clip is from a checkpoint loader node your checkpoint does not contain a valid clip or text encoder model.")

        tokens = clip.tokenize(text)
        return_encoding = clip.encode_from_tokens_scheduled(tokens)

        return return_encoding

class WarpedSingleEncoder(ComfyNodeABC):
    @classmethod
    def INPUT_TYPES(self):
        return {"required": { "clip": ("CLIP", ),
                              "input_text": (IO.STRING, {"multiline": True, "dynamicPrompts": True, "tooltip": "The prompt to be encoded."}),
                            },
               }
    RETURN_TYPES = (IO.CONDITIONING, "STRING", )
    RETURN_NAMES = ("conditioning", "prompt")
    FUNCTION = "process"
    CATEGORY = "Warped/General/Conditioning"
    DESCRIPTION = "Encodes a single prompt."

    def process(self, clip, input_text=""):
        print("WarpedSingleEncoder: Loading clip model to device: {}".format(clip.patcher.load_device))
        clip.patcher.model.to(device=clip.patcher.load_device)

        print("WarpedSingleEncoder: Encoding Prompt...")
        output_conditioning = self.encode(clip, input_text)
        print("WarpedSingleEncoder: Encoding Prompt...Done.")

        print("WarpedSingleEncoder: Unloading clip model to device: {}".format(get_offload_device()))
        clip.patcher.model.to(device=get_offload_device())

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        gc.collect()
        time.sleep(1)

        return (output_conditioning, input_text, )

    def encode(self, clip, text):

        if clip is None:
            raise RuntimeError("ERROR: clip input is invalid: None\n\nIf the clip is from a checkpoint loader node your checkpoint does not contain a valid clip or text encoder model.")

        tokens = clip.tokenize(text)
        return_encoding = clip.encode_from_tokens_scheduled(tokens)

        return return_encoding

class WarpedCLIPLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "clip_name": (folder_paths.get_filename_list("text_encoders"), ),
                              "clip_type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace"], ),
                              },
                "optional": {
                              "device": (get_available_devices(), {"default": "cpu"}),
                             }}
    RETURN_TYPES = ("CLIP",)
    FUNCTION = "load_clip"

    CATEGORY = "Warped/General/Loaders"

    DESCRIPTION = "[Recipes]\n\nstable_diffusion: clip-l\nstable_cascade: clip-g\nsd3: t5 xxl/ clip-g / clip-l\nstable_audio: t5 base\nmochi: t5 xxl\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\n hidream: llama-3.1 (Recommend) or t5"

    def load_clip(self, clip_name, clip_type="stable_diffusion", device="cpu"):
        clip_type_attr = getattr(comfy.sd.CLIPType, clip_type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)

        print("WarpedCLIPLoader: Clip Type: {}  |  {}".format(clip_type, clip_type_attr))

        model_options = {}
        if device == "cpu":
            model_options["load_device"] = model_options["offload_device"] = torch.device("cpu")
        else:
            model_options["load_device"] = model_options["offload_device"] = torch.device("cpu")

            if torch.cuda.is_available():
                if device == "cuda":
                    model_options["load_device"] = torch.cuda.current_device()
                else:
                    temp_device = device.split(':')
                    device_number = int(temp_device[len(temp_device) - 1])
                    model_options["load_device"] = torch.device(device_number)

        print("WarpedCLIPLoader: {}".format(model_options))

        clip_path = folder_paths.get_full_path_or_raise("text_encoders", clip_name)
        clip = self.sd_load_clip(ckpt_paths=[clip_path], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=clip_type_attr, model_options=model_options)
        return (clip,)

    def sd_load_clip(self, ckpt_paths, embedding_directory=None, clip_type=comfy.sd.CLIPType.STABLE_DIFFUSION, model_options={"offload_device": get_offload_device()}):
        checkpoint_temp = []
        for p in ckpt_paths:
            print("Reading: {}".format(p))
            with open(p, "rb") as file:
                content = file.read()

            checkpoint_temp.append(content)

        clip_data = []

        for p in checkpoint_temp:
            print("Loading Clip...")
            clip_data.append(warped_load_torch_file(p))
            print("Loading Clip...Done")

        return_clip = comfy.sd.load_text_encoder_state_dicts(clip_data, embedding_directory=embedding_directory, clip_type=clip_type, model_options=model_options)

        checkpoint_temp = None
        clip_data = None

        return return_clip

    def encode(self, clip, text):

        if clip is None:
            raise RuntimeError("ERROR: clip input is invalid: None\n\nIf the clip is from a checkpoint loader node your checkpoint does not contain a valid clip or text encoder model.")

        tokens = clip.tokenize(text)
        return_encoding = clip.encode_from_tokens_scheduled(tokens)

        return return_encoding

class WarpedDualCLIPLoader:
    @classmethod
    def INPUT_TYPES(self):
        return {"required": { "clip_name1": (folder_paths.get_filename_list("text_encoders"), ),
                              "clip_name2": (folder_paths.get_filename_list("text_encoders"), ),
                              "type": (["sdxl", "sd3", "flux", "hunyuan_video", "hunyuan_video_15"], ),
                              "device": (get_available_devices(), {"default": "cuda"}),
                            },
               }
    RETURN_TYPES = ("CLIP", )
    RETURN_NAMES = ("clip", )
    FUNCTION = "load_clip"

    CATEGORY = "Warped/General/Loaders"

    DESCRIPTION = "[Recipes]\n\nsdxl: clip-l, clip-g\nsd3: clip-l, clip-g / clip-l, t5 / clip-g, t5\nflux: clip-l, t5"

    def load_clip(self, clip_name1, clip_name2, type, device="cpu"):
        clip_path1 = folder_paths.get_full_path_or_raise("text_encoders", clip_name1)
        clip_path2 = folder_paths.get_full_path_or_raise("text_encoders", clip_name2)
        if type == "sdxl":
            clip_type = comfy.sd.CLIPType.STABLE_DIFFUSION
        elif type == "sd3":
            clip_type = comfy.sd.CLIPType.SD3
        elif type == "flux":
            clip_type = comfy.sd.CLIPType.FLUX
        elif type == "hunyuan_video":
            clip_type = comfy.sd.CLIPType.HUNYUAN_VIDEO
        elif type == "hunyuan_video_15":
            clip_type = comfy.sd.CLIPType.HUNYUAN_VIDEO_15

        model_options = {}
        if device == "cpu":
            model_options["load_device"] = model_options["offload_device"] = torch.device("cpu")
        else:
            model_options["load_device"] = model_options["offload_device"] = torch.device("cpu")

            if torch.cuda.is_available():
                if device == "cuda":
                    model_options["load_device"] = torch.cuda.current_device()
                else:
                    temp_device = device.split(':')
                    device_number = int(temp_device[len(temp_device) - 1])
                    model_options["load_device"] = torch.device(device_number)

        print("WarpedDualCLIPLoader: {}".format(model_options))

        clip = self.sd_load_clip(ckpt_paths=[clip_path1, clip_path2], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=clip_type, model_options=model_options)

        return (clip, )

    def sd_load_clip(self, ckpt_paths, embedding_directory=None, clip_type=comfy.sd.CLIPType.STABLE_DIFFUSION, model_options={"offload_device": get_offload_device()}):
        checkpoint_temp = []
        for p in ckpt_paths:
            print("Reading: {}".format(p))
            with open(p, "rb") as file:
                content = file.read()

            checkpoint_temp.append(content)

        clip_data = []
        i = 1
        for p in checkpoint_temp:
            print("Loading Clip: {}...".format(i))
            clip_data.append(warped_load_torch_file(p))
            print("Loading Clip: {}...Done".format(i))
            i += 1

        return_clip = comfy.sd.load_text_encoder_state_dicts(clip_data, embedding_directory=embedding_directory, clip_type=clip_type, model_options=model_options)

        checkpoint_temp = None
        clip_data = None

        return return_clip

    def encode(self, clip, text):

        if clip is None:
            raise RuntimeError("ERROR: clip input is invalid: None\n\nIf the clip is from a checkpoint loader node your checkpoint does not contain a valid clip or text encoder model.")

        tokens = clip.tokenize(text)
        return_encoding = clip.encode_from_tokens_scheduled(tokens)

        return return_encoding

script_directory = os.path.dirname(os.path.abspath(__file__))

class WarpedLoadFramePackModel:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": (folder_paths.get_filename_list("diffusion_models"), {"tooltip": "These models are loaded from the 'ComfyUI/models/diffusion_models' -folder",}),

            "base_precision": (["fp32", "bf16", "fp16"], {"default": "bf16"}),
            "quantization": (['disabled', 'fp8_e4m3fn', 'fp8_e4m3fn_fast', 'fp8_e5m2'], {"default": 'disabled', "tooltip": "optional quantization method"}),
            "load_device": (get_available_devices(), {"default": "cuda", "tooltip": "set load device."}),
            },
            "optional": {
                "attention_mode": ([
                    "sdpa",
                    "flash_attn",
                    "sageattn",
                    ], {"default": "sdpa"}),
                "compile_args": ("FRAMEPACKCOMPILEARGS", ),
                "lora": ("FPLORA", {"default": None, "tooltip": "LORA model to load"}),
            }
        }

    RETURN_TYPES = ("FramePackMODEL",)
    RETURN_NAMES = ("model", )
    FUNCTION = "loadmodel"
    CATEGORY = "Warped/Framepack/Loaders"

    def loadmodel(self, model, base_precision, quantization,
                  compile_args=None, attention_mode="sdpa", lora=None, load_device="main_device"):
        from .diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModel
        from accelerate.utils import set_module_tensor_to_device
        from accelerate import init_empty_weights
        from .diffusers_helper.memory import DynamicSwapInstaller, move_model_to_device_with_memory_preservation

        mm.unload_all_models()
        gc.collect()
        time.sleep(1)

        base_dtype = {"fp8_e4m3fn": torch.float8_e4m3fn, "fp8_e4m3fn_fast": torch.float8_e4m3fn, "bf16": torch.bfloat16, "fp16": torch.float16, "fp16_fast": torch.float16, "fp32": torch.float32}[base_precision]

        offload_device = get_offload_device()

        if load_device == "cuda":
            device = mm.get_torch_device()
        elif load_device == "cpu":
            device = offload_device
        else:
            temp_strings = load_device.split(':')
            device = torch.device(int(temp_strings[len(temp_strings) - 1]))

        transformer_load_device = offload_device

        model_path = folder_paths.get_full_path_or_raise("diffusion_models", model)
        model_config_path = os.path.join(script_directory, "transformer_config.json")
        import json
        with open(model_config_path, "r") as f:
            config = json.load(f)

        print("Reading: {}".format(model_path))
        with open(model_path, "rb") as file:
            checkpoint_temp = file.read()

        print("Loading Checkpoint...")
        sd = warped_load_torch_file(checkpoint_temp)
        print("Loading Checkpoint...Done")
        checkpoint_temp = None

        gc.collect()
        time.sleep(1)

        # for key in sd.keys():
        #     print("Framepack Model Key: {}".format(key))

        # sd = load_torch_file(model_path, device=offload_device, safe_load=True)
        model_weight_dtype = sd['single_transformer_blocks.0.attn.to_k.weight'].dtype

        with init_empty_weights():
            transformer = HunyuanVideoTransformer3DModel(**config, attention_mode=attention_mode)

        params_to_keep = {"norm", "bias", "time_in", "vector_in", "guidance_in", "txt_in", "img_in"}
        if quantization == "fp8_e4m3fn" or quantization == "fp8_e4m3fn_fast" or quantization == "fp8_scaled":
            dtype = torch.float8_e4m3fn
        elif quantization == "fp8_e5m2":
            dtype = torch.float8_e5m2
        else:
            dtype = base_dtype

        if lora is not None:
            after_lora_dtype = dtype
            dtype = base_dtype

        print("Using accelerate to load and assign model weights to device...")
        param_count = sum(1 for _ in transformer.named_parameters())
        for name, param in tqdm(transformer.named_parameters(),
                desc=f"Loading transformer parameters to {transformer_load_device}",
                total=param_count,
                leave=True):
            dtype_to_use = base_dtype if any(keyword in name for keyword in params_to_keep) else dtype

            set_module_tensor_to_device(transformer, name, device=transformer_load_device, dtype=dtype_to_use, value=sd[name])

        if lora is not None:
            from diffusers.loaders.lora_conversion_utils import _convert_hunyuan_video_lora_to_diffusers

            adapter_list = []
            adapter_weights = []

            lora_number = 1
            for l in lora:
                fuse = True if l["fuse_lora"] else False

                blocktypes = None

                if "blocks" in l.keys():
                    blocktypes = l["blocks"]

                print("Reading LORA: {}".format(l["path"]))
                with open(l["path"], "rb") as file:
                    lora_temp = file.read()

                print("Loading Lora: {}...".format(lora_number))
                lora_sd = warped_load_torch_file(lora_temp)
                print("Loading Lora: {}...Done".format(lora_number))
                lora_number += 1

                # Convert Lora To diffusion_model format
                lora_sd = convert_lora(lora_sd, convert_to="diffusion_model")

                # for key in lora_sd:
                #     print("Lora Key: {}".format(key))

                # If blocks value exist in lora metadata, then select block types
                if not blocktypes is None:
                    lora_sd = filter_lora_keys(lora_sd, blocktypes)

                if "lora_unet_single_transformer_blocks_0_attn_to_k.lora_up.weight" in lora_sd:
                    from .utils import convert_to_diffusers
                    lora_sd = convert_to_diffusers("lora_unet_", lora_sd)

                if not "transformer.single_transformer_blocks.0.attn_to.k.lora_A.weight" in lora_sd:
                    print(f"Converting LoRA weights from {l['path']} to diffusers format...")
                    lora_sd = _convert_hunyuan_video_lora_to_diffusers(lora_sd)

                lora_rank = None
                for key, val in lora_sd.items():
                    if "lora_B" in key or "lora_up" in key:
                        lora_rank = val.shape[1]
                        break
                if lora_rank is not None:
                    print(f"Merging rank {lora_rank} LoRA weights from {l['path']} with strength {l['strength']}")
                    adapter_name = l['path'].split("/")[-1].split(".")[0]
                    adapter_weight = l['strength']
                    transformer.load_lora_adapter(lora_sd, weight_name=l['path'].split("/")[-1], lora_rank=lora_rank, adapter_name=adapter_name)

                    adapter_list.append(adapter_name)
                    adapter_weights.append(adapter_weight)

                del lora_sd
                mm.soft_empty_cache()

            if adapter_list:
                transformer.set_adapters(adapter_list, weights=adapter_weights)
                if fuse:
                    if model_weight_dtype not in [torch.float32, torch.float16, torch.bfloat16]:
                        raise ValueError("Fusing LoRA doesn't work well with fp8 model weights. Please use a bf16 model file, or disable LoRA fusing.")
                    lora_scale = 1
                    transformer.fuse_lora(lora_scale=lora_scale)
                    transformer.delete_adapters(adapter_list)

            if quantization == "fp8_e4m3fn" or quantization == "fp8_e4m3fn_fast" or quantization == "fp8_e5m2":
                params_to_keep = {"norm", "bias", "time_in", "vector_in", "guidance_in", "txt_in", "img_in"}
                for name, param in transformer.named_parameters():
                    # Make sure to not cast the LoRA weights to fp8.
                    if not any(keyword in name for keyword in params_to_keep) and not 'lora' in name:
                        param.data = param.data.to(after_lora_dtype)

        if quantization == "fp8_e4m3fn_fast":
            from .fp8_optimization import convert_fp8_linear
            convert_fp8_linear(transformer, base_dtype, params_to_keep=params_to_keep)

        DynamicSwapInstaller.install_model(transformer, device=device)

        if compile_args is not None:
            if compile_args["compile_single_blocks"]:
                for i, block in enumerate(transformer.single_transformer_blocks):
                    transformer.single_transformer_blocks[i] = torch.compile(block, fullgraph=compile_args["fullgraph"], dynamic=compile_args["dynamic"], backend=compile_args["backend"], mode=compile_args["mode"])
            if compile_args["compile_double_blocks"]:
                for i, block in enumerate(transformer.transformer_blocks):
                    transformer.transformer_blocks[i] = torch.compile(block, fullgraph=compile_args["fullgraph"], dynamic=compile_args["dynamic"], backend=compile_args["backend"], mode=compile_args["mode"])

        pipe = {
            "transformer": transformer.eval(),
            "dtype": base_dtype,
        }
        return (pipe, )

def warped_clip_vision_load(ckpt_path):
    print("Reading: {}".format(ckpt_path))
    with open(ckpt_path, "rb") as file:
        clip_vision_temp = file.read()

    print("Loading Clip Vision Model...")
    sd = warped_load_torch_file(clip_vision_temp)
    print("Loading Clip Vision Model...Done")

    clip_vision_temp = None
    gc.collect()
    time.sleep(1)

    # sd = load_torch_file(ckpt_path)
    if "visual.transformer.resblocks.0.attn.in_proj_weight" in sd:
        return comfy.clip_vision.load_clipvision_from_sd(sd, prefix="visual.", convert_keys=True)
    else:
        return comfy.clip_vision.load_clipvision_from_sd(sd)

class WarpedCLIPVisionLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "clip_name": (folder_paths.get_filename_list("clip_vision"), ),
                             }}
    RETURN_TYPES = ("CLIP_VISION",)
    RETURN_NAMES = ("clip_vision",)
    FUNCTION = "load_clip"
    CATEGORY = "Warped/General/Loaders"

    def load_clip(self, clip_name):
        clip_path = folder_paths.get_full_path_or_raise("clip_vision", clip_name)
        clip_vision = warped_clip_vision_load(clip_path)
        return (clip_vision,)

class WarpedVAELoader:
    @staticmethod
    def vae_list():
        vaes = folder_paths.get_filename_list("vae")
        approx_vaes = folder_paths.get_filename_list("vae_approx")
        sdxl_taesd_enc = False
        sdxl_taesd_dec = False
        sd1_taesd_enc = False
        sd1_taesd_dec = False
        sd3_taesd_enc = False
        sd3_taesd_dec = False
        f1_taesd_enc = False
        f1_taesd_dec = False

        for v in approx_vaes:
            if v.startswith("taesd_decoder."):
                sd1_taesd_dec = True
            elif v.startswith("taesd_encoder."):
                sd1_taesd_enc = True
            elif v.startswith("taesdxl_decoder."):
                sdxl_taesd_dec = True
            elif v.startswith("taesdxl_encoder."):
                sdxl_taesd_enc = True
            elif v.startswith("taesd3_decoder."):
                sd3_taesd_dec = True
            elif v.startswith("taesd3_encoder."):
                sd3_taesd_enc = True
            elif v.startswith("taef1_encoder."):
                f1_taesd_dec = True
            elif v.startswith("taef1_decoder."):
                f1_taesd_enc = True
        if sd1_taesd_dec and sd1_taesd_enc:
            vaes.append("taesd")
        if sdxl_taesd_dec and sdxl_taesd_enc:
            vaes.append("taesdxl")
        if sd3_taesd_dec and sd3_taesd_enc:
            vaes.append("taesd3")
        if f1_taesd_dec and f1_taesd_enc:
            vaes.append("taef1")
        return vaes

    @staticmethod
    def load_taesd(name):
        sd = {}
        approx_vaes = folder_paths.get_filename_list("vae_approx")

        encoder = next(filter(lambda a: a.startswith("{}_encoder.".format(name)), approx_vaes))
        decoder = next(filter(lambda a: a.startswith("{}_decoder.".format(name)), approx_vaes))

        encoder_path = folder_paths.get_full_path_or_raise("vae_approx", encoder)

        print("Reading vae_approx encoder: {}".format(encoder_path))

        with open(encoder_path, "rb") as file:
            vae_temp = file.read()

        enc = warped_load_torch_file(vae_temp)
        vae_temp = None

        for k in enc:
            sd["taesd_encoder.{}".format(k)] = enc[k]

        decoder_path = folder_paths.get_full_path_or_raise("vae_approx", decoder)

        print("Reading vae_approx encoder: {}".format(decoder_path))

        with open(encoder_path, "rb") as file:
            vae_temp = file.read()

        dec = warped_load_torch_file(vae_temp)
        vae_temp = None

        for k in dec:
            sd["taesd_decoder.{}".format(k)] = dec[k]

        if name == "taesd":
            sd["vae_scale"] = torch.tensor(0.18215)
            sd["vae_shift"] = torch.tensor(0.0)
        elif name == "taesdxl":
            sd["vae_scale"] = torch.tensor(0.13025)
            sd["vae_shift"] = torch.tensor(0.0)
        elif name == "taesd3":
            sd["vae_scale"] = torch.tensor(1.5305)
            sd["vae_shift"] = torch.tensor(0.0609)
        elif name == "taef1":
            sd["vae_scale"] = torch.tensor(0.3611)
            sd["vae_shift"] = torch.tensor(0.1159)
        return sd

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "vae_name": (s.vae_list(), )}}
    RETURN_TYPES = ("VAE", "DICT", )
    RETURN_NAMES = ("vae", "vae_state_dict", )
    FUNCTION = "load_vae"

    CATEGORY = "Warped/General/Loaders"

    #TODO: scale factor?
    def load_vae(self, vae_name):
        if vae_name in ["taesd", "taesdxl", "taesd3", "taef1"]:
            sd = self.load_taesd(vae_name)
        else:
            vae_path = folder_paths.get_full_path_or_raise("vae", vae_name)

            print("Reading VAE: {}".format(vae_path))
            with open(vae_path, "rb") as file:
                vae_temp = file.read()

            sd = warped_load_torch_file(vae_temp)
            vae_temp = None

        vae = comfy.sd.VAE(sd=sd)

        gc.collect()
        time.sleep(1)

        vae.throw_exception_if_invalid()
        return (vae, sd, )

class WarpedNumericalConversion:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "optional": {
                        "int_value": ("INT", {"default": None, "forceInput": True}),
                        "float_value": ("FLOAT", {"default": None, "forceInput": True}),
                        "number_value": ("NUMBER", {"default": None, "forceInput": True}),
                        "bool_value": ("BOOL", {"default": None, "forceInput": True}),
            }
        }

    RETURN_TYPES = ("INT","FLOAT", "NUMBER", "STRING", )
    RETURN_NAMES = ("int","float", "number", "string", )
    OUTPUT_NODE = True
    FUNCTION = "int_to_number"
    CATEGORY = "Warped/General/Utils"

    def int_to_number(self, int_value=None, float_value=None, number_value=None, bool_value=None):
        value = None

        if (int_value == None) and (float_value == None) and (number_value == None) and (bool_value == None):
            raise ValueError("WarpedNumericalConversion: All inputs are None. Nothing to convert.")

        i = 0

        if not int_value is None:
            i += 1
            value = int_value

        if not float_value is None:
            i += 1
            value = float_value

        if not number_value is None:
            i += 1
            value = number_value

        if not bool_value is None:
            i += 1
            value = bool_value

        if i > 1:
            raise ValueError("WarpedNumericalConversion: More than one type of value inputs simultaneously is not supported.")

        if not value is None:
            return {"ui": {"string": ["{}".format(value),]}, "result": (int(value), float(value), value, "{}".format(value),), }

        return {"ui": {"string": ["{}".format(0),]}, "result": (int(0), float(0), value, "0",), }

class WarpedLoraReSave:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "lora_name": (folder_paths.get_filename_list("loras"),),
                "save_path": ("STRING", {"default": get_default_gen_output_path()}),
                "trigger_words": ("STRING", {"default": ""}),
                "additional_info": ("STRING", {"default": ""}),
            },
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    OUTPUT_IS_LIST = (True,)
    FUNCTION = "load_lora"
    CATEGORY = "Warped/General/Lora"
    DESCRIPTION = "Re-saves LORA with new metadata."

    def get_metadata(self, lora_path, trigger_words, additional_info):
        model_config_path = os.path.join(script_directory, "your_metadata.json")
        model_config_path = os.path.abspath(model_config_path)

        print("metadata config file path: {}".format(model_config_path))

        try:
            with open(model_config_path, "r") as f:
                metadata = json.load(f)

            if len(trigger_words) > 0:
                metadata["trigger_words"] = trigger_words

            if len(additional_info) > 0:
                metadata["additional_info"] = additional_info

            return metadata
        except Exception as e:
            print("Exception attempting to read config file: {}".format(model_config_path))
            raise e

        return {"metadata": "No Metadata" }

    def load_lora(self, lora_name: str, save_path: str, trigger_words: str, additional_info: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
        if not lora_name:
            return {"ui": {"tags": ["Nothing here to Re-save"]}}

        # Get the full path to the LoRA file
        lora_path = folder_paths.get_full_path("loras", lora_name)
        if not os.path.exists(lora_path):
            raise ValueError(f"LoRA file not found: {lora_path}")

        lora_weights = comfy.utils.load_torch_file(lora_path)

        for key in lora_weights:
            lora_weights[key] = lora_weights[key].to(dtype=torch.bfloat16)

        metadata = self.get_metadata(lora_path, trigger_words, additional_info)

        temp_strings = save_path.split('\\')
        del temp_strings[len(temp_strings) - 1]

        print("Save Path: {}".format(save_path))

        save_folder = ""
        for temp_string in temp_strings:
            if len(save_folder) > 0:
                save_folder = "{}\\{}".format(save_folder, temp_string)
            else:
                save_folder = temp_string

        print("Save_folder: {}".format(save_folder))
        os.makedirs(save_folder, exist_ok = True)

        utils.save_torch_file(lora_weights, save_path, metadata=metadata)

        save_message = "Weights Saved To: {}".format(save_path)

        return {"ui": {"tags": [save_message]}}

class WarpedReverseImageBatch:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE", {"tooltip": "Image batch to reverse."}),
            }
        }
    RETURN_TYPES = ("IMAGE", )
    RETURN_NAMES = ("images", )
    FUNCTION = "reverse_images"
    CATEGORY = "Warped/General/Image"
    DESCRIPTION = "Reverses the order of the batch of images input into the node."

    def reverse_images(self, images=None):
        if images is None:
            raise ValueError("images input cannot be None.")

        images_tuple = torch.split(images, 1, dim=0)
        images_split = [item for item in images_tuple]

        images_split.reverse()

        resulting_images = None
        for temp_image in images_split:
            if not resulting_images is None:
                resulting_images = torch.cat((resulting_images, temp_image), 0)
            else:
                resulting_images = temp_image

        return (resulting_images, )

class WarpedCheckpointLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "ckpt_name": (folder_paths.get_filename_list("checkpoints"), {"tooltip": "The name of the checkpoint (model) to load."}),
            }
        }
    RETURN_TYPES = ("MODEL", "CLIP", "VAE", )
    RETURN_NAMES = ("model", "clip", "vae", )
    OUTPUT_TOOLTIPS = ("The model used for denoising latents.",
                       "The CLIP model used for encoding text prompts.",
                       "The VAE model used for encoding and decoding images to and from latent space.")
    FUNCTION = "load_checkpoint"
    CATEGORY = "Warped/General/Loaders"
    DESCRIPTION = "Loads a diffusion model checkpoint, diffusion models are used to denoise latents."

    def load_checkpoint(self, ckpt_name):
        ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
        out = self.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
        return out[:3]

    def load_checkpoint_guess_config(self, ckpt_path, output_vae=True, output_clip=True, embedding_directory=None):
        print("Reading: {}".format(ckpt_path))
        with open(ckpt_path, "rb") as file:
            checkpoint_temp = file.read()

        print("Loading Checkpoint...")
        ckpt = warped_load_torch_file(checkpoint_temp)
        print("Loading Checkpoint...Done")
        checkpoint_temp = None

        gc.collect()
        time.sleep(1)

        out = comfy.sd.load_state_dict_guess_config(ckpt, output_vae, output_clip, embedding_directory)
        if out is None:
            raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path))

        return out

class WarpedDatabase:
    def __init__(self, the_dictionary):
        self.data = the_dictionary

    def catExists(self, category):
        return category in self.data

    def keyExists(self, category, key):
        return category in self.data and key in self.data[category]

    def insert(self, category, key, value):
        if not isinstance(category, str) or not isinstance(key, str):
            cstr("Category and key must be strings").error.print()
            return

        if category not in self.data:
            self.data[category] = {}
        self.data[category][key] = value
        self._save()

    def update(self, category, key, value):
        if category in self.data and key in self.data[category]:
            self.data[category][key] = value
            self._save()

    def updateCat(self, category, dictionary):
        self.data[category].update(dictionary)
        self._save()

    def get(self, category, key):
        return self.data.get(category, {}).get(key, None)

    def getDB(self):
        return self.data

    def insertCat(self, category):
        if not isinstance(category, str):
            cstr("Category must be a string").error.print()
            return

        if category in self.data:
            cstr(f"The database category '{category}' already exists!").error.print()
            return
        self.data[category] = {}
        self._save()

    def getDict(self, category):
        if category not in self.data:
            cstr(f"The database category '{category}' does not exist!").error.print()
            return {}
        return self.data[category]

warped_DB = WarpedDatabase({"custom_tokens": {}})

class TextTokens:
    def __init__(self):
        self.WDB = warped_DB
        if not self.WDB.getDB().__contains__('custom_tokens'):
            self.WDB.insertCat('custom_tokens')
        self.custom_tokens = self.WDB.getDict('custom_tokens')

        self.tokens = {
            '[time]': str(time.time()).replace('.','_'),
            '[hostname]': socket.gethostname(),
            '[cuda_device]': str(comfy.model_management.get_torch_device()),
            '[cuda_name]': str(comfy.model_management.get_torch_device_name(device=comfy.model_management.get_torch_device())),
        }

        if '.' in self.tokens['[time]']:
            self.tokens['[time]'] = self.tokens['[time]'].split('.')[0]

        try:
            self.tokens['[user]'] = os.getlogin() if os.getlogin() else 'null'
        except Exception:
            self.tokens['[user]'] = 'null'

    def addToken(self, name, value):
        self.custom_tokens.update({name: value})
        self._update()

    def removeToken(self, name):
        self.custom_tokens.pop(name)
        self._update()

    def format_time(self, format_code):
        return time.strftime(format_code, time.localtime(time.time()))

    def parseTokens(self, text):
        tokens = self.tokens.copy()
        if self.custom_tokens:
            tokens.update(self.custom_tokens)

        # Update time
        tokens['[time]'] = str(time.time())
        if '.' in tokens['[time]']:
            tokens['[time]'] = tokens['[time]'].split('.')[0]

        for token, value in tokens.items():
            if token.startswith('[time('):
                continue
            pattern = re.compile(re.escape(token))
            text = pattern.sub(value, text)

        def replace_custom_time(match):
            format_code = match.group(1)
            return self.format_time(format_code)

        text = re.sub(r'\[time\((.*?)\)\]', replace_custom_time, text)

        return text

    def _update(self):
        self.WDB.updateCat('custom_tokens', self.custom_tokens)

class WarpedSamplerScriptsBase:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        arg_dict = {
            "required": {
            },
            "optional": {
            }
        }

        arg_dict["optional"]["dummy"] = ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True})

        for i in range(self.num_batchs):
            arg_dict["optional"][f"batch_{i}"] = ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True})

        return arg_dict

    RETURN_TYPES = ("WARPEDSCRIPTS",)
    RETURN_NAMES = ("scripts", )
    FUNCTION = "do_scripts"

    CATEGORY = "Warped/General/Scripts"

    def do_scripts(self, **kwargs):
        scripts = {}

        temp_text = kwargs.get(f"dummy")

        if len(temp_text) > 0:
            scripts["dummy"] = self.get_tokens(temp_text)

        for i in range(12):
            temp_text = kwargs.get(f"batch_{i}")

            if len(temp_text) > 0:
                scripts["{}".format(i)] = self.get_tokens(temp_text)

        return (scripts, )

    def get_tokens(self, text):
        new_text = []
        for line in io.StringIO(text):
            if not line.strip().startswith('#'):
                new_text.append(line.replace("\n", ''))
        new_text = "\n".join(new_text)

        tokens = TextTokens()
        new_text = tokens.parseTokens(new_text)

        return new_text

class WarpedSamplerScripts5(WarpedSamplerScriptsBase):
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(self):
        arg_dict = {
            "required": {
            },
            "optional": {
            }
        }

        arg_dict["optional"]["dummy"] = ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True})

        for i in range(5):
            arg_dict["optional"][f"batch_{i}"] = ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True})

        return arg_dict

    def do_scripts(self, **kwargs):
        scripts = {}

        temp_text = kwargs.get(f"dummy")

        if len(temp_text) > 0:
            scripts["dummy"] = self.get_tokens(temp_text)

        for i in range(5):
            temp_text = kwargs.get(f"batch_{i}")

            if len(temp_text) > 0:
                scripts["{}".format(i)] = self.get_tokens(temp_text)

        return (scripts, )

class WarpedSamplerScripts8(WarpedSamplerScriptsBase):
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(self):
        arg_dict = {
            "required": {
            },
            "optional": {
            }
        }

        arg_dict["optional"]["dummy"] = ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True})

        for i in range(8):
            arg_dict["optional"][f"batch_{i}"] = ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True})

        return arg_dict

    def do_scripts(self, **kwargs):
        scripts = {}

        temp_text = kwargs.get(f"dummy")

        if len(temp_text) > 0:
            scripts["dummy"] = self.get_tokens(temp_text)

        for i in range(8):
            temp_text = kwargs.get(f"batch_{i}")

            if len(temp_text) > 0:
                scripts["{}".format(i)] = self.get_tokens(temp_text)

        return (scripts, )

class WarpedSamplerScripts12(WarpedSamplerScriptsBase):
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(self):
        arg_dict = {
            "required": {
            },
            "optional": {
            }
        }

        arg_dict["optional"]["dummy"] = ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True})

        for i in range(12):
            arg_dict["optional"][f"batch_{i}"] = ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True})

        return arg_dict

    def do_scripts(self, **kwargs):
        scripts = {}

        temp_text = kwargs.get(f"dummy")

        if len(temp_text) > 0:
            scripts["dummy"] = self.get_tokens(temp_text)

        for i in range(12):
            temp_text = kwargs.get(f"batch_{i}")

            if len(temp_text) > 0:
                scripts["{}".format(i)] = self.get_tokens(temp_text)

        return (scripts, )

class WarpedSamplerScripts16(WarpedSamplerScriptsBase):
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(self):
        arg_dict = {
            "required": {
            },
            "optional": {
            }
        }

        arg_dict["optional"]["dummy"] = ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True})

        for i in range(16):
            arg_dict["optional"][f"batch_{i}"] = ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True})

        return arg_dict

    def do_scripts(self, **kwargs):
        scripts = {}

        temp_text = kwargs.get(f"dummy")

        if len(temp_text) > 0:
            scripts["dummy"] = self.get_tokens(temp_text)

        for i in range(16):
            temp_text = kwargs.get(f"batch_{i}")

            if len(temp_text) > 0:
                scripts["{}".format(i)] = self.get_tokens(temp_text)

        return (scripts, )

class WarpedSamplerScripts20(WarpedSamplerScriptsBase):
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(self):
        arg_dict = {
            "required": {
            },
            "optional": {
            }
        }

        arg_dict["optional"]["dummy"] = ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True})

        for i in range(20):
            arg_dict["optional"][f"batch_{i}"] = ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True})

        return arg_dict

    def do_scripts(self, **kwargs):
        scripts = {}

        temp_text = kwargs.get(f"dummy")

        if len(temp_text) > 0:
            scripts["dummy"] = self.get_tokens(temp_text)

        for i in range(20):
            temp_text = kwargs.get(f"batch_{i}")

            if len(temp_text) > 0:
                scripts["{}".format(i)] = self.get_tokens(temp_text)

        return (scripts, )

class WarpedSamplerScripts30(WarpedSamplerScriptsBase):
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(self):
        arg_dict = {
            "required": {
            },
            "optional": {
            }
        }

        arg_dict["optional"]["dummy"] = ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True})

        for i in range(30):
            arg_dict["optional"][f"batch_{i}"] = ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True})

        return arg_dict

    def do_scripts(self, **kwargs):
        scripts = {}

        temp_text = kwargs.get(f"dummy")

        if len(temp_text) > 0:
            scripts["dummy"] = self.get_tokens(temp_text)

        for i in range(30):
            temp_text = kwargs.get(f"batch_{i}")

            if len(temp_text) > 0:
                scripts["{}".format(i)] = self.get_tokens(temp_text)

        return (scripts, )

class WarpedSamplerScripts40(WarpedSamplerScriptsBase):
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(self):
        arg_dict = {
            "required": {
            },
            "optional": {
            }
        }

        arg_dict["optional"]["dummy"] = ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True})

        for i in range(40):
            arg_dict["optional"][f"batch_{i}"] = ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True})

        return arg_dict

    def do_scripts(self, **kwargs):
        scripts = {}

        temp_text = kwargs.get(f"dummy")

        if len(temp_text) > 0:
            scripts["dummy"] = self.get_tokens(temp_text)

        for i in range(40):
            temp_text = kwargs.get(f"batch_{i}")

            if len(temp_text) > 0:
                scripts["{}".format(i)] = self.get_tokens(temp_text)

        return (scripts, )

class WarpedLoadImages:
    def __init__(self, index=0):
        self.index = index
        self.output_dir = folder_paths.get_output_directory()
        self.previous_path = ""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "mode": (["incremental_image", "random"], {"default": "incremental_image"}),
                "index": ("INT", {"default": 0, "min": 0, "max": 150000, "step": 1}),
                "label": ("STRING", {"default": 'Batch 001', "multiline": False}),
                "path": ("STRING", {"default": '', "multiline": False}),
                "periodic_sleep": ("BOOLEAN", {"default": False})
            },
            "optional": {
                "suffix": ("STRING", {"default": '', "multiline": False})
            }
        }

    RETURN_TYPES = ("STRING", "INT", "STRING", "STRING", "STRING")
    RETURN_NAMES = ("filename", "index", "prefixorg", "prefixseg", "prefixmsk")

    FUNCTION = "load_batch_images"

    CATEGORY = "Warped/General/Image"

    def load_batch_images(self, path, index=0, mode="incremental_image", label='Batch 001', periodic_sleep=False, suffix=""):
        if not os.path.exists(path):
            return (None, )

        if path != self.previous_path:
            self.index = index
            self.previous_path = path

        if periodic_sleep:
            if (self.index > 0) and (self.index % 10 == 0):
                time.sleep(2)

        fl = self.BatchImageLoader(path, label, '*', index)
        new_paths = fl.image_paths

        retry = True

        try:
            filename = fl.image_paths[self.index]
        except:
            if retry:
                retry = False
                self.index = 0
                filename = fl.image_paths[self.index]

        print("Filename: {}".format(filename))

        temp_strings = filename.split(".")
        temp_fileend = temp_strings[len(temp_strings) - 1]

        temp_filename = ""

        i = 0
        while i < (len(temp_strings) - 1):
            temp_filename = "{}{}".format(temp_filename, temp_strings[i])
            i += 1

        tempFilename = "{}/{}_{}.{}"
        tempFilenamesuffix = "{}/{}_{}_{}.{}"

        tempStrings1 = temp_filename.split("\\")
        temp_filename = tempStrings1[len(tempStrings1) - 1]

        if len(suffix) < 1:
            prefixorg = tempFilename.format(self.output_dir, temp_filename, "org", temp_fileend)
            prefixseg = tempFilename.format(self.output_dir, temp_filename, "seg", temp_fileend)
            prefixmsk = tempFilename.format(self.output_dir, temp_filename, "msk", temp_fileend)
        else:
            prefixorg = tempFilenamesuffix.format(self.output_dir, temp_filename, "org", suffix, temp_fileend)
            prefixseg = tempFilenamesuffix.format(self.output_dir, temp_filename, "seg", suffix, temp_fileend)
            prefixmsk = tempFilenamesuffix.format(self.output_dir, temp_filename, "msk", suffix, temp_fileend)

        self.index += 1

        if self.index >= len(fl.image_paths):
            self.index = 0

        return filename, self.index - 1, prefixorg, prefixseg, prefixmsk,

    class BatchImageLoader:
        def __init__(self, directory_path, label, pattern, index):
            self.image_paths = []
            self.load_images(directory_path, pattern)
            self.image_paths.sort()

            self.index = index
            self.label = label

        def load_images(self, directory_path, pattern):
            for file_name in glob.glob(os.path.join(directory_path, pattern), recursive=True):
                if file_name.lower().endswith(ALLOWED_EXT):
                    abs_file_path = os.path.abspath(file_name)
                    self.image_paths.append(abs_file_path)

        def get_image_by_id(self, image_id):
            if image_id < 0 or image_id >= len(self.image_paths):
                cstr(f"Invalid image index `{image_id}`").error.print()
                return

            return self.image_paths[image_id]

        def get_next_image(self):
            if self.index >= len(self.image_paths):
                self.index = 0

            image_path = self.image_paths[self.index]
            self.index += 1

            if self.index == len(self.image_paths):
                self.index = 0

            cstr(f'{cstr.color.YELLOW}{self.label}{cstr.color.END} Index: {self.index}').msg.print()

            return image_path

        def get_current_image(self):
            if self.index >= len(self.image_paths):
                self.index = 0
            image_path = self.image_paths[self.index]

            return image_path

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return float("NaN")

class WarpedSaveImageCaption:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"image_path": ("STRING", {"forceInput": True}),
                     "caption": ("STRING", {"forceInput": True}),
                    },
                }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("caption",)
    OUTPUT_NODE = True
    FUNCTION = "save_caption"

    CATEGORY = "Warped/General/Captioning"

    def save_caption(self, image_path, caption):
        tempStrings = image_path.split('.')
        caption_path = tempStrings[0]

        i = 1

        while i < (len(tempStrings) - 1):
            caption_path = "{}.{}".format(caption_path, tempStrings[i])
            i += 1

        caption_path = "{}.txt".format(caption_path)

        with open(caption_path, 'w') as fp:
            fp.write(caption)

        return {"ui": {"string": [caption,]}, "result": (caption,)}

class WarpedModifyCaptionFile:
    def __init__(self, index=0):
        self.index = index
        self.previous_path = ""

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"mode": (["incremental_caption", "random"], {"default": "incremental_caption"}),
                     "index": ("INT", {"default": 0, "min": 0, "max": 150000, "step": 1}),
                     "caption_path": ("STRING", {"forceInput": False}),
                     "find_text": ("STRING", {"forceInput": False}),
                     "replace_text": ("STRING", {"forceInput": False}),
                    },
                }

    RETURN_TYPES = ("STRING", "STRING", "STRING", )
    RETURN_NAMES = ("filename", "original", "modified", )
    OUTPUT_NODE = True
    FUNCTION = "load_batch_captions"

    CATEGORY = "Warped/General/Captioning"

    def load_batch_captions(self, mode="incremental_caption", index=0, caption_path="", find_text="", replace_text=""):
        if (len(caption_path) == 0) or (len(find_text) == 0) or (not os.path.exists(caption_path)):
            return ("None", "None", "None")

        if caption_path != self.previous_path:
            self.index = index
            self.previous_path = caption_path

        fl = self.BatchCaptionLoader(caption_path, '*', index)

        retry = True

        try:
            filename = fl.caption_paths[self.index]
        except:
            if retry:
                retry = False
                self.index = 0
                filename = fl.caption_paths[self.index]

        print("Filename: {}".format(filename))

        original, modified = self.find_replace(filename, find_text, replace_text)

        self.index += 1

        if self.index >= len(fl.caption_paths):
            self.index = 0

        return {"ui": {"string": [filename, original, modified]}, "result": (filename, original, modified)}

    def find_replace(self, filename, find_text, replace_text):
        with open(filename, "r") as fp:
            original = fp.read()

        print("\nOriginal Text: {}".format(original))

        modified = original.replace(find_text.strip(), replace_text.strip())

        print("Modified Text: {}".format(modified))

        with open(filename, "w") as fp:
            fp.write(modified)

        return (original, modified)

    class BatchCaptionLoader:
        def __init__(self, directory_path, pattern, index):
            self.caption_paths = []
            self.load_captions(directory_path, pattern)
            self.caption_paths.sort()
            self.index = index

        def load_captions(self, directory_path, pattern):
            for file_name in glob.glob(os.path.join(directory_path, pattern), recursive=True):
                if file_name.lower().endswith(ALLOWED_CAPTION_EXT):
                    abs_file_path = os.path.abspath(file_name)
                    self.caption_paths.append(abs_file_path)

        def get_caption_by_id(self, caption_id):
            if caption_id < 0 or caption_id >= len(self.caption_paths):
                cstr(f"Invalid caption index `{caption_id}`").error.print()
                return

            return self.caption_paths[caption_id]

        def get_next_caption(self):
            if self.index >= len(self.caption_paths):
                self.index = 0

            caption_path = self.caption_paths[self.index]
            self.index += 1

            if self.index == len(self.caption_paths):
                self.index = 0

            cstr(f'{cstr.color.YELLOW}{self.label}{cstr.color.END} Index: {self.index}').msg.print()

            return caption_path

        def get_current_caption(self):
            if self.index >= len(self.caption_paths):
                self.index = 0
            caption_path = self.caption_paths[self.index]

            return caption_path

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return float("NaN")

class WarpedAddToCaption:
    def __init__(self, index=0):
        self.index = index
        self.previous_path = ""

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"mode": (["incremental_caption", "random"], {"default": "incremental_caption"}),
                     "index": ("INT", {"default": 0, "min": 0, "max": 150000, "step": 1}),
                     "caption_path": ("STRING", {"forceInput": False}),
                     "ignore_if_contains": ("STRING", {"forceInput": False}),
                     "add_text": ("STRING", {"forceInput": False}),
                     "add_mode": (["prefix", "postfix"], {"default": "prefix"}),
                    },
                }

    RETURN_TYPES = ("STRING", "STRING", "STRING", )
    RETURN_NAMES = ("filename", "original", "modified", )
    OUTPUT_NODE = True
    FUNCTION = "load_batch_captions"

    CATEGORY = "Warped/General/Captioning"

    def load_batch_captions(self, mode="incremental_caption", index=0, caption_path="", ignore_if_contains="", add_text="", add_mode="prefix"):
        if (len(caption_path) == 0) or (len(add_text) == 0) or (not os.path.exists(caption_path)):
            return ("None", "None", "None")

        if caption_path != self.previous_path:
            self.index = index
            self.previous_path = caption_path

        fl = self.BatchCaptionLoader(caption_path, '*', index)

        retry = True

        try:
            filename = fl.caption_paths[self.index]
        except:
            if retry:
                retry = False
                self.index = 0
                filename = fl.caption_paths[self.index]

        print("Filename: {}".format(filename))

        original, modified = self.process_prompt(filename, ignore_if_contains, add_text, add_mode)

        self.index += 1

        if self.index >= len(fl.caption_paths):
            self.index = 0

        return {"ui": {"string": [filename, original, modified]}, "result": (filename, original, modified)}

    def process_prompt(self, filename, ignore_if_contains, add_text, add_mode):
        with open(filename, "r") as fp:
            original = fp.read()

        print("\nOriginal Text: {}".format(original))

        if len(ignore_if_contains) > 0:
            if original.find(ignore_if_contains) != -1:
                return (original, original)

        if add_mode == "prefix":
            modified = "{}, {}".format(add_text, original)
        else:
            modified = "{}, {}".format(original, add_text)

        print("Modified Text: {}".format(modified))

        with open(filename, "w") as fp:
            fp.write(modified)

        return (original, modified)

    class BatchCaptionLoader:
        def __init__(self, directory_path, pattern, index):
            self.caption_paths = []
            self.load_captions(directory_path, pattern)
            self.caption_paths.sort()
            self.index = index

        def load_captions(self, directory_path, pattern):
            for file_name in glob.glob(os.path.join(directory_path, pattern), recursive=True):
                if file_name.lower().endswith(ALLOWED_CAPTION_EXT):
                    abs_file_path = os.path.abspath(file_name)
                    self.caption_paths.append(abs_file_path)

        def get_caption_by_id(self, caption_id):
            if caption_id < 0 or caption_id >= len(self.caption_paths):
                cstr(f"Invalid caption index `{caption_id}`").error.print()
                return

            return self.caption_paths[caption_id]

        def get_next_caption(self):
            if self.index >= len(self.caption_paths):
                self.index = 0

            caption_path = self.caption_paths[self.index]
            self.index += 1

            if self.index == len(self.caption_paths):
                self.index = 0

            cstr(f'{cstr.color.YELLOW}{self.label}{cstr.color.END} Index: {self.index}').msg.print()

            return caption_path

        def get_current_caption(self):
            if self.index >= len(self.caption_paths):
                self.index = 0
            caption_path = self.caption_paths[self.index]

            return caption_path

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return float("NaN")

class WarpedPromptConcat:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"prompt1": ("STRING", {"default": ""}),
                     "prompt2": ("STRING", {"default": "", "forceInput": True}),
                     "delimiter": ([",", "."], {"default": ","}),
                    },
                }

    RETURN_TYPES = ("STRING", )
    RETURN_NAMES = ("prompt", )
    FUNCTION = "do_concat"

    CATEGORY = "Warped/General/Utils"

    def do_concat(self, prompt1="", prompt2="", delimiter=","):
        temp_prompt1 = prompt1.strip().strip("\n").strip("\0").strip("\r")
        temp_prompt2 = prompt2.strip().strip("\n").strip("\0").strip("\r")

        output_prompt = "{}{} {}".format(temp_prompt1, delimiter, temp_prompt2)

        print("output_prompt: {}".format(output_prompt))

        return (output_prompt,)

class WarpedPromptConcatExt:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"prompt1": ("STRING", {"default": ""}),
                     "prompt2": ("STRING", {"default": "", "forceInput": True}),
                     "delimiter": ([",", "."], {"default": ","}),
                     "filter_json_file": ("STRING", {"default": ""})
                    },
                }

    RETURN_TYPES = ("STRING", )
    RETURN_NAMES = ("prompt", )
    FUNCTION = "do_concat"

    CATEGORY = "Warped/General/Utils"

    def do_concat(self, prompt1="", prompt2="", delimiter=",", filter_json_file=""):
        if len(filter_json_file) < 1:
            filter_json_file = None

        temp_prompt1 = prompt1.strip().strip("\n").strip("\0").strip("\r")
        temp_prompt2 = prompt2.strip().strip("\n").strip("\0").strip("\r")

        output_prompt = "{}{} {}".format(temp_prompt1, delimiter, temp_prompt2)

        filter_list = []
        if not filter_json_file is None:
            if not os.path.exists(filter_json_file):
                raise ValueError("{} Not Found!".format(filter_json_file))

            # try:
            with open(filter_json_file) as fp:
                filter_list = json.load(fp)
            # except:
            #     print("Error Encountered reading filter_json_file. Returning unfiltered prompt.")
            #     pass

            for entry in filter_list:
                output_prompt = output_prompt.replace(entry["find"], entry["replace"])

        print("output_prompt: {}".format(output_prompt))

        return (output_prompt,)

class WarpedImageLossCalc:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"predicted": ("IMAGE", ),
                     "target": ("IMAGE", ),
                    },
                }

    RETURN_TYPES = ("STRING", )
    RETURN_NAMES = ("loss_calc", )
    FUNCTION = "do_calc"

    CATEGORY = "Warped/General/Utils"

    def do_calc(self, predicted, target):
        loss = torch.nn.functional.mse_loss(predicted.to(dtype=torch.bfloat16), target.to(dtype=torch.bfloat16), reduction="none")
        loss = loss.mean()
        loss_calc = "{:5.4f}".format(loss)

        return (loss_calc,)

class WarpedUNETLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "unet_name": (folder_paths.get_filename_list("diffusion_models"), ),
                              "weight_dtype": (["default", "f16", "bf16", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"],)
                             }}
    RETURN_TYPES = ("MODEL",)
    FUNCTION = "load_unet"

    CATEGORY = "Warped/General/Loaders"

    def load_unet(self, unet_name, weight_dtype):
        model_options = {}
        if weight_dtype == "fp8_e4m3fn":
            model_options["dtype"] = torch.float8_e4m3fn
        elif weight_dtype == "fp8_e4m3fn_fast":
            model_options["dtype"] = torch.float8_e4m3fn
            model_options["fp8_optimizations"] = True
        elif weight_dtype == "fp8_e5m2":
            model_options["dtype"] = torch.float8_e5m2
        elif weight_dtype == "bf16":
            model_options["dtype"] = torch.bfloat16
        elif weight_dtype == "f16":
            model_options["dtype"] = torch.float16

        unet_path = folder_paths.get_full_path_or_raise("diffusion_models", unet_name)

        print("Reading: {} . . .".format(unet_path))
        with open(unet_path, "rb") as file:
            content = file.read()
        print("Reading: {} . . . Done!".format(unet_path))

        chkpt, metadata = warped_load_torch_file(content, return_metadata=True)

        for key in chkpt:
            print("Diffusiion Model Key: {}  |  Shape: {}".format(key, chkpt[key].shape))

        model = comfy.sd.load_diffusion_model_state_dict(chkpt, model_options=model_options, metadata=metadata)

        if model is None:
            logging.error("ERROR UNSUPPORTED DIFFUSION MODEL {}".format(unet_path))
            raise RuntimeError("ERROR: Could not detect model type of: {}\n{}".format(unet_path, model_detection_error_hint(unet_path, chkpt)))

        # model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options)
        return (model,)

class WarpedUNETConverter:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "unet_name": (folder_paths.get_filename_list("diffusion_models"), ),
                              "weight_dtype": (["default", "fp16", "bf16", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"],)
                             }}
    RETURN_TYPES = ()
    OUTPUT_NODE = True
    OUTPUT_IS_LIST = (True,)
    FUNCTION = "convert_model"

    CATEGORY = "Warped/General/Conversion"

    def load_unet(self, unet_name, weight_dtype):
        model_options = {}
        if weight_dtype == "fp8_e4m3fn":
            model_options["dtype"] = torch.float8_e4m3fn
        elif weight_dtype == "fp8_e4m3fn_fast":
            model_options["dtype"] = torch.float8_e4m3fn
            model_options["fp8_optimizations"] = True
        elif weight_dtype == "fp8_e5m2":
            model_options["dtype"] = torch.float8_e5m2
        elif weight_dtype == "bf16":
            model_options["dtype"] = torch.bfloat16
        elif weight_dtype == "fp16":
            model_options["dtype"] = torch.float16

        unet_path = folder_paths.get_full_path_or_raise("diffusion_models", unet_name)

        print("Reading: {} . . .".format(unet_path))
        with open(unet_path, "rb") as file:
            content = file.read()
        print("Reading: {} . . . Done!".format(unet_path))

        chkpt, metadata = warped_load_torch_file(content, return_metadata=True)
        content = None
        new_chkpt = {}

        for key in chkpt:
            print("Diffusiion Model Key: {}  |  Shape: {}".format(key, chkpt[key].shape))
            new_chkpt[key] = chkpt[key].to(dtype=model_options["dtype"])

        return new_chkpt, metadata, unet_path

    def convert_model(self, unet_name, weight_dtype):
        unet_data, metadata, unet_path = self.load_unet(unet_name, weight_dtype)

        metadata = {"original_metadata": "{}".format(metadata)}

        save_path = "{}_{}{}".format(os.path.splitext(unet_path)[0], weight_dtype, os.path.splitext(unet_path)[1])

        utils.save_torch_file(unet_data, save_path, metadata=metadata)

        unet_data = None
        mm.soft_empty_cache()
        gc.collect()
        time.sleep(1)

        save_message = "Weights Saved To: {}".format(save_path)

        return {"ui": {"tags": [save_message]}}

class WarpedHunyuanLoraDoubleBlocksSwap:
    @classmethod
    def INPUT_TYPES(s):
        arg_dict = {
            "required": {
                "lora_model_1": (['None'] + folder_paths.get_filename_list("loras"),),
                "lora_model_2": (['None'] + folder_paths.get_filename_list("loras"),),
                "mainstrength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "display": "number"}),
                "save_path": ("STRING", {"default": get_default_output_path()}),
                "save_new_lora": ("BOOLEAN", {"default": False}),
                "return_state_only": ("BOOLEAN", {"default": False}),
            },
            "optional": {
                "model": ("MODEL", {"default": None}),
                "state_dictionary": ("WARPEDSTATEDICT", {"default": None}),
                "metadata_dict": ("WARPEDMETADICT", {"default": None}),
                "metadata_flush": ("BOOLEAN", {"default": False}),
                "verbose_messaging": ("BOOLEAN", {"default": False}),
            }
        }

        for i in range(20):
            arg_dict["required"][f"double_blocks.{i}."] = ("BOOLEAN", {"default": False})

        return arg_dict

    RETURN_TYPES = ("MODEL", "STRING", "WARPEDSTATEDICT", "WARPEDMETADICT",)
    RETURN_NAMES = ("model", "metadata", "state_dict", "metadata_dict",)
    FUNCTION = "load_lora"
    CATEGORY = "Warped/Hunyuan/Mixers"
    OUTPUT_NODE = False
    DESCRIPTION = "LoRA, single blocks double blocks"

    def load_lora(self, lora_model_1, lora_model_2, mainstrength, save_path, save_new_lora=False, return_state_only=False, model=None, state_dictionary=None, metadata_dict=None, metadata_flush=False, verbose_messaging=False, **kwargs):
        if ((lora_model_1 is None) and (state_dictionary)) or (lora_model_2 is None):
            ValueError("Both LORA models must be a valid selection, or LORA 1 has to be replaced by optional state dictionary input.")

        from comfy.utils import load_torch_file, save_torch_file
        from comfy.sd import load_lora_for_models
        from comfy.lora import load_lora

        new_metadata = None
        metadata = None
        torch.cuda.empty_cache()
        gc.collect()
        time.sleep(0.25)

        if state_dictionary is None:
            lora_1, metadata = warped_load_lora_weights(lora_model_1, return_metadata=True)
        else:
            lora_1 = state_dictionary

            if not metadata_dict is None:
                metadata = metadata_dict
            else:
                metadata = {}

        if metadata_flush or (metadata is None):
            metadata = {}

        new_metadata = metadata.copy()

        lora_2 = warped_load_lora_weights(lora_model_2)

        diffusers_lora_1 = convert_lora(lora_1, convert_to="diffusion_model")
        filtered_lora_1 = filter_lora_keys(diffusers_lora_1, "all")

        diffusers_lora_2 = convert_lora(lora_2, convert_to="diffusion_model")
        filtered_lora_2 = filter_lora_keys(diffusers_lora_2, "all")

        max_dimension = 0

        for key in lora_1:
            if lora_1[key].shape[0] < lora_1[key].shape[1]:
                temp_dimension = lora_1[key].shape[0]
            else:
                temp_dimension = lora_1[key].shape[1]

            if temp_dimension > max_dimension:
                max_dimension = temp_dimension

            break

        for key in lora_2:
            if lora_2[key].shape[0] < lora_2[key].shape[1]:
                temp_dimension = lora_2[key].shape[0]
            else:
                temp_dimension = lora_2[key].shape[1]

            if temp_dimension > max_dimension:
                max_dimension = temp_dimension

            break

        filtered_lora_1 = convert_lora_dimensions(max_dimension, filtered_lora_1)
        filtered_lora_2 = convert_lora_dimensions(max_dimension, filtered_lora_2)

        block_settings = {
            **kwargs,
        }

        swap_block_settings = {k: v for k, v in block_settings.items() if v is True}
        swap_blocks = []

        for key in swap_block_settings:
            if swap_block_settings[key]:
                temp_strings = key.split('.')

                if len(temp_strings) > 1:
                    swap_blocks.append(int(temp_strings[1]))

        for entry in swap_blocks:
            temp_search_val = ".{}.".format(entry)

            for key in filtered_lora_1:
                if (temp_search_val in key) and (key in filtered_lora_2):
                    filtered_lora_1[key] = filtered_lora_2[key]

        if len(swap_blocks) > 0:
            if (not "modified_loras" in new_metadata) and (not "tune_counter" in new_metadata):
                new_metadata = {"modified_loras": "{} and {}".format(lora_model_1, lora_model_2)}

                swap_block_data = ""

                for entry in swap_blocks:
                    swap_block_data = "{} {}".format(swap_block_data, entry).strip()

                new_metadata["replaced_blocks"] = swap_block_data
            else:
                new_metadata["modified_loras"] = "{} | {} and {}".format(new_metadata["modified_loras"], lora_model_1, lora_model_2)

                swap_block_data = ""

                for entry in swap_blocks:
                    swap_block_data = "{} {}".format(swap_block_data, entry).strip()

                if "replaced_blocks" in new_metadata:
                    new_metadata["replaced_blocks"] = "{}  |  {}".format(new_metadata["replaced_blocks"], swap_block_data)
                else:
                    new_metadata["replaced_blocks"] = "{}".format(swap_block_data)

        if verbose_messaging:
            print("\nLora Metadata: {}".format(new_metadata))

        if save_new_lora and (len(swap_blocks) > 0):
            os.makedirs(os.path.dirname(save_path), exist_ok=True)

            for key in filtered_lora_1:
                filtered_lora_1[key] = filtered_lora_1[key].to(dtype=torch.bfloat16)

            print("Saving Model To: {}...".format(save_path))
            save_torch_file(filtered_lora_1, save_path, metadata=new_metadata)
            print("Saving Model To: {}...Done.".format(save_path))

        if (not return_state_only) and (not model is None):
            new_model, _ = load_lora_for_models(model, None, filtered_lora_1, mainstrength, 0)
            if new_model is not None:
                return (new_model, new_metadata, None, None, )

        return (model, new_metadata, filtered_lora_1, new_metadata, )

    @classmethod
    def IS_CHANGED(s, lora_model_1, lora_model_2, mainstrength, save_path, save_new_lora, return_state_only=False, model=None, state_dictionary=None, metadata_dict=None, metadata_flush=False, verbose_messaging=False, **kwargs):
        return f"{lora_model_1}_{lora_model_2}_{mainstrength}_{state_dictionary}_{metadata_dict}"

def get_base_layer_type(layer_key):
    layer_key = layer_key.replace(".weight", "")

    if "single" in layer_key:
        layer_key = layer_key.replace("diffusion_model.single_blocks.", "")
    else:
        layer_key = layer_key.replace("diffusion_model.double_blocks.", "")

    if "img_" in layer_key:
        temp_strings = layer_key.split("img_")
        layer_key = "img_{}".format(temp_strings[len(temp_strings) - 1])
    else:
        temp_strings = layer_key.split("txt_")
        layer_key = "txt_{}".format(temp_strings[len(temp_strings) - 1])

    return layer_key

class WarpedHunyuanLoraDoubleBlocksLayersBlend:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "lora_model_1": (['None'] + folder_paths.get_filename_list("loras"),),
                "lora_model_2": (['None'] + folder_paths.get_filename_list("loras"),),
                "mainstrength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "display": "number"}),
                "save_path": ("STRING", {"default": get_default_output_path()}),
                "save_new_lora": ("BOOLEAN", {"default": False}),
                "return_state_only": ("BOOLEAN", {"default": False}),
                "img_attn_proj_lora_A": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "img_attn_proj_lora_B": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "img_attn_qkv_lora_A": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "img_attn_qkv_lora_B": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "img_mlp_fc1_lora_A": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "img_mlp_fc1_lora_B": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "img_mlp_fc2_lora_A": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "img_mlp_fc2_lora_B": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "img_mod_linear_lora_A": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "img_mod_linear_lora_B": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "txt_attn_proj_lora_A": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "txt_attn_proj_lora_B": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "txt_attn_qkv_lora_A": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "txt_attn_qkv_lora_B": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "txt_mlp_fc1_lora_A": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "txt_mlp_fc1_lora_B": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "txt_mlp_fc2_lora_A": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "txt_mlp_fc2_lora_B": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "txt_mod_linear_lora_A": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "txt_mod_linear_lora_B": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "for_all_img_layers": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
                "for_all_txt_layers": ("FLOAT", {"default": 0.000, "min": 0.000, "max": 1.000, "step": 0.001}),
            },
            "optional": {
                "model": ("MODEL", {"default": None}),
                "state_dictionary": ("WARPEDSTATEDICT", {"default": None}),
                "metadata_dict": ("WARPEDMETADICT", {"default": None}),
                "metadata_flush": (["tuner_only", "full", "none"], {"default": "tuner_only"}),
                "discard_single_blocks": ("BOOLEAN", {"default": True}),
                "verbose_messaging": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ("MODEL", "STRING", "WARPEDSTATEDICT", "WARPEDMETADICT",)
    RETURN_NAMES = ("model", "metadata", "state_dict", "metadata_dict",)
    FUNCTION = "load_lora"
    CATEGORY = "Warped/Hunyuan/Mixers"
    OUTPUT_NODE = False
    DESCRIPTION = "LoRA, single blocks double blocks"

    def load_lora(self, lora_model_1, lora_model_2, mainstrength, save_path, save_new_lora=False, return_state_only=False, model=None, state_dictionary=None, metadata_dict=None, metadata_flush="tuner_only", discard_single_blocks=True, verbose_messaging=False,
                img_attn_proj_lora_A=0.00, img_attn_proj_lora_B=0.00, img_attn_qkv_lora_A=0.00, img_attn_qkv_lora_B=0.00, img_mlp_fc1_lora_A=0.00,
                img_mlp_fc1_lora_B=0.00, img_mlp_fc2_lora_A=0.00, img_mlp_fc2_lora_B=0.00, img_mod_linear_lora_A=0.00, img_mod_linear_lora_B=0.00, txt_attn_proj_lora_A=0.00,
                txt_attn_proj_lora_B=0.00, txt_attn_qkv_lora_A=0.00, txt_attn_qkv_lora_B=0.00, txt_mlp_fc1_lora_A=0.00, txt_mlp_fc1_lora_B=0.00, txt_mlp_fc2_lora_A=0.00,
                txt_mlp_fc2_lora_B=0.00, txt_mod_linear_lora_A=0.00, txt_mod_linear_lora_B=0.00, for_all_img_layers=0.00, for_all_txt_layers=0.00):

        if lora_model_1 is None or lora_model_2 is None:
            ValueError("Bother LORA models must be valid selections.")

        from comfy.utils import load_torch_file, save_torch_file
        from comfy.sd import load_lora_for_models
        from comfy.lora import load_lora

        use_for_all_img = False

        if Decimal(for_all_img_layers).compare(Decimal(0.0)) != 0:
            use_for_all_img = True

        if use_for_all_img:
            print("Using For All Img.")
        else:
            print("Not Using For All Img.")

        use_for_all_txt = False

        if Decimal(for_all_txt_layers).compare(Decimal(0.0)) != 0:
            use_for_all_txt = True

        if use_for_all_txt:
            print("Using For All Txt.")
        else:
            print("Not Using For All Txt.")

        if state_dictionary is None:
            lora_1, metadata = warped_load_lora_weights(lora_model_1, return_metadata=True)
            print("++++ Metadata: ++++\n{}".format(metadata))
        else:
            lora_1 = state_dictionary

            if not metadata_dict is None:
                metadata = metadata_dict
            else:
                metadata = {}

            # print("---- Metadata: ----\n{}".format(metadata))

        lora_2 = warped_load_lora_weights(lora_model_2)

        diffusers_lora_1 = convert_lora(lora_1, convert_to="diffusion_model")
        filtered_lora_1 = filter_lora_keys(diffusers_lora_1, "all")

        diffusers_lora_2 = convert_lora(lora_2, convert_to="diffusion_model")
        filtered_lora_2 = filter_lora_keys(diffusers_lora_2, "all")

        max_dimension = 0

        for key in lora_1:
            if lora_1[key].shape[0] < lora_1[key].shape[1]:
                temp_dimension = lora_1[key].shape[0]
            else:
                temp_dimension = lora_1[key].shape[1]

            if temp_dimension > max_dimension:
                max_dimension = temp_dimension

            break

        for key in lora_2:
            if lora_2[key].shape[0] < lora_2[key].shape[1]:
                temp_dimension = lora_2[key].shape[0]
            else:
                temp_dimension = lora_2[key].shape[1]

            if temp_dimension > max_dimension:
                max_dimension = temp_dimension

            break

        filtered_lora_1 = convert_lora_dimensions(max_dimension, filtered_lora_1)
        filtered_lora_2 = convert_lora_dimensions(max_dimension, filtered_lora_2)

        tune_data = {}
        sub_key = ""

        dummy_lora = filtered_lora_1.copy()

        for key in dummy_lora:
            if "single_blocks" in key:
                if discard_single_blocks:
                    del filtered_lora_1[key]

                continue

            do_mod = False

            if not use_for_all_img and "img_" in key:
                if "img_attn_proj.lora_A" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - img_attn_proj_lora_A
                        perc_2 = img_attn_proj_lora_A

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["img_attn_proj.lora_A"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "img_attn_proj.lora_B" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - img_attn_proj_lora_B
                        perc_2 = img_attn_proj_lora_B

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["img_attn_proj.lora_B"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "img_attn_qkv.lora_A" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - img_attn_qkv_lora_A
                        perc_2 = img_attn_qkv_lora_A

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["img_attn_qkv.lora_A"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "img_attn_qkv.lora_B" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - img_attn_qkv_lora_B
                        perc_2 = img_attn_qkv_lora_B

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["img_attn_qkv.lora_B"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "img_mlp.fc1.lora_A" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - img_mlp_fc1_lora_A
                        perc_2 = img_mlp_fc1_lora_A

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["img_mlp.fc1.lora_A"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "img_mlp.fc1.lora_B" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - img_mlp_fc1_lora_B
                        perc_2 = img_mlp_fc1_lora_B

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["img_mlp.fc1.lora_B"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "img_mlp.fc2.lora_A" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - img_mlp_fc2_lora_A
                        perc_2 = img_mlp_fc2_lora_A

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["img_mlp.fc2.lora_A"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "img_mlp.fc2.lora_B" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - img_mlp_fc2_lora_B
                        perc_2 = img_mlp_fc2_lora_B

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["img_mlp.fc2.lora_B"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "img_mod.linear.lora_A" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - img_mod_linear_lora_A
                        perc_2 = img_mod_linear_lora_A

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["img_mod.linear.lora_A"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "img_mod.linear.lora_B" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - img_mod_linear_lora_B
                        perc_2 = img_mod_linear_lora_B

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["img_mod.linear.lora_B"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
            elif "img_" in key:
                sub_key = key.replace(".weight", "")
                sub_key = sub_key.replace("diffusion_model.double_blocks.", "")

                temp_strings = sub_key.split('img_')
                sub_key = "img_{}".format(temp_strings[len(temp_strings) - 1])

                perc_1 = 1.0 - for_all_img_layers
                perc_2 = for_all_img_layers

                tune_data[sub_key] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                do_mod = True

            if not use_for_all_txt and "txt_" in key:
                if "txt_attn_proj.lora_A" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - txt_attn_proj_lora_A
                        perc_2 = txt_attn_proj_lora_A

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["txt_attn_proj.lora_A"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "txt_attn_proj.lora_B" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - txt_attn_proj_lora_B
                        perc_2 = txt_attn_proj_lora_B

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["txt_attn_proj.lora_B"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "txt_attn_qkv.lora_A" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - txt_attn_qkv_lora_A
                        perc_2 = txt_attn_qkv_lora_A

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["txt_attn_qkv.lora_A"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "txt_attn_qkv.lora_B" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - txt_attn_qkv_lora_B
                        perc_2 = txt_attn_qkv_lora_B

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["txt_attn_qkv.lora_B"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "txt_mlp.fc1.lora_A" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - txt_mlp_fc1_lora_A
                        perc_2 = txt_mlp_fc1_lora_A

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["txt_mlp.fc1.lora_A"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "txt_mlp.fc1.lora_B" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - txt_mlp_fc1_lora_B
                        perc_2 = txt_mlp_fc1_lora_B

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["txt_mlp.fc1.lora_B"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "txt_mlp.fc2.lora_A" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - txt_mlp_fc2_lora_A
                        perc_2 = txt_mlp_fc2_lora_A

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["txt_mlp.fc2.lora_A"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "txt_mlp.fc2.lora_B" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - txt_mlp_fc2_lora_B
                        perc_2 = txt_mlp_fc2_lora_B

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["txt_mlp.fc2.lora_B"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "txt_mod.linear.lora_A" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - txt_mod_linear_lora_A
                        perc_2 = txt_mod_linear_lora_A

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["txt_mod.linear.lora_A"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
                elif "txt_mod.linear.lora_B" in key:
                    if key in filtered_lora_2:
                        perc_1 = 1.0 - txt_mod_linear_lora_B
                        perc_2 = txt_mod_linear_lora_B

                        if Decimal(perc_1).compare(Decimal(1.0)) != 0:
                            tune_data["txt_mod.linear.lora_B"] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                            do_mod = True
            elif "txt_" in key:
                sub_key = key.replace(".weight", "")
                sub_key = sub_key.replace("diffusion_model.double_blocks.", "")

                temp_strings = sub_key.split('txt_')
                sub_key = "txt_{}".format(temp_strings[len(temp_strings) - 1])

                perc_1 = 1.0 - for_all_txt_layers
                perc_2 = for_all_txt_layers

                tune_data[sub_key] = "perc_1: {} | perc_2: {}".format(perc_1, perc_2)
                do_mod = True

            if (key in filtered_lora_1) and (key in filtered_lora_2) and do_mod:
                mod_1 = torch.mul(filtered_lora_1[key], perc_1)
                mod_2 = torch.mul(filtered_lora_2[key], perc_2)

                new_lora = torch.add(mod_1, mod_2)
            else:
                continue

        if (metadata_flush == "full") or (metadata is None):
            metadata = {}

        print("**** Metadata: ****\n{}".format(metadata))

        if ((not "tune_counter_1" in metadata) and (not "modified_loras" in metadata)) or (metadata_flush == "full"):
            metadata = {"modified_loras": "{} and {}".format(lora_model_1, lora_model_2)}
        else:
            metadata["modified_loras"] = "{}  |  {} and {}".format(metadata["modified_loras"], lora_model_1, lora_model_2)

        if len(tune_data.keys()) > 0:
            if (not "tune_counter" in metadata) or (not metadata_flush == "none"):
                metadata["tune_counter"] = "{}".format(1)
                metadata[f'{tune_data}_1'] = "{}".format(tune_data)
            else:
                tune_counter = int(metadata["tune_counter"]) + 1
                metadata["tune_counter"] = "{}".format(tune_counter)
                metadata[f'{tune_data}_{tune_counter}'] = "{}".format(tune_data)

        if verbose_messaging:
            print("\nLora Metadata: {}".format(metadata))

        if save_new_lora: # and (len(swap_blocks) > 0):
            os.makedirs(os.path.dirname(save_path), exist_ok=True)

            for key in filtered_lora_1:
                filtered_lora_1[key] = filtered_lora_1[key].to(dtype=torch.bfloat16)

            print("Saving Model To: {}...".format(save_path))
            save_torch_file(filtered_lora_1, save_path, metadata=metadata)
            print("Saving Model To: {}...Done.".format(save_path))

        if (not return_state_only) and (not model is None):
            new_model, _ = load_lora_for_models(model, None, filtered_lora_1, mainstrength, 0)
            if new_model is not None:
                return (new_model, metadata, None, None, )

        return (model, metadata, filtered_lora_1, metadata, )

    @classmethod
    def IS_CHANGED(s, lora_model_1, lora_model_2, mainstrength, save_path, save_new_lora=False, return_state_only=False, model=None, state_dictionary=None, metadata_dict=None, metadata_flush="tuner_only", discard_single_blocks=True, verbose_messaging=False,
                img_attn_proj_lora_A=0.00, img_attn_proj_lora_B=0.00, img_attn_qkv_lora_A=0.00, img_attn_qkv_lora_B=0.00, img_mlp_fc1_lora_A=0.00, img_mlp_fc1_lora_B=0.00, img_mlp_fc2_lora_A=0.00,
                img_mlp_fc2_lora_B=0.00, img_mod_linear_lora_A=0.00, img_mod_linear_lora_B=0.00, txt_attn_proj_lora_A=0.00, txt_attn_proj_lora_B=0.00, txt_attn_qkv_lora_A=0.00, txt_attn_qkv_lora_B=0.00, txt_mlp_fc1_lora_A=0.00, txt_mlp_fc1_lora_B=0.00,
                txt_mlp_fc2_lora_A=0.00, txt_mlp_fc2_lora_B=0.00, txt_mod_linear_lora_A=0.00, txt_mod_linear_lora_B=0.00, for_all_img_layers=0.00, for_all_txt_layers=0.00):
        return f"{lora_model_1}_{lora_model_2}_{mainstrength}"

class WarpedHunyuanLoraDoubleBlocksRemoveLinear:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "lora_model": (['None'] + folder_paths.get_filename_list("loras"),),
                "mainstrength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "display": "number"}),
                "save_path": ("STRING", {"default": get_default_output_path()}),
                "save_new_lora": ("BOOLEAN", {"default": False}),
                "return_state_only": ("BOOLEAN", {"default": False}),
                "max_dimension": ([32, 64, 128], {"default": 128}),
            },
            "optional": {
                "model": ("MODEL", {"default": None}),
                "state_dictionary": ("WARPEDSTATEDICT", {"default": None}),
                "metadata_dict": ("WARPEDMETADICT", {"default": None}),
                "metadata_flush": (["deletes_only", "full", "none"], {"default": "deletes_only"}),
                "discard_single_blocks": ("BOOLEAN", {"default": True}),
                "verbose_messaging": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ("MODEL", "STRING", "WARPEDSTATEDICT", "WARPEDMETADICT",)
    RETURN_NAMES = ("model", "metadata", "state_dict", "metadata_dict",)
    FUNCTION = "load_lora"
    CATEGORY = "Warped/Hunyuan/Mixers"
    OUTPUT_NODE = False
    DESCRIPTION = "LoRA, single blocks double blocks"

    def load_lora(self, lora_model, mainstrength, save_path, save_new_lora=False, return_state_only=False, max_dimension=128, model=None, state_dictionary=None, metadata_dict=None, metadata_flush="deletes_only",
                discard_single_blocks=True, verbose_messaging=False):

        if lora_model is None and state_dictionary is None:
            ValueError("Either lora_model or state_dictionary input must be valid selections")

        from comfy.utils import load_torch_file, save_torch_file
        from comfy.sd import load_lora_for_models
        from comfy.lora import load_lora

        if state_dictionary is None:
            lora, metadata = warped_load_lora_weights(lora_model, return_metadata=True)
        else:
            lora = state_dictionary

            if not metadata_dict is None:
                metadata = metadata_dict
            else:
                metadata = {}

        diffusers_lora = convert_lora(lora, convert_to="diffusion_model")
        filtered_lora = filter_lora_keys(diffusers_lora, "all")
        filtered_lora = convert_lora_dimensions(max_dimension, filtered_lora)

        tune_data = {}
        sub_key = ""

        dummy_lora = filtered_lora.copy()
        deleted_keys = {}

        for key in dummy_lora:
            if "single_blocks" in key:
                if discard_single_blocks:
                    temp_key = get_base_layer_type(key)

                    deleted_keys[temp_key] = "Deleted All"
                    del filtered_lora[key]
                elif "linear" in key:
                    temp_key = get_base_layer_type(key)

                    deleted_keys[temp_key] = "Deleted All"
                    del filtered_lora[key]

                continue

            if not "linear" in key:
                continue

            temp_key = get_base_layer_type(key)

            deleted_keys[temp_key] = "Deleted All"

            del filtered_lora[key]

        if metadata_flush == "full":
            metadata = {}

        if ((not "deleted_keys_1" in metadata) and (not "modified_loras" in metadata)) or (metadata_flush == "full"):
            metadata = {"modified_loras": "{}".format(lora_model)}
        else:
            metadata["modified_loras"] = "{}  |  {}".format(metadata["modified_loras"], lora_model)

        if len(deleted_keys.keys()) > 0:
            if (not "deleted_counter" in metadata) or (not metadata_flush == "none"):
                metadata["deleted_counter"] = "{}".format(1)
                metadata[f'deleted_data_1'] = "{}".format(deleted_keys)
            else:
                deleted_counter = int(metadata["deleted_counter"]) + 1
                metadata["deleted_counter"] = "{}".format(deleted_counter)
                metadata[f'deleted_data_{deleted_counter}'] = "{}".format(deleted_keys)

        if verbose_messaging:
            print("\nLora Metadata: {}".format(metadata))

        if save_new_lora: # and (len(swap_blocks) > 0):
            os.makedirs(os.path.dirname(save_path), exist_ok=True)

            for key in filtered_lora:
                filtered_lora[key] = filtered_lora[key].to(dtype=torch.bfloat16)

            print(metadata)

            print("Saving Model To: {}...".format(save_path))
            save_torch_file(filtered_lora, save_path, metadata=metadata)
            print("Saving Model To: {}...Done.".format(save_path))

        if (not return_state_only) and (not model is None):
            new_model, _ = load_lora_for_models(model, None, filtered_lora, mainstrength, 0)
            if new_model is not None:
                return (new_model, metadata, None, None, )

        return (model, metadata, filtered_lora, metadata, )

    @classmethod
    def IS_CHANGED(s, lora_model, mainstrength, save_path, save_new_lora=False, return_state_only=False, max_dimension=128, model=None, state_dictionary=None, metadata_dict=None, metadata_flush="tuner_only", discard_single_blocks=True, verbose_messaging=False):
        return f"{lora_model}_{mainstrength}"

# class WarpedHunyuanLoraDoubleBlocksModifySegment:
#     @classmethod
#     def INPUT_TYPES(s):
#         return {
#             "required": {
#                 "lora_model": (['None'] + folder_paths.get_filename_list("loras"),),
#                 "source_lora": (['None'] + folder_paths.get_filename_list("loras"),),
#                 "mainstrength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "display": "number"}),
#                 "save_path": ("STRING", {"default": get_default_output_path()}),
#                 "save_new_lora": ("BOOLEAN", {"default": False}),
#                 "return_state_only": ("BOOLEAN", {"default": False}),
#                 "segment_number": ("INT", {"default": 2, "min": 1, "max": 127, "step": 1}),
#                 "test_mode": (["zero_all", "perc_all", "by_block_num", "random_noise", "use_source"], {"default": "perc_all"}),
#                 "max_dimension": ([32, 64, 128], {"default": 128}),
#             },
#             "optional": {
#                 "model": ("MODEL", {"default": None}),
#                 "state_dictionary": ("WARPEDSTATEDICT", {"default": None}),
#                 "metadata_dict": ("WARPEDMETADICT", {"default": None}),
#                 "block_number": ([0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19], {"default": 0}),
#                 "percentage": ("FLOAT", {"default": 1.000, "min": 0.000, "max": 5.000, "step": 0.001}),
#                 "layer_type": (["all", "img", "txt"], {"default": "all"}),
#                 "discard_single_blocks": ("BOOLEAN", {"default": True}),
#                 "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
#                 "verbose_messaging": ("BOOLEAN", {"default": False}),
#             }
#         }
#
#     RETURN_TYPES = ("MODEL", "STRING", "WARPEDSTATEDICT", "WARPEDMETADICT",)
#     RETURN_NAMES = ("model", "metadata", "state_dict", "metadata_dict",)
#     FUNCTION = "load_lora"
#     CATEGORY = "Warped/Hunyuan/Mixers/Experimental"
#     OUTPUT_NODE = False
#     DESCRIPTION = "LoRA, single blocks double blocks"
#
#     def load_lora(self, lora_model, source_lora, mainstrength, save_path, save_new_lora=False, return_state_only=False, segment_number=0, test_mode="perc_all", max_dimension=128, model=None, state_dictionary=None, metadata_dict=None,
#                 block_number=0, percentage=1.000, layer_type="all", discard_single_blocks=True, seed=0, verbose_messaging=False):
#
#         if lora_model is None and state_dictionary is None:
#             raise ValueError("Either lora_model or state_dictionary input must be valid selections")
#
#         if segment_number > (max_dimension - 1):
#             raise ValueError("segment_number cannot be greater than max_dimension - 1.")
#
#         from comfy.utils import save_torch_file
#         from comfy.sd import load_lora_for_models
#         from comfy.lora import load_lora
#
#         if state_dictionary is None:
#
#             lora, metadata = warped_load_lora_weights(lora_model, return_metadata=True)
#         else:
#             lora = state_dictionary
#
#             if not metadata_dict is None:
#                 metadata = metadata_dict
#             else:
#                 metadata = {}
#
#         if not "modified_loras" in metadata:
#             metadata["modified_loras"] = "{} and {}".format(lora_model, source_lora)
#         else:
#             metadata["modified_loras"] = "{}  |  {} and {}".format(metadata["modified_loras"], lora_model, source_lora)
#
#         diffusers_lora = convert_lora(lora, convert_to="diffusion_model")
#
#         if discard_single_blocks:
#             filtered_lora = filter_lora_keys(diffusers_lora, "double_blocks")
#         else:
#             filtered_lora = filter_lora_keys(diffusers_lora, "all")
#
#         filtered_lora = convert_lora_dimensions(max_dimension, filtered_lora)
#
#         source_filtered_lora = None
#
#         if (test_mode == "use_source") and (not source_lora == "None"):
#             source_lora_path = folder_paths.get_full_path("loras", source_lora)
#
#             if not os.path.exists(source_lora_path):
#                 raise Exception(f"Lora {source_lora} not found at {source_lora_path}")
#
#             temp_source_lora = warped_load_lora_weights(source_lora)
#
#             source_diffusers_lora = convert_lora(temp_source_lora, convert_to="diffusion_model")
#             source_filtered_lora = filter_lora_keys(source_diffusers_lora, "double_blocks")
#             source_filtered_lora = convert_lora_dimensions(max_dimension, source_filtered_lora)
#
#         if segment_number > max_dimension:
#             segment_number = max_dimension
#         elif segment_number < 0:
#             segment_number = 0
#
#         block_filter = ""
#
#         if test_mode == "by_block_num":
#             block_filter = ".{}.".format(block_number)
#
#         if not "slice_manipulation_counter" in metadata:
#             metadata["slice_manipulation_counter"] = "1"
#             metadata_key = "slice_manipulation_1"
#         else:
#             metadata["slice_manipulation_counter"] = "{}".format(int(metadata["slice_manipulation_counter"]) + 1)
#             metadata_key = f'slice_manipulation_{"{}".format(metadata["slice_manipulation_counter"])}'
#
#         metadata[metadata_key] = {"seed": "{}".format(seed)}
#         metadata[metadata_key]["lora_model"] = lora_model
#         metadata[metadata_key]["test_mode"] = test_mode
#         metadata[metadata_key]["segment_number"] = "{}".format(segment_number)
#         metadata[metadata_key]["max_dimension"] = "{}".format(max_dimension)
#         metadata[metadata_key]["block_number"] = "{}".format(block_number)
#         metadata[metadata_key]["layer_type"] = "{}".format(layer_type)
#         metadata[metadata_key]["percentage"] = "{}".format(percentage)
#         metadata[metadata_key]["discard_single_blocks"] = "{}".format(discard_single_blocks)
#
#         if test_mode == "use_source":
#             metadata[metadata_key]["source_model"] = source_lora
#         else:
#             metadata[metadata_key]["source_model"] = None
#
#         metadata[metadata_key] = "{}".format(metadata[metadata_key])
#
#         for key in filtered_lora:
#             if "single_blocks" in key:
#                 continue
#
#             if ("layer_type" == "img") and ("txt_" in key):
#                 continue
#
#             if ("layer_type" == "txt") and ("img_" in key):
#                 continue
#
#             if filtered_lora[key].shape[0] < filtered_lora[key].shape[1]:
#                 use_length = filtered_lora[key].shape[0]
#                 test_length = 1
#             else:
#                 use_length = int(int(filtered_lora[key].shape[0]) // int(filtered_lora[key].shape[1]))
#                 test_length = int(filtered_lora[key].shape[0])
#
#             temp_tensor = torch.zeros_like(filtered_lora[key])
#
#             if (test_mode == "zero_all") or ((len(block_filter) > 0) and (block_filter in key)):
#                 if use_length == test_length:
#                     if (not segment_number == 0) and (not segment_number == max_dimension):
#                         temp_tensor[:segment_number,:] = filtered_lora[key][:segment_number,:]
#                         temp_tensor[segment_number + 1:,:] = filtered_lora[key][segment_number + 1:,:]
#                     elif segment_number == 0:
#                         temp_tensor[segment_number + 1:,:] = filtered_lora[key][segment_number + 1:,:]
#                     else:
#                         temp_tensor[:segment_number,:] = filtered_lora[key][:segment_number,:]
#                 else:
#                     if (not segment_number == 0) and (not segment_number == max_dimension):
#                         temp_tensor[:segment_number * test_length,:] = filtered_lora[key][:segment_number * test_length,:]
#                         temp_tensor[segment_number + test_length:,:] = filtered_lora[key][segment_number + test_length:,:]
#                     elif segment_number == 0:
#                         temp_tensor[segment_number + test_length:,:] = filtered_lora[key][segment_number + test_length:,:]
#                     else:
#                         temp_tensor[:segment_number * test_length,:] = filtered_lora[key][:segment_number * test_length,:]
#             elif test_mode == "perc_all":
#                 temp_perc_tensor = torch.zeros_like(filtered_lora[key])
#
#                 if use_length == test_length:
#                     if (not segment_number == 0) and (not segment_number == max_dimension):
#                         temp_tensor[:segment_number,:] = filtered_lora[key][:segment_number,:]
#                         temp_tensor[segment_number + 1:,:] = filtered_lora[key][segment_number + 1:,:]
#
#                         temp_perc_tensor[segment_number - 1:segment_number,:] = filtered_lora[key][segment_number - 1:segment_number,:] * percentage
#                     elif segment_number == 0:
#                         temp_tensor[segment_number + 1:,:] = filtered_lora[key][segment_number + 1:,:]
#
#                         temp_perc_tensor[:1,:] = filtered_lora[key][:1,:] * percentage
#                     else:
#                         temp_tensor[:segment_number,:] = filtered_lora[key][:segment_number,:]
#
#                         temp_perc_tensor[segment_number:segment_number + 1,:] = filtered_lora[key][segment_number:segment_number + 1,:] * percentage
#                 else:
#                     if (not segment_number == 0) and (not segment_number == max_dimension):
#                         temp_tensor[:segment_number * test_length,:] = filtered_lora[key][:segment_number * test_length,:]
#                         temp_tensor[segment_number + test_length:,:] = filtered_lora[key][segment_number + test_length:,:]
#
#                         temp_perc_tensor[segment_number - test_length:segment_number,:] = filtered_lora[key][segment_number - test_length:segment_number,:] * percentage
#                     elif segment_number == 0:
#                         temp_tensor[segment_number + test_length:,:] = filtered_lora[key][segment_number + test_length:,:]
#
#                         temp_perc_tensor[:test_length,:] = filtered_lora[key][:test_length,:] * percentage
#                     else:
#                         temp_tensor[:segment_number * test_length,:] = filtered_lora[key][:segment_number * test_length,:]
#
#                         temp_perc_tensor[(segment_number * test_length) - test_length:(segment_number * test_length) + test_length,:] = filtered_lora[key][(segment_number * test_length) - test_length:(segment_number * test_length) + test_length,:] * percentage
#             elif test_mode == "random_noise":
#                 random_noise_latent = warped_prepare_noise(torch.zeros_like(filtered_lora[key]), seed)
#                 temp_perc_tensor = torch.zeros_like(filtered_lora[key])
#
#                 if use_length == test_length:
#                     if (not segment_number == 0) and (not segment_number == max_dimension):
#                         temp_tensor[:segment_number,:] = filtered_lora[key][:segment_number,:]
#                         temp_tensor[segment_number + 1:,:] = filtered_lora[key][segment_number + 1:,:]
#
#                         temp_perc_tensor[segment_number - 1:segment_number,:] = random_noise_latent[segment_number - 1:segment_number,:] * percentage
#                     elif segment_number == 0:
#                         temp_tensor[segment_number + 1:,:] = filtered_lora[key][segment_number + 1:,:]
#
#                         temp_perc_tensor[:1,:] = random_noise_latent[:1,:] * percentage
#                     else:
#                         temp_tensor[:segment_number,:] = filtered_lora[key][:segment_number,:]
#
#                         temp_perc_tensor[segment_number:segment_number + 1,:] = random_noise_latent[segment_number:segment_number + 1,:] * percentage
#                 else:
#                     if (not segment_number == 0) and (not segment_number == max_dimension):
#                         temp_tensor[:segment_number * test_length,:] = filtered_lora[key][:segment_number * test_length,:]
#                         temp_tensor[segment_number + test_length:,:] = filtered_lora[key][segment_number + test_length:,:]
#
#                         temp_perc_tensor[segment_number - test_length:segment_number,:] = random_noise_latent[segment_number - test_length:segment_number,:] * percentage
#                     elif segment_number == 0:
#                         temp_tensor[segment_number + test_length:,:] = filtered_lora[key][segment_number + test_length:,:]
#
#                         temp_perc_tensor[:test_length,:] = random_noise_latent[:test_length,:] * percentage
#                     else:
#                         temp_tensor[:segment_number * test_length,:] = filtered_lora[key][:segment_number * test_length,:]
#
#                         temp_perc_tensor[(segment_number * test_length) - test_length:(segment_number * test_length) + test_length,:] = random_noise_latent[(segment_number * test_length) - test_length:(segment_number * test_length) + test_length,:] * percentage
#
#                 temp_tensor = torch.add(temp_tensor, temp_perc_tensor)
#             elif test_mode == "use_source":
#                 if not source_filtered_lora is None:
#                     if not key in source_filtered_lora:
#                         continue
#
#                     temp_perc_tensor = torch.zeros_like(filtered_lora[key])
#
#                     if use_length == test_length:
#                         if (not segment_number == 0) and (not segment_number == max_dimension):
#                             temp_tensor[:segment_number,:] = filtered_lora[key][:segment_number,:]
#                             temp_tensor[segment_number + 1:,:] = filtered_lora[key][segment_number + 1:,:]
#                             temp_perc_tensor[segment_number - 1:segment_number,:] = source_filtered_lora[key][segment_number - 1:segment_number,:] * percentage
#                         elif segment_number == 0:
#                             temp_tensor[segment_number + 1:,:] = filtered_lora[key][segment_number + 1:,:]
#                             temp_perc_tensor[:1,:] = source_filtered_lora[key][:1,:] * percentage
#                         else:
#                             temp_tensor[:segment_number,:] = filtered_lora[key][:segment_number,:]
#                             temp_perc_tensor[segment_number:segment_number + 1,:] = source_filtered_lora[key][segment_number:segment_number + 1,:] * percentage
#                     else:
#                         if (not segment_number == 0) and (not segment_number == max_dimension):
#                             temp_tensor[:segment_number * test_length,:] = filtered_lora[key][:segment_number * test_length,:]
#                             temp_tensor[segment_number + test_length:,:] = filtered_lora[key][segment_number + test_length:,:]
#                             temp_perc_tensor[segment_number - test_length:segment_number,:] = source_filtered_lora[key][segment_number - test_length:segment_number,:] * percentage
#                         elif segment_number == 0:
#                             temp_tensor[segment_number + test_length:,:] = filtered_lora[key][segment_number + test_length:,:]
#                             temp_perc_tensor[:test_length,:] = source_filtered_lora[key][:test_length,:] * percentage
#                         else:
#                             temp_tensor[:segment_number * test_length,:] = filtered_lora[key][:segment_number * test_length,:]
#                             temp_perc_tensor[(segment_number * test_length) - test_length:(segment_number * test_length) + test_length,:] = source_filtered_lora[key][(segment_number * test_length) - test_length:(segment_number * test_length) + test_length,:] * percentage
#
#                     temp_tensor = torch.add(temp_tensor, temp_perc_tensor)
#                 else:
#                     print("**** No Source LORA Provided. Unable to make mosification. ****")
#
#             filtered_lora[key] = temp_tensor.to(torch.bfloat16)
#
#         if verbose_messaging:
#             print("Tester Metadata: {}".format(metadata))
#
#         if save_new_lora: # and (len(swap_blocks) > 0):
#             os.makedirs(os.path.dirname(save_path), exist_ok=True)
#
#             for key in filtered_lora:
#                 filtered_lora[key] = filtered_lora[key].to(dtype=torch.bfloat16)
#
#             print("Saving Model To: {}...".format(save_path))
#             save_torch_file(filtered_lora, save_path, metadata=metadata)
#             print("Saving Model To: {}...Done.".format(save_path))
#
#         if (not return_state_only) and (not model is None):
#             new_model, _ = load_lora_for_models(model, None, filtered_lora, mainstrength, 0)
#             if new_model is not None:
#                 return (new_model, metadata, None, None, )
#
#         return (model, metadata, filtered_lora, metadata, )
#
#
#     @classmethod
#     def IS_CHANGED(s, lora_model, source_lora, mainstrength, save_path, save_new_lora=False, return_state_only=False, segment_number=0, test_mode="perc_all", max_dimension=128, model=None, state_dictionary=None, metadata_dict=None, block_number=0, percentage=1.000, layer_type="all", discard_single_blocks=True, seed=0, verbose_messaging=False):
#         return f"{lora_model}_{mainstrength}"

class WarpedHunyuanLoraDoubleBlocksModifyMultipleSegments:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "lora_model": (['None'] + folder_paths.get_filename_list("loras"),),
                "source_lora": (['None'] + folder_paths.get_filename_list("loras"),),
                "mainstrength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "display": "number"}),
                "save_path": ("STRING", {"default": get_default_output_path()}),
                "save_new_lora": ("BOOLEAN", {"default": False}),
                "return_state_only": ("BOOLEAN", {"default": False}),
                "segment_numbers": ("STRING", {"default": "2,23"}),
                "test_mode": (["zero_all", "perc_all", "by_block_num", "random_noise", "use_source", "add_source", "subtract_source"], {"default": "perc_all"}),
                "max_dimension": ([32, 64, 128], {"default": 128}),
            },
            "optional": {
                "model": ("MODEL", {"default": None}),
                "state_dictionary": ("WARPEDSTATEDICT", {"default": None}),
                "metadata_dict": ("WARPEDMETADICT", {"default": None}),
                "block_number": ([0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19], {"default": 0}),
                "percentage": ("FLOAT", {"default": 1.000, "min": 0.000, "max": 5.000, "step": 0.001}),
                "layer_type": (["all", "img", "txt"], {"default": "all"}),
                "discard_single_blocks": ("BOOLEAN", {"default": True}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                "verbose_messaging": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ("MODEL", "STRING", "WARPEDSTATEDICT", "WARPEDMETADICT",)
    RETURN_NAMES = ("model", "metadata", "state_dict", "metadata_dict",)
    FUNCTION = "load_lora"
    CATEGORY = "Warped/Hunyuan/Mixers/Experimental"
    OUTPUT_NODE = False
    DESCRIPTION = "LoRA, single blocks double blocks"

    def load_lora(self, lora_model, source_lora, mainstrength, save_path, save_new_lora=False, return_state_only=False, segment_numbers="2,23", test_mode="perc_all", max_dimension=128, model=None, state_dictionary=None, metadata_dict=None,
                block_number=0, percentage=1.000, layer_type="all", discard_single_blocks=True, seed=0, verbose_messaging=False):

        if lora_model is None and state_dictionary is None:
            raise ValueError("Either lora_model or state_dictionary input must be valid selections")

        if len(segment_numbers) < 1:
            raise ValueError("segment_numbers cannot be empty.")

        temp_seg_split = segment_numbers.split(',')

        use_segment_numbers = []

        for entry in temp_seg_split:
            # if range of segments
            if "-" in entry.strip():
                temp_seg_split_2 = entry.split("-")

                start = int(temp_seg_split_2[0].strip())
                end = int(temp_seg_split_2[1].strip())

                if start > end:
                    raise ValueError("Beginning of segment range cannot be greater than the end of the segment range.")

                for i in range(start, end):
                    use_segment_numbers.append(i)

                continue

            use_segment_numbers.append(int(entry.strip()))

        for segment_number in use_segment_numbers:
            if segment_number > (max_dimension - 1):
                raise ValueError("segment_number cannot be greater than max_dimension - 1.")

            if segment_number < 1:
                raise ValueError("segment_number cannot be less than 1.")

        from comfy.utils import save_torch_file
        from comfy.sd import load_lora_for_models
        from comfy.lora import load_lora

        if state_dictionary is None:

            lora, metadata = warped_load_lora_weights(lora_model, return_metadata=True)
        else:
            lora = state_dictionary

            if not metadata_dict is None:
                metadata = metadata_dict
            else:
                metadata = {}

        if not "modified_loras" in metadata:
            metadata["modified_loras"] = "{} and {}".format(lora_model, source_lora)
        else:
            metadata["modified_loras"] = "{}  |  {} and {}".format(metadata["modified_loras"], lora_model, source_lora)

        diffusers_lora = convert_lora(lora, convert_to="diffusion_model")

        if discard_single_blocks:
            filtered_lora = filter_lora_keys(diffusers_lora, "double_blocks")
        else:
            filtered_lora = filter_lora_keys(diffusers_lora, "all")

        filtered_lora = convert_lora_dimensions(max_dimension, filtered_lora)

        source_filtered_lora = None

        if ((test_mode == "use_source") or (test_mode == "add_source") or (test_mode == "subtract_source")) and (not source_lora == "None"):
            source_lora_path = folder_paths.get_full_path("loras", source_lora)

            if not os.path.exists(source_lora_path):
                raise Exception(f"Lora {source_lora} not found at {source_lora_path}")

            temp_source_lora = warped_load_lora_weights(source_lora)

            source_diffusers_lora = convert_lora(temp_source_lora, convert_to="diffusion_model")
            source_filtered_lora = filter_lora_keys(source_diffusers_lora, "double_blocks")
            source_filtered_lora = convert_lora_dimensions(max_dimension, source_filtered_lora)

        block_filter = ""

        if test_mode == "by_block_num":
            block_filter = ".{}.".format(block_number)

        if not "slice_manipulation_counter" in metadata:
            metadata["slice_manipulation_counter"] = "1"
            metadata_key = "slice_manipulation_1"
        else:
            metadata["slice_manipulation_counter"] = "{}".format(int(metadata["slice_manipulation_counter"]) + 1)
            metadata_key = f'slice_manipulation_{"{}".format(metadata["slice_manipulation_counter"])}'

        metadata[metadata_key] = {"seed": "{}".format(seed)}
        metadata[metadata_key]["lora_model"] = lora_model
        metadata[metadata_key]["test_mode"] = test_mode
        metadata[metadata_key]["segment_numbers"] = "{}".format(segment_numbers)
        metadata[metadata_key]["max_dimension"] = "{}".format(max_dimension)
        metadata[metadata_key]["block_number"] = "{}".format(block_number)
        metadata[metadata_key]["layer_type"] = "{}".format(layer_type)
        metadata[metadata_key]["percentage"] = "{}".format(percentage)
        metadata[metadata_key]["discard_single_blocks"] = "{}".format(discard_single_blocks)

        if (test_mode == "use_source") or (test_mode == "add_source") or (test_mode == "subtract_source"):
            metadata[metadata_key]["source_model"] = source_lora
        else:
            metadata[metadata_key]["source_model"] = None

        metadata[metadata_key] = "{}".format(metadata[metadata_key])

        for key in filtered_lora:
            if "single_blocks" in key:
                continue

            if ("layer_type" == "img") and ("txt_" in key):
                continue

            if ("layer_type" == "txt") and ("img_" in key):
                continue

            if filtered_lora[key].shape[0] < filtered_lora[key].shape[1]:
                use_length = filtered_lora[key].shape[0]
                test_length = 1
            else:
                use_length = int(int(filtered_lora[key].shape[0]) // int(filtered_lora[key].shape[1]))
                test_length = int(filtered_lora[key].shape[0])

            temp_tensor = torch.zeros_like(filtered_lora[key])

            if (test_mode == "zero_all") or ((len(block_filter) > 0) and (block_filter in key)):
                for segment_number in use_segment_numbers:
                    if use_length == test_length:
                        if (not segment_number == 0) and (not segment_number == max_dimension):
                            temp_tensor[:segment_number,:] = filtered_lora[key][:segment_number,:]
                            temp_tensor[segment_number + 1:,:] = filtered_lora[key][segment_number + 1:,:]
                        elif segment_number == 0:
                            temp_tensor[segment_number + 1:,:] = filtered_lora[key][segment_number + 1:,:]
                        else:
                            temp_tensor[:segment_number,:] = filtered_lora[key][:segment_number,:]
                    else:
                        if (not segment_number == 0) and (not segment_number == max_dimension):
                            temp_tensor[:segment_number * test_length,:] = filtered_lora[key][:segment_number * test_length,:]
                            temp_tensor[segment_number + test_length:,:] = filtered_lora[key][segment_number + test_length:,:]
                        elif segment_number == 0:
                            temp_tensor[segment_number + test_length:,:] = filtered_lora[key][segment_number + test_length:,:]
                        else:
                            temp_tensor[:segment_number * test_length,:] = filtered_lora[key][:segment_number * test_length,:]

                    filtered_lora[key] = temp_tensor.to(torch.bfloat16)
            elif test_mode == "perc_all":
                temp_perc_tensor = torch.zeros_like(filtered_lora[key])

                for segment_number in use_segment_numbers:
                    if use_length == test_length:
                        if (not segment_number == 0) and (not segment_number == max_dimension):
                            temp_tensor[:segment_number,:] = filtered_lora[key][:segment_number,:]
                            temp_tensor[segment_number + 1:,:] = filtered_lora[key][segment_number + 1:,:]

                            temp_perc_tensor[segment_number - 1:segment_number,:] = filtered_lora[key][segment_number - 1:segment_number,:] * percentage
                        elif segment_number == 0:
                            temp_tensor[segment_number + 1:,:] = filtered_lora[key][segment_number + 1:,:]

                            temp_perc_tensor[:1,:] = filtered_lora[key][:1,:] * percentage
                        else:
                            temp_tensor[:segment_number,:] = filtered_lora[key][:segment_number,:]

                            temp_perc_tensor[segment_number:segment_number + 1,:] = filtered_lora[key][segment_number:segment_number + 1,:] * percentage
                    else:
                        if (not segment_number == 0) and (not segment_number == max_dimension):
                            temp_tensor[:segment_number * test_length,:] = filtered_lora[key][:segment_number * test_length,:]
                            temp_tensor[segment_number + test_length:,:] = filtered_lora[key][segment_number + test_length:,:]

                            temp_perc_tensor[segment_number - test_length:segment_number,:] = filtered_lora[key][segment_number - test_length:segment_number,:] * percentage
                        elif segment_number == 0:
                            temp_tensor[segment_number + test_length:,:] = filtered_lora[key][segment_number + test_length:,:]

                            temp_perc_tensor[:test_length,:] = filtered_lora[key][:test_length,:] * percentage
                        else:
                            temp_tensor[:segment_number * test_length,:] = filtered_lora[key][:segment_number * test_length,:]

                            temp_perc_tensor[(segment_number * test_length) - test_length:(segment_number * test_length) + test_length,:] = filtered_lora[key][(segment_number * test_length) - test_length:(segment_number * test_length) + test_length,:] * percentage

                    filtered_lora[key] = temp_tensor.to(torch.bfloat16)
            elif test_mode == "random_noise":
                rnd = torch.manual_seed(seed)

                random_noise_latent = warped_prepare_noise(torch.zeros_like(filtered_lora[key]), seed, generator=rnd)
                temp_perc_tensor = torch.zeros_like(filtered_lora[key])

                for segment_number in use_segment_numbers:
                    if use_length == test_length:
                        if (not segment_number == 0) and (not segment_number == max_dimension):
                            temp_tensor[:segment_number,:] = filtered_lora[key][:segment_number,:]
                            temp_tensor[segment_number + 1:,:] = filtered_lora[key][segment_number + 1:,:]

                            temp_perc_tensor[segment_number - 1:segment_number,:] = random_noise_latent[segment_number - 1:segment_number,:] * percentage
                        elif segment_number == 0:
                            temp_tensor[segment_number + 1:,:] = filtered_lora[key][segment_number + 1:,:]

                            temp_perc_tensor[:1,:] = random_noise_latent[:1,:] * percentage
                        else:
                            temp_tensor[:segment_number,:] = filtered_lora[key][:segment_number,:]

                            temp_perc_tensor[segment_number:segment_number + 1,:] = random_noise_latent[segment_number:segment_number + 1,:] * percentage
                    else:
                        if (not segment_number == 0) and (not segment_number == max_dimension):
                            temp_tensor[:segment_number * test_length,:] = filtered_lora[key][:segment_number * test_length,:]
                            temp_tensor[segment_number + test_length:,:] = filtered_lora[key][segment_number + test_length:,:]

                            temp_perc_tensor[segment_number - test_length:segment_number,:] = random_noise_latent[segment_number - test_length:segment_number,:] * percentage
                        elif segment_number == 0:
                            temp_tensor[segment_number + test_length:,:] = filtered_lora[key][segment_number + test_length:,:]

                            temp_perc_tensor[:test_length,:] = random_noise_latent[:test_length,:] * percentage
                        else:
                            temp_tensor[:segment_number * test_length,:] = filtered_lora[key][:segment_number * test_length,:]

                            temp_perc_tensor[(segment_number * test_length) - test_length:(segment_number * test_length) + test_length,:] = random_noise_latent[(segment_number * test_length) - test_length:(segment_number * test_length) + test_length,:] * percentage

                    temp_tensor = torch.add(temp_tensor, temp_perc_tensor)
                    filtered_lora[key] = temp_tensor.to(torch.bfloat16)
            elif test_mode == "use_source":
                if not source_filtered_lora is None:
                    if not key in source_filtered_lora:
                        continue

                    for segment_number in use_segment_numbers:
                        temp_perc_tensor = torch.zeros_like(filtered_lora[key])

                        if use_length == test_length:
                            if (not segment_number == 0) and (not segment_number == max_dimension):
                                temp_tensor[:segment_number,:] = filtered_lora[key][:segment_number,:]
                                temp_tensor[segment_number + 1:,:] = filtered_lora[key][segment_number + 1:,:]
                                temp_perc_tensor[segment_number - 1:segment_number,:] = source_filtered_lora[key][segment_number - 1:segment_number,:] * percentage
                            elif segment_number == 0:
                                temp_tensor[segment_number + 1:,:] = filtered_lora[key][segment_number + 1:,:]
                                temp_perc_tensor[:1,:] = source_filtered_lora[key][:1,:] * percentage
                            else:
                                temp_tensor[:segment_number,:] = filtered_lora[key][:segment_number,:]
                                temp_perc_tensor[segment_number:segment_number + 1,:] = source_filtered_lora[key][segment_number:segment_number + 1,:] * percentage
                        else:
                            if (not segment_number == 0) and (not segment_number == max_dimension):
                                temp_tensor[:segment_number * test_length,:] = filtered_lora[key][:segment_number * test_length,:]
                                temp_tensor[segment_number + test_length:,:] = filtered_lora[key][segment_number + test_length:,:]
                                temp_perc_tensor[segment_number - test_length:segment_number,:] = source_filtered_lora[key][segment_number - test_length:segment_number,:] * percentage
                            elif segment_number == 0:
                                temp_tensor[segment_number + test_length:,:] = filtered_lora[key][segment_number + test_length:,:]
                                temp_perc_tensor[:test_length,:] = source_filtered_lora[key][:test_length,:] * percentage
                            else:
                                temp_tensor[:segment_number * test_length,:] = filtered_lora[key][:segment_number * test_length,:]
                                temp_perc_tensor[(segment_number * test_length) - test_length:(segment_number * test_length) + test_length,:] = source_filtered_lora[key][(segment_number * test_length) - test_length:(segment_number * test_length) + test_length,:] * percentage

                        temp_tensor = torch.add(temp_tensor, temp_perc_tensor)
                        filtered_lora[key] = temp_tensor.to(torch.bfloat16)
            elif test_mode == "add_source":
                if not source_filtered_lora is None:
                    if not key in source_filtered_lora:
                        continue

                    for segment_number in use_segment_numbers:
                        temp_perc_tensor = torch.zeros_like(filtered_lora[key])

                        if use_length == test_length:
                            if (not segment_number == 0) and (not segment_number == max_dimension):
                                temp_perc_tensor[segment_number - 1:segment_number,:] = source_filtered_lora[key][segment_number - 1:segment_number,:] * percentage
                            elif segment_number == 0:
                                temp_perc_tensor[:1,:] = source_filtered_lora[key][:1,:] * percentage
                            else:
                                temp_perc_tensor[segment_number:segment_number + 1,:] = source_filtered_lora[key][segment_number:segment_number + 1,:] * percentage
                        else:
                            if (not segment_number == 0) and (not segment_number == max_dimension):
                                temp_perc_tensor[segment_number - test_length:segment_number,:] = source_filtered_lora[key][segment_number - test_length:segment_number,:] * percentage
                            elif segment_number == 0:
                                temp_perc_tensor[:test_length,:] = source_filtered_lora[key][:test_length,:] * percentage
                            else:
                                temp_perc_tensor[(segment_number * test_length) - test_length:(segment_number * test_length) + test_length,:] = source_filtered_lora[key][(segment_number * test_length) - test_length:(segment_number * test_length) + test_length,:] * percentage

                        temp_tensor = torch.add(filtered_lora[key], temp_perc_tensor)
                        filtered_lora[key] = temp_tensor.to(torch.bfloat16)
            elif test_mode == "subtract_source":
                if not source_filtered_lora is None:
                    if not key in source_filtered_lora:
                        continue

                    for segment_number in use_segment_numbers:
                        temp_perc_tensor = torch.zeros_like(filtered_lora[key])

                        if use_length == test_length:
                            if (not segment_number == 0) and (not segment_number == max_dimension):
                                temp_perc_tensor[segment_number - 1:segment_number,:] = source_filtered_lora[key][segment_number - 1:segment_number,:] * percentage
                            elif segment_number == 0:
                                temp_perc_tensor[:1,:] = source_filtered_lora[key][:1,:] * percentage
                            else:
                                temp_perc_tensor[segment_number:segment_number + 1,:] = source_filtered_lora[key][segment_number:segment_number + 1,:] * percentage
                        else:
                            if (not segment_number == 0) and (not segment_number == max_dimension):
                                temp_perc_tensor[segment_number - test_length:segment_number,:] = source_filtered_lora[key][segment_number - test_length:segment_number,:] * percentage
                            elif segment_number == 0:
                                temp_perc_tensor[:test_length,:] = source_filtered_lora[key][:test_length,:] * percentage
                            else:
                                temp_perc_tensor[(segment_number * test_length) - test_length:(segment_number * test_length) + test_length,:] = source_filtered_lora[key][(segment_number * test_length) - test_length:(segment_number * test_length) + test_length,:] * percentage

                        temp_tensor = torch.sub(filtered_lora[key], temp_perc_tensor)
                        filtered_lora[key] = temp_tensor.to(torch.bfloat16)
                else:
                    print("**** No Source LORA Provided. Unable to make mosification. ****")

        if verbose_messaging:
            print("Tester Metadata: {}".format(metadata))

        if save_new_lora: # and (len(swap_blocks) > 0):
            os.makedirs(os.path.dirname(save_path), exist_ok=True)

            for key in filtered_lora:
                filtered_lora[key] = filtered_lora[key].to(dtype=torch.bfloat16)

            print("Saving Model To: {}...".format(save_path))
            save_torch_file(filtered_lora, save_path, metadata=metadata)
            print("Saving Model To: {}...Done.".format(save_path))

        if (not return_state_only) and (not model is None):
            new_model, _ = load_lora_for_models(model, None, filtered_lora, mainstrength, 0)
            if new_model is not None:
                return (new_model, metadata, None, None, )

        return (model, metadata, filtered_lora, metadata, )


    @classmethod
    def IS_CHANGED(s, lora_model, source_lora, mainstrength, save_path, save_new_lora=False, return_state_only=False, segment_numbers="2,23", test_mode="perc_all", max_dimension=128, model=None, state_dictionary=None, metadata_dict=None, block_number=0, percentage=1.000, layer_type="all", discard_single_blocks=True, seed=0, verbose_messaging=False):
        return f"{lora_model}_{mainstrength}"

class WarpedLoadHunyuanLoraWeightsByPrefix:
    def __init__(self):
        self.index = 0
        self.lora_dir = ""
        self.last_prefix = ""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "lora_dir": (get_lora_directories(), ),
                "lora_prefix": ("STRING", {"default": '', "multiline": False}),
            },
        }

    RETURN_TYPES = ("WARPEDSTATEDICT", "WARPEDMETADICT", )
    RETURN_NAMES = ("state_dict", "metadata_dict", )
    FUNCTION = "load_batch_loras"

    CATEGORY = "Warped/Hunyuan/Lora"

    def load_batch_loras(self, lora_dir, lora_prefix):
        self.lora_dir = lora_dir
        path = lora_dir
        print(path)

        if not os.path.exists(path):
            return ("", "", )

        if not (self.last_prefix == lora_prefix):
            self.last_prefix = lora_prefix
            self.index = 0

        retry = False
        index = 0

        try:
            filename, full_filename = self.do_the_load(path, lora_prefix, index)
            print("WarpedLoadHunyuanLoraWeightsByPrefix: Filename: {}  |  Full File Path: {}".format(filename, full_filename))

            lora, metadata = warped_load_lora_weights(filename, return_metadata=True)

            return (lora, metadata,)
        except:
            self.index = 0
            retry = True

        if retry:
            retry = False
            filename, full_filename = self.do_the_load(path, lora_prefix, index)
            print("WarpedLoadHunyuanLoraWeightsByPrefix: Retrying: Filename: {}  |  Full File Path: {}".format(filename, full_filename))

            lora, metadata = warped_load_lora_weights(filename, return_metadata=True)

            return (lora, metadata,)

        return (None, None, )

    def do_the_load(self, path, prefix, index):
        prefix = prefix.strip(' ')

        if (len(prefix) == 1) and (prefix == '*'):
            fl = self.BatchLoraLoader(path, '*', index)
        else:
            prefix = prefix.strip('*')
            fl = self.BatchLoraLoader(path, "{}*".format(prefix), index)

        new_paths = fl.lora_paths

        filename = fl.lora_paths[self.index]

        # filename = os.path.join(self.sub_folder, filename)
        full_filename = os.path.join(path, filename)
        base_dir, lora_name = get_lora_path_parts(full_filename)

        self.index += 1

        if self.index >= len(fl.lora_paths):
            self.index = 0

        return lora_name, full_filename


    class BatchLoraLoader:
        def __init__(self, directory_path, pattern, index):
            self.lora_paths = []
            self.load_loras(directory_path, pattern)
            self.lora_paths.sort()

            self.index = index

        def load_loras(self, directory_path, pattern):
            for file_name in glob.glob(os.path.join(directory_path, pattern), recursive=True):
                temp_strings = file_name.split('\\')
                file_name = temp_strings[len(temp_strings) - 1]

                if file_name.lower().endswith("safetensors"):
                    self.lora_paths.append(file_name)

        def get_lora_by_id(self, lora_id):
            if lora_id < 0 or lora_id >= len(self.lora_paths):
                cstr(f"WarpedLoadLorasBatchByPrefix: Invalid lora index `{lora_id}`").error.print()
                return

            return self.lora_paths[lora_id]

        def get_next_lora(self):
            if self.index >= len(self.lora_paths):
                self.index = 0

            lora_path = self.lora_paths[self.index]
            self.index += 1

            if self.index == len(self.lora_paths):
                self.index = 0

            cstr(f'{cstr.color.YELLOW}{cstr.color.END} Index: {self.index}').msg.print()

            return lora_path

        def get_current_lora(self):
            if self.index >= len(self.lora_paths):
                self.index = 0
            lora_path = self.lora_paths[self.index]

            return lora_path

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return float("NaN")

class WarpedAnyType(str):
    """A special type that always compares equal to any value."""

    def __ne__(self, __value: object) -> bool:
        return False

WARPEDANYTYPE = WarpedAnyType("*")

class WarpedContinueWorkflowManual:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "any": (WARPEDANYTYPE,),
                "continue_workflow": ("BOOLEAN", {"default": True})
            },
            "hidden": {
                "id": "UNIQUE_ID",
            },
        }

    RETURN_TYPES = (WARPEDANYTYPE, )
    RETURN_NAMES = ( "any", )
    FUNCTION = "execute"
    CATEGORY = "Warped/General/Utils"
    OUTPUT_NODE = True

    def execute(self, any=None, continue_workflow=True, id=None):
        if not continue_workflow:
            # print(f"Cancelled workflow for {id}")
            raise comfy.model_management.InterruptProcessingException()

        return {"result": (any,)}

class WarpedContinueWorkflowAuto:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "any": (WARPEDANYTYPE,),
                "continue_workflow": ("BOOL", {"default": True, "forceInput": True})
            },
            "hidden": {
                "id": "UNIQUE_ID",
            },
        }

    RETURN_TYPES = (WARPEDANYTYPE, )
    RETURN_NAMES = ( "any", )
    FUNCTION = "execute"
    CATEGORY = "Warped/General/Utils"
    OUTPUT_NODE = True

    def execute(self, any=None, continue_workflow=True, id=None):
        if not continue_workflow:
            # print(f"Cancelled workflow for {id}")
            raise comfy.model_management.InterruptProcessingException()

        return {"result": (any,)}