401 lines
16 KiB
Python
401 lines
16 KiB
Python
"""
|
|
TensorPrism Intelligent Tensor Selector
|
|
=======================================
|
|
|
|
Advanced tensor-by-tensor selection that evaluates multiple models and chooses
|
|
the best tensor based on quality metrics and configurable weights.
|
|
|
|
Inspired by Hyphoria's merge methodology - evaluates each tensor individually
|
|
and selects the best one based on multiple quality criteria.
|
|
|
|
Author: Arctenox
|
|
Version: 1.0.1
|
|
License: GPL-3.0
|
|
"""
|
|
|
|
import torch
|
|
import numpy as np
|
|
from typing import Dict, List, Tuple, Optional
|
|
import gc
|
|
import comfy.model_management
|
|
|
|
class TensorPrism_IntelligentTensorSelector:
|
|
"""
|
|
Evaluates tensors from multiple models and intelligently selects the best
|
|
ones based on configurable quality metrics and weights.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_A": ("MODEL",),
|
|
"model_B": ("MODEL",),
|
|
"selection_mode": ([
|
|
"best_quality",
|
|
"balanced",
|
|
"conservative",
|
|
"aggressive"
|
|
], {"default": "best_quality"}),
|
|
|
|
# Quality metric weights
|
|
"magnitude_weight": ("FLOAT", {
|
|
"default": 0.3, "min": 0.0, "max": 2.0, "step": 0.05,
|
|
"tooltip": "Weight for tensor magnitude (strength)"
|
|
}),
|
|
"variance_weight": ("FLOAT", {
|
|
"default": 0.25, "min": 0.0, "max": 2.0, "step": 0.05,
|
|
"tooltip": "Weight for tensor variance (detail preservation)"
|
|
}),
|
|
"spectral_weight": ("FLOAT", {
|
|
"default": 0.25, "min": 0.0, "max": 2.0, "step": 0.05,
|
|
"tooltip": "Weight for spectral content (frequency richness)"
|
|
}),
|
|
"stability_weight": ("FLOAT", {
|
|
"default": 0.2, "min": 0.0, "max": 2.0, "step": 0.05,
|
|
"tooltip": "Weight for numerical stability"
|
|
}),
|
|
|
|
# Bias controls
|
|
"model_A_bias": ("FLOAT", {
|
|
"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.05,
|
|
"tooltip": "Bias toward Model A (-1 = strong B preference, +1 = strong A preference)"
|
|
}),
|
|
"attention_preference": ("FLOAT", {
|
|
"default": 0.0, "min": -0.5, "max": 0.5, "step": 0.05,
|
|
"tooltip": "Extra weight for attention layers"
|
|
}),
|
|
"output_preference": ("FLOAT", {
|
|
"default": 0.0, "min": -0.5, "max": 0.5, "step": 0.05,
|
|
"tooltip": "Extra weight for output layers"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"model_C": ("MODEL",),
|
|
"model_D": ("MODEL",),
|
|
"enable_block_analysis": ("BOOLEAN", {"default": True}),
|
|
"memory_efficient": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "STRING")
|
|
RETURN_NAMES = ("merged_model", "selection_report")
|
|
FUNCTION = "select_best_tensors"
|
|
CATEGORY = "Tensor_Prism/Advanced"
|
|
|
|
def __init__(self):
|
|
self.device = comfy.model_management.get_torch_device()
|
|
|
|
def _to_device(self, tensor: torch.Tensor) -> torch.Tensor:
|
|
"""Safely move tensor to working device."""
|
|
if tensor.device != self.device:
|
|
return tensor.to(self.device)
|
|
return tensor
|
|
|
|
def _evaluate_tensor_quality(self, tensor: torch.Tensor,
|
|
magnitude_w: float, variance_w: float,
|
|
spectral_w: float, stability_w: float) -> float:
|
|
"""
|
|
Evaluate the quality of a tensor based on multiple metrics.
|
|
Returns a quality score (higher is better).
|
|
"""
|
|
try:
|
|
# Move to working device
|
|
tensor = self._to_device(tensor)
|
|
|
|
# Magnitude score - strength of the tensor
|
|
magnitude = torch.norm(tensor).item()
|
|
magnitude_score = np.log1p(magnitude) # Log scale to handle large values
|
|
|
|
# Variance score - detail preservation
|
|
variance = torch.var(tensor).item()
|
|
variance_score = np.log1p(variance)
|
|
|
|
# Spectral score - frequency content richness
|
|
if tensor.numel() > 16 and tensor.dim() >= 2:
|
|
try:
|
|
flat = tensor.flatten()[:1024] # Limit for efficiency
|
|
fft = torch.fft.fft(flat.float())
|
|
spectral_energy = torch.abs(fft).mean().item()
|
|
spectral_score = np.log1p(spectral_energy)
|
|
except:
|
|
spectral_score = 0.0
|
|
else:
|
|
spectral_score = 0.0
|
|
|
|
# Stability score - inverse of numerical extremes
|
|
max_val = torch.abs(tensor).max().item()
|
|
if max_val > 1e-8:
|
|
stability_score = 1.0 / (1.0 + np.log1p(max_val / (torch.abs(tensor).mean().item() + 1e-8)))
|
|
else:
|
|
stability_score = 1.0
|
|
|
|
# Weighted combination
|
|
total_score = (
|
|
magnitude_score * magnitude_w +
|
|
variance_score * variance_w +
|
|
spectral_score * spectral_w +
|
|
stability_score * stability_w
|
|
)
|
|
|
|
return total_score
|
|
|
|
except Exception as e:
|
|
print(f"Warning: Tensor evaluation failed: {e}")
|
|
return 0.0
|
|
|
|
def _get_layer_type(self, key: str) -> str:
|
|
"""Determine the type of layer from its key."""
|
|
key_lower = key.lower()
|
|
if any(term in key_lower for term in ['attn', 'attention']):
|
|
return 'attention'
|
|
elif any(term in key_lower for term in ['output_blocks', 'out.']):
|
|
return 'output'
|
|
elif any(term in key_lower for term in ['input_blocks']):
|
|
return 'input'
|
|
elif any(term in key_lower for term in ['middle_block']):
|
|
return 'middle'
|
|
else:
|
|
return 'other'
|
|
|
|
def _apply_layer_bias(self, base_score: float, layer_type: str,
|
|
attention_pref: float, output_pref: float) -> float:
|
|
"""Apply layer-specific bias to the quality score."""
|
|
if layer_type == 'attention':
|
|
return base_score * (1.0 + attention_pref)
|
|
elif layer_type == 'output':
|
|
return base_score * (1.0 + output_pref)
|
|
return base_score
|
|
|
|
def _select_best_model(self, models: List[Tuple[str, Dict]], key: str,
|
|
weights: Dict, biases: Dict) -> Tuple[str, torch.Tensor]:
|
|
"""
|
|
Select the best tensor from available models for a specific key.
|
|
Returns (model_name, selected_tensor).
|
|
"""
|
|
if len(models) == 0:
|
|
raise ValueError("No models provided")
|
|
|
|
# Find models that have this key
|
|
candidates = [(name, sd[key]) for name, sd in models if key in sd]
|
|
|
|
if len(candidates) == 0:
|
|
# Return from first model if key not in others
|
|
return models[0][0], models[0][1][key]
|
|
|
|
if len(candidates) == 1:
|
|
return candidates[0][0], candidates[0][1]
|
|
|
|
# Evaluate each candidate
|
|
scores = []
|
|
layer_type = self._get_layer_type(key)
|
|
|
|
for name, tensor in candidates:
|
|
# Base quality score
|
|
quality = self._evaluate_tensor_quality(
|
|
tensor,
|
|
weights['magnitude'],
|
|
weights['variance'],
|
|
weights['spectral'],
|
|
weights['stability']
|
|
)
|
|
|
|
# Apply layer-specific bias
|
|
quality = self._apply_layer_bias(
|
|
quality, layer_type,
|
|
biases['attention'],
|
|
biases['output']
|
|
)
|
|
|
|
# Apply model bias
|
|
if name == 'A':
|
|
quality *= (1.0 + biases['model_A'])
|
|
elif name == 'B':
|
|
quality *= (1.0 - biases['model_A'])
|
|
|
|
scores.append((name, tensor, quality))
|
|
|
|
# Select best
|
|
best = max(scores, key=lambda x: x[2])
|
|
return best[0], best[1]
|
|
|
|
def select_best_tensors(self, model_A, model_B, selection_mode: str,
|
|
magnitude_weight: float, variance_weight: float,
|
|
spectral_weight: float, stability_weight: float,
|
|
model_A_bias: float, attention_preference: float,
|
|
output_preference: float,
|
|
model_C=None, model_D=None,
|
|
enable_block_analysis: bool = True,
|
|
memory_efficient: bool = True):
|
|
"""
|
|
Main selection function - evaluates and selects best tensors.
|
|
"""
|
|
|
|
print("\n" + "="*70)
|
|
print("🎯 INTELLIGENT TENSOR SELECTOR (Tensor Prism)")
|
|
print("="*70)
|
|
|
|
# Prepare weights based on mode
|
|
if selection_mode == "best_quality":
|
|
pass # Use provided weights
|
|
elif selection_mode == "balanced":
|
|
magnitude_weight = variance_weight = spectral_weight = stability_weight = 0.25
|
|
elif selection_mode == "conservative":
|
|
magnitude_weight *= 0.8
|
|
stability_weight *= 1.5
|
|
elif selection_mode == "aggressive":
|
|
magnitude_weight *= 1.3
|
|
variance_weight *= 1.2
|
|
|
|
weights = {
|
|
'magnitude': magnitude_weight,
|
|
'variance': variance_weight,
|
|
'spectral': spectral_weight,
|
|
'stability': stability_weight
|
|
}
|
|
|
|
biases = {
|
|
'model_A': model_A_bias,
|
|
'attention': attention_preference,
|
|
'output': output_preference
|
|
}
|
|
|
|
print(f"\n⚙️ Configuration:")
|
|
print(f" Selection Mode: {selection_mode}")
|
|
print(f" Weights - Mag: {magnitude_weight:.2f}, Var: {variance_weight:.2f}, "
|
|
f"Spec: {spectral_weight:.2f}, Stab: {stability_weight:.2f}")
|
|
print(f" Model A Bias: {model_A_bias:+.2f}")
|
|
|
|
# Collect models - keep state dicts in CPU to avoid memory issues
|
|
models = [('A', model_A.model.state_dict()), ('B', model_B.model.state_dict())]
|
|
if model_C is not None:
|
|
models.append(('C', model_C.model.state_dict()))
|
|
if model_D is not None:
|
|
models.append(('D', model_D.model.state_dict()))
|
|
|
|
print(f" Evaluating {len(models)} models")
|
|
|
|
# Get all keys from first model
|
|
base_state_dict = models[0][1]
|
|
all_keys = list(base_state_dict.keys())
|
|
|
|
print(f"\n🔍 Processing {len(all_keys)} tensors...")
|
|
|
|
# Selection tracking
|
|
selection_stats = {name: 0 for name, _ in models}
|
|
layer_stats = {
|
|
'attention': {name: 0 for name, _ in models},
|
|
'output': {name: 0 for name, _ in models},
|
|
'input': {name: 0 for name, _ in models},
|
|
'middle': {name: 0 for name, _ in models},
|
|
'other': {name: 0 for name, _ in models}
|
|
}
|
|
|
|
# Create patches
|
|
patches = {}
|
|
processed = 0
|
|
|
|
for key in all_keys:
|
|
try:
|
|
# Select best tensor
|
|
selected_model, selected_tensor = self._select_best_model(
|
|
models, key, weights, biases
|
|
)
|
|
|
|
# Track selection
|
|
selection_stats[selected_model] += 1
|
|
layer_type = self._get_layer_type(key)
|
|
layer_stats[layer_type][selected_model] += 1
|
|
|
|
# Create patch if not from model A
|
|
if selected_model != 'A':
|
|
original = base_state_dict[key]
|
|
# Ensure tensors are on same device before subtraction
|
|
selected_tensor = self._to_device(selected_tensor)
|
|
original = self._to_device(original)
|
|
diff = selected_tensor - original
|
|
|
|
if torch.abs(diff).max() > 1e-8:
|
|
# Move diff back to CPU for patches
|
|
patches[key] = (diff.cpu(),)
|
|
|
|
processed += 1
|
|
|
|
# Memory cleanup
|
|
if memory_efficient and processed % 100 == 0:
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ Warning: Failed to process {key}: {e}")
|
|
continue
|
|
|
|
# Apply patches
|
|
print(f"\n📦 Applying {len(patches)} patches...")
|
|
merged_model = model_A.clone()
|
|
if patches:
|
|
merged_model.add_patches(patches, 1.0)
|
|
|
|
# Generate report
|
|
report = self._generate_report(selection_stats, layer_stats, len(all_keys),
|
|
selection_mode, weights, biases)
|
|
|
|
print("\n✅ Selection complete!")
|
|
print("="*70 + "\n")
|
|
|
|
return (merged_model, report)
|
|
|
|
def _generate_report(self, selection_stats: Dict, layer_stats: Dict,
|
|
total_tensors: int, mode: str, weights: Dict,
|
|
biases: Dict) -> str:
|
|
"""Generate detailed selection report."""
|
|
|
|
report = "="*60 + "\n"
|
|
report += "INTELLIGENT TENSOR SELECTION REPORT\n"
|
|
report += "="*60 + "\n\n"
|
|
|
|
report += f"Selection Mode: {mode}\n"
|
|
report += f"Total Tensors Processed: {total_tensors}\n\n"
|
|
|
|
report += "OVERALL SELECTION STATISTICS:\n"
|
|
report += "-"*40 + "\n"
|
|
for model, count in sorted(selection_stats.items()):
|
|
percentage = (count / total_tensors * 100) if total_tensors > 0 else 0
|
|
report += f"Model {model}: {count} tensors ({percentage:.1f}%)\n"
|
|
|
|
report += "\n\nLAYER-SPECIFIC BREAKDOWN:\n"
|
|
report += "-"*40 + "\n"
|
|
for layer_type, stats in layer_stats.items():
|
|
layer_total = sum(stats.values())
|
|
if layer_total > 0:
|
|
report += f"\n{layer_type.upper()}:\n"
|
|
for model, count in sorted(stats.items()):
|
|
percentage = (count / layer_total * 100) if layer_total > 0 else 0
|
|
report += f" Model {model}: {count} ({percentage:.1f}%)\n"
|
|
|
|
report += "\n\nQUALITY METRIC WEIGHTS:\n"
|
|
report += "-"*40 + "\n"
|
|
report += f"Magnitude: {weights['magnitude']:.2f}\n"
|
|
report += f"Variance: {weights['variance']:.2f}\n"
|
|
report += f"Spectral: {weights['spectral']:.2f}\n"
|
|
report += f"Stability: {weights['stability']:.2f}\n"
|
|
|
|
report += "\n\nBIAS SETTINGS:\n"
|
|
report += "-"*40 + "\n"
|
|
report += f"Model A Bias: {biases['model_A']:+.2f}\n"
|
|
report += f"Attention Preference: {biases['attention']:+.2f}\n"
|
|
report += f"Output Preference: {biases['output']:+.2f}\n"
|
|
|
|
report += "\n" + "="*60 + "\n"
|
|
|
|
return report
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"TensorPrism_IntelligentTensorSelector": TensorPrism_IntelligentTensorSelector
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"TensorPrism_IntelligentTensorSelector": "Intelligent Tensor Selector (Tensor Prism)"
|
|
} |