346 lines
13 KiB
Python
346 lines
13 KiB
Python
import torch
|
|
import numpy as np
|
|
import re
|
|
|
|
class TensorPrism_ModelMaskGenerator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mask_type": (["layer_based", "block_based", "attention_only", "feedforward_only", "custom_pattern", "random_sparse", "depth_gradient"],),
|
|
"intensity": ("FLOAT", {
|
|
"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05,
|
|
"help": "Overall mask intensity"
|
|
}),
|
|
"reference_model": ("MODEL", {
|
|
"help": "Model to analyze structure from"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"layer_start": ("INT", {
|
|
"default": 0, "min": 0, "max": 50, "step": 1,
|
|
"help": "Starting layer for range-based masks"
|
|
}),
|
|
"layer_end": ("INT", {
|
|
"default": -1, "min": -1, "max": 50, "step": 1,
|
|
"help": "Ending layer (-1 for all)"
|
|
}),
|
|
"gradient_direction": (["shallow_to_deep", "deep_to_shallow", "center_out", "edges_in"],),
|
|
"sparsity": ("FLOAT", {
|
|
"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
|
|
"help": "Sparsity for random masks (0=dense, 1=sparse)"
|
|
}),
|
|
"custom_pattern": ("STRING", {
|
|
"default": "attn,mlp.fc1",
|
|
"help": "Comma-separated layer name patterns"
|
|
}),
|
|
"falloff": ("FLOAT", {
|
|
"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.05,
|
|
"help": "Gradient falloff smoothness"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
RETURN_NAMES = ("mask",)
|
|
FUNCTION = "generate_mask"
|
|
CATEGORY = "Tensor Prism/Mask"
|
|
|
|
def generate_mask(self, mask_type, intensity, reference_model,
|
|
layer_start=0, layer_end=-1, gradient_direction="shallow_to_deep",
|
|
sparsity=0.5, custom_pattern="attn,mlp", falloff=0.1):
|
|
"""
|
|
Generate masks that work with actual model layer structure for merging.
|
|
"""
|
|
|
|
# Get model structure
|
|
state_dict = reference_model.model.state_dict()
|
|
layer_info = self._analyze_model_structure(state_dict)
|
|
|
|
# Generate mask based on type
|
|
if mask_type == "layer_based":
|
|
mask = self._create_layer_range_mask(layer_info, layer_start, layer_end, intensity)
|
|
|
|
elif mask_type == "block_based":
|
|
mask = self._create_block_mask(layer_info, layer_start, layer_end, intensity)
|
|
|
|
elif mask_type == "attention_only":
|
|
mask = self._create_component_mask(layer_info, ["attn", "attention", "self_attn"], intensity)
|
|
|
|
elif mask_type == "feedforward_only":
|
|
mask = self._create_component_mask(layer_info, ["mlp", "fc", "feedforward", "ffn"], intensity)
|
|
|
|
elif mask_type == "custom_pattern":
|
|
patterns = [p.strip() for p in custom_pattern.split(",")]
|
|
mask = self._create_pattern_mask(layer_info, patterns, intensity)
|
|
|
|
elif mask_type == "random_sparse":
|
|
mask = self._create_random_mask(layer_info, sparsity, intensity)
|
|
|
|
elif mask_type == "depth_gradient":
|
|
mask = self._create_depth_gradient_mask(layer_info, gradient_direction, intensity, falloff)
|
|
|
|
else:
|
|
mask = self._create_uniform_mask(layer_info, intensity)
|
|
|
|
# Package as MASK type
|
|
mask = {
|
|
"mask_dict": mask,
|
|
"mask_type": mask_type,
|
|
"intensity": intensity,
|
|
"layer_info": layer_info
|
|
}
|
|
|
|
return (mask,)
|
|
|
|
def _analyze_model_structure(self, state_dict):
|
|
"""Analyze model structure to understand layers and components"""
|
|
layer_info = {
|
|
"layers": {},
|
|
"total_layers": 0,
|
|
"layer_names": list(state_dict.keys())
|
|
}
|
|
|
|
# Pattern matching for different model architectures
|
|
layer_patterns = [
|
|
r"layers\.(\d+)\.", # Transformer layers
|
|
r"blocks\.(\d+)\.", # Vision transformer blocks
|
|
r"h\.(\d+)\.", # GPT-style layers
|
|
r"layer\.(\d+)\.", # BERT-style layers
|
|
r"encoder\.layer\.(\d+)\.", # BERT encoder
|
|
r"decoder\.layer\.(\d+)\.", # BERT decoder
|
|
]
|
|
|
|
for name in state_dict.keys():
|
|
layer_num = None
|
|
|
|
# Try to extract layer number
|
|
for pattern in layer_patterns:
|
|
match = re.search(pattern, name)
|
|
if match:
|
|
layer_num = int(match.group(1))
|
|
break
|
|
|
|
if layer_num is not None:
|
|
if layer_num not in layer_info["layers"]:
|
|
layer_info["layers"][layer_num] = []
|
|
layer_info["layers"][layer_num].append(name)
|
|
layer_info["total_layers"] = max(layer_info["total_layers"], layer_num + 1)
|
|
|
|
return layer_info
|
|
|
|
def _create_layer_range_mask(self, layer_info, start, end, intensity):
|
|
"""Create mask for specific layer range"""
|
|
mask = {}
|
|
|
|
if end == -1:
|
|
end = layer_info["total_layers"]
|
|
|
|
for name in layer_info["layer_names"]:
|
|
# Find which layer this parameter belongs to
|
|
layer_num = self._get_layer_number(name)
|
|
|
|
if layer_num is not None and start <= layer_num < end:
|
|
mask[name] = intensity
|
|
else:
|
|
mask[name] = 0.0
|
|
|
|
return mask
|
|
|
|
def _create_block_mask(self, layer_info, start, end, intensity):
|
|
"""Create mask for transformer blocks with smooth transitions"""
|
|
mask = {}
|
|
|
|
if end == -1:
|
|
end = layer_info["total_layers"]
|
|
|
|
total_range = max(end - start, 1)
|
|
|
|
for name in layer_info["layer_names"]:
|
|
layer_num = self._get_layer_number(name)
|
|
|
|
if layer_num is not None and start <= layer_num < end:
|
|
# Create smooth gradient within the range
|
|
progress = (layer_num - start) / total_range
|
|
mask_value = intensity * (0.5 + 0.5 * np.cos(progress * np.pi))
|
|
mask[name] = mask_value
|
|
else:
|
|
mask[name] = 0.0
|
|
|
|
return mask
|
|
|
|
def _create_component_mask(self, layer_info, component_patterns, intensity):
|
|
"""Create mask for specific components (attention, MLP, etc.)"""
|
|
mask = {}
|
|
|
|
for name in layer_info["layer_names"]:
|
|
should_mask = False
|
|
|
|
for pattern in component_patterns:
|
|
if pattern.lower() in name.lower():
|
|
should_mask = True
|
|
break
|
|
|
|
mask[name] = intensity if should_mask else 0.0
|
|
|
|
return mask
|
|
|
|
def _create_pattern_mask(self, layer_info, patterns, intensity):
|
|
"""Create mask based on custom patterns"""
|
|
mask = {}
|
|
|
|
for name in layer_info["layer_names"]:
|
|
should_mask = False
|
|
|
|
for pattern in patterns:
|
|
if pattern.lower() in name.lower():
|
|
should_mask = True
|
|
break
|
|
|
|
mask[name] = intensity if should_mask else 0.0
|
|
|
|
return mask
|
|
|
|
def _create_random_mask(self, layer_info, sparsity, intensity):
|
|
"""Create random sparse mask"""
|
|
mask = {}
|
|
np.random.seed(42) # For reproducibility
|
|
|
|
for name in layer_info["layer_names"]:
|
|
if np.random.random() > sparsity:
|
|
mask[name] = intensity
|
|
else:
|
|
mask[name] = 0.0
|
|
|
|
return mask
|
|
|
|
def _create_depth_gradient_mask(self, layer_info, direction, intensity, falloff):
|
|
"""Create gradient mask based on model depth"""
|
|
mask = {}
|
|
total_layers = max(layer_info["total_layers"], 1)
|
|
|
|
for name in layer_info["layer_names"]:
|
|
layer_num = self._get_layer_number(name)
|
|
|
|
if layer_num is not None:
|
|
# Calculate position (0 to 1)
|
|
position = layer_num / (total_layers - 1) if total_layers > 1 else 0.5
|
|
|
|
if direction == "shallow_to_deep":
|
|
mask_value = position
|
|
elif direction == "deep_to_shallow":
|
|
mask_value = 1.0 - position
|
|
elif direction == "center_out":
|
|
mask_value = 1.0 - 2.0 * abs(position - 0.5)
|
|
else: # edges_in
|
|
mask_value = 2.0 * abs(position - 0.5)
|
|
|
|
# Apply falloff
|
|
if falloff > 0:
|
|
mask_value = np.power(mask_value, 1.0 / max(falloff, 0.01))
|
|
|
|
mask[name] = mask_value * intensity
|
|
else:
|
|
# For parameters not in layers (embeddings, etc.)
|
|
mask[name] = intensity * 0.1
|
|
|
|
return mask
|
|
|
|
def _create_uniform_mask(self, layer_info, intensity):
|
|
"""Create uniform mask for all parameters"""
|
|
mask = {}
|
|
for name in layer_info["layer_names"]:
|
|
mask[name] = intensity
|
|
return mask
|
|
|
|
def _get_layer_number(self, param_name):
|
|
"""Extract layer number from parameter name"""
|
|
layer_patterns = [
|
|
r"layers\.(\d+)\.",
|
|
r"blocks\.(\d+)\.",
|
|
r"h\.(\d+)\.",
|
|
r"layer\.(\d+)\.",
|
|
r"encoder\.layer\.(\d+)\.",
|
|
r"decoder\.layer\.(\d+)\.",
|
|
]
|
|
|
|
for pattern in layer_patterns:
|
|
match = re.search(pattern, param_name)
|
|
if match:
|
|
return int(match.group(1))
|
|
|
|
return None
|
|
|
|
|
|
# Updated merge node to work with MASK
|
|
class TensorPrism_WeightedMaskMerge:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model_A": ("MODEL",),
|
|
"model_B": ("MODEL",),
|
|
"mask": ("MASK",),
|
|
"merge_ratio": ("FLOAT", {
|
|
"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01,
|
|
"help": "Global multiplier for mask intensity"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("merged_model",)
|
|
FUNCTION = "merge_models"
|
|
CATEGORY = "Tensor Prism/Mask"
|
|
|
|
def merge_models(self, model_A, model_B, mask, merge_ratio):
|
|
"""
|
|
Merge models using structure-aware masks
|
|
"""
|
|
|
|
# Clone model A as base
|
|
merged_model = model_A.clone()
|
|
|
|
# Get state dicts
|
|
state_dict_A = model_A.model.state_dict()
|
|
state_dict_B = model_B.model.state_dict()
|
|
|
|
# Get mask dictionary
|
|
mask_dict = mask["mask_dict"]
|
|
|
|
# Create patches
|
|
patches = {}
|
|
|
|
for key in state_dict_A.keys():
|
|
if key in state_dict_B and key in mask_dict:
|
|
weight_A = state_dict_A[key]
|
|
weight_B = state_dict_B[key]
|
|
|
|
if isinstance(weight_A, torch.Tensor) and isinstance(weight_B, torch.Tensor):
|
|
if weight_A.shape == weight_B.shape:
|
|
# Get mask value for this parameter
|
|
mask_value = mask_dict[key] * merge_ratio
|
|
|
|
# Calculate merged weights
|
|
merged_weight = weight_A * (1 - mask_value) + weight_B * mask_value
|
|
|
|
# Store as patch
|
|
if mask_value > 0:
|
|
patches[key] = (merged_weight - weight_A,)
|
|
|
|
# Apply patches
|
|
if patches:
|
|
merged_model.add_patches(patches, 1.0)
|
|
|
|
return (merged_model,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"TensorPrism_ModelMaskGenerator": TensorPrism_ModelMaskGenerator,
|
|
"TensorPrism_WeightedMaskMerge": TensorPrism_WeightedMaskMerge
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"TensorPrism_ModelMaskGenerator": "Model Mask Generator (Tensor Prism)",
|
|
"TensorPrism_WeightedMaskMerge": "Weighted Mask Merge (Tensor Prism)"
|
|
} |