12336 lines
540 KiB
Python
12336 lines
540 KiB
Python
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,)}
|