Creative Update
@@ -1,58 +1,136 @@
|
||||
## 🛠️ Easy LoRA Merger: Complexity Made Simple
|
||||
# 🛠️ Easy LoRA Merger
|
||||
*Complexity Made Simple*
|
||||
|
||||
Models are evolving fast—between Z-Image, Klein 4B/9B, Flux 2 Dev, and SDXL, it's a lot to keep track of. I built this suite of nodes because I wanted a way to merge these different architectures without worrying about tensor dimensions or scaling math.
|
||||
Models are evolving fast — between Z-Image, Klein 4B/9B, Flux 2 Dev, and SDXL, it's a lot to keep track of.
|
||||
This suite of nodes lets you merge different architectures without worrying about tensor dimensions, sparsity mismatches, or scaling math.
|
||||
|
||||
The goal isn't to be the most "complex" tool, but the most helpful. It does the heavy lifting in the background so you can just focus on the results.
|
||||
The goal isn’t to be the most complex tool — it’s to be the most helpful.
|
||||
|
||||
---
|
||||
|
||||
## 🧪 Proof in the Results: Identity Blending
|
||||
# 🔍 Example Workflow
|
||||
|
||||
I tested this by merging characters with very different training backgrounds. By balancing the weights, you can achieve perfect hybrids that keep the identity of both. (workflow included)
|
||||
This is the basic "Identity Hybrid" setup used in the examples below.
|
||||
|
||||
✅ Z-Image (Turbo + Base Mix)
|
||||
✅ Flux Klein 9B Character Blends
|
||||
✅ SDXL Style + Character Mixes
|
||||

|
||||
|
||||
> Drag the workflow image into ComfyUI to load it instantly.
|
||||
|
||||
---
|
||||
|
||||
## 📦 The Toolbox (Node List)
|
||||
# 🧠 Main Node Interface
|
||||
|
||||
- **Easy LoRA Merger**: The main engine. Connect your LoRAs, hit merge, and preview instantly.
|
||||
- **Easy LoRA-Only Merger**: Perfect for chaining multiple merges or keeping your workspace clean.
|
||||
- **🎨 Triple Merger (Experimental)**: Mix a character, an outfit, and an art style in one node.
|
||||
- **🔄 Musubi LoRA Converter**: A specialized bridge for LoRAs trained in the Musubi tuner environment.
|
||||
- **🔄 Z-Image Normalizer**: Fixes weight balancing when mixing Z-Image Turbo and Base models.
|
||||
- **🔥 Base Model Baker (Experimental)**: Bakes your merge directly into the model for instant testing.
|
||||
The core node with blend modes, diagnostics, and merge methods.
|
||||
|
||||

|
||||
|
||||
Includes:
|
||||
- Blend mode selection
|
||||
- Advanced merge methods
|
||||
- Fun Mode toggles
|
||||
- Layer-by-layer diagnostics
|
||||
- Compatibility checks
|
||||
|
||||
---
|
||||
|
||||
## 🚧 Status: Beta - Help Wanted!
|
||||
# 🧪 Identity Blending Example
|
||||
|
||||
This project works great for my workflows, but I need your help to push it further:
|
||||
### Source LoRAs
|
||||
|
||||
- **Flux 2 Dev Testing**: Feedback on newest Flux models
|
||||
- **High-VRAM Users (24GB+)**: Try baking 9B models and report performance
|
||||
- **SDXL Users**: Test style + character mixes
|
||||
| LoRA A | LoRA B |
|
||||
|--------|--------|
|
||||
|  |  |
|
||||
|
||||
### Merged Result
|
||||
|
||||

|
||||
|
||||
By carefully balancing weights, both identities remain visible without overpowering each other.
|
||||
|
||||
Tested With:
|
||||
- ✅ Z-Image (Turbo + Base Mix)
|
||||
- ✅ Flux Klein 9B Character Blends
|
||||
- ✅ SDXL Style + Character Mixes
|
||||
|
||||
---
|
||||
|
||||
## 📝 Smart Diagnostics
|
||||
# 🎮 Blend Modes
|
||||
|
||||
The console status window shows:
|
||||
|
||||
✅ Green checkmarks when math is aligned
|
||||
📊 Layer-by-layer statistics
|
||||
⚠️ Clear warnings if something needs attention
|
||||
| Mode | Description |
|
||||
|------|-------------|
|
||||
| **auto** | Smart selection based on sparsity patterns |
|
||||
| **balanced** | Preserves unique activation structure (great for mixed trainers) |
|
||||
| **dense** | Traditional dense tensor blending |
|
||||
| **fun_mode** | Experimental math for creative effects |
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Get Started
|
||||
## 🎲 Fun Mode Effects by Method
|
||||
|
||||
- **Download**: https://github.com/Terpentinas/EasyLoRAMerger
|
||||
- **Workflow**: Drag the attached image into ComfyUI to see the Identity Hybrid workflow
|
||||
- **Questions?**: Open an issue on GitHub or comment below!
|
||||
| Method | Fun Mode Effect |
|
||||
|--------|-----------------|
|
||||
| magnitude | "chaos mode" – Random weighted selection |
|
||||
| subtract | "glitch mode" – Random add/subtract behavior |
|
||||
| dare_lite | "gambler mode" – Random amplification |
|
||||
| ties_strict | "disagreement mode" – Keeps opposing tensor signs |
|
||||
| ties_gentle | "drama mode" – Amplifies disagreements |
|
||||
| dare_rescale | "lottery mode" – Random jackpot scaling |
|
||||
| svd_preserve | "patchwork mode" – Block-based swapping |
|
||||
| noise_aware | "static mode" – Structured noise injection |
|
||||
| gradient_alignment | "mood swing mode" – Random alignment bias |
|
||||
|
||||
---
|
||||
|
||||
*Built because merging different LoRA formats shouldn't require a PhD in tensor math* 💪
|
||||
# 📦 Node Toolbox
|
||||
|
||||
- **Easy LoRA Merger** — Main engine for merging and previewing
|
||||
- **Easy LoRA-Only Merger** — Clean chaining of multiple merges
|
||||
- **🎨 Triple Merger (Experimental)** — Character + outfit + style in one pass
|
||||
- **🔄 Musubi LoRA Converter** — Compatibility bridge for Musubi-trained LoRAs
|
||||
- **🔄 Z-Image Normalizer** — Fix weight imbalance between Turbo and Base
|
||||
- **🔥 Base Model Baker (Experimental)** — Bake merged LoRA directly into model
|
||||
|
||||
---
|
||||
|
||||
# 📝 Smart Diagnostics
|
||||
|
||||

|
||||
|
||||
Console output includes:
|
||||
|
||||
- ✅ Alignment verification
|
||||
- 📊 Layer-by-layer statistics
|
||||
- ⚠️ Clear warnings for mismatches
|
||||
- 🔍 Sparsity + scaling reports
|
||||
|
||||
---
|
||||
|
||||
# 🧪 Resources Used in Examples
|
||||
|
||||
These are the LoRAs used in the demonstration images above.
|
||||
|
||||
- **Jinx Arcane (Z-Image Turbo)** by Tekemo — [https://civitai.com/models/2198444/jinx-arcane-z-image-turbo-lora?modelVersionId=2475322]
|
||||
- **Izzy AI Character (Flux 2)** by berts_eyebrow — [https://civitai.com/models/2194957/izzy-ai-character-flux2-9b-qwen-2512-z-image-xl-wan?modelVersionId=2650293]
|
||||
|
||||
---
|
||||
|
||||
# 🚧 Beta Status – Feedback Wanted
|
||||
|
||||
Looking for:
|
||||
|
||||
- **Flux 2 Dev testers**
|
||||
- **24GB+ VRAM users** to test baking 9B models
|
||||
- **SDXL users** to try style + character merges
|
||||
|
||||
---
|
||||
|
||||
# 🚀 Get Started
|
||||
|
||||
- Download: https://github.com/Terpentinas/EasyLoRAMerger
|
||||
- Load workflow into ComfyUI
|
||||
- Experiment with blend modes
|
||||
- Open an issue for feedback
|
||||
|
||||
---
|
||||
|
||||
*Built because merging different LoRA formats shouldn’t require a PhD in tensor math.* 💪
|
||||
|
After Width: | Height: | Size: 2.2 MiB |
|
After Width: | Height: | Size: 1.1 MiB |
|
After Width: | Height: | Size: 1.2 MiB |
|
After Width: | Height: | Size: 30 KiB |
@@ -73,6 +73,7 @@ class MergeConfig:
|
||||
uniqueness: float = 0.7 # For feature_mix
|
||||
threshold: float = 0.0 # For subtract
|
||||
blend: float = 0.5 # For magnitude
|
||||
blend_mode: str = "auto" # "auto", "balanced", "dense", "fun_mode"
|
||||
|
||||
@classmethod
|
||||
def from_inputs(cls, **kwargs) -> 'MergeConfig':
|
||||
@@ -1296,13 +1297,39 @@ class StreamingMergeEngine:
|
||||
|
||||
# --- END DEBUG INFO ---
|
||||
|
||||
# After scaling and rank adjustment, before creating method_args
|
||||
if self.config.blend_mode == "auto":
|
||||
# Calculate sparsity (percentage of zeros)
|
||||
sparsity_a = (t_a == 0).float().mean().item()
|
||||
sparsity_b = (t_b == 0).float().mean().item()
|
||||
|
||||
print(f" Sparsity: A={1-sparsity_a:.1%} active, B={1-sparsity_b:.1%} active")
|
||||
|
||||
# If sparsity differs significantly, use active blending
|
||||
if abs(sparsity_a - sparsity_b) > 0.2: # 20% difference threshold
|
||||
active_mode = "active"
|
||||
#print(f" 🔄 Auto blend: detected different sparsity patterns ({sparsity_a:.1%} vs {sparsity_b:.1%})")
|
||||
else:
|
||||
active_mode = "standard"
|
||||
#print(f" 🔄 Auto blend: similar sparsity patterns ({sparsity_a:.1%} vs {sparsity_b:.1%})")
|
||||
else:
|
||||
# Map user selection to internal mode
|
||||
mode_map = {
|
||||
"balanced": "active",
|
||||
"dense": "standard",
|
||||
"fun_mode": "crazy_mode"
|
||||
}
|
||||
active_mode = mode_map.get(self.config.blend_mode, "active")
|
||||
|
||||
# Get merge method
|
||||
merge_func = MergeMethodRegistry.get_method(self.config.method)
|
||||
|
||||
# Prepare arguments with SCALED tensors
|
||||
method_args = {
|
||||
'a': t_a, 'b': t_b,
|
||||
'wa': self.config.weight_a, 'wb': self.config.weight_b
|
||||
'wa': self.config.weight_a,
|
||||
'wb': self.config.weight_b,
|
||||
'blend_mode': active_mode # ← Use the variable
|
||||
}
|
||||
|
||||
# Add method-specific parameters from self.config
|
||||
@@ -1685,6 +1712,10 @@ class EasyLoRAmergerNode:
|
||||
"tooltip": "Tensors per batch (lower = less VRAM)"}),
|
||||
"streaming": ("BOOLEAN", {"default": True,
|
||||
"tooltip": "Stream tensors to save VRAM"}),
|
||||
"blend_mode": (["auto", "balanced", "dense", "fun_mode"], {
|
||||
"default": "auto",
|
||||
"tooltip": "auto: Smart choice based on LoRAs | balanced: Preserve patterns | dense: Traditional mix | fun_mode: math a bit incorrect"
|
||||
}),
|
||||
# METHOD INFO DISPLAY (read-only)
|
||||
"method_info": ("STRING", {
|
||||
"default": "Select a method to see details here...",
|
||||
@@ -1704,7 +1735,7 @@ class EasyLoRAmergerNode:
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
"""Cache implementation for ComfyUI change detection."""
|
||||
# Create cache key from relevant inputs
|
||||
# Create cache key from relevant inputs - MUST match merge() parameters
|
||||
cache_key = (
|
||||
kwargs.get('lora_a', 'None'),
|
||||
kwargs.get('lora_b', 'None'),
|
||||
@@ -1714,6 +1745,14 @@ class EasyLoRAmergerNode:
|
||||
kwargs.get('density', 1.0),
|
||||
kwargs.get('save_trigger', False),
|
||||
kwargs.get('streaming', True),
|
||||
kwargs.get('uniqueness', 0.7), # Used by merge_feature_mix
|
||||
kwargs.get('threshold', 0.0), # Used by merge_subtract, merge_ties_gentle
|
||||
kwargs.get('blend', 0.5), # Used by merge_magnitude
|
||||
kwargs.get('blend_mode', 'auto'), # Affects all merges
|
||||
kwargs.get('precision', 'auto'), # Affects computation
|
||||
kwargs.get('device', 'auto'), # Affects computation
|
||||
kwargs.get('batch_size', 32), # Affects streaming
|
||||
kwargs.get('metadata_mode', 'merge_basic'), # Affects output metadata
|
||||
)
|
||||
|
||||
# Generate hash
|
||||
@@ -1735,6 +1774,7 @@ class EasyLoRAmergerNode:
|
||||
save_trigger=False, save_folder="", filename="merged_lora",
|
||||
device="auto", precision="auto", metadata_mode="merge_basic",
|
||||
batch_size=32, streaming=True,
|
||||
blend_mode="auto",
|
||||
method_info=None):
|
||||
|
||||
print("\n" + "="*50)
|
||||
@@ -1752,6 +1792,7 @@ class EasyLoRAmergerNode:
|
||||
blend=blend,
|
||||
device_type=device,
|
||||
precision=precision,
|
||||
blend_mode=blend_mode,
|
||||
metadata_mode=metadata_mode,
|
||||
batch_size=batch_size,
|
||||
streaming=streaming
|
||||
@@ -1973,6 +2014,10 @@ class EasyLoRAonlyMerger:
|
||||
"tooltip": "Tensors per batch (lower = less VRAM)"}),
|
||||
"streaming": ("BOOLEAN", {"default": True,
|
||||
"tooltip": "Stream tensors to save VRAM"}),
|
||||
"blend_mode": (["auto", "balanced", "dense", "fun_mode"], {
|
||||
"default": "auto",
|
||||
"tooltip": "auto: Smart choice based on LoRAs | balanced: Preserve patterns | dense: Traditional mix | fun_mode: math a bit incorrect"
|
||||
}),
|
||||
# METHOD INFO DISPLAY (read-only)
|
||||
"method_info": ("STRING", {
|
||||
"default": "Select a method to see details here...",
|
||||
@@ -1997,7 +2042,7 @@ class EasyLoRAonlyMerger:
|
||||
uniqueness=0.7, threshold=0.0, blend=0.5,
|
||||
save_trigger=False, save_folder="", filename="merged_lora",
|
||||
device="auto", precision="auto", metadata_mode="merge_basic",
|
||||
batch_size=32, streaming=True,
|
||||
batch_size=32, streaming=True, blend_mode="auto",
|
||||
method_info=None, node_id=None):
|
||||
|
||||
print("\n" + "="*50)
|
||||
@@ -2013,9 +2058,9 @@ class EasyLoRAonlyMerger:
|
||||
uniqueness=uniqueness,
|
||||
threshold=threshold,
|
||||
blend=blend,
|
||||
# REMOVED: attn_weight, mlp_weight, confidence_a, confidence_b
|
||||
device_type=device,
|
||||
precision=precision,
|
||||
blend_mode=blend_mode,
|
||||
metadata_mode=metadata_mode,
|
||||
batch_size=batch_size,
|
||||
streaming=streaming
|
||||
@@ -2155,6 +2200,10 @@ class EasyLoRATripleMerger:
|
||||
"device": (["auto", "cuda", "cpu"], {"default": "auto"}),
|
||||
"precision": (["auto", "float32", "bfloat16", "float16"], {"default": "auto"}),
|
||||
"batch_size": ("INT", {"default": 32, "min": 1, "max": 256, "step": 8}),
|
||||
"blend_mode": (["auto", "balanced", "dense", "fun_mode"], {
|
||||
"default": "auto",
|
||||
"tooltip": "auto: Smart choice based on LoRAs | balanced: Preserve patterns | dense: Traditional mix | fun_mode: math a bit incorrect"
|
||||
}),
|
||||
"method_info": ("STRING", {
|
||||
"default": "Select a method to see details...",
|
||||
"multiline": True,
|
||||
@@ -2172,7 +2221,7 @@ class EasyLoRATripleMerger:
|
||||
lora_a="None", lora_b="None", lora_c="None",
|
||||
lora_data_a=None, lora_data_b=None, lora_data_c=None,
|
||||
weight_a=1.0, weight_b=1.0, weight_c=1.0,
|
||||
uniqueness=0.7, threshold=0.0, blend=0.5,
|
||||
uniqueness=0.7, threshold=0.0, blend=0.5, blend_mode="auto",
|
||||
save_trigger=False, save_folder="", filename="triple_merged",
|
||||
device="auto", precision="auto", batch_size=32,
|
||||
method_info=None):
|
||||
@@ -2190,6 +2239,7 @@ class EasyLoRATripleMerger:
|
||||
uniqueness=uniqueness,
|
||||
threshold=threshold,
|
||||
blend=blend,
|
||||
blend_mode=blend_mode,
|
||||
device_type=device,
|
||||
precision=precision,
|
||||
batch_size=batch_size,
|
||||
@@ -2244,7 +2294,7 @@ class EasyLoRATripleMerger:
|
||||
merged_dict = self._merge_triple_method(
|
||||
normalized_sds,
|
||||
[weight_a, weight_b, weight_c],
|
||||
method, density, uniqueness, threshold, blend
|
||||
method, density, uniqueness, threshold, blend, blend_mode=blend_mode
|
||||
)
|
||||
|
||||
# Finalize and save
|
||||
@@ -2270,7 +2320,7 @@ class EasyLoRATripleMerger:
|
||||
|
||||
return (model_lora, clip_lora, str(output_path), lora_tuple)
|
||||
|
||||
def _merge_triple_method(self, sds, weights, method, density, uniqueness, threshold, blend):
|
||||
def _merge_triple_method(self, sds, weights, method, density, uniqueness, threshold, blend, blend_mode="auto"):
|
||||
"""Core triple merge logic with rank adjustment"""
|
||||
all_keys = set()
|
||||
for sd in sds:
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
"""
|
||||
LoRA Merging Methods for EasyLoRAMerger - Version 1.2.0
|
||||
Includes Adaptive Shape Safety for 4B/9B Architecture Mixing
|
||||
LoRA Merging Methods for EasyLoRAMerger - Version 1.3.0
|
||||
Includes Active Region Blending for better sparse LoRA compatibility
|
||||
"""
|
||||
|
||||
import torch
|
||||
@@ -43,98 +43,504 @@ def universal_merge_executor(method_fn, a, b, wa, wb, **kwargs):
|
||||
a_safe, b_safe = ensure_shape_match(a, b)
|
||||
return method_fn(a_safe, b_safe, wa, wb, **kwargs)
|
||||
|
||||
# ==================== CORE METHODS ====================
|
||||
# ==================== ACTIVE REGION BLENDING HELPER ====================
|
||||
|
||||
def merge_linear(a, b, wa, wb, **kwargs):
|
||||
"""Simple weighted average: a*wa + b*wb"""
|
||||
return a * wa + b * wb
|
||||
def apply_active_blending(method_fn, a, b, wa, wb, blend_mode="active", **kwargs):
|
||||
"""
|
||||
Wrapper that applies merge methods only where tensors are active (non-zero)
|
||||
This preserves sparsity patterns and improves cross-trainer blending
|
||||
"""
|
||||
dtype = a.dtype # Store original dtype
|
||||
|
||||
if blend_mode != "active":
|
||||
# Standard mode - apply to all values
|
||||
return method_fn(a, b, wa, wb, **kwargs)
|
||||
|
||||
# Find active regions (non-zero)
|
||||
a_active = (a != 0)
|
||||
b_active = (b != 0)
|
||||
|
||||
both_active = a_active & b_active
|
||||
only_a = a_active & ~b_active
|
||||
only_b = b_active & ~a_active
|
||||
neither = ~a_active & ~b_active
|
||||
|
||||
# Initialize result
|
||||
result = torch.zeros_like(a)
|
||||
|
||||
# Where both are active: apply the merge method
|
||||
if both_active.any():
|
||||
# Ensure the method returns tensors with correct dtype
|
||||
merged_values = method_fn(
|
||||
a[both_active], b[both_active],
|
||||
wa, wb, **kwargs
|
||||
)
|
||||
# Convert to original dtype if needed
|
||||
if merged_values.dtype != dtype:
|
||||
merged_values = merged_values.to(dtype)
|
||||
result[both_active] = merged_values
|
||||
|
||||
# Where only A is active: just take A (weighted)
|
||||
if only_a.any():
|
||||
result[only_a] = a[only_a] * wa
|
||||
|
||||
# Where only B is active: just take B (weighted)
|
||||
if only_b.any():
|
||||
result[only_b] = b[only_b] * wb
|
||||
|
||||
# Where neither is active: already zero
|
||||
|
||||
return result
|
||||
|
||||
# ==================== CORE METHODS (UPDATED WITH ACTIVE BLENDING) ====================
|
||||
|
||||
def merge_linear(a, b, wa, wb, blend_mode="active", **kwargs):
|
||||
"""Simple weighted average with active region preservation"""
|
||||
dtype = a.dtype
|
||||
|
||||
if blend_mode == "active":
|
||||
a_active = (a != 0)
|
||||
b_active = (b != 0)
|
||||
both_active = a_active & b_active
|
||||
only_a = a_active & ~b_active
|
||||
only_b = b_active & ~a_active
|
||||
|
||||
result = torch.zeros_like(a)
|
||||
|
||||
if both_active.any():
|
||||
result[both_active] = a[both_active] * wa + b[both_active] * wb
|
||||
if only_a.any():
|
||||
result[only_a] = a[only_a] * wa
|
||||
if only_b.any():
|
||||
result[only_b] = b[only_b] * wb
|
||||
|
||||
return result
|
||||
|
||||
if blend_mode == "crazy_mode":
|
||||
a_w = a * wa
|
||||
b_w = b * wb
|
||||
|
||||
# Safe cross term - use absolute values for sqrt
|
||||
cross = (a_w.abs() * b_w.abs()).sqrt() * 0.3
|
||||
|
||||
# Preserve signs from dominant tensor
|
||||
a_dominant = a_w.abs() > b_w.abs()
|
||||
sign = torch.where(a_dominant, torch.sign(a_w), torch.sign(b_w))
|
||||
|
||||
result = a_w + b_w + cross * sign
|
||||
return result
|
||||
|
||||
else:
|
||||
return a * wa + b * wb
|
||||
|
||||
# ==================== TIES METHODS ====================
|
||||
|
||||
def merge_ties_strict(a, b, wa, wb, **kwargs):
|
||||
def merge_ties_strict(a, b, wa, wb, blend_mode="active", **kwargs):
|
||||
"""TIES merging: only keep weights where signs agree."""
|
||||
a_weighted = a * wa
|
||||
b_weighted = b * wb
|
||||
sign_agreement = torch.sign(a_weighted) == torch.sign(b_weighted)
|
||||
return (a_weighted + b_weighted) * sign_agreement
|
||||
|
||||
def merge_ties_gentle(a, b, wa, wb, agreement_threshold=0.3, **kwargs):
|
||||
"""Gentle TIES: Only apply when strong disagreement."""
|
||||
a_weighted = a * wa
|
||||
b_weighted = b * wb
|
||||
dtype = a.dtype
|
||||
|
||||
a_flat = a_weighted.flatten().float()
|
||||
b_flat = b_weighted.flatten().float()
|
||||
if blend_mode == "active":
|
||||
a_active = (a != 0)
|
||||
b_active = (b != 0)
|
||||
both_active = a_active & b_active
|
||||
only_a = a_active & ~b_active
|
||||
only_b = b_active & ~a_active
|
||||
|
||||
result = torch.zeros_like(a)
|
||||
|
||||
if both_active.any():
|
||||
a_w = a[both_active] * wa
|
||||
b_w = b[both_active] * wb
|
||||
sign_agreement = (torch.sign(a_w) == torch.sign(b_w)).to(dtype)
|
||||
result[both_active] = (a_w + b_w) * sign_agreement
|
||||
|
||||
if only_a.any():
|
||||
result[only_a] = a[only_a] * wa
|
||||
|
||||
if only_b.any():
|
||||
result[only_b] = b[only_b] * wb
|
||||
|
||||
return result
|
||||
|
||||
norm_prod = (torch.norm(a_flat) * torch.norm(b_flat) + 1e-8)
|
||||
cos_sim = torch.dot(a_flat, b_flat) / norm_prod
|
||||
if blend_mode == "crazy_mode":
|
||||
a_w = a * wa
|
||||
b_w = b * wb
|
||||
|
||||
# Convert to float32 for bitwise operations, then back
|
||||
a_sign = torch.sign(a_w).float()
|
||||
b_sign = torch.sign(b_w).float()
|
||||
|
||||
# Where they disagree (signs different)
|
||||
sign_disagreement = (a_sign != b_sign).to(dtype)
|
||||
|
||||
# Random keep - generate in float32 then convert
|
||||
random_keep = (torch.rand_like(a_w, dtype=torch.float32) > 0.7).to(dtype)
|
||||
|
||||
# Combine using logical_or (works with any dtype)
|
||||
keep_mask = torch.logical_or(sign_disagreement, random_keep).to(dtype)
|
||||
|
||||
result = (a_w + b_w) * keep_mask
|
||||
return result
|
||||
|
||||
if cos_sim < -agreement_threshold:
|
||||
sign_agreement = torch.sign(a_weighted) == torch.sign(b_weighted)
|
||||
else:
|
||||
a_weighted = a * wa
|
||||
b_weighted = b * wb
|
||||
sign_agreement = (torch.sign(a_weighted) == torch.sign(b_weighted)).to(dtype)
|
||||
return (a_weighted + b_weighted) * sign_agreement
|
||||
return a_weighted + b_weighted
|
||||
|
||||
|
||||
def merge_ties_gentle(a, b, wa, wb, agreement_threshold=0.3, blend_mode="active", **kwargs):
|
||||
"""Gentle TIES: Only apply when strong disagreement."""
|
||||
dtype = a.dtype # Store original dtype
|
||||
|
||||
if blend_mode == "active":
|
||||
a_active = (a != 0)
|
||||
b_active = (b != 0)
|
||||
both_active = a_active & b_active
|
||||
only_a = a_active & ~b_active
|
||||
only_b = b_active & ~a_active
|
||||
|
||||
result = torch.zeros_like(a)
|
||||
|
||||
if both_active.any():
|
||||
a_w = a[both_active] * wa
|
||||
b_w = b[both_active] * wb
|
||||
|
||||
# Keep operations in original dtype
|
||||
a_flat = a_w.flatten()
|
||||
b_flat = b_w.flatten()
|
||||
|
||||
norm_prod = (torch.norm(a_flat) * torch.norm(b_flat) + 1e-8)
|
||||
cos_sim = torch.dot(a_flat, b_flat) / norm_prod
|
||||
|
||||
if cos_sim < -agreement_threshold:
|
||||
sign_agreement = (torch.sign(a_w) == torch.sign(b_w)).to(dtype)
|
||||
result[both_active] = (a_w + b_w) * sign_agreement
|
||||
else:
|
||||
result[both_active] = a_w + b_w
|
||||
|
||||
if only_a.any():
|
||||
result[only_a] = a[only_a] * wa
|
||||
|
||||
if only_b.any():
|
||||
result[only_b] = b[only_b] * wb
|
||||
|
||||
return result
|
||||
|
||||
if blend_mode == "crazy_mode":
|
||||
a_w = a * wa
|
||||
b_w = b * wb
|
||||
|
||||
# Amplify disagreements, ignore agreements
|
||||
a_sign = torch.sign(a_w).float()
|
||||
b_sign = torch.sign(b_w).float()
|
||||
|
||||
# Where they disagree, boost the values
|
||||
disagreement = (a_sign != b_sign).to(dtype)
|
||||
agreement = (a_sign == b_sign).to(dtype)
|
||||
|
||||
# Disagreement areas get amplified, agreement areas get muted
|
||||
result = (a_w + b_w) * (1.0 + disagreement) * 0.5
|
||||
return result
|
||||
|
||||
else:
|
||||
a_weighted = a * wa
|
||||
b_weighted = b * wb
|
||||
|
||||
# Convert to float32 only for the cosine similarity calculation
|
||||
a_flat = a_weighted.flatten().float()
|
||||
b_flat = b_weighted.flatten().float()
|
||||
|
||||
norm_prod = (torch.norm(a_flat) * torch.norm(b_flat) + 1e-8)
|
||||
cos_sim = torch.dot(a_flat, b_flat) / norm_prod
|
||||
|
||||
if cos_sim < -agreement_threshold:
|
||||
sign_agreement = (torch.sign(a_weighted) == torch.sign(b_weighted)).to(dtype)
|
||||
return (a_weighted + b_weighted) * sign_agreement
|
||||
return a_weighted + b_weighted
|
||||
|
||||
|
||||
# ==================== DARE METHODS ====================
|
||||
|
||||
def merge_dare_lite(a, b, wa, wb, drop_rate=0.1, **kwargs):
|
||||
def merge_dare_lite(a, b, wa, wb, drop_rate=0.1, blend_mode="active", **kwargs):
|
||||
"""DARE: Random dropout without rescaling."""
|
||||
merged = a * wa + b * wb
|
||||
mask = (torch.rand_like(merged) > drop_rate).float()
|
||||
return merged * mask
|
||||
|
||||
def merge_dare_rescale(a, b, wa, wb, drop_rate=0.1, **kwargs):
|
||||
"""DARE with proper rescaling."""
|
||||
drop_rate = max(0.0, min(0.9, drop_rate))
|
||||
mask_a = (torch.rand_like(a) > drop_rate).float()
|
||||
mask_b = (torch.rand_like(b) > drop_rate).float()
|
||||
rescale = 1.0 / (1.0 - drop_rate) if drop_rate < 1.0 else 1.0
|
||||
return (a * mask_a * rescale) * wa + (b * mask_b * rescale) * wb
|
||||
|
||||
# ==================== NEW INTUITIVE METHODS ====================
|
||||
|
||||
def merge_subtract(a, b, wa, wb, threshold=0.0, **kwargs):
|
||||
"""Subtract B from A - useful for removing unwanted styles."""
|
||||
a_weighted = a * wa
|
||||
b_weighted = b * wb
|
||||
dtype = a.dtype
|
||||
|
||||
if blend_mode == "active":
|
||||
a_active = (a != 0)
|
||||
b_active = (b != 0)
|
||||
both_active = a_active & b_active
|
||||
only_a = a_active & ~b_active
|
||||
only_b = b_active & ~a_active
|
||||
|
||||
result = torch.zeros_like(a)
|
||||
|
||||
if both_active.any():
|
||||
merged = a[both_active] * wa + b[both_active] * wb
|
||||
# Generate mask in same dtype
|
||||
mask = (torch.rand_like(merged, dtype=torch.float32) > drop_rate).to(dtype)
|
||||
result[both_active] = merged * mask
|
||||
|
||||
if only_a.any():
|
||||
result[only_a] = a[only_a] * wa
|
||||
|
||||
if only_b.any():
|
||||
result[only_b] = b[only_b] * wb
|
||||
|
||||
return result
|
||||
|
||||
if blend_mode == "crazy_mode":
|
||||
# Instead of random dropout, random amplification
|
||||
a_w = a * wa
|
||||
b_w = b * wb
|
||||
|
||||
# Randomly amplify some positions by up to 3x
|
||||
amp_a = 1.0 + torch.rand_like(a_w) * 2.0
|
||||
amp_b = 1.0 + torch.rand_like(b_w) * 2.0
|
||||
|
||||
# Randomly choose which to amplify where
|
||||
choose_a = torch.rand_like(a_w) > 0.5
|
||||
|
||||
result = torch.where(choose_a, a_w * amp_a, b_w * amp_b)
|
||||
return result
|
||||
|
||||
if threshold > 0:
|
||||
a_magnitude = a_weighted.abs()
|
||||
b_magnitude = b_weighted.abs()
|
||||
b_significant = b_magnitude > (threshold * a_magnitude + 1e-8)
|
||||
b_to_subtract = b_weighted * b_significant.float()
|
||||
return a_weighted - b_to_subtract
|
||||
else:
|
||||
return a_weighted - b_weighted
|
||||
merged = a * wa + b * wb
|
||||
# Generate mask in same dtype
|
||||
mask = (torch.rand_like(merged, dtype=torch.float32) > drop_rate).to(dtype)
|
||||
return merged * mask
|
||||
|
||||
def merge_magnitude(a, b, wa, wb, blend=0.5, **kwargs):
|
||||
"""Keep larger magnitude from either LoRA."""
|
||||
a_w = a * wa
|
||||
b_w = b * wb
|
||||
a_larger = a_w.abs() > b_w.abs()
|
||||
selected = torch.where(a_larger, a_w, b_w)
|
||||
averaged = (a_w + b_w) / 2
|
||||
return selected * blend + averaged * (1 - blend)
|
||||
|
||||
def merge_feature_mix(a, b, wa, wb, uniqueness=0.7, **kwargs):
|
||||
"""Preserve features unique to each LoRA."""
|
||||
a_w = a * wa
|
||||
b_w = b * wb
|
||||
ratio = a_w.abs() / (b_w.abs() + 1e-8)
|
||||
a_unique = (ratio > (1/uniqueness)).float()
|
||||
b_unique = (ratio < uniqueness).float()
|
||||
shared = torch.clamp(1 - a_unique - b_unique, 0, 1)
|
||||
def merge_dare_rescale(a, b, wa, wb, drop_rate=0.1, blend_mode="active", **kwargs):
|
||||
"""DARE with proper rescaling."""
|
||||
dtype = a.dtype
|
||||
drop_rate = max(0.0, min(0.9, drop_rate))
|
||||
|
||||
result = a_w * a_unique + b_w * b_unique
|
||||
result += ((a_w + b_w) / 2) * shared
|
||||
return result
|
||||
if blend_mode == "active":
|
||||
a_active = (a != 0)
|
||||
b_active = (b != 0)
|
||||
both_active = a_active & b_active
|
||||
only_a = a_active & ~b_active
|
||||
only_b = b_active & ~a_active
|
||||
|
||||
result = torch.zeros_like(a)
|
||||
rescale = 1.0 / (1.0 - drop_rate) if drop_rate < 1.0 else 1.0
|
||||
|
||||
if both_active.any():
|
||||
# Generate masks in float32 then convert to original dtype
|
||||
mask_a = (torch.rand_like(a[both_active], dtype=torch.float32) > drop_rate).to(dtype)
|
||||
mask_b = (torch.rand_like(b[both_active], dtype=torch.float32) > drop_rate).to(dtype)
|
||||
result[both_active] = (a[both_active] * mask_a * rescale) * wa + (b[both_active] * mask_b * rescale) * wb
|
||||
|
||||
if only_a.any():
|
||||
mask_a = (torch.rand_like(a[only_a], dtype=torch.float32) > drop_rate).to(dtype)
|
||||
result[only_a] = (a[only_a] * mask_a * rescale) * wa
|
||||
|
||||
if only_b.any():
|
||||
mask_b = (torch.rand_like(b[only_b], dtype=torch.float32) > drop_rate).to(dtype)
|
||||
result[only_b] = (b[only_b] * mask_b * rescale) * wb
|
||||
|
||||
return result
|
||||
|
||||
if blend_mode == "crazy_mode":
|
||||
a_w = a * wa
|
||||
b_w = b * wb
|
||||
|
||||
# Instead of random dropout, random jackpots!
|
||||
jackpot_a = (torch.rand_like(a_w, dtype=torch.float32) > 0.9).to(dtype)
|
||||
jackpot_b = (torch.rand_like(b_w, dtype=torch.float32) > 0.9).to(dtype)
|
||||
|
||||
# Winners get 10x, losers get 0.1x
|
||||
a_boosted = a_w * torch.where(jackpot_a > 0, 10.0, 0.1)
|
||||
b_boosted = b_w * torch.where(jackpot_b > 0, 10.0, 0.1)
|
||||
|
||||
result = a_boosted + b_boosted
|
||||
return result
|
||||
|
||||
else:
|
||||
# Generate masks in float32 then convert to original dtype
|
||||
mask_a = (torch.rand_like(a, dtype=torch.float32) > drop_rate).to(dtype)
|
||||
mask_b = (torch.rand_like(b, dtype=torch.float32) > drop_rate).to(dtype)
|
||||
rescale = 1.0 / (1.0 - drop_rate) if drop_rate < 1.0 else 1.0
|
||||
return (a * mask_a * rescale) * wa + (b * mask_b * rescale) * wb
|
||||
|
||||
|
||||
# ==================== INTUITIVE METHODS ====================
|
||||
|
||||
def merge_subtract(a, b, wa, wb, threshold=0.2, blend_mode="active", **kwargs):
|
||||
"""Subtract B from A where B is significant."""
|
||||
dtype = a.dtype
|
||||
|
||||
if blend_mode == "active":
|
||||
a_active = (a != 0)
|
||||
b_active = (b != 0)
|
||||
both_active = a_active & b_active
|
||||
only_a = a_active & ~b_active
|
||||
only_b = b_active & ~a_active
|
||||
|
||||
result = torch.zeros_like(a)
|
||||
|
||||
if both_active.any():
|
||||
a_w = a[both_active] * wa
|
||||
b_w = b[both_active] * wb
|
||||
mask = (b_w.abs() > (threshold * a_w.abs())).to(dtype)
|
||||
result[both_active] = (a_w - b_w) * mask + (a_w * (1 - mask))
|
||||
|
||||
if only_a.any():
|
||||
result[only_a] = a[only_a] * wa
|
||||
|
||||
if only_b.any():
|
||||
result[only_b] = -b[only_b] * wb
|
||||
|
||||
return result
|
||||
|
||||
if blend_mode == "crazy_mode":
|
||||
a_w = a * wa
|
||||
b_w = b * wb
|
||||
|
||||
# Instead of subtract where significant, subtract randomly
|
||||
# but keep the sign pattern
|
||||
random_mask = (torch.rand_like(a_w) > 0.5).to(dtype)
|
||||
|
||||
# Sometimes add, sometimes subtract, sometimes ignore
|
||||
result = a_w + b_w * (random_mask * 2 - 1) * 0.5
|
||||
return result
|
||||
|
||||
else:
|
||||
a_w = a * wa
|
||||
b_w = b * wb
|
||||
mask = (b_w.abs() > (threshold * a_w.abs())).to(dtype)
|
||||
return (a_w - b_w) * mask + (a_w * (1 - mask))
|
||||
|
||||
|
||||
def merge_magnitude(a, b, wa, wb, blend=0.5, blend_mode="active", **kwargs):
|
||||
"""Keep larger magnitude from either LoRA."""
|
||||
dtype = a.dtype
|
||||
|
||||
if blend_mode == "active":
|
||||
a_active = (a != 0)
|
||||
b_active = (b != 0)
|
||||
both_active = a_active & b_active
|
||||
only_a = a_active & ~b_active
|
||||
only_b = b_active & ~a_active
|
||||
|
||||
result = torch.zeros_like(a)
|
||||
|
||||
if both_active.any():
|
||||
a_w = a[both_active] * wa
|
||||
b_w = b[both_active] * wb
|
||||
a_larger = a_w.abs() > b_w.abs()
|
||||
selected = torch.where(a_larger, a_w, b_w)
|
||||
averaged = (a_w + b_w) / 2
|
||||
result[both_active] = selected * blend + averaged * (1 - blend)
|
||||
|
||||
if only_a.any():
|
||||
result[only_a] = a[only_a] * wa
|
||||
|
||||
if only_b.any():
|
||||
result[only_b] = b[only_b] * wb
|
||||
|
||||
return result
|
||||
|
||||
if blend_mode == "crazy_mode":
|
||||
a_w = a * wa
|
||||
b_w = b * wb
|
||||
|
||||
# Instead of taking the stronger, randomly choose with bias
|
||||
a_mag = a_w.abs()
|
||||
b_mag = b_w.abs()
|
||||
total_mag = a_mag + b_mag + 1e-8
|
||||
|
||||
# Create probability based on relative strength
|
||||
a_prob = a_mag / total_mag
|
||||
b_prob = b_mag / total_mag
|
||||
|
||||
# Random choice weighted by magnitude
|
||||
random_val = torch.rand_like(a_mag)
|
||||
take_a = random_val < a_prob
|
||||
|
||||
result = torch.where(take_a, a_w, b_w)
|
||||
return result
|
||||
|
||||
else:
|
||||
a_w = a * wa
|
||||
b_w = b * wb
|
||||
a_larger = a_w.abs() > b_w.abs()
|
||||
selected = torch.where(a_larger, a_w, b_w)
|
||||
averaged = (a_w + b_w) / 2
|
||||
return selected * blend + averaged * (1 - blend)
|
||||
|
||||
|
||||
def merge_feature_mix(a, b, wa, wb, uniqueness=0.7, blend_mode="active", **kwargs):
|
||||
"""Preserve unique features from each LoRA."""
|
||||
dtype = a.dtype
|
||||
|
||||
if blend_mode == "active":
|
||||
a_active = (a != 0)
|
||||
b_active = (b != 0)
|
||||
both_active = a_active & b_active
|
||||
only_a = a_active & ~b_active
|
||||
only_b = b_active & ~a_active
|
||||
|
||||
result = torch.zeros_like(a)
|
||||
|
||||
if both_active.any():
|
||||
a_w = a[both_active] * wa
|
||||
b_w = b[both_active] * wb
|
||||
total = a_w.abs() + b_w.abs() + 1e-8
|
||||
a_share = a_w.abs() / total
|
||||
|
||||
mask_a = (a_share > uniqueness).to(dtype)
|
||||
mask_b = (a_share < (1.0 - uniqueness)).to(dtype)
|
||||
mask_shared = (1.0 - mask_a - mask_b).to(dtype)
|
||||
|
||||
result[both_active] = (a_w * mask_a) + (b_w * mask_b) + ((a_w + b_w) / 2) * mask_shared
|
||||
|
||||
if only_a.any():
|
||||
result[only_a] = a[only_a] * wa
|
||||
|
||||
if only_b.any():
|
||||
result[only_b] = b[only_b] * wb
|
||||
|
||||
return result
|
||||
|
||||
if blend_mode == "crazy_mode":
|
||||
# Old style for backward compatibility
|
||||
a_w = a * wa
|
||||
b_w = b * wb
|
||||
ratio = a_w.abs() / (b_w.abs() + 1e-8)
|
||||
a_unique = (ratio > (1/uniqueness)).to(dtype)
|
||||
b_unique = (ratio < uniqueness).to(dtype)
|
||||
shared = torch.clamp(1 - a_unique - b_unique, 0, 1).to(dtype)
|
||||
|
||||
result = a_w * a_unique + b_w * b_unique
|
||||
result += ((a_w + b_w) / 2) * shared
|
||||
return result
|
||||
|
||||
else:
|
||||
a_w = a * wa
|
||||
b_w = b * wb
|
||||
total = a_w.abs() + b_w.abs() + 1e-8
|
||||
a_share = a_w.abs() / total
|
||||
|
||||
mask_a = (a_share > uniqueness).to(dtype)
|
||||
mask_b = (a_share < (1.0 - uniqueness)).to(dtype)
|
||||
mask_shared = (1.0 - mask_a - mask_b).to(dtype)
|
||||
|
||||
return (a_w * mask_a) + (b_w * mask_b) + ((a_w + b_w) / 2) * mask_shared
|
||||
|
||||
|
||||
# ==================== ADVANCED METHODS ====================
|
||||
|
||||
def merge_svd_preserve(a, b, wa, wb, preserve_ratio=0.8, **kwargs):
|
||||
def merge_svd_preserve(a, b, wa, wb, preserve_ratio=0.8, blend_mode="active", **kwargs):
|
||||
"""SVD-based merge with rank reduction."""
|
||||
if a.numel() < 100:
|
||||
return a * wa + b * wb
|
||||
dtype = a.dtype
|
||||
|
||||
if a.numel() < 100:
|
||||
return merge_linear(a, b, wa, wb, blend_mode=blend_mode)
|
||||
|
||||
# For SVD, we need to work with full tensors
|
||||
# Convert to float32 for SVD, then back to original dtype
|
||||
combined = a.float() * wa + b.float() * wb
|
||||
if min(combined.shape) > 1:
|
||||
try:
|
||||
@@ -146,34 +552,167 @@ def merge_svd_preserve(a, b, wa, wb, preserve_ratio=0.8, **kwargs):
|
||||
result = U[:, :k] @ torch.diag(S[:k]) @ Vh[:k, :]
|
||||
except:
|
||||
result = combined
|
||||
|
||||
if blend_mode == "crazy_mode":
|
||||
# Simple but visible: randomly swap chunks between A and B
|
||||
a_w = a * wa
|
||||
b_w = b * wb
|
||||
|
||||
# Create random blocks (works for 2D tensors)
|
||||
if len(a_w.shape) == 2:
|
||||
h, w = a_w.shape
|
||||
block_h, block_w = max(1, h//8), max(1, w//8)
|
||||
|
||||
# Create block pattern
|
||||
result = a_w.clone()
|
||||
for i in range(0, h, block_h):
|
||||
for j in range(0, w, block_w):
|
||||
if torch.rand(1).item() > 0.5:
|
||||
# Swap this block from B
|
||||
i_end = min(i + block_h, h)
|
||||
j_end = min(j + block_w, w)
|
||||
result[i:i_end, j:j_end] = b_w[i:i_end, j:j_end]
|
||||
return result
|
||||
else:
|
||||
# Fallback for non-2D tensors
|
||||
return a_w * 0.7 + b_w * 0.3
|
||||
|
||||
else:
|
||||
result = combined
|
||||
return result.to(a.dtype)
|
||||
return result.to(dtype)
|
||||
|
||||
def merge_noise_aware(a, b, wa, wb, noise_threshold=0.01, **kwargs):
|
||||
|
||||
def merge_noise_aware(a, b, wa, wb, noise_threshold=0.01, blend_mode="active", **kwargs):
|
||||
"""Reduce noise before merging."""
|
||||
threshold_a = noise_threshold * a.abs().max()
|
||||
threshold_b = noise_threshold * b.abs().max()
|
||||
a_clean = torch.where(a.abs() < threshold_a, a * 0.1, a)
|
||||
b_clean = torch.where(b.abs() < threshold_b, b * 0.1, b)
|
||||
return a_clean * wa + b_clean * wb
|
||||
dtype = a.dtype
|
||||
|
||||
if blend_mode == "active":
|
||||
a_active = (a != 0)
|
||||
b_active = (b != 0)
|
||||
both_active = a_active & b_active
|
||||
only_a = a_active & ~b_active
|
||||
only_b = b_active & ~a_active
|
||||
|
||||
result = torch.zeros_like(a)
|
||||
|
||||
if both_active.any():
|
||||
threshold_a = noise_threshold * a[both_active].abs().max()
|
||||
threshold_b = noise_threshold * b[both_active].abs().max()
|
||||
a_clean = torch.where(a[both_active].abs() < threshold_a, a[both_active] * 0.1, a[both_active])
|
||||
b_clean = torch.where(b[both_active].abs() < threshold_b, b[both_active] * 0.1, b[both_active])
|
||||
result[both_active] = a_clean * wa + b_clean * wb
|
||||
|
||||
if only_a.any():
|
||||
threshold_a = noise_threshold * a[only_a].abs().max()
|
||||
a_clean = torch.where(a[only_a].abs() < threshold_a, a[only_a] * 0.1, a[only_a])
|
||||
result[only_a] = a_clean * wa
|
||||
|
||||
if only_b.any():
|
||||
threshold_b = noise_threshold * b[only_b].abs().max()
|
||||
b_clean = torch.where(b[only_b].abs() < threshold_b, b[only_b] * 0.1, b[only_b])
|
||||
result[only_b] = b_clean * wb
|
||||
|
||||
return result
|
||||
|
||||
if blend_mode == "crazy_mode":
|
||||
# Add noise, but make it structured
|
||||
a_w = a * wa
|
||||
b_w = b * wb
|
||||
|
||||
# Create patterned noise (grid-like)
|
||||
h, w = a_w.shape[-2:] if len(a_w.shape) > 2 else (1, 1)
|
||||
noise_pattern = (torch.arange(h).to(a_w.device) % 2).float()
|
||||
noise_pattern = noise_pattern.view(-1, 1).repeat(1, w)
|
||||
|
||||
# Add noise only to certain positions
|
||||
noise = (torch.rand_like(a_w) - 0.5) * noise_pattern * 0.1
|
||||
|
||||
result = a_w + b_w + noise
|
||||
return result
|
||||
|
||||
else:
|
||||
threshold_a = noise_threshold * a.abs().max()
|
||||
threshold_b = noise_threshold * b.abs().max()
|
||||
a_clean = torch.where(a.abs() < threshold_a, a * 0.1, a)
|
||||
b_clean = torch.where(b.abs() < threshold_b, b * 0.1, b)
|
||||
return a_clean * wa + b_clean * wb
|
||||
|
||||
def merge_gradient_alignment(a, b, wa, wb, **kwargs):
|
||||
|
||||
def merge_gradient_alignment(a, b, wa, wb, blend_mode="active", **kwargs):
|
||||
"""Merge based on directional alignment."""
|
||||
a_flat, b_flat = a.flatten().float(), b.flatten().float()
|
||||
norm_prod = (torch.norm(a_flat) * torch.norm(b_flat) + 1e-8)
|
||||
cos_sim = torch.dot(a_flat, b_flat) / norm_prod
|
||||
alignment = (cos_sim + 1) / 2
|
||||
return a * wa * alignment + b * wb * alignment
|
||||
|
||||
# ==================== UTILITY FUNCTIONS ====================
|
||||
dtype = a.dtype
|
||||
|
||||
if blend_mode == "active":
|
||||
a_active = (a != 0)
|
||||
b_active = (b != 0)
|
||||
both_active = a_active & b_active
|
||||
only_a = a_active & ~b_active
|
||||
only_b = b_active & ~a_active
|
||||
|
||||
result = torch.zeros_like(a)
|
||||
|
||||
if both_active.any():
|
||||
# Use float32 only for the cosine similarity calculation
|
||||
a_flat = a[both_active].flatten().float()
|
||||
b_flat = b[both_active].flatten().float()
|
||||
norm_prod = (torch.norm(a_flat) * torch.norm(b_flat) + 1e-8)
|
||||
cos_sim = torch.dot(a_flat, b_flat) / norm_prod
|
||||
alignment = ((cos_sim + 1) / 2).to(dtype)
|
||||
# Convert alignment back to original dtype
|
||||
alignment = torch.tensor(alignment, dtype=dtype, device=a.device)
|
||||
result[both_active] = a[both_active] * wa * alignment + b[both_active] * wb * alignment
|
||||
|
||||
if only_a.any():
|
||||
result[only_a] = a[only_a] * wa
|
||||
|
||||
if only_b.any():
|
||||
result[only_b] = b[only_b] * wb
|
||||
|
||||
return result
|
||||
|
||||
if blend_mode == "crazy_mode":
|
||||
a_w = a * wa
|
||||
b_w = b * wb
|
||||
|
||||
# Calculate alignment but add random mood swings
|
||||
a_flat = a_w.flatten().float()
|
||||
b_flat = b_w.flatten().float()
|
||||
|
||||
cos_sim = torch.dot(a_flat, b_flat) / (torch.norm(a_flat) * torch.norm(b_flat) + 1e-8)
|
||||
|
||||
# Add random mood (changes per layer)
|
||||
mood = torch.randn(1, device=a.device).item() * 0.3
|
||||
|
||||
# Blend based on mood-influenced alignment
|
||||
alignment = torch.clamp(cos_sim + mood, -1, 1)
|
||||
alignment = ((alignment + 1) / 2).to(dtype)
|
||||
|
||||
result = a_w * alignment + b_w * (1 - alignment)
|
||||
return result
|
||||
|
||||
else:
|
||||
# Use float32 only for the cosine similarity calculation
|
||||
a_flat = a.flatten().float()
|
||||
b_flat = b.flatten().float()
|
||||
norm_prod = (torch.norm(a_flat) * torch.norm(b_flat) + 1e-8)
|
||||
cos_sim = torch.dot(a_flat, b_flat) / norm_prod
|
||||
alignment = ((cos_sim + 1) / 2).to(dtype)
|
||||
# Convert alignment back to original dtype
|
||||
alignment = torch.tensor(alignment, dtype=dtype, device=a.device)
|
||||
return a * wa * alignment + b * wb * alignment
|
||||
|
||||
def apply_density(tensor, density):
|
||||
"""Keep top 'density' percentage of weights."""
|
||||
dtype = tensor.dtype
|
||||
|
||||
if density >= 1:
|
||||
return tensor
|
||||
flat = tensor.abs().flatten()
|
||||
|
||||
# Use float32 for the threshold calculation to avoid dtype issues
|
||||
flat = tensor.abs().float().flatten()
|
||||
k = max(1, int(flat.numel() * density))
|
||||
threshold = torch.topk(flat, k).values.min()
|
||||
mask = tensor.abs() >= threshold
|
||||
|
||||
# Create mask in original dtype
|
||||
mask = (tensor.abs() >= threshold).to(dtype)
|
||||
return tensor * mask
|
||||
|
After Width: | Height: | Size: 1.2 MiB |
|
After Width: | Height: | Size: 273 KiB |
|
After Width: | Height: | Size: 268 KiB |
|
After Width: | Height: | Size: 1.1 MiB |