714 lines
28 KiB
Python
714 lines
28 KiB
Python
"""
|
|
TensorPrism Core Merge Operations
|
|
==================================
|
|
Consolidates: MainMerge, LayeredBlend, and Prism (spectral merging)
|
|
|
|
Author: Arctenox
|
|
Version: 1.6.5
|
|
License: GPL-3.0
|
|
"""
|
|
|
|
import torch
|
|
import math
|
|
import torch.fft
|
|
import comfy.model_management
|
|
|
|
# ==================== MAIN MERGE NODE ====================
|
|
|
|
class TensorPrism_MainMerge:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_A": ("MODEL",),
|
|
"model_B": ("MODEL",),
|
|
"merge_ratio": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}),
|
|
"method": (["linear", "slerp", "cosine", "directional", "frequency", "stochastic"],),
|
|
},
|
|
"optional": {
|
|
"random_seed": ("INT", {"default": 42, "min": 0}),
|
|
"stochastic_prob": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0}),
|
|
"freq_low_ratio": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}),
|
|
"freq_high_ratio": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("merged_model",)
|
|
FUNCTION = "merge_models"
|
|
CATEGORY = "Tensor_Prism/Core"
|
|
|
|
def __init__(self):
|
|
self.device = self._get_optimal_device()
|
|
|
|
def _get_optimal_device(self):
|
|
try:
|
|
return comfy.model_management.get_torch_device()
|
|
except:
|
|
if torch.cuda.is_available():
|
|
return torch.device("cuda")
|
|
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
|
return torch.device("mps")
|
|
else:
|
|
return torch.device("cpu")
|
|
|
|
def _to_device(self, tensor: torch.Tensor, target_device: torch.device = None) -> torch.Tensor:
|
|
if target_device is None:
|
|
target_device = self.device
|
|
if tensor.device != target_device:
|
|
return tensor.to(target_device)
|
|
return tensor
|
|
|
|
@staticmethod
|
|
def _blend_linear(t1, t2, alpha):
|
|
device = t1.device
|
|
t2 = t2.to(device)
|
|
return (1 - alpha) * t1 + alpha * t2
|
|
|
|
def _blend_slerp(self, t1, t2, alpha):
|
|
if t1.numel() == 0 or t2.numel() == 0:
|
|
return t1
|
|
|
|
device = t1.device
|
|
t2 = t2.to(device)
|
|
|
|
t1_flat, t2_flat = t1.view(-1), t2.view(-1)
|
|
t1_norm = torch.nn.functional.normalize(t1_flat, dim=0, eps=1e-8)
|
|
t2_norm = torch.nn.functional.normalize(t2_flat, dim=0, eps=1e-8)
|
|
|
|
dot = torch.clamp(torch.dot(t1_norm, t2_norm), -1.0 + 1e-7, 1.0 - 1e-7)
|
|
|
|
if abs(dot.item()) > 0.9995:
|
|
return self._blend_linear(t1, t2, alpha)
|
|
|
|
theta = torch.acos(dot)
|
|
sin_theta = torch.sin(theta)
|
|
|
|
if sin_theta.abs() < 1e-6:
|
|
return self._blend_linear(t1, t2, alpha)
|
|
|
|
w1 = torch.sin((1 - alpha) * theta) / sin_theta
|
|
w2 = torch.sin(alpha * theta) / sin_theta
|
|
|
|
result = w1 * t1_norm + w2 * t2_norm
|
|
original_norm = torch.norm(t1_flat) * (1 - alpha) + torch.norm(t2_flat) * alpha
|
|
result = result * original_norm
|
|
|
|
return result.view(t1.shape).to(device)
|
|
|
|
@staticmethod
|
|
def _blend_cosine(t1, t2, alpha):
|
|
device = t1.device
|
|
t2 = t2.to(device)
|
|
mu = (1 - math.cos(alpha * math.pi)) / 2
|
|
return (1 - mu) * t1 + mu * t2
|
|
|
|
@staticmethod
|
|
def _blend_directional(t1, t2, alpha):
|
|
device = t1.device
|
|
t2 = t2.to(device)
|
|
return t1 + (t2 - t1) * alpha
|
|
|
|
def _blend_frequency(self, t1, t2, low_ratio, high_ratio):
|
|
device = t1.device
|
|
t2 = t2.to(device)
|
|
|
|
if t1.numel() < 4 or t2.numel() < 4:
|
|
return self._blend_linear(t1, t2, (low_ratio + high_ratio) / 2)
|
|
|
|
try:
|
|
orig_shape = t1.shape
|
|
if t1.dim() < 2:
|
|
t1_fft = t1.view(-1, 1)
|
|
t2_fft = t2.view(-1, 1)
|
|
else:
|
|
t1_fft = t1
|
|
t2_fft = t2
|
|
|
|
fA = torch.fft.fft2(t1_fft.float())
|
|
fB = torch.fft.fft2(t2_fft.float())
|
|
|
|
h, w = fA.shape[-2:]
|
|
cy, cx = h // 2, w // 2
|
|
|
|
y, x = torch.meshgrid(torch.arange(h, device=device), torch.arange(w, device=device), indexing='ij')
|
|
|
|
dist = torch.sqrt((y - cy).float()**2 + (x - cx).float()**2)
|
|
max_dist = min(cy, cx)
|
|
|
|
if max_dist > 0:
|
|
normalized_dist = dist / max_dist
|
|
low_freq_mask = (normalized_dist < 0.3).float()
|
|
high_freq_mask = (normalized_dist >= 0.3).float()
|
|
else:
|
|
low_freq_mask = torch.ones_like(dist)
|
|
high_freq_mask = torch.zeros_like(dist)
|
|
|
|
fA_low = fA * low_freq_mask.unsqueeze(0) if fA.dim() > 2 else fA * low_freq_mask
|
|
fB_low = fB * low_freq_mask.unsqueeze(0) if fB.dim() > 2 else fB * low_freq_mask
|
|
fA_high = fA * high_freq_mask.unsqueeze(0) if fA.dim() > 2 else fA * high_freq_mask
|
|
fB_high = fB * high_freq_mask.unsqueeze(0) if fB.dim() > 2 else fB * high_freq_mask
|
|
|
|
merged_low = fA_low * (1 - low_ratio) + fB_low * low_ratio
|
|
merged_high = fA_high * (1 - high_ratio) + fB_high * high_ratio
|
|
merged_freq = merged_low + merged_high
|
|
|
|
result = torch.fft.ifft2(merged_freq).real
|
|
|
|
return result.view(orig_shape).to(t1.dtype).to(device)
|
|
|
|
except Exception as e:
|
|
return self._blend_linear(t1, t2, (low_ratio + high_ratio) / 2)
|
|
|
|
def _blend_stochastic(self, t1, t2, alpha, prob, seed):
|
|
device = t1.device
|
|
t2 = t2.to(device)
|
|
|
|
generator = torch.Generator(device=device)
|
|
generator.manual_seed(seed)
|
|
|
|
mask = torch.rand(t1.shape, generator=generator, device=device) < prob
|
|
|
|
base_blend = self._blend_linear(t1, t2, alpha)
|
|
stochastic_component = torch.where(mask, t2, t1)
|
|
|
|
return (base_blend * (1 - prob) + stochastic_component * prob).to(device)
|
|
|
|
def merge_models(self, model_A, model_B, merge_ratio, method,
|
|
random_seed=42, stochastic_prob=0.1,
|
|
freq_low_ratio=0.5, freq_high_ratio=0.5):
|
|
|
|
merged_model = model_A.clone()
|
|
|
|
state_dict_A = model_A.model.state_dict()
|
|
state_dict_B = model_B.model.state_dict()
|
|
|
|
patches = {}
|
|
|
|
for key in state_dict_A.keys():
|
|
if key in state_dict_B:
|
|
t1 = state_dict_A[key]
|
|
t2 = state_dict_B[key]
|
|
|
|
if isinstance(t1, torch.Tensor) and isinstance(t2, torch.Tensor):
|
|
if t1.shape == t2.shape:
|
|
try:
|
|
device = t1.device
|
|
t2 = t2.to(device)
|
|
|
|
if method == "linear":
|
|
merged_tensor = self._blend_linear(t1, t2, merge_ratio)
|
|
elif method == "slerp":
|
|
merged_tensor = self._blend_slerp(t1, t2, merge_ratio)
|
|
elif method == "cosine":
|
|
merged_tensor = self._blend_cosine(t1, t2, merge_ratio)
|
|
elif method == "directional":
|
|
merged_tensor = self._blend_directional(t1, t2, merge_ratio)
|
|
elif method == "frequency":
|
|
merged_tensor = self._blend_frequency(t1, t2, freq_low_ratio, freq_high_ratio)
|
|
elif method == "stochastic":
|
|
merged_tensor = self._blend_stochastic(t1, t2, merge_ratio, stochastic_prob, random_seed)
|
|
else:
|
|
merged_tensor = self._blend_linear(t1, t2, merge_ratio)
|
|
|
|
merged_tensor = merged_tensor.to(device)
|
|
|
|
diff = merged_tensor - t1
|
|
if torch.abs(diff).max() > 1e-8:
|
|
patches[key] = (diff.cpu(),)
|
|
|
|
except Exception as e:
|
|
print(f"Warning: Failed to merge parameter {key} with method {method}: {e}")
|
|
continue
|
|
|
|
if patches:
|
|
merged_model.add_patches(patches, 1.0)
|
|
|
|
return (merged_model,)
|
|
|
|
|
|
# ==================== LAYERED BLEND NODE ====================
|
|
|
|
class TensorPrism_LayeredBlend:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_A": ("MODEL",),
|
|
"model_B": ("MODEL",),
|
|
"text_encoder_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}),
|
|
"text_encoder_method": (["linear", "slerp", "cosine"],),
|
|
"unet_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}),
|
|
"unet_method": (["linear", "slerp", "cosine"],),
|
|
"time_embed_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}),
|
|
"input_blocks_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}),
|
|
"middle_block_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}),
|
|
"output_blocks_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}),
|
|
"out_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}),
|
|
},
|
|
"optional": {
|
|
"vae_A": ("VAE",),
|
|
"vae_B": ("VAE",),
|
|
"vae_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}),
|
|
"vae_method": (["linear", "slerp", "cosine"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "VAE",)
|
|
RETURN_NAMES = ("merged_model", "merged_vae",)
|
|
FUNCTION = "blend_models"
|
|
CATEGORY = "Tensor_Prism/Core"
|
|
|
|
def __init__(self):
|
|
self.device = self._get_optimal_device()
|
|
|
|
def _get_optimal_device(self):
|
|
try:
|
|
return comfy.model_management.get_torch_device()
|
|
except:
|
|
if torch.cuda.is_available():
|
|
return torch.device("cuda")
|
|
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
|
return torch.device("mps")
|
|
else:
|
|
return torch.device("cpu")
|
|
|
|
@staticmethod
|
|
def _blend_linear(t1: torch.Tensor, t2: torch.Tensor, alpha: float) -> torch.Tensor:
|
|
device = t1.device
|
|
t2 = t2.to(device)
|
|
return (1 - alpha) * t1 + alpha * t2
|
|
|
|
@staticmethod
|
|
def _blend_slerp(t1: torch.Tensor, t2: torch.Tensor, alpha: float) -> torch.Tensor:
|
|
if t1.numel() == 0 or t2.numel() == 0:
|
|
return t1
|
|
|
|
device = t1.device
|
|
t2 = t2.to(device)
|
|
|
|
t1_flat = t1.view(-1)
|
|
t2_flat = t2.view(-1)
|
|
|
|
t1_norm = torch.nn.functional.normalize(t1_flat, dim=0, eps=1e-8)
|
|
t2_norm = torch.nn.functional.normalize(t2_flat, dim=0, eps=1e-8)
|
|
|
|
dot = torch.clamp(torch.dot(t1_norm, t2_norm), -1.0 + 1e-7, 1.0 - 1e-7)
|
|
|
|
if abs(dot.item()) > 0.9995:
|
|
return TensorPrism_LayeredBlend._blend_linear(t1, t2, alpha)
|
|
|
|
theta = torch.acos(dot)
|
|
sin_theta = torch.sin(theta)
|
|
|
|
if sin_theta.abs() < 1e-6:
|
|
return TensorPrism_LayeredBlend._blend_linear(t1, t2, alpha)
|
|
|
|
w1 = torch.sin((1 - alpha) * theta) / sin_theta
|
|
w2 = torch.sin(alpha * theta) / sin_theta
|
|
|
|
result = w1 * t1_norm + w2 * t2_norm
|
|
|
|
original_norm = torch.norm(t1_flat) * (1 - alpha) + torch.norm(t2_flat) * alpha
|
|
result = result * original_norm
|
|
|
|
return result.view(t1.shape).to(device)
|
|
|
|
@staticmethod
|
|
def _blend_cosine(t1: torch.Tensor, t2: torch.Tensor, alpha: float) -> torch.Tensor:
|
|
device = t1.device
|
|
t2 = t2.to(device)
|
|
mu = (1 - math.cos(alpha * math.pi)) / 2
|
|
return (1 - mu) * t1 + mu * t2
|
|
|
|
def _get_component_type(self, param_name: str) -> str:
|
|
param_lower = param_name.lower()
|
|
|
|
if any(pattern in param_lower for pattern in ['cond_stage_model', 'text_encoder', 'clip', 'transformer.text_model']):
|
|
return 'text_encoder'
|
|
|
|
if any(pattern in param_lower for pattern in ['model.diffusion_model', 'unet']):
|
|
if 'time_embed' in param_lower:
|
|
return 'time_embed'
|
|
elif 'input_blocks' in param_lower:
|
|
return 'input_blocks'
|
|
elif 'middle_block' in param_lower:
|
|
return 'middle_block'
|
|
elif 'output_blocks' in param_lower:
|
|
return 'output_blocks'
|
|
elif param_lower.endswith('.out.weight') or param_lower.endswith('.out.bias'):
|
|
return 'out'
|
|
else:
|
|
return 'unet'
|
|
|
|
if any(pattern in param_lower for pattern in ['first_stage_model', 'vae', 'decoder', 'encoder']):
|
|
return 'vae'
|
|
|
|
return 'unet'
|
|
|
|
def _apply_blend_method(self, t1: torch.Tensor, t2: torch.Tensor, strength: float, method: str) -> torch.Tensor:
|
|
device = t1.device
|
|
t2 = t2.to(device)
|
|
|
|
if method == "linear":
|
|
return self._blend_linear(t1, t2, strength)
|
|
elif method == "slerp":
|
|
return self._blend_slerp(t1, t2, strength)
|
|
elif method == "cosine":
|
|
return self._blend_cosine(t1, t2, strength)
|
|
else:
|
|
return self._blend_linear(t1, t2, strength)
|
|
|
|
def blend_models(self, model_A, model_B,
|
|
text_encoder_strength, text_encoder_method,
|
|
unet_strength, unet_method,
|
|
time_embed_strength, input_blocks_strength,
|
|
middle_block_strength, output_blocks_strength, out_strength,
|
|
vae_A=None, vae_B=None,
|
|
vae_strength=0.5, vae_method="linear"):
|
|
|
|
merged_model = model_A.clone()
|
|
|
|
state_dict_A = model_A.model.state_dict()
|
|
state_dict_B = model_B.model.state_dict()
|
|
|
|
patches = {}
|
|
|
|
component_strengths = {
|
|
'text_encoder': (text_encoder_strength, text_encoder_method),
|
|
'time_embed': (time_embed_strength, unet_method),
|
|
'input_blocks': (input_blocks_strength, unet_method),
|
|
'middle_block': (middle_block_strength, unet_method),
|
|
'output_blocks': (output_blocks_strength, unet_method),
|
|
'out': (out_strength, unet_method),
|
|
'unet': (unet_strength, unet_method),
|
|
'vae': (0.0, "linear")
|
|
}
|
|
|
|
for key in state_dict_A.keys():
|
|
if key in state_dict_B:
|
|
t1 = state_dict_A[key]
|
|
t2 = state_dict_B[key]
|
|
|
|
if isinstance(t1, torch.Tensor) and isinstance(t2, torch.Tensor):
|
|
if t1.shape == t2.shape:
|
|
try:
|
|
device = t1.device
|
|
t2 = t2.to(device)
|
|
|
|
component_type = self._get_component_type(key)
|
|
|
|
strength, method = component_strengths.get(component_type, (unet_strength, unet_method))
|
|
|
|
if strength == 0:
|
|
continue
|
|
|
|
merged_tensor = self._apply_blend_method(t1, t2, strength, method)
|
|
merged_tensor = merged_tensor.to(device)
|
|
|
|
diff = merged_tensor - t1
|
|
if torch.abs(diff).max() > 1e-8:
|
|
patches[key] = (diff.cpu(),)
|
|
|
|
except Exception as e:
|
|
print(f"Warning: Failed to blend parameter {key}: {e}")
|
|
continue
|
|
|
|
if patches:
|
|
merged_model.add_patches(patches, 1.0)
|
|
|
|
merged_vae = None
|
|
if vae_A is not None and vae_B is not None:
|
|
try:
|
|
merged_vae = self._blend_vaes(vae_A, vae_B, vae_strength, vae_method)
|
|
except Exception as e:
|
|
print(f"Warning: Failed to blend VAEs: {e}")
|
|
merged_vae = vae_A
|
|
|
|
return (merged_model, merged_vae)
|
|
|
|
def _blend_vaes(self, vae_A, vae_B, strength: float, method: str):
|
|
if strength == 0:
|
|
return vae_A
|
|
if strength == 1:
|
|
return vae_B
|
|
|
|
merged_vae = vae_A.clone() if hasattr(vae_A, 'clone') else vae_A
|
|
|
|
try:
|
|
if hasattr(vae_A, 'first_stage_model') and hasattr(vae_B, 'first_stage_model'):
|
|
state_dict_A = vae_A.first_stage_model.state_dict()
|
|
state_dict_B = vae_B.first_stage_model.state_dict()
|
|
|
|
vae_patches = {}
|
|
|
|
for key in state_dict_A.keys():
|
|
if key in state_dict_B:
|
|
t1 = state_dict_A[key]
|
|
t2 = state_dict_B[key]
|
|
|
|
if isinstance(t1, torch.Tensor) and isinstance(t2, torch.Tensor):
|
|
if t1.shape == t2.shape:
|
|
device = t1.device
|
|
t2 = t2.to(device)
|
|
|
|
merged_tensor = self._apply_blend_method(t1, t2, strength, method)
|
|
merged_tensor = merged_tensor.to(device)
|
|
|
|
diff = merged_tensor - t1
|
|
|
|
if torch.abs(diff).max() > 1e-8:
|
|
vae_patches[key] = (diff.cpu(),)
|
|
|
|
if vae_patches and hasattr(merged_vae, 'add_patches'):
|
|
merged_vae.add_patches(vae_patches, 1.0)
|
|
|
|
except Exception as e:
|
|
print(f"VAE blending failed, using VAE A: {e}")
|
|
return vae_A
|
|
|
|
return merged_vae
|
|
|
|
|
|
# ==================== PRISM (SPECTRAL) NODE ====================
|
|
|
|
class TensorPrism_FastPrism:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_A": ("MODEL",),
|
|
"model_B": ("MODEL",),
|
|
"merge_method": ([
|
|
"spectral_blend", "frequency_bands", "magnitude_weighted",
|
|
"adaptive_mix", "harmonic_merge"
|
|
],),
|
|
"strength": ("FLOAT", {
|
|
"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01
|
|
}),
|
|
},
|
|
"optional": {
|
|
"low_freq_bias": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.05}),
|
|
"high_freq_bias": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.05}),
|
|
"spectral_precision": (["fast", "balanced", "precise"], {"default": "fast"}),
|
|
"layer_selective": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("merged_model",)
|
|
FUNCTION = "fast_prism_merge"
|
|
CATEGORY = "Tensor_Prism/Core"
|
|
|
|
def fast_prism_merge(self, model_A, model_B, merge_method, strength,
|
|
low_freq_bias=0.7, high_freq_bias=0.3,
|
|
spectral_precision="fast", layer_selective=False):
|
|
|
|
merged_model = model_A.clone()
|
|
|
|
state_dict_A = model_A.model.state_dict()
|
|
state_dict_B = model_B.model.state_dict()
|
|
|
|
patches = {}
|
|
|
|
for key in state_dict_A.keys():
|
|
if key not in state_dict_B:
|
|
continue
|
|
|
|
tensor_A = state_dict_A[key]
|
|
tensor_B = state_dict_B[key]
|
|
|
|
if not (isinstance(tensor_A, torch.Tensor) and isinstance(tensor_B, torch.Tensor)):
|
|
continue
|
|
|
|
if tensor_A.shape != tensor_B.shape:
|
|
continue
|
|
|
|
if tensor_A.numel() < 10:
|
|
continue
|
|
|
|
device = tensor_A.device
|
|
tensor_B = tensor_B.to(device)
|
|
|
|
if layer_selective:
|
|
layer_strength = self._get_layer_strength(key, strength)
|
|
else:
|
|
layer_strength = strength
|
|
|
|
try:
|
|
merged_tensor = self._fast_spectral_merge(
|
|
tensor_A, tensor_B, merge_method, layer_strength,
|
|
low_freq_bias, high_freq_bias, spectral_precision
|
|
)
|
|
|
|
merged_tensor = merged_tensor.to(device)
|
|
|
|
diff = merged_tensor - tensor_A
|
|
if torch.abs(diff).max() > 1e-6:
|
|
patches[key] = (diff.cpu(),)
|
|
|
|
except Exception as e:
|
|
merged_tensor = tensor_A * (1 - layer_strength) + tensor_B * layer_strength
|
|
merged_tensor = merged_tensor.to(device)
|
|
diff = merged_tensor - tensor_A
|
|
if torch.abs(diff).max() > 1e-6:
|
|
patches[key] = (diff.cpu(),)
|
|
|
|
if patches:
|
|
merged_model.add_patches(patches, 1.0)
|
|
|
|
return (merged_model,)
|
|
|
|
def _get_layer_strength(self, layer_name, base_strength):
|
|
layer_name = layer_name.lower()
|
|
|
|
if any(x in layer_name for x in ['attention', 'attn']):
|
|
return base_strength * 0.8
|
|
elif any(x in layer_name for x in ['norm', 'ln']):
|
|
return base_strength * 0.6
|
|
elif any(x in layer_name for x in ['embed', 'pos']):
|
|
return base_strength * 0.4
|
|
elif any(x in layer_name for x in ['output', 'head']):
|
|
return base_strength * 0.9
|
|
else:
|
|
return base_strength
|
|
|
|
def _fast_spectral_merge(self, tensor_A, tensor_B, method, strength,
|
|
low_bias, high_bias, precision):
|
|
|
|
device = tensor_A.device
|
|
tensor_B = tensor_B.to(device)
|
|
|
|
if method == "spectral_blend":
|
|
return self._spectral_blend_fast(tensor_A, tensor_B, strength, low_bias, high_bias)
|
|
elif method == "frequency_bands":
|
|
return self._frequency_bands_fast(tensor_A, tensor_B, strength, low_bias, high_bias, precision)
|
|
elif method == "magnitude_weighted":
|
|
return self._magnitude_weighted_fast(tensor_A, tensor_B, strength)
|
|
elif method == "adaptive_mix":
|
|
return self._adaptive_mix_fast(tensor_A, tensor_B, strength)
|
|
elif method == "harmonic_merge":
|
|
return self._harmonic_merge_fast(tensor_A, tensor_B, strength, low_bias, high_bias)
|
|
else:
|
|
return tensor_A * (1 - strength) + tensor_B * strength
|
|
|
|
def _spectral_blend_fast(self, tensor_A, tensor_B, strength, low_bias, high_bias):
|
|
device = tensor_A.device
|
|
tensor_B = tensor_B.to(device)
|
|
|
|
mag_A = torch.abs(tensor_A)
|
|
mag_B = torch.abs(tensor_B)
|
|
|
|
threshold = (mag_A.mean() + mag_B.mean()) / 2
|
|
|
|
low_freq_mask = (mag_A + mag_B) > threshold
|
|
high_freq_mask = ~low_freq_mask
|
|
|
|
result = torch.zeros_like(tensor_A, device=device)
|
|
result[low_freq_mask] = (tensor_A[low_freq_mask] * (1 - strength * low_bias) +
|
|
tensor_B[low_freq_mask] * (strength * low_bias))
|
|
result[high_freq_mask] = (tensor_A[high_freq_mask] * (1 - strength * high_bias) +
|
|
tensor_B[high_freq_mask] * (strength * high_bias))
|
|
|
|
return result.to(device)
|
|
|
|
def _frequency_bands_fast(self, tensor_A, tensor_B, strength, low_bias, high_bias, precision):
|
|
device = tensor_A.device
|
|
tensor_B = tensor_B.to(device)
|
|
|
|
if tensor_A.dim() < 2 or precision == "fast":
|
|
var_A = torch.var(tensor_A, dim=-1, keepdim=True)
|
|
var_B = torch.var(tensor_B, dim=-1, keepdim=True)
|
|
|
|
freq_weight = torch.sigmoid((var_A + var_B) - (var_A.mean() + var_B.mean()))
|
|
|
|
effective_bias = low_bias * (1 - freq_weight) + high_bias * freq_weight
|
|
effective_strength = strength * effective_bias
|
|
|
|
return (tensor_A * (1 - effective_strength) + tensor_B * effective_strength).to(device)
|
|
|
|
else:
|
|
try:
|
|
flat_A = tensor_A.flatten()
|
|
flat_B = tensor_B.flatten()
|
|
|
|
fft_A = torch.fft.fft(flat_A.float())
|
|
fft_B = torch.fft.fft(flat_B.float())
|
|
|
|
n = len(fft_A)
|
|
low_cutoff = n // 4
|
|
high_cutoff = 3 * n // 4
|
|
|
|
fft_merged = fft_A.clone()
|
|
fft_merged[:low_cutoff] = (fft_A[:low_cutoff] * (1 - strength * low_bias) +
|
|
fft_B[:low_cutoff] * (strength * low_bias))
|
|
fft_merged[high_cutoff:] = (fft_A[high_cutoff:] * (1 - strength * high_bias) +
|
|
fft_B[high_cutoff:] * (strength * high_bias))
|
|
fft_merged[low_cutoff:high_cutoff] = (fft_A[low_cutoff:high_cutoff] * (1 - strength) +
|
|
fft_B[low_cutoff:high_cutoff] * strength)
|
|
|
|
merged_flat = torch.fft.ifft(fft_merged).real.to(tensor_A.dtype)
|
|
|
|
return merged_flat.view(tensor_A.shape).to(device)
|
|
|
|
except Exception as e:
|
|
return self._spectral_blend_fast(tensor_A, tensor_B, strength, low_bias, high_bias)
|
|
|
|
def _magnitude_weighted_fast(self, tensor_A, tensor_B, strength):
|
|
device = tensor_A.device
|
|
tensor_B = tensor_B.to(device)
|
|
|
|
mag_A = torch.norm(tensor_A)
|
|
mag_B = torch.norm(tensor_B)
|
|
|
|
if mag_A + mag_B > 1e-8:
|
|
weight_B = mag_B / (mag_A + mag_B)
|
|
effective_strength = strength * weight_B + (1 - strength) * 0.5
|
|
else:
|
|
effective_strength = strength
|
|
|
|
return (tensor_A * (1 - effective_strength) + tensor_B * effective_strength).to(device)
|
|
|
|
def _adaptive_mix_fast(self, tensor_A, tensor_B, strength):
|
|
device = tensor_A.device
|
|
tensor_B = tensor_B.to(device)
|
|
|
|
flat_A = tensor_A.flatten()
|
|
flat_B = tensor_B.flatten()
|
|
|
|
similarity = torch.cosine_similarity(flat_A.unsqueeze(0), flat_B.unsqueeze(0), dim=1)
|
|
|
|
adaptive_strength = strength * (1 - 0.3 * torch.abs(similarity))
|
|
|
|
return (tensor_A * (1 - adaptive_strength) + tensor_B * adaptive_strength).to(device)
|
|
|
|
def _harmonic_merge_fast(self, tensor_A, tensor_B, strength, low_bias, high_bias):
|
|
device = tensor_A.device
|
|
tensor_B = tensor_B.to(device)
|
|
|
|
sign_A = torch.sign(tensor_A)
|
|
sign_B = torch.sign(tensor_B)
|
|
|
|
alignment = (sign_A * sign_B) > 0
|
|
|
|
effective_bias = torch.where(alignment,
|
|
torch.tensor(low_bias, device=device),
|
|
torch.tensor(high_bias, device=device))
|
|
effective_strength = strength * effective_bias
|
|
|
|
return (tensor_A * (1 - effective_strength) + tensor_B * effective_strength).to(device)
|
|
|
|
|
|
# ==================== NODE REGISTRATION ====================
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"TensorPrism_MainMerge": TensorPrism_MainMerge,
|
|
"TensorPrism_LayeredBlend": TensorPrism_LayeredBlend,
|
|
"TensorPrism_Prism": TensorPrism_FastPrism
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"TensorPrism_MainMerge": "Main Merge (Tensor Prism)",
|
|
"TensorPrism_LayeredBlend": "Layered Blend (Tensor Prism)",
|
|
"TensorPrism_Prism": "Prism (Tensor Prism)"
|
|
} |