124 lines
4.7 KiB
Python
124 lines
4.7 KiB
Python
import torch
|
|
|
|
class ZImageAdvancedConditioning:
|
|
"""
|
|
FINAL PRODUCTION VERSION.
|
|
- Fixes 'Drift' bug (No more in-place modification of cached inputs).
|
|
- Removes debug prints for speed.
|
|
- Safely handles Z-Image/Qwen/Llama architectures.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"conditioning_to": ("CONDITIONING", ),
|
|
"conditioning_from": ("CONDITIONING", ),
|
|
"operation": (["mix_slerp", "purge_ortho", "add_perpendicular"], ),
|
|
"strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
RETURN_NAMES = ("conditioning",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "RES4LYF/conditioning"
|
|
|
|
def process(self, conditioning_to, conditioning_from, operation, strength):
|
|
results = []
|
|
|
|
# Match list lengths
|
|
min_len = min(len(conditioning_to), len(conditioning_from))
|
|
|
|
for i in range(min_len):
|
|
# --- 1. PROCESS MAIN TENSOR (t[0]) ---
|
|
t0_input = conditioning_to[i][0]
|
|
t0_ref = conditioning_from[i][0]
|
|
|
|
# Create a SAFE COPY of the main tensor to avoid modifying the input cache
|
|
t0_output = t0_input.clone()
|
|
|
|
# Handle Shape Mismatch (Slice to common length)
|
|
common_len = min(t0_input.shape[1], t0_ref.shape[1])
|
|
|
|
# Extract working slices
|
|
v0 = t0_input[:, :common_len, :].clone()
|
|
v1 = t0_ref[:, :common_len, :].clone()
|
|
|
|
# Apply Math
|
|
v_processed = self.apply_math(v0, v1, operation, strength)
|
|
|
|
# Write into the NEW output tensor (not the input!)
|
|
t0_output[:, :common_len, :] = v_processed
|
|
|
|
# --- 2. PROCESS DICTIONARY ---
|
|
# Create a shallow copy of the dict, but we will replace values with new tensors
|
|
new_dict = conditioning_to[i][1].copy()
|
|
ref_dict = conditioning_from[i][1]
|
|
|
|
target_keys = ["conditioning_llama3", "llama_embeds", "pooled_output"]
|
|
|
|
for key in new_dict.keys():
|
|
if key in target_keys and key in ref_dict:
|
|
val_to = new_dict[key]
|
|
val_from = ref_dict[key]
|
|
|
|
if val_to is not None and val_from is not None and isinstance(val_to, torch.Tensor):
|
|
# Ensure we only process if shapes align
|
|
if val_to.shape == val_from.shape:
|
|
# Apply Math
|
|
# apply_math returns a new tensor, so this is safe
|
|
new_dict[key] = self.apply_math(val_to, val_from, operation, strength)
|
|
|
|
results.append([t0_output, new_dict])
|
|
|
|
return (results,)
|
|
|
|
def apply_math(self, v0, v1, operation, strength):
|
|
"""Helper for vector operations"""
|
|
# Epsilon for stability
|
|
eps = 1e-8
|
|
|
|
if operation == "mix_slerp":
|
|
# Normalize
|
|
v0_n = v0 / (v0.norm(dim=-1, keepdim=True) + eps)
|
|
v1_n = v1 / (v1.norm(dim=-1, keepdim=True) + eps)
|
|
|
|
dot = (v0_n * v1_n).sum(dim=-1, keepdim=True)
|
|
dot = torch.clamp(dot, -0.9995, 0.9995)
|
|
|
|
theta = torch.acos(dot)
|
|
sin_theta = torch.sin(theta) + eps
|
|
|
|
w0 = torch.sin((1.0 - strength) * theta) / sin_theta
|
|
w1 = torch.sin(strength * theta) / sin_theta
|
|
|
|
return (w0 * v0 + w1 * v1).contiguous()
|
|
|
|
elif operation == "purge_ortho":
|
|
# Project v0 onto v1
|
|
dot_v0_v1 = (v0 * v1).sum(dim=-1, keepdim=True)
|
|
dot_v1_v1 = (v1 * v1).sum(dim=-1, keepdim=True)
|
|
|
|
proj = (dot_v0_v1 / (dot_v1_v1 + eps)) * v1
|
|
return (v0 - (proj * strength)).contiguous()
|
|
|
|
elif operation == "add_perpendicular":
|
|
# Find part of v1 orthogonal to v0
|
|
dot_v1_v0 = (v1 * v0).sum(dim=-1, keepdim=True)
|
|
dot_v0_v0 = (v0 * v0).sum(dim=-1, keepdim=True)
|
|
|
|
proj = (dot_v1_v0 / (dot_v0_v0 + eps)) * v0
|
|
ortho_v1 = v1 - proj
|
|
|
|
return (v0 + (ortho_v1 * strength)).contiguous()
|
|
|
|
return v0.contiguous()
|
|
|
|
# Register
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ZImageAdvancedConditioning": ZImageAdvancedConditioning
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ZImageAdvancedConditioning": "Z-Image Advanced Mixing"
|
|
} |