chore: untrack pre-existing plan docs swept in by git add -A
These docs/plans/2026-03/04-*.md were untracked local notes from earlier sessions that an over-broad 'git add -A' committed. Untrack them (kept on disk) to restore the prior state; this session's 2026-06-22 design doc stays. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Ethanfel
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
34ac112041
commit
588d6b0243
@@ -1,59 +0,0 @@
|
||||
# Per-Prefix Adaptive Merge Strategy
|
||||
|
||||
## Problem
|
||||
|
||||
The optimizer currently computes per-prefix conflict stats in Pass 1 but averages them into a single global conflict ratio. One global decision (TIES vs weighted_average) is applied to all prefixes. This loses important information:
|
||||
|
||||
- Two LoRAs may only overlap in certain blocks (e.g., attention layers 4-7). Non-overlapping prefixes have zero conflict but get treated the same as high-conflict ones.
|
||||
- When the global average falls below the TIES threshold, genuinely conflicting prefixes get simple averaging — conflicts go unresolved.
|
||||
- When the global average is above the threshold, non-overlapping prefixes get TIES trimming — their weights are unnecessarily reduced.
|
||||
|
||||
## Solution
|
||||
|
||||
Per-prefix adaptive strategy: each prefix picks its own merge mode based on its own conflict data.
|
||||
|
||||
### Decision table (per prefix)
|
||||
|
||||
| Condition | Strategy |
|
||||
|-----------|----------|
|
||||
| Only 1 LoRA touches this prefix | `weighted_sum` — full strength, no dilution |
|
||||
| 2+ LoRAs, conflict <= 25% | `weighted_average` — compatible, simple merge |
|
||||
| 2+ LoRAs, conflict > 25% | `ties` — resolve conflicts with trim/elect/merge |
|
||||
|
||||
Density and sign method are also computed per-prefix when TIES is selected.
|
||||
|
||||
### Toggle
|
||||
|
||||
New input `optimization_mode` with values `per_prefix` (default) and `global` (original behavior).
|
||||
|
||||
### Architecture
|
||||
|
||||
```
|
||||
Pass 1 → per-prefix stats stored in dict
|
||||
↓
|
||||
global stats still computed for report + global mode fallback
|
||||
↓
|
||||
Pass 2 → if per_prefix: each prefix looks up its own stats, picks strategy
|
||||
if global: same mode/density/sign_method for all (original behavior)
|
||||
```
|
||||
|
||||
### What stays the same
|
||||
|
||||
- Two-pass streaming architecture (no extra memory)
|
||||
- Auto-strength (LoRA-level, not per-prefix)
|
||||
- Cache mechanism (add optimization_mode to cache key)
|
||||
- All merge modes unchanged
|
||||
- Report format (additive sections)
|
||||
|
||||
### Report additions
|
||||
|
||||
```
|
||||
--- Per-Prefix Strategy ---
|
||||
weighted_sum (single LoRA): 120 prefixes
|
||||
weighted_average (low conflict): 32 prefixes
|
||||
ties (high conflict): 40 prefixes
|
||||
```
|
||||
|
||||
### Files modified
|
||||
|
||||
Only `lora_optimizer.py`.
|
||||
@@ -1,419 +0,0 @@
|
||||
# Per-Prefix Adaptive Merge Implementation Plan
|
||||
|
||||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||||
|
||||
**Goal:** Make the optimizer pick merge strategy (TIES vs weighted_average vs weighted_sum) per-prefix instead of globally, so non-overlapping LoRA regions keep full strength and only genuinely conflicting regions use TIES.
|
||||
|
||||
**Architecture:** Pass 1 already computes per-prefix conflict/magnitude data. Instead of averaging into one global decision, store per-prefix stats. Pass 2 looks up each prefix's stats to pick its own strategy. A toggle (`optimization_mode`) lets users switch between `per_prefix` (new default) and `global` (original behavior).
|
||||
|
||||
**Tech Stack:** Python, PyTorch (existing dependencies only)
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Add `optimization_mode` input and update cache key
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:426-447` (INPUT_TYPES)
|
||||
- Modify: `lora_optimizer.py:455-475` (_compute_cache_key)
|
||||
- Modify: `lora_optimizer.py:477-482` (IS_CHANGED)
|
||||
|
||||
**Step 1: Add `optimization_mode` to INPUT_TYPES**
|
||||
|
||||
In `INPUT_TYPES` (line 435), add to the `optional` dict after `free_vram_between_passes`:
|
||||
|
||||
```python
|
||||
"optimization_mode": (["per_prefix", "global"], {
|
||||
"default": "per_prefix",
|
||||
"tooltip": "per_prefix: each weight group picks its own merge strategy based on local conflict. global: single strategy for all (original behavior)."
|
||||
}),
|
||||
```
|
||||
|
||||
**Step 2: Add `optimization_mode` to cache key**
|
||||
|
||||
In `_compute_cache_key` (line 474), change the hash suffix to include the new param:
|
||||
|
||||
```python
|
||||
h.update(f"|os={output_strength}|csm={clip_strength_multiplier}|as={auto_strength}|om={optimization_mode}".encode())
|
||||
```
|
||||
|
||||
Add `optimization_mode` as a parameter to the method signature:
|
||||
|
||||
```python
|
||||
@staticmethod
|
||||
def _compute_cache_key(lora_stack, output_strength, clip_strength_multiplier, auto_strength, optimization_mode="per_prefix"):
|
||||
```
|
||||
|
||||
**Step 3: Update IS_CHANGED to pass optimization_mode**
|
||||
|
||||
Update IS_CHANGED signature and call (lines 477-482):
|
||||
|
||||
```python
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, model, clip, lora_stack, output_strength,
|
||||
clip_strength_multiplier=1.0, auto_strength="disabled",
|
||||
free_vram_between_passes="disabled", optimization_mode="per_prefix"):
|
||||
return cls._compute_cache_key(lora_stack, output_strength,
|
||||
clip_strength_multiplier, auto_strength,
|
||||
optimization_mode)
|
||||
```
|
||||
|
||||
**Step 4: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: add optimization_mode input (per_prefix/global toggle)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: Store per-prefix stats during Pass 1
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:1117-1166` (Pass 1 accumulators and _collect_analysis_result)
|
||||
|
||||
**Step 1: Add per-prefix stats accumulator**
|
||||
|
||||
After the existing accumulators (line 1132), add:
|
||||
|
||||
```python
|
||||
prefix_stats = {} # prefix -> {conflict_ratio, n_loras, magnitude_samples, magnitude_ratio}
|
||||
```
|
||||
|
||||
**Step 2: Update `_collect_analysis_result` to store per-prefix data**
|
||||
|
||||
After the existing accumulation logic (line 1150), add per-prefix stat storage. The function already has `partial_stats` (which loras contribute), `pair_conflicts` (conflict data), and `mag_samples`. Store these per-prefix:
|
||||
|
||||
```python
|
||||
# Store per-prefix stats for per_prefix optimization mode
|
||||
if len(partial_stats) > 0:
|
||||
# Number of LoRAs contributing to this prefix
|
||||
n_contributing = len(partial_stats)
|
||||
|
||||
# Per-prefix conflict ratio
|
||||
pf_overlap = sum(ov for ov, _ in pair_conflicts.values())
|
||||
pf_conflict = sum(conf for _, conf in pair_conflicts.values())
|
||||
pf_conflict_ratio = pf_conflict / pf_overlap if pf_overlap > 0 else 0.0
|
||||
|
||||
# Per-prefix magnitude ratio (max/min L2 among contributing LoRAs)
|
||||
pf_l2s = [l2 for _, _, l2 in partial_stats if l2 > 0]
|
||||
if len(pf_l2s) >= 2:
|
||||
pf_mag_ratio = max(pf_l2s) / min(pf_l2s)
|
||||
else:
|
||||
pf_mag_ratio = 1.0
|
||||
|
||||
prefix_stats[prefix] = {
|
||||
"n_loras": n_contributing,
|
||||
"conflict_ratio": pf_conflict_ratio,
|
||||
"magnitude_ratio": pf_mag_ratio,
|
||||
"magnitude_samples": list(mag_samples), # copy, not reference
|
||||
}
|
||||
```
|
||||
|
||||
**Step 3: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: store per-prefix conflict and magnitude stats during Pass 1"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Per-prefix decision logic in Pass 2
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:1282-1398` (Pass 2 merge logic)
|
||||
|
||||
**Step 1: Update optimize_merge signature**
|
||||
|
||||
Add `optimization_mode="per_prefix"` to the `optimize_merge` method signature (line 1018):
|
||||
|
||||
```python
|
||||
def optimize_merge(self, model, clip, lora_stack, output_strength, clip_strength_multiplier=1.0, auto_strength="disabled", free_vram_between_passes="disabled", optimization_mode="per_prefix"):
|
||||
```
|
||||
|
||||
**Step 2: Update cache_key call in optimize_merge**
|
||||
|
||||
Update the cache_key computation (line 1066) to include optimization_mode:
|
||||
|
||||
```python
|
||||
cache_key = self._compute_cache_key(lora_stack, output_strength,
|
||||
clip_strength_multiplier, auto_strength,
|
||||
optimization_mode)
|
||||
```
|
||||
|
||||
**Step 3: Modify `_merge_one_prefix` to use per-prefix decisions**
|
||||
|
||||
The inner function `_merge_one_prefix` (line 1292) currently uses closure variables `mode`, `density`, `sign_method`. Modify it to accept and use per-prefix params. Replace the merge call section (lines 1367-1371):
|
||||
|
||||
Before (current code):
|
||||
```python
|
||||
merged_diff = self._merge_diffs(
|
||||
diffs_list, mode,
|
||||
density=density, majority_sign_method=sign_method,
|
||||
compute_device=compute_device
|
||||
)
|
||||
```
|
||||
|
||||
After:
|
||||
```python
|
||||
# Determine strategy for this prefix
|
||||
if optimization_mode == "per_prefix" and lora_prefix in prefix_stats:
|
||||
pf = prefix_stats[lora_prefix]
|
||||
if pf["n_loras"] <= 1 or len(diffs_list) <= 1:
|
||||
# Single LoRA on this prefix: weighted_sum, full strength
|
||||
pf_mode = "weighted_sum"
|
||||
pf_density = 0.5
|
||||
pf_sign = "frequency"
|
||||
else:
|
||||
pf_mode, pf_density, pf_sign, _ = self._auto_select_params(
|
||||
pf["conflict_ratio"], pf["magnitude_ratio"],
|
||||
magnitude_samples=pf.get("magnitude_samples")
|
||||
)
|
||||
else:
|
||||
pf_mode = mode
|
||||
pf_density = density
|
||||
pf_sign = sign_method
|
||||
|
||||
merged_diff = self._merge_diffs(
|
||||
diffs_list, pf_mode,
|
||||
density=pf_density, majority_sign_method=pf_sign,
|
||||
compute_device=compute_device
|
||||
)
|
||||
```
|
||||
|
||||
**Step 4: Track per-prefix strategy counts for reporting**
|
||||
|
||||
Add counters before the Pass 2 loop (after line 1290):
|
||||
|
||||
```python
|
||||
strategy_counts = {"weighted_sum": 0, "weighted_average": 0, "ties": 0}
|
||||
```
|
||||
|
||||
Update `_merge_one_prefix` to return the chosen mode. Change the return (line 1374) to:
|
||||
|
||||
```python
|
||||
return (target_key, is_clip_key, merged_diff, pf_mode)
|
||||
```
|
||||
|
||||
Update `_collect_merge_result` (line 1376) to count strategies:
|
||||
|
||||
```python
|
||||
def _collect_merge_result(result):
|
||||
nonlocal processed_keys
|
||||
if result is None:
|
||||
return
|
||||
target_key, is_clip_key, merged_diff, used_mode = result
|
||||
if is_clip_key:
|
||||
clip_patches[target_key] = ("diff", (merged_diff,))
|
||||
else:
|
||||
model_patches[target_key] = ("diff", (merged_diff,))
|
||||
processed_keys += 1
|
||||
strategy_counts[used_mode] = strategy_counts.get(used_mode, 0) + 1
|
||||
```
|
||||
|
||||
**Step 5: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: per-prefix adaptive merge strategy in Pass 2"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Update report to show per-prefix strategy breakdown
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:938-1016` (_build_report)
|
||||
- Modify: `lora_optimizer.py:1424-1437` (report call site)
|
||||
|
||||
**Step 1: Add strategy_counts parameter to _build_report**
|
||||
|
||||
Update `_build_report` signature (line 938) to accept `strategy_counts` and `optimization_mode`:
|
||||
|
||||
```python
|
||||
def _build_report(self, lora_stats, pairwise_conflicts, collection_stats,
|
||||
mode, density, sign_method, reasoning, merge_summary,
|
||||
auto_strength_info=None, strategy_counts=None, optimization_mode="global"):
|
||||
```
|
||||
|
||||
**Step 2: Add per-prefix strategy section to report**
|
||||
|
||||
After the "Auto-Selected Parameters" section (after line 995), add:
|
||||
|
||||
```python
|
||||
# Per-Prefix Strategy breakdown (only in per_prefix mode)
|
||||
if optimization_mode == "per_prefix" and strategy_counts:
|
||||
lines.append("")
|
||||
lines.append("--- Per-Prefix Strategy ---")
|
||||
total_pf = sum(strategy_counts.values())
|
||||
if strategy_counts.get("weighted_sum", 0) > 0:
|
||||
n = strategy_counts["weighted_sum"]
|
||||
lines.append(f" weighted_sum (single LoRA): {n:>4} prefixes ({n/total_pf:.0%})")
|
||||
if strategy_counts.get("weighted_average", 0) > 0:
|
||||
n = strategy_counts["weighted_average"]
|
||||
lines.append(f" weighted_average (low conflict): {n:>4} prefixes ({n/total_pf:.0%})")
|
||||
if strategy_counts.get("ties", 0) > 0:
|
||||
n = strategy_counts["ties"]
|
||||
lines.append(f" ties (high conflict): {n:>4} prefixes ({n/total_pf:.0%})")
|
||||
lines.append(f" Total: {total_pf:>4} prefixes")
|
||||
```
|
||||
|
||||
**Step 3: Update the Auto-Selected Parameters section**
|
||||
|
||||
When in `per_prefix` mode, the "Auto-Selected Parameters" section should clarify that these are the *global fallback* parameters. Add a note:
|
||||
|
||||
After line 995 (the sign_method line), add:
|
||||
|
||||
```python
|
||||
if optimization_mode == "per_prefix":
|
||||
lines.append(" (global fallback — each prefix uses its own parameters)")
|
||||
```
|
||||
|
||||
**Step 4: Update report call site**
|
||||
|
||||
Update the `_build_report` call (line 1433) to pass the new params:
|
||||
|
||||
```python
|
||||
report = self._build_report(
|
||||
lora_stats, pairwise_conflicts, collection_stats,
|
||||
mode, density, sign_method, reasoning, merge_summary,
|
||||
auto_strength_info=auto_strength_info,
|
||||
strategy_counts=strategy_counts if optimization_mode == "per_prefix" else None,
|
||||
optimization_mode=optimization_mode
|
||||
)
|
||||
```
|
||||
|
||||
**Step 5: Update saved report params**
|
||||
|
||||
Update the `selected_params` dict (line 1444) to include the mode:
|
||||
|
||||
```python
|
||||
selected_params = {"mode": mode, "density": density, "sign_method": sign_method, "optimization_mode": optimization_mode}
|
||||
if optimization_mode == "per_prefix":
|
||||
selected_params["strategy_counts"] = dict(strategy_counts)
|
||||
```
|
||||
|
||||
**Step 6: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 7: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: add per-prefix strategy breakdown to analysis report"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: Add logging for per-prefix mode
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` (Pass 2 logging section, around line 1400)
|
||||
|
||||
**Step 1: Add per-prefix strategy logging**
|
||||
|
||||
After the Pass 2 patches log line (line 1400), add:
|
||||
|
||||
```python
|
||||
if optimization_mode == "per_prefix":
|
||||
logging.info(f"[LoRA Optimizer] Per-prefix strategies: "
|
||||
f"{strategy_counts.get('weighted_sum', 0)} weighted_sum, "
|
||||
f"{strategy_counts.get('weighted_average', 0)} weighted_average, "
|
||||
f"{strategy_counts.get('ties', 0)} ties")
|
||||
```
|
||||
|
||||
Also update the Pass 2 start log (line 1285) to indicate mode:
|
||||
|
||||
```python
|
||||
logging.info(f"[LoRA Optimizer] Pass 2: Merging {len(all_key_targets)} keys "
|
||||
f"({optimization_mode} strategy, "
|
||||
f"{'sequential' if use_gpu else 'threaded'})...")
|
||||
```
|
||||
|
||||
**Step 2: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: add per-prefix strategy logging"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 6: Update DESCRIPTION and README
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:453` (DESCRIPTION)
|
||||
- Modify: `README.md`
|
||||
|
||||
**Step 1: Update node DESCRIPTION**
|
||||
|
||||
Change line 453:
|
||||
|
||||
```python
|
||||
DESCRIPTION = "Auto-analyzes LoRA stack and selects optimal merge strategy per weight group. Outputs merged model + analysis report."
|
||||
```
|
||||
|
||||
**Step 2: Update README decision table**
|
||||
|
||||
In the "How it decides" section, add a note about per-prefix mode. After the existing decision table, add:
|
||||
|
||||
```markdown
|
||||
**Per-prefix mode** (default): Each weight prefix picks its own strategy based on local conflict data. Non-overlapping prefixes use `weighted_sum` at full strength. Only genuinely conflicting prefixes use TIES. Set `optimization_mode` to `global` for the original single-strategy behavior.
|
||||
```
|
||||
|
||||
**Step 3: Update the Example Report in README**
|
||||
|
||||
Add the per-prefix strategy section to the example report.
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py README.md
|
||||
git commit -m "docs: update README and node description for per-prefix mode"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 7: Final verification
|
||||
|
||||
**Step 1: Full syntax check**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 2: Verify node registration unchanged**
|
||||
|
||||
Run: `python -c "exec(open('lora_optimizer.py').read()); print(list(NODE_CLASS_MAPPINGS.keys()))"`
|
||||
Expected: `['LoRAStack', 'LoRAOptimizer']`
|
||||
|
||||
**Step 3: Verify INPUT_TYPES has optimization_mode**
|
||||
|
||||
Run: `python -c "exec(open('lora_optimizer.py').read()); print([k for k in LoRAOptimizer.INPUT_TYPES()['optional'].keys()])"`
|
||||
Expected: includes `optimization_mode`
|
||||
|
||||
**Step 4: Verify default is per_prefix**
|
||||
|
||||
Run: `python -c "exec(open('lora_optimizer.py').read()); om = LoRAOptimizer.INPUT_TYPES()['optional']['optimization_mode']; print(om[1]['default'])"`
|
||||
Expected: `per_prefix`
|
||||
@@ -1,109 +0,0 @@
|
||||
# Architecture-Aware Key Normalization Design
|
||||
|
||||
## Problem
|
||||
|
||||
The LoRA Optimizer merges multiple LoRAs by grouping their weights by key prefix. When two LoRAs use different key formats for the same model layer (e.g., one trained with Kohya, another with AI-Toolkit), their keys don't align — the optimizer treats them as separate prefixes and they never interact during conflict analysis or merging. This is a silent failure.
|
||||
|
||||
The most severe case is Z-Image Turbo, which uses fused QKV attention (`attention.qkv [11520, 3840]`) instead of separate `to_q`/`to_k`/`to_v`. A LoRA trained on diffusers format has separate Q/K/V keys that cannot be directly merged with a LoRA that targets Z-Image's native fused format.
|
||||
|
||||
## Solution
|
||||
|
||||
Add architecture detection and per-architecture key normalization as a pre-processing step inside `_normalize_stack()`, gated by a new `normalize_keys` boolean input (default: disabled).
|
||||
|
||||
When enabled:
|
||||
1. Auto-detect architecture from LoRA key patterns
|
||||
2. Normalize all LoRAs to a canonical key format for that architecture
|
||||
3. The rest of the pipeline (prefix grouping, analysis, merge) works unchanged
|
||||
4. For Z-Image: re-fuse split Q/K/V back to `qkv` at patch application time
|
||||
|
||||
## Architecture Detection
|
||||
|
||||
New `_detect_architecture(lora_sd)` method on `_LoRAMergeBase`. Returns `"zimage"`, `"flux"`, `"wan"`, `"sdxl"`, or `"unknown"`.
|
||||
|
||||
Detection order (first match wins):
|
||||
- **Z-Image**: `diffusion_model.layers.N` or `single_transformer_blocks.N` (without `transformer.` prefix) or `lora_unet_layers_`
|
||||
- **FLUX**: `double_blocks`/`single_blocks` or `transformer.transformer_blocks`/`transformer.single_transformer_blocks`
|
||||
- **Wan**: `blocks.N` with `self_attn`/`cross_attn`/`ffn`
|
||||
- **SDXL**: `lora_te1_`/`lora_te2_` or `input_blocks`/`output_blocks`/`down_blocks`/`up_blocks`
|
||||
|
||||
## Per-Architecture Normalization
|
||||
|
||||
New `_normalize_keys(lora_sd, architecture)` method. Returns a new dict with normalized keys.
|
||||
|
||||
### Z-Image (Lumina2)
|
||||
1. **QKV splitting**: Fused `attention.qkv.lora_A/B` split into `attention.to_q`, `attention.to_k`, `attention.to_v` (slice along dim=0, /3). Alpha copied to all three.
|
||||
2. **Output remap**: `attention.out` -> `attention.to_out.0`
|
||||
3. **Prefix standardization**: Ensure `diffusion_model.layers.N.` prefix
|
||||
4. **Musubi Tuner format**: `lora_unet_layers_N_attention_...` -> `diffusion_model.layers.N.attention...`
|
||||
|
||||
### FLUX
|
||||
1. **AI-Toolkit**: `transformer.transformer_blocks.N` -> `diffusion_model.double_blocks.N`; `transformer.single_transformer_blocks.N` -> `diffusion_model.single_blocks.N`
|
||||
2. **Kohya**: `lora_transformer_double_blocks_N` -> `diffusion_model.double_blocks.N`
|
||||
3. **Standard**: Ensure `diffusion_model.` prefix
|
||||
|
||||
### Wan
|
||||
Port existing `_wan_standardize_lora_key_format` logic:
|
||||
1. LyCORIS/diffusers/Fun LoRA/finetrainer format normalization
|
||||
2. RS-LoRA alpha compensation
|
||||
|
||||
### SDXL
|
||||
1. Ensure consistent `lora_unet_`/`lora_te1_`/`lora_te2_` prefixes
|
||||
2. Handle diffusers-format block naming variants
|
||||
|
||||
## Z-Image Re-Fusion at Apply Time
|
||||
|
||||
After merging at the `to_q`/`to_k`/`to_v` level, the merged patches must be re-fused to `qkv` before applying to a Z-Image model. This happens in `_process_prefix()`:
|
||||
- Detect groups of `to_q`/`to_k`/`to_v` patches for the same layer
|
||||
- Concatenate into single `qkv` patch
|
||||
- Remap `to_out.0` -> `out`
|
||||
|
||||
## Data Flow
|
||||
|
||||
```
|
||||
LoRA stack input
|
||||
|
|
||||
_normalize_stack()
|
||||
|-- Load/unpack each LoRA
|
||||
|-- if normalize_keys == "enabled":
|
||||
| |-- _detect_architecture(first_lora_sd)
|
||||
| +-- For each LoRA: _normalize_keys(lora_sd, arch)
|
||||
+-- Return normalized dicts
|
||||
|
|
||||
Pass 1: Analysis (unchanged)
|
||||
|
|
||||
Pass 2: Merge (unchanged)
|
||||
|
|
||||
_process_prefix() / patch application
|
||||
|-- if arch == "zimage": re-fuse to_q/to_k/to_v -> qkv
|
||||
+-- else: apply normally
|
||||
|
|
||||
Patched model output
|
||||
```
|
||||
|
||||
## Node Input Change
|
||||
|
||||
Add to `LoRAOptimizer.INPUT_TYPES` optional:
|
||||
```python
|
||||
"normalize_keys": (["disabled", "enabled"], {"default": "disabled"})
|
||||
```
|
||||
|
||||
## Report Change
|
||||
|
||||
New line at top of analysis report:
|
||||
```
|
||||
Architecture: Z-Image Turbo (auto-detected)
|
||||
Key normalization: enabled (3 LoRAs normalized)
|
||||
```
|
||||
|
||||
## Files Changed
|
||||
|
||||
| File | Change |
|
||||
|------|--------|
|
||||
| `lora_optimizer.py` | Add `_detect_architecture()`, `_normalize_keys()` to `_LoRAMergeBase`. Modify `_normalize_stack()`. Modify `_process_prefix()` for Z-Image re-fusion. Add `normalize_keys` input. |
|
||||
|
||||
## Design Decisions
|
||||
|
||||
1. **Disabled by default**: Normalization is opt-in to avoid changing behavior for existing users
|
||||
2. **Split QKV, not fuse**: Splitting fused QKV into separate Q/K/V preserves per-component granularity for conflict analysis. Re-fuse only at apply time.
|
||||
3. **Single file change**: All changes in `lora_optimizer.py` — no new files needed
|
||||
4. **Shared on `_LoRAMergeBase`**: Detection and normalization methods available to both LoRAOptimizer and WanVideoLoRAOptimizer
|
||||
@@ -1,907 +0,0 @@
|
||||
# Architecture-Aware Key Normalization Implementation Plan
|
||||
|
||||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||||
|
||||
**Goal:** Add architecture detection and per-architecture key normalization to `LoRAOptimizer` so LoRAs trained by different tools get their keys aligned before merging — including Z-Image QKV tensor splitting/re-fusion.
|
||||
|
||||
**Architecture:** A `normalize_keys` boolean input gates the feature. When enabled, `_normalize_stack()` auto-detects architecture from LoRA keys, then normalizes each LoRA's state dict to a canonical format. For Z-Image, fused QKV tensors are split into separate Q/K/V for per-component conflict analysis, then re-fused at patch application time.
|
||||
|
||||
**Tech Stack:** Python, PyTorch, ComfyUI node API
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Add `_detect_architecture()` to `_LoRAMergeBase`
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:65-94` (after `_get_compute_device`, before `_load_lora`)
|
||||
|
||||
**Step 1: Add `import re` at top of file**
|
||||
|
||||
The file currently has no `import re`. Add it after `import time` (line 13):
|
||||
|
||||
```python
|
||||
import re
|
||||
```
|
||||
|
||||
Check first — if `import re` already exists, skip this step.
|
||||
|
||||
**Step 2: Add `_detect_architecture` static method**
|
||||
|
||||
Insert after `_get_compute_device()` (after line 83), before `_load_lora()`:
|
||||
|
||||
```python
|
||||
@staticmethod
|
||||
def _detect_architecture(lora_sd):
|
||||
"""
|
||||
Detect model architecture from LoRA key patterns.
|
||||
Returns: 'zimage', 'flux', 'wan', 'sdxl', 'ltx', 'qwen_image', or 'unknown'.
|
||||
"""
|
||||
keys = list(lora_sd.keys())
|
||||
keys_str = ' '.join(k.lower() for k in keys)
|
||||
|
||||
# Z-Image Turbo (Lumina2): layers.N with attention patterns
|
||||
# Handles: diffusion_model.layers.N, single_transformer_blocks.N (non-FLUX),
|
||||
# lora_unet_layers_N (Musubi Tuner)
|
||||
if any('diffusion_model.layers.' in k and ('attention' in k or 'adaln' in k.lower())
|
||||
for k in keys):
|
||||
return 'zimage'
|
||||
if any('lora_unet_layers_' in k and 'attention' in k.lower() for k in keys):
|
||||
return 'zimage'
|
||||
# single_transformer_blocks WITHOUT transformer. prefix = Z-Image
|
||||
if any('single_transformer_blocks' in k and 'transformer.single_transformer_blocks' not in k
|
||||
for k in keys):
|
||||
return 'zimage'
|
||||
|
||||
# FLUX: double/single blocks in various trainer formats
|
||||
if any('transformer.single_transformer_blocks' in k or 'transformer.transformer_blocks' in k
|
||||
for k in keys):
|
||||
return 'flux'
|
||||
if any('transformer_single_transformer_blocks' in k or 'transformer_double_blocks' in k
|
||||
for k in keys):
|
||||
return 'flux'
|
||||
if any('double_blocks' in k or 'single_blocks' in k for k in keys):
|
||||
return 'flux'
|
||||
|
||||
# Wan: blocks.N with self_attn/cross_attn/ffn
|
||||
if any(('blocks.' in k or 'blocks_' in k) and
|
||||
any(x in k for x in ['self_attn', 'cross_attn', 'ffn'])
|
||||
for k in keys):
|
||||
return 'wan'
|
||||
|
||||
# LTX Video: transformer_blocks with attn1/attn2 and adaln_single
|
||||
if any('adaln_single' in k for k in keys):
|
||||
return 'ltx'
|
||||
if any('transformer_blocks' in k and ('attn1' in k or 'attn2' in k)
|
||||
and not any('transformer_blocks' in k2 and 'img_mlp' in k2 for k2 in keys)
|
||||
for k in keys):
|
||||
# transformer_blocks with attn1/attn2 but not Qwen-Image patterns
|
||||
return 'ltx'
|
||||
|
||||
# Qwen-Image: transformer_blocks with img_mlp/txt_mlp/img_mod/txt_mod
|
||||
if any('transformer_blocks' in k and
|
||||
any(x in k for x in ['img_mlp', 'txt_mlp', 'img_mod', 'txt_mod', 'add_q_proj'])
|
||||
for k in keys):
|
||||
return 'qwen_image'
|
||||
|
||||
# SDXL: text encoders or UNet block patterns
|
||||
if 'lora_te1_' in keys_str or 'lora_te2_' in keys_str:
|
||||
return 'sdxl'
|
||||
if any('input_blocks' in k or 'output_blocks' in k for k in keys):
|
||||
return 'sdxl'
|
||||
if any('down_blocks' in k or 'up_blocks' in k for k in keys):
|
||||
return 'sdxl'
|
||||
|
||||
return 'unknown'
|
||||
```
|
||||
|
||||
**Step 3: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: add _detect_architecture() to _LoRAMergeBase"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: Add `_normalize_keys()` for Z-Image
|
||||
|
||||
This is the most complex normalizer — handles QKV tensor splitting and output projection remapping.
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` (add method to `_LoRAMergeBase`, after `_detect_architecture`)
|
||||
|
||||
**Step 1: Add `_normalize_keys_zimage` static method**
|
||||
|
||||
Insert after `_detect_architecture()`:
|
||||
|
||||
```python
|
||||
@staticmethod
|
||||
def _normalize_keys_zimage(lora_sd):
|
||||
"""
|
||||
Normalize Z-Image Turbo (Lumina2) LoRA keys to a canonical format.
|
||||
|
||||
Handles:
|
||||
1. Split fused QKV (attention.qkv) into separate to_q/to_k/to_v
|
||||
for per-component conflict analysis during merge.
|
||||
2. Remap attention.out -> attention.to_out.0 (diffusers convention).
|
||||
3. Standardize Musubi Tuner format (lora_unet_layers_N_...).
|
||||
4. Ensure diffusion_model.layers.N prefix.
|
||||
|
||||
Returns new dict with normalized keys. Original dict is not modified.
|
||||
"""
|
||||
normalized = {}
|
||||
processed = set()
|
||||
|
||||
# First pass: standardize prefixes (Musubi Tuner -> canonical)
|
||||
prefix_fixed = {}
|
||||
for k, v in lora_sd.items():
|
||||
new_k = k
|
||||
# Musubi Tuner: lora_unet_layers_N_attention_... -> diffusion_model.layers.N.attention...
|
||||
if new_k.startswith('lora_unet_'):
|
||||
new_k = new_k.replace('lora_unet_', 'diffusion_model.')
|
||||
# Convert underscore-separated to dot-separated for known patterns
|
||||
new_k = re.sub(r'layers_(\d+)_', r'layers.\1.', new_k)
|
||||
new_k = re.sub(r'attention_', 'attention.', new_k)
|
||||
new_k = re.sub(r'feed_forward_', 'feed_forward.', new_k)
|
||||
# Ensure diffusion_model. prefix
|
||||
if new_k.startswith('layers.'):
|
||||
new_k = 'diffusion_model.' + new_k
|
||||
prefix_fixed[new_k] = v
|
||||
|
||||
# Second pass: split fused QKV and remap output projection
|
||||
# Find all layer indices
|
||||
layer_pattern = re.compile(r'((?:diffusion_model\.)?layers\.(\d+)\.attention)\.')
|
||||
layers_seen = set()
|
||||
for k in prefix_fixed:
|
||||
m = layer_pattern.search(k)
|
||||
if m:
|
||||
layers_seen.add((m.group(1), int(m.group(2))))
|
||||
|
||||
for base, layer_idx in layers_seen:
|
||||
# --- Split fused QKV ---
|
||||
for lora_fmt in [('.lora_A.weight', '.lora_B.weight'),
|
||||
('.lora_up.weight', '.lora_down.weight'),
|
||||
('.lora_B.weight', '.lora_A.weight'),
|
||||
('.lora.up.weight', '.lora.down.weight')]:
|
||||
# Try each LoRA format for the fused QKV key
|
||||
down_suffix, up_suffix = lora_fmt[0], lora_fmt[1]
|
||||
qkv_down_key = f"{base}.qkv{down_suffix}"
|
||||
qkv_up_key = f"{base}.qkv{up_suffix}"
|
||||
|
||||
if qkv_down_key in prefix_fixed and qkv_up_key in prefix_fixed:
|
||||
qkv_down = prefix_fixed[qkv_down_key] # [rank*3, dim] or [rank, dim]
|
||||
qkv_up = prefix_fixed[qkv_up_key] # [dim*3, rank] or [dim, rank]
|
||||
|
||||
# Determine which is A (down) and B (up) by shape
|
||||
# lora_A/down: [rank, in_features] or [rank*3, in_features]
|
||||
# lora_B/up: [out_features, rank] or [out_features*3, rank]
|
||||
rank_down = qkv_down.shape[0]
|
||||
rank_up = qkv_up.shape[1] if qkv_up.dim() == 2 else qkv_up.shape[0]
|
||||
|
||||
# Split down matrix (rank*3 -> 3 x rank, or just chunk into 3)
|
||||
if rank_down % 3 == 0:
|
||||
q_down, k_down, v_down = torch.chunk(qkv_down, 3, dim=0)
|
||||
else:
|
||||
# Can't split evenly — skip
|
||||
break
|
||||
|
||||
# Split up matrix (out_features*3 -> 3 x out_features)
|
||||
out_dim = qkv_up.shape[0]
|
||||
if out_dim % 3 == 0:
|
||||
q_up, k_up, v_up = torch.chunk(qkv_up, 3, dim=0)
|
||||
else:
|
||||
break
|
||||
|
||||
for comp, comp_down, comp_up in [('to_q', q_down, q_up),
|
||||
('to_k', k_down, k_up),
|
||||
('to_v', v_down, v_up)]:
|
||||
normalized[f"{base}.{comp}{down_suffix}"] = comp_down
|
||||
normalized[f"{base}.{comp}{up_suffix}"] = comp_up
|
||||
|
||||
# Copy alpha (same for all three components)
|
||||
alpha_key = f"{base}.qkv.alpha"
|
||||
if alpha_key in prefix_fixed:
|
||||
for comp in ('to_q', 'to_k', 'to_v'):
|
||||
normalized[f"{base}.{comp}.alpha"] = prefix_fixed[alpha_key]
|
||||
processed.add(alpha_key)
|
||||
|
||||
processed.add(qkv_down_key)
|
||||
processed.add(qkv_up_key)
|
||||
break # Found QKV format, don't try others
|
||||
|
||||
# --- Remap output projection: attention.out -> attention.to_out.0 ---
|
||||
for lora_fmt in [('.lora_A.weight', '.lora_B.weight'),
|
||||
('.lora_up.weight', '.lora_down.weight'),
|
||||
('.lora_B.weight', '.lora_A.weight'),
|
||||
('.lora.up.weight', '.lora.down.weight')]:
|
||||
sfx_a, sfx_b = lora_fmt
|
||||
out_a = f"{base}.out{sfx_a}"
|
||||
out_b = f"{base}.out{sfx_b}"
|
||||
if out_a in prefix_fixed and out_b in prefix_fixed:
|
||||
normalized[f"{base}.to_out.0{sfx_a}"] = prefix_fixed[out_a]
|
||||
normalized[f"{base}.to_out.0{sfx_b}"] = prefix_fixed[out_b]
|
||||
processed.add(out_a)
|
||||
processed.add(out_b)
|
||||
|
||||
out_alpha = f"{base}.out.alpha"
|
||||
if out_alpha in prefix_fixed:
|
||||
normalized[f"{base}.to_out.0.alpha"] = prefix_fixed[out_alpha]
|
||||
processed.add(out_alpha)
|
||||
break
|
||||
|
||||
# Pass through all unprocessed keys
|
||||
for k, v in prefix_fixed.items():
|
||||
if k not in processed:
|
||||
normalized[k] = v
|
||||
|
||||
return normalized
|
||||
```
|
||||
|
||||
**Step 2: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: add _normalize_keys_zimage() for QKV splitting and key remap"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Add `_normalize_keys()` for FLUX, Wan, SDXL, LTX, Qwen-Image
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` (add methods to `_LoRAMergeBase`, after `_normalize_keys_zimage`)
|
||||
|
||||
**Step 1: Add `_normalize_keys_flux` static method**
|
||||
|
||||
```python
|
||||
@staticmethod
|
||||
def _normalize_keys_flux(lora_sd):
|
||||
"""
|
||||
Normalize FLUX LoRA keys from various trainer formats to canonical format.
|
||||
|
||||
Canonical format: diffusion_model.double_blocks.N.* / diffusion_model.single_blocks.N.*
|
||||
|
||||
Handles:
|
||||
- AI-Toolkit: transformer.transformer_blocks.N -> double_blocks.N
|
||||
transformer.single_transformer_blocks.N -> single_blocks.N
|
||||
- Kohya: lora_transformer_double_blocks_N -> double_blocks.N
|
||||
lora_transformer_single_transformer_blocks_N -> single_blocks.N
|
||||
- Standard: double_blocks.N / single_blocks.N (ensure prefix)
|
||||
"""
|
||||
normalized = {}
|
||||
for k, v in lora_sd.items():
|
||||
new_k = k
|
||||
|
||||
# AI-Toolkit format
|
||||
# transformer.single_transformer_blocks.N -> diffusion_model.single_blocks.N
|
||||
new_k = re.sub(
|
||||
r'^transformer\.single_transformer_blocks\.(\d+)\.',
|
||||
r'diffusion_model.single_blocks.\1.', new_k)
|
||||
# transformer.transformer_blocks.N -> diffusion_model.double_blocks.N
|
||||
new_k = re.sub(
|
||||
r'^transformer\.transformer_blocks\.(\d+)\.',
|
||||
r'diffusion_model.double_blocks.\1.', new_k)
|
||||
|
||||
# Kohya underscore format
|
||||
# lora_transformer_single_transformer_blocks_N_... -> diffusion_model.single_blocks.N....
|
||||
m = re.match(r'^lora_transformer_single_transformer_blocks_(\d+)_(.*)', new_k)
|
||||
if m:
|
||||
block_num = m.group(1)
|
||||
rest = m.group(2).replace('_', '.')
|
||||
new_k = f"diffusion_model.single_blocks.{block_num}.{rest}"
|
||||
m = re.match(r'^lora_transformer_double_blocks_(\d+)_(.*)', new_k)
|
||||
if m:
|
||||
block_num = m.group(1)
|
||||
rest = m.group(2).replace('_', '.')
|
||||
new_k = f"diffusion_model.double_blocks.{block_num}.{rest}"
|
||||
|
||||
# Ensure diffusion_model. prefix for standard format
|
||||
if new_k.startswith('double_blocks.') or new_k.startswith('single_blocks.'):
|
||||
new_k = 'diffusion_model.' + new_k
|
||||
|
||||
# Generic transformer. prefix -> diffusion_model.
|
||||
if new_k.startswith('transformer.'):
|
||||
new_k = new_k.replace('transformer.', 'diffusion_model.', 1)
|
||||
|
||||
normalized[new_k] = v
|
||||
return normalized
|
||||
```
|
||||
|
||||
**Step 2: Add `_normalize_keys_wan` static method**
|
||||
|
||||
```python
|
||||
@staticmethod
|
||||
def _normalize_keys_wan(lora_sd):
|
||||
"""
|
||||
Normalize Wan LoRA keys from various trainer formats to canonical format.
|
||||
|
||||
Canonical format: diffusion_model.blocks.N.{self_attn,cross_attn,ffn}.*
|
||||
|
||||
Handles LyCORIS, diffusers, Fun LoRA, finetrainer formats.
|
||||
Also applies RS-LoRA alpha compensation if detected.
|
||||
"""
|
||||
normalized = {}
|
||||
for k, v in lora_sd.items():
|
||||
new_k = k
|
||||
|
||||
# LyCORIS/aitoolkit format
|
||||
if new_k.startswith('lycoris_blocks_'):
|
||||
new_k = new_k.replace('lycoris_blocks_', 'blocks.')
|
||||
new_k = new_k.replace('_cross_attn_', '.cross_attn.')
|
||||
new_k = new_k.replace('_self_attn_', '.self_attn.')
|
||||
new_k = new_k.replace('_ffn_net_0_proj', '.ffn.0')
|
||||
new_k = new_k.replace('_ffn_net_2', '.ffn.2')
|
||||
new_k = new_k.replace('to_out_0', 'o')
|
||||
|
||||
# Diffusers format prefixes
|
||||
if new_k.startswith('transformer.'):
|
||||
new_k = new_k.replace('transformer.', 'diffusion_model.', 1)
|
||||
if new_k.startswith('pipe.dit.'):
|
||||
new_k = new_k.replace('pipe.dit.', 'diffusion_model.', 1)
|
||||
if new_k.startswith('blocks.'):
|
||||
new_k = 'diffusion_model.' + new_k
|
||||
if new_k.startswith('vace_blocks.'):
|
||||
new_k = 'diffusion_model.' + new_k
|
||||
|
||||
# Common diffusers cleanup
|
||||
new_k = new_k.replace('.default.', '.')
|
||||
new_k = new_k.replace('.diff_m', '.modulation.diff')
|
||||
new_k = new_k.replace('base_model.model.', 'diffusion_model.')
|
||||
|
||||
# Fun LoRA format: lora_unet__blocks_N_...
|
||||
if new_k.startswith('lora_unet__'):
|
||||
parts = new_k.split('.')
|
||||
main_part = parts[0]
|
||||
weight_type = '.'.join(parts[1:]) if len(parts) > 1 else None
|
||||
|
||||
if 'blocks_' in main_part:
|
||||
components = main_part[len('lora_unet__'):].split('_')
|
||||
rebuilt = 'diffusion_model'
|
||||
|
||||
if components[0] == 'blocks':
|
||||
rebuilt += f".blocks.{components[1]}"
|
||||
idx = 2
|
||||
if idx < len(components):
|
||||
if (components[idx] == 'self' and idx + 1 < len(components)
|
||||
and components[idx + 1] == 'attn'):
|
||||
rebuilt += '.self_attn'
|
||||
idx += 2
|
||||
elif (components[idx] == 'cross' and idx + 1 < len(components)
|
||||
and components[idx + 1] == 'attn'):
|
||||
rebuilt += '.cross_attn'
|
||||
idx += 2
|
||||
elif components[idx] == 'ffn':
|
||||
rebuilt += '.ffn'
|
||||
idx += 1
|
||||
if idx < len(components):
|
||||
component = components[idx]
|
||||
idx += 1
|
||||
if idx < len(components) and components[idx] == 'img':
|
||||
component += '_img'
|
||||
idx += 1
|
||||
rebuilt += f'.{component}'
|
||||
|
||||
if weight_type:
|
||||
if weight_type == 'alpha':
|
||||
rebuilt += '.alpha'
|
||||
elif weight_type in ('lora_down.weight', 'lora_down'):
|
||||
rebuilt += '.lora_A.weight'
|
||||
elif weight_type in ('lora_up.weight', 'lora_up'):
|
||||
rebuilt += '.lora_B.weight'
|
||||
else:
|
||||
rebuilt += f'.{weight_type}'
|
||||
if not rebuilt.endswith('.weight'):
|
||||
rebuilt += '.weight'
|
||||
new_k = rebuilt
|
||||
else:
|
||||
new_k = main_part.replace('lora_unet__', 'diffusion_model.')
|
||||
new_k = new_k.replace('_', '.')
|
||||
if weight_type:
|
||||
if weight_type == 'alpha':
|
||||
new_k += '.alpha'
|
||||
elif weight_type in ('lora_down.weight', 'lora_down'):
|
||||
new_k += '.lora_A.weight'
|
||||
elif weight_type in ('lora_up.weight', 'lora_up'):
|
||||
new_k += '.lora_B.weight'
|
||||
else:
|
||||
new_k += f'.{weight_type}'
|
||||
if not new_k.endswith('.weight'):
|
||||
new_k += '.weight'
|
||||
|
||||
# Finetrainer format
|
||||
new_k = new_k.replace('.attn1.to_q.', '.self_attn.q.')
|
||||
new_k = new_k.replace('.attn1.to_k.', '.self_attn.k.')
|
||||
new_k = new_k.replace('.attn1.to_v.', '.self_attn.v.')
|
||||
new_k = new_k.replace('.attn1.to_out.0.', '.self_attn.o.')
|
||||
new_k = new_k.replace('.attn2.to_q.', '.cross_attn.q.')
|
||||
new_k = new_k.replace('.attn2.to_k.', '.cross_attn.k.')
|
||||
new_k = new_k.replace('.attn2.to_v.', '.cross_attn.v.')
|
||||
new_k = new_k.replace('.attn2.to_out.0.', '.cross_attn.o.')
|
||||
|
||||
normalized[new_k] = v
|
||||
|
||||
# RS-LoRA compensation: detect and fix alpha scaling
|
||||
rs_marker = 'diffusion_model.blocks.0.cross_attn.k.lora_A.weight'
|
||||
if rs_marker in normalized:
|
||||
rank = normalized[rs_marker].shape[0]
|
||||
import math
|
||||
corrected_alpha = torch.tensor(rank * (rank ** 0.5))
|
||||
for nk in list(normalized.keys()):
|
||||
if nk.endswith('.lora_A.weight'):
|
||||
alpha_key = nk.replace('.lora_A.weight', '.alpha')
|
||||
if alpha_key not in normalized:
|
||||
normalized[alpha_key] = corrected_alpha
|
||||
|
||||
return normalized
|
||||
```
|
||||
|
||||
**Step 3: Add `_normalize_keys_sdxl` static method**
|
||||
|
||||
```python
|
||||
@staticmethod
|
||||
def _normalize_keys_sdxl(lora_sd):
|
||||
"""
|
||||
Normalize SDXL LoRA keys to canonical format.
|
||||
|
||||
Canonical format: lora_unet_* / lora_te1_* / lora_te2_* (Kohya convention).
|
||||
|
||||
Handles diffusers-format keys (down_blocks.N, up_blocks.N, mid_block).
|
||||
"""
|
||||
normalized = {}
|
||||
for k, v in lora_sd.items():
|
||||
new_k = k
|
||||
|
||||
# Diffusers format: text_encoder.* -> lora_te1_*, text_encoder_2.* -> lora_te2_*
|
||||
if new_k.startswith('text_encoder_2.'):
|
||||
new_k = 'lora_te2_' + new_k[len('text_encoder_2.'):].replace('.', '_')
|
||||
elif new_k.startswith('text_encoder.'):
|
||||
new_k = 'lora_te1_' + new_k[len('text_encoder.'):].replace('.', '_')
|
||||
|
||||
# Diffusers UNet: unet.* -> lora_unet_*
|
||||
if new_k.startswith('unet.'):
|
||||
new_k = 'lora_unet_' + new_k[len('unet.'):].replace('.', '_')
|
||||
|
||||
# base_model.model -> strip prefix
|
||||
new_k = new_k.replace('base_model.model.', '')
|
||||
|
||||
normalized[new_k] = v
|
||||
return normalized
|
||||
```
|
||||
|
||||
**Step 4: Add `_normalize_keys_ltx` static method**
|
||||
|
||||
```python
|
||||
@staticmethod
|
||||
def _normalize_keys_ltx(lora_sd):
|
||||
"""
|
||||
Normalize LTX Video LoRA keys to canonical format.
|
||||
|
||||
Canonical format: diffusion_model.transformer_blocks.N.attn1/attn2.to_q/to_k/to_v.*
|
||||
|
||||
LTX uses standard separate Q/K/V — only prefix standardization needed.
|
||||
"""
|
||||
normalized = {}
|
||||
for k, v in lora_sd.items():
|
||||
new_k = k
|
||||
|
||||
# Kohya format: lora_unet_transformer_blocks_N_... -> diffusion_model.transformer_blocks.N...
|
||||
if new_k.startswith('lora_unet_'):
|
||||
new_k = new_k.replace('lora_unet_', 'diffusion_model.')
|
||||
# Convert underscores back to dots for known structural segments
|
||||
new_k = re.sub(r'transformer_blocks_(\d+)_', r'transformer_blocks.\1.', new_k)
|
||||
new_k = re.sub(r'attn(\d)_', r'attn\1.', new_k)
|
||||
new_k = re.sub(r'to_(q|k|v|out)_', r'to_\1.', new_k)
|
||||
new_k = re.sub(r'ff_net_(\d+)_', r'ff.net.\1.', new_k)
|
||||
|
||||
# Diffusers format: unet.* -> diffusion_model.*
|
||||
if new_k.startswith('unet.'):
|
||||
new_k = new_k.replace('unet.', 'diffusion_model.', 1)
|
||||
|
||||
# Ensure diffusion_model. prefix
|
||||
if new_k.startswith('transformer_blocks.'):
|
||||
new_k = 'diffusion_model.' + new_k
|
||||
|
||||
# base_model.model prefix
|
||||
new_k = new_k.replace('base_model.model.', 'diffusion_model.')
|
||||
|
||||
normalized[new_k] = v
|
||||
return normalized
|
||||
```
|
||||
|
||||
**Step 5: Add `_normalize_keys_qwen_image` static method**
|
||||
|
||||
```python
|
||||
@staticmethod
|
||||
def _normalize_keys_qwen_image(lora_sd):
|
||||
"""
|
||||
Normalize Qwen-Image LoRA keys to canonical format.
|
||||
|
||||
Canonical format: diffusion_model.transformer_blocks.N.*
|
||||
|
||||
Qwen-Image uses separate Q/K/V with dual-stream (image+text) attention.
|
||||
Supports transformer.*, lycoris_*, and direct key formats.
|
||||
"""
|
||||
normalized = {}
|
||||
for k, v in lora_sd.items():
|
||||
new_k = k
|
||||
|
||||
# LyCORIS format: lycoris_transformer_blocks_N_... -> diffusion_model.transformer_blocks.N...
|
||||
if new_k.startswith('lycoris_'):
|
||||
new_k = new_k.replace('lycoris_', 'diffusion_model.')
|
||||
new_k = re.sub(r'transformer_blocks_(\d+)_', r'transformer_blocks.\1.', new_k)
|
||||
# Restore dots for known component names
|
||||
for comp in ['to_q', 'to_k', 'to_v', 'add_q_proj', 'add_k_proj', 'add_v_proj',
|
||||
'img_mlp', 'txt_mlp', 'img_mod', 'txt_mod', 'img_norm1', 'img_norm2',
|
||||
'txt_norm1', 'txt_norm2']:
|
||||
new_k = new_k.replace(f'_{comp}_', f'.{comp}.')
|
||||
if new_k.endswith(f'_{comp}'):
|
||||
new_k = new_k[:-len(f'_{comp}')] + f'.{comp}'
|
||||
|
||||
# transformer.* -> diffusion_model.*
|
||||
if new_k.startswith('transformer.'):
|
||||
new_k = new_k.replace('transformer.', 'diffusion_model.', 1)
|
||||
|
||||
# Ensure diffusion_model. prefix
|
||||
if new_k.startswith('transformer_blocks.'):
|
||||
new_k = 'diffusion_model.' + new_k
|
||||
|
||||
# base_model.model prefix
|
||||
new_k = new_k.replace('base_model.model.', 'diffusion_model.')
|
||||
|
||||
normalized[new_k] = v
|
||||
return normalized
|
||||
```
|
||||
|
||||
**Step 6: Add dispatcher `_normalize_keys` method**
|
||||
|
||||
```python
|
||||
@classmethod
|
||||
def _normalize_keys(cls, lora_sd, architecture):
|
||||
"""
|
||||
Dispatch to architecture-specific key normalizer.
|
||||
Returns a new dict with normalized keys.
|
||||
"""
|
||||
if architecture == 'zimage':
|
||||
return cls._normalize_keys_zimage(lora_sd)
|
||||
elif architecture == 'flux':
|
||||
return cls._normalize_keys_flux(lora_sd)
|
||||
elif architecture == 'wan':
|
||||
return cls._normalize_keys_wan(lora_sd)
|
||||
elif architecture == 'sdxl':
|
||||
return cls._normalize_keys_sdxl(lora_sd)
|
||||
elif architecture == 'ltx':
|
||||
return cls._normalize_keys_ltx(lora_sd)
|
||||
elif architecture == 'qwen_image':
|
||||
return cls._normalize_keys_qwen_image(lora_sd)
|
||||
return lora_sd # unknown — pass through unchanged
|
||||
```
|
||||
|
||||
**Step 7: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 8: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: add key normalizers for FLUX, Wan, SDXL, LTX, Qwen-Image and dispatcher"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Wire normalization into `_normalize_stack()` and add `normalize_keys` input
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` — `LoRAOptimizer.INPUT_TYPES` (~line 538), `_normalize_stack` (~line 463), `optimize_merge` (~line 1272), `_compute_cache_key` (~line 576), `IS_CHANGED` (~line 598)
|
||||
|
||||
**Step 1: Add `normalize_keys` to INPUT_TYPES optional inputs**
|
||||
|
||||
In `LoRAOptimizer.INPUT_TYPES`, add after the `svd_device` entry (before the closing `}` of optional):
|
||||
|
||||
```python
|
||||
"normalize_keys": (["disabled", "enabled"], {
|
||||
"default": "disabled",
|
||||
"tooltip": "Auto-detect LoRA architecture and normalize keys before merging. Enable when mixing LoRAs from different trainers (e.g., Kohya + AI-Toolkit) or for Z-Image QKV fusion."
|
||||
}),
|
||||
```
|
||||
|
||||
**Step 2: Add `normalize_keys` to `_compute_cache_key` signature and hash**
|
||||
|
||||
Modify `_compute_cache_key` to include `normalize_keys` in the hash:
|
||||
|
||||
Change the hash update line from:
|
||||
```python
|
||||
h.update(f"|os={output_strength}|csm={clip_strength_multiplier}|as={auto_strength}|om={optimization_mode}|cp={compress_patches}|sd={svd_device}".encode())
|
||||
```
|
||||
To:
|
||||
```python
|
||||
h.update(f"|os={output_strength}|csm={clip_strength_multiplier}|as={auto_strength}|om={optimization_mode}|cp={compress_patches}|sd={svd_device}|nk={normalize_keys}".encode())
|
||||
```
|
||||
|
||||
Add `normalize_keys="disabled"` parameter to the `_compute_cache_key` signature.
|
||||
|
||||
**Step 3: Add `normalize_keys` to `IS_CHANGED` signature**
|
||||
|
||||
Add `normalize_keys="disabled"` parameter and pass it through to `_compute_cache_key`.
|
||||
|
||||
**Step 4: Modify `_normalize_stack` to accept and use `normalize_keys`**
|
||||
|
||||
Add `normalize_keys="disabled"` parameter to `_normalize_stack`. After building the normalized list (both tuple and dict paths), add before the `return`:
|
||||
|
||||
```python
|
||||
# Architecture-aware key normalization
|
||||
if normalize_keys == "enabled" and len(normalized) > 0:
|
||||
first_sd = normalized[0]["lora"]
|
||||
arch = self._detect_architecture(first_sd)
|
||||
if arch != "unknown":
|
||||
logging.info(f"[LoRA Optimizer] Architecture detected: {arch}")
|
||||
logging.info(f"[LoRA Optimizer] Normalizing keys for {len(normalized)} LoRAs...")
|
||||
for item in normalized:
|
||||
item["lora"] = self._normalize_keys(item["lora"], arch)
|
||||
self._detected_arch = arch
|
||||
else:
|
||||
logging.info("[LoRA Optimizer] Architecture: unknown (no key normalization applied)")
|
||||
self._detected_arch = None
|
||||
else:
|
||||
self._detected_arch = None
|
||||
```
|
||||
|
||||
Note: We store the detected architecture on `self._detected_arch` for use in Z-Image re-fusion later.
|
||||
|
||||
**Step 5: Pass `normalize_keys` through `optimize_merge`**
|
||||
|
||||
Add `normalize_keys="disabled"` to `optimize_merge` signature. Pass it to:
|
||||
- `self._normalize_stack(lora_stack, normalize_keys=normalize_keys)`
|
||||
- `self._compute_cache_key(..., normalize_keys=normalize_keys)`
|
||||
|
||||
**Step 6: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 7: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: wire normalize_keys input through stack normalization and caching"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: Add Z-Image re-fusion at patch application time
|
||||
|
||||
When the detected architecture is `zimage`, merged patches for `to_q`/`to_k`/`to_v` must be re-fused to `qkv` before applying to the model, because the model expects fused QKV.
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` — add `_refuse_zimage_patches` method, call it in `optimize_merge` before `model.add_patches()`
|
||||
|
||||
**Step 1: Add `_refuse_zimage_patches` method to `_LoRAMergeBase`**
|
||||
|
||||
Insert after `_normalize_keys`:
|
||||
|
||||
```python
|
||||
@staticmethod
|
||||
def _refuse_zimage_patches(patches):
|
||||
"""
|
||||
Re-fuse split to_q/to_k/to_v patches back into fused QKV patches
|
||||
for Z-Image Turbo models. Also remaps to_out.0 -> out.
|
||||
|
||||
Called after merging, before applying patches to the model.
|
||||
Returns a new dict with fused patches.
|
||||
"""
|
||||
fused = {}
|
||||
qkv_groups = {} # base -> {comp: (key, patch)}
|
||||
out_remap = {} # base -> (old_key, patch)
|
||||
|
||||
for key, patch in patches.items():
|
||||
# Handle both string keys and tuple keys
|
||||
if isinstance(key, tuple):
|
||||
key_str = key[0]
|
||||
else:
|
||||
key_str = key
|
||||
|
||||
# Detect to_q/to_k/to_v patterns in the key
|
||||
m = re.search(r'(layers\.\d+\.attention)\.to_(q|k|v)(?:\.|$)', key_str)
|
||||
if m:
|
||||
base = m.group(1)
|
||||
comp = m.group(2)
|
||||
if base not in qkv_groups:
|
||||
qkv_groups[base] = {}
|
||||
qkv_groups[base][comp] = (key, patch)
|
||||
continue
|
||||
|
||||
# Detect to_out.0 -> out remap
|
||||
m_out = re.search(r'(layers\.\d+\.attention)\.to_out\.0(?:\.|$)', key_str)
|
||||
if m_out:
|
||||
base_out = m_out.group(1)
|
||||
new_key_str = key_str.replace('.to_out.0', '.out')
|
||||
if isinstance(key, tuple):
|
||||
new_key = (new_key_str,) + key[1:]
|
||||
else:
|
||||
new_key = new_key_str
|
||||
fused[new_key] = patch
|
||||
continue
|
||||
|
||||
# Not a QKV or out key — pass through
|
||||
fused[key] = patch
|
||||
|
||||
# Fuse QKV groups
|
||||
for base, comps in qkv_groups.items():
|
||||
if len(comps) == 3 and 'q' in comps and 'k' in comps and 'v' in comps:
|
||||
# All three components present — fuse
|
||||
q_key, q_patch = comps['q']
|
||||
k_key, k_patch = comps['k']
|
||||
v_key, v_patch = comps['v']
|
||||
|
||||
# Build the fused key name
|
||||
if isinstance(q_key, tuple):
|
||||
fused_key_str = re.sub(r'\.to_q(?=\.|$)', '.qkv', q_key[0])
|
||||
fused_key = (fused_key_str,) + q_key[1:]
|
||||
else:
|
||||
fused_key_str = re.sub(r'\.to_q(?=\.|$)', '.qkv', q_key)
|
||||
fused_key = fused_key_str
|
||||
|
||||
# Handle different patch formats
|
||||
if isinstance(q_patch, tuple) and q_patch[0] == "diff":
|
||||
# Full-rank diff patch: ("diff", (tensor,))
|
||||
q_diff = q_patch[1][0]
|
||||
k_diff = k_patch[1][0]
|
||||
v_diff = v_patch[1][0]
|
||||
fused_diff = torch.cat([q_diff, k_diff, v_diff], dim=0)
|
||||
fused[fused_key] = ("diff", (fused_diff,))
|
||||
elif isinstance(q_patch, LoRAAdapter):
|
||||
# Low-rank patch — concatenate A and B matrices
|
||||
q_data = q_patch.lora_data
|
||||
k_data = k_patch.lora_data
|
||||
v_data = v_patch.lora_data
|
||||
# (mat_up, mat_down, alpha, mid, dora_scale, bias)
|
||||
fused_up = torch.cat([q_data[0], k_data[0], v_data[0]], dim=0)
|
||||
fused_down = torch.cat([q_data[1], k_data[1], v_data[1]], dim=0)
|
||||
fused_alpha = q_data[2] # alpha is the same for all components
|
||||
fused_mid = None # Mid not expected for attention
|
||||
fused_patch = LoRAAdapter(set(), (fused_up, fused_down, fused_alpha, fused_mid, None, None))
|
||||
fused[fused_key] = fused_patch
|
||||
else:
|
||||
# Unknown patch format — pass through unfused
|
||||
for comp_key, comp_patch in [(q_key, q_patch), (k_key, k_patch), (v_key, v_patch)]:
|
||||
fused[comp_key] = comp_patch
|
||||
else:
|
||||
# Incomplete QKV group — pass through individual components
|
||||
for comp, (comp_key, comp_patch) in comps.items():
|
||||
fused[comp_key] = comp_patch
|
||||
|
||||
return fused
|
||||
```
|
||||
|
||||
**Step 2: Call re-fusion in `optimize_merge` before `add_patches`**
|
||||
|
||||
In `optimize_merge`, after Pass 2 merge is done and before the "Apply patches" section (around line 1798), add:
|
||||
|
||||
```python
|
||||
# Re-fuse Z-Image QKV patches if architecture normalization was used
|
||||
if getattr(self, '_detected_arch', None) == 'zimage':
|
||||
if len(model_patches) > 0:
|
||||
model_patches = self._refuse_zimage_patches(model_patches)
|
||||
logging.info(f"[LoRA Optimizer] Re-fused Z-Image QKV patches ({len(model_patches)} model patches)")
|
||||
```
|
||||
|
||||
**Step 3: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: add Z-Image QKV re-fusion at patch application time"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 6: Add architecture info to report and update report signature
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` — `_build_report`, `optimize_merge`
|
||||
|
||||
**Step 1: Add architecture info to `_build_report`**
|
||||
|
||||
Add `detected_arch=None, normalize_keys="disabled"` parameters to `_build_report`. After the header lines (after line 1119 `"=" * 50`), add:
|
||||
|
||||
```python
|
||||
# Architecture info (when normalization is enabled)
|
||||
if normalize_keys == "enabled" and detected_arch:
|
||||
arch_names = {
|
||||
'zimage': 'Z-Image Turbo (Lumina2)',
|
||||
'flux': 'FLUX',
|
||||
'wan': 'Wan 2.1/2.2',
|
||||
'sdxl': 'SDXL',
|
||||
'ltx': 'LTX Video',
|
||||
'qwen_image': 'Qwen-Image',
|
||||
}
|
||||
lines.append(f"Architecture: {arch_names.get(detected_arch, detected_arch)} (auto-detected)")
|
||||
lines.append(f"Key normalization: enabled")
|
||||
lines.append("")
|
||||
```
|
||||
|
||||
**Step 2: Pass architecture info from `optimize_merge` to `_build_report`**
|
||||
|
||||
In the `_build_report` call in `optimize_merge`, add:
|
||||
|
||||
```python
|
||||
detected_arch=getattr(self, '_detected_arch', None),
|
||||
normalize_keys=normalize_keys,
|
||||
```
|
||||
|
||||
**Step 3: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: add architecture info to analysis report"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 7: Final verification
|
||||
|
||||
**Step 1: Full syntax check**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 2: Verify all node classes are registered**
|
||||
|
||||
Run: `python -c "from lora_optimizer import NODE_CLASS_MAPPINGS; print(list(NODE_CLASS_MAPPINGS.keys()))"`
|
||||
Expected: `['LoRAStack', 'LoRAOptimizer']`
|
||||
|
||||
**Step 3: Verify the new methods exist on `_LoRAMergeBase`**
|
||||
|
||||
Run:
|
||||
```python
|
||||
python -c "
|
||||
from lora_optimizer import _LoRAMergeBase
|
||||
for m in ['_detect_architecture', '_normalize_keys', '_normalize_keys_zimage',
|
||||
'_normalize_keys_flux', '_normalize_keys_wan', '_normalize_keys_sdxl',
|
||||
'_normalize_keys_ltx', '_normalize_keys_qwen_image',
|
||||
'_refuse_zimage_patches']:
|
||||
assert hasattr(_LoRAMergeBase, m), f'Missing {m}'
|
||||
print('All methods present')
|
||||
"
|
||||
```
|
||||
Expected: `All methods present`
|
||||
|
||||
**Step 4: Verify `normalize_keys` appears in INPUT_TYPES**
|
||||
|
||||
Run:
|
||||
```python
|
||||
python -c "
|
||||
from lora_optimizer import LoRAOptimizer
|
||||
inputs = LoRAOptimizer.INPUT_TYPES()
|
||||
assert 'normalize_keys' in inputs['optional'], 'normalize_keys not in optional inputs'
|
||||
print('normalize_keys input present')
|
||||
"
|
||||
```
|
||||
Expected: `normalize_keys input present`
|
||||
|
||||
**Step 5: Commit any final cleanup**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "chore: final verification pass"
|
||||
```
|
||||
@@ -1,58 +0,0 @@
|
||||
# Design: Base Model Filter for LoRAStackDynamic
|
||||
|
||||
## Problem
|
||||
|
||||
Users with large LoRA collections (mixed SDXL, Flux, Wan, etc.) see all LoRAs in the dropdown regardless of which model they're using. ComfyUI-Lora-Manager already categorizes LoRAs by base model type — we can leverage its API to filter the dropdown.
|
||||
|
||||
## Solution
|
||||
|
||||
Add a `base_model_filter` dropdown to `LoRAStackDynamic` that queries Lora Manager's API to filter LoRA dropdowns. Hidden if Lora Manager is not installed.
|
||||
|
||||
### Data flow
|
||||
|
||||
1. Node created → JS calls `GET /api/lm/loras/base-models` → populates filter dropdown with `["All"] + model types`
|
||||
2. User selects a model type (e.g., "Wan") → JS calls `GET /api/lm/loras/list?base_model=Wan&page_size=10000` → gets filtered LoRA filenames
|
||||
3. JS updates all `lora_name_N` COMBO widget options with filtered list
|
||||
4. User selects "All" → JS restores original full LoRA list
|
||||
|
||||
### Graceful degradation
|
||||
|
||||
- Lora Manager not installed: API call fails → filter widget hidden, dropdowns show all LoRAs
|
||||
- API call fails mid-session: revert to full LoRA list, no error
|
||||
- Filter set to "All": show full LoRA list
|
||||
|
||||
## Files to modify
|
||||
|
||||
- `lora_optimizer.py` — add `base_model_filter` widget to `LoRAStackDynamic.INPUT_TYPES()` (hidden input, not used in execution)
|
||||
- `js/lora_stack_dynamic.js` — add filter logic: detect Lora Manager, populate filter, intercept changes, update LoRA dropdowns
|
||||
|
||||
## Implementation details
|
||||
|
||||
### Python side
|
||||
|
||||
Add to `LoRAStackDynamic.INPUT_TYPES()` optional inputs:
|
||||
```python
|
||||
"base_model_filter": (["All"], {
|
||||
"default": "All",
|
||||
"tooltip": "Filter LoRA list by base model type. Requires ComfyUI-Lora-Manager."
|
||||
})
|
||||
```
|
||||
|
||||
The `build_stack` method ignores this widget — it's purely for UI filtering. The options list starts with just `["All"]` and is dynamically populated by the JS extension.
|
||||
|
||||
### JS side (in lora_stack_dynamic.js)
|
||||
|
||||
On `LoRAStackDynamic` node creation:
|
||||
|
||||
1. **Detect Lora Manager**: `fetch("/api/lm/loras/base-models")`
|
||||
- Success → populate `base_model_filter` widget options with `["All", ...names]`
|
||||
- Failure → hide the `base_model_filter` widget
|
||||
|
||||
2. **Intercept filter changes**: same `Object.defineProperty` pattern already used for `mode` and `lora_count`
|
||||
|
||||
3. **On filter change**:
|
||||
- If "All" → restore original LoRA list (cached from initial widget options)
|
||||
- Otherwise → `fetch("/api/lm/loras/list?base_model=X&page_size=10000")` → extract filenames → update all `lora_name_N` widget options
|
||||
- Preserve currently selected values if they exist in the filtered list
|
||||
|
||||
4. **Cache**: store the full LoRA list on first load so "All" restoration is instant
|
||||
@@ -1,102 +0,0 @@
|
||||
# LoRA Conflict Editor Node — Design
|
||||
|
||||
## Overview
|
||||
|
||||
An advanced node that sits between the stacker and optimizer. It analyzes pairwise conflicts between LoRAs, auto-suggests per-LoRA conflict modes, and lets users manually override both conflict modes and merge strategy. Pass-through design: enriches the LORA_STACK with resolved settings.
|
||||
|
||||
## Position in Workflow
|
||||
|
||||
Stacker → **LoRA Conflict Editor** → Optimizer
|
||||
|
||||
## Inputs
|
||||
|
||||
| Input | Type | Default | Notes |
|
||||
|-------|------|---------|-------|
|
||||
| `lora_stack` | LORA_STACK | required | From any stacker node |
|
||||
| `merge_strategy` | dropdown | `auto` | `auto` / `ties` / `weighted_average` / `weighted_sum` |
|
||||
| `conflict_mode_1..10` | dropdown | `auto` | `auto` / `all` / `low_conflict` / `high_conflict` |
|
||||
|
||||
`conflict_mode_i` maps to LoRAs by stack order (1st LoRA = `conflict_mode_1`).
|
||||
|
||||
## Outputs
|
||||
|
||||
| Output | Type | Notes |
|
||||
|--------|------|-------|
|
||||
| `lora_stack` | LORA_STACK | Enriched with resolved `conflict_mode` per entry |
|
||||
| `analysis_report` | STRING | Pairwise conflicts, per-LoRA stats, applied settings |
|
||||
| `merge_strategy` | STRING | Resolved strategy for optimizer to consume |
|
||||
|
||||
## Override Logic (Manual-First)
|
||||
|
||||
- `auto`: node computes suggestion each run, writes to **output stack only**, never modifies the widget
|
||||
- `all` / `low_conflict` / `high_conflict`: user's explicit choice, passed through unchanged
|
||||
- Widget values are never overwritten — no re-run surprises
|
||||
|
||||
Same for `merge_strategy`: `auto` runs heuristic, explicit values pass through.
|
||||
|
||||
## Execution Logic
|
||||
|
||||
1. Normalize incoming stack (load LoRAs, cache)
|
||||
2. Run pairwise conflict analysis (sign conflicts, cosine similarity) — reuse optimizer's `_sample_conflict` logic
|
||||
3. For each LoRA position `i`:
|
||||
- If `conflict_mode_i == "auto"`: compute suggestion from conflict stats, bake into output
|
||||
- Otherwise: use the explicit value
|
||||
4. For `merge_strategy`:
|
||||
- If `auto`: auto-select based on avg conflict ratio (same heuristic as optimizer)
|
||||
- Otherwise: pass through
|
||||
5. Build analysis report
|
||||
6. Output enriched stack + report + resolved strategy
|
||||
|
||||
## Auto-Suggestion Heuristic
|
||||
|
||||
For each LoRA, compute average pairwise sign conflict ratio with all others:
|
||||
- < 15% avg conflict → suggest `all` (compatible, no filtering needed)
|
||||
- 15–40% avg conflict → suggest `low_conflict` (conservative, safe)
|
||||
- > 40% avg conflict → suggest `high_conflict` (very different, let it dominate contested regions)
|
||||
|
||||
Report shows reasoning for each suggestion.
|
||||
|
||||
## Optimizer Integration
|
||||
|
||||
Add optional `merge_strategy_override` input (STRING) to the optimizer. When connected and non-empty, it overrides auto-detection. When not connected, existing behavior is unchanged.
|
||||
|
||||
## JS Widget Visibility
|
||||
|
||||
Show all 10 `conflict_mode_i` dropdowns. Hide unused ones. Since stack size is only known at execution time, two options:
|
||||
- Show all 10 always (unused ones are harmless)
|
||||
- After first run, use PromptServer to hide extras (nice but complex)
|
||||
|
||||
Start with showing all 10. Can refine later.
|
||||
|
||||
## Analysis Report Format
|
||||
|
||||
```
|
||||
============ LORA CONFLICT EDITOR ============
|
||||
|
||||
--- Stack (3 LoRAs) ---
|
||||
1. style_lora.safetensors (strength: 1.0)
|
||||
2. character_lora.safetensors (strength: 0.8)
|
||||
3. lighting_lora.safetensors (strength: 0.6)
|
||||
|
||||
--- Pairwise Conflicts ---
|
||||
style × character:
|
||||
Overlap: 12,450 positions
|
||||
Sign conflicts: 3,112 (25.0%)
|
||||
Cosine similarity: 0.62
|
||||
style × lighting:
|
||||
Overlap: 8,200 positions
|
||||
Sign conflicts: 820 (10.0%)
|
||||
Cosine similarity: 0.85
|
||||
character × lighting:
|
||||
Overlap: 9,100 positions
|
||||
Sign conflicts: 3,640 (40.0%)
|
||||
Cosine similarity: 0.31
|
||||
|
||||
--- Applied Conflict Modes ---
|
||||
1. style_lora: all (auto-suggested — low avg conflict 17.5%)
|
||||
2. character_lora: high_conflict (manual override)
|
||||
3. lighting_lora: low_conflict (auto-suggested — high avg conflict 25.0%)
|
||||
|
||||
--- Merge Strategy ---
|
||||
Applied: ties (auto — avg conflict 25.0% > threshold)
|
||||
```
|
||||
@@ -1,511 +0,0 @@
|
||||
# LoRA Conflict Editor — Implementation Plan
|
||||
|
||||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||||
|
||||
**Goal:** Add a LoRAConflictEditor node that analyzes pairwise LoRA conflicts, auto-suggests per-LoRA conflict modes, and lets users manually override conflict modes and merge strategy before the optimizer runs.
|
||||
|
||||
**Architecture:** Pass-through node between stacker and optimizer. Takes LORA_STACK in, runs pairwise conflict analysis (reusing optimizer internals), resolves conflict_mode per LoRA (auto-suggest or manual), and outputs an enriched LORA_STACK. Inherits from `_LoRAMergeBase` to reuse LoRA loading, normalization, and conflict analysis methods.
|
||||
|
||||
**Tech Stack:** Python (ComfyUI node), JS (widget visibility)
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Define LoRAConflictEditor class skeleton
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:3013-3025` (NODE_CLASS_MAPPINGS registration)
|
||||
- Modify: `lora_optimizer.py` (new class, insert before NODE_CLASS_MAPPINGS)
|
||||
|
||||
**Step 1: Add the class with INPUT_TYPES**
|
||||
|
||||
Insert before `NODE_CLASS_MAPPINGS` (line 3013):
|
||||
|
||||
```python
|
||||
class LoRAConflictEditor(_LoRAMergeBase):
|
||||
"""
|
||||
Analyzes pairwise conflicts between LoRAs and lets users control
|
||||
per-LoRA conflict modes and merge strategy. Sits between a stacker
|
||||
and the optimizer.
|
||||
"""
|
||||
|
||||
MAX_LORAS = 10
|
||||
|
||||
def __init__(self):
|
||||
self.loaded_loras = {}
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
inputs = {
|
||||
"required": {
|
||||
"lora_stack": ("LORA_STACK", {
|
||||
"tooltip": "Connect a LoRA Stack node here. The editor will analyze conflicts between these LoRAs."
|
||||
}),
|
||||
"merge_strategy": (["auto", "ties", "weighted_average", "weighted_sum"], {
|
||||
"default": "auto",
|
||||
"tooltip": "Merge strategy to pass to the optimizer. "
|
||||
"'auto': let the optimizer decide based on conflict analysis. "
|
||||
"'ties': force TIES merging (good for high-conflict stacks). "
|
||||
"'weighted_average': force simple averaging (good for compatible LoRAs). "
|
||||
"'weighted_sum': force direct addition (preserves all weights exactly)."
|
||||
}),
|
||||
}
|
||||
}
|
||||
for i in range(1, cls.MAX_LORAS + 1):
|
||||
inputs["required"][f"conflict_mode_{i}"] = (["auto", "all", "low_conflict", "high_conflict"], {
|
||||
"default": "auto",
|
||||
"tooltip": f"LoRA #{i} conflict filter. "
|
||||
f"'auto': node suggests based on analysis. "
|
||||
f"'all': apply everywhere. "
|
||||
f"'low_conflict': only where this LoRA agrees with the majority. "
|
||||
f"'high_conflict': only where this LoRA disagrees."
|
||||
})
|
||||
return inputs
|
||||
|
||||
RETURN_TYPES = ("LORA_STACK", "STRING", "STRING")
|
||||
RETURN_NAMES = ("lora_stack", "analysis_report", "merge_strategy")
|
||||
FUNCTION = "analyze_and_enrich"
|
||||
CATEGORY = "loaders/lora"
|
||||
DESCRIPTION = "Analyzes LoRA conflicts and lets you control per-LoRA conflict modes and merge strategy. Connect between a LoRA Stack and the LoRA Optimizer."
|
||||
```
|
||||
|
||||
**Step 2: Register in NODE_CLASS_MAPPINGS**
|
||||
|
||||
Add to `NODE_CLASS_MAPPINGS` (line 3013):
|
||||
```python
|
||||
"LoRAConflictEditor": LoRAConflictEditor,
|
||||
```
|
||||
|
||||
Add to `NODE_DISPLAY_NAME_MAPPINGS` (line 3020):
|
||||
```python
|
||||
"LoRAConflictEditor": "LoRA Conflict Editor",
|
||||
```
|
||||
|
||||
**Step 3: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 4: Commit**
|
||||
```
|
||||
feat: add LoRAConflictEditor class skeleton
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: Implement analyze_and_enrich method
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` (LoRAConflictEditor class)
|
||||
|
||||
**Step 1: Add the analysis and enrichment logic**
|
||||
|
||||
Add to LoRAConflictEditor class:
|
||||
|
||||
```python
|
||||
def analyze_and_enrich(self, lora_stack, merge_strategy="auto", **kwargs):
|
||||
"""
|
||||
Analyze pairwise conflicts, resolve conflict modes, enrich the stack.
|
||||
"""
|
||||
if not lora_stack or len(lora_stack) == 0:
|
||||
return (lora_stack, "No LoRAs in stack.", merge_strategy)
|
||||
|
||||
# Normalize to consistent dict format
|
||||
normalized = self._normalize_stack(lora_stack)
|
||||
active_loras = [item for item in normalized if item["strength"] != 0]
|
||||
|
||||
if len(active_loras) == 0:
|
||||
return (lora_stack, "No active LoRAs (all zero strength).", merge_strategy)
|
||||
|
||||
n_loras = len(active_loras)
|
||||
|
||||
# --- Pairwise conflict analysis ---
|
||||
compute_device = self._get_compute_device()
|
||||
use_gpu = (compute_device is not None and compute_device.type == "cuda")
|
||||
|
||||
# Collect all unique prefixes across all LoRAs
|
||||
all_prefixes = set()
|
||||
for item in active_loras:
|
||||
for key in item["lora"].keys():
|
||||
prefix = key.rsplit(".", 1)[0] if "." in key else key
|
||||
# Only track LoRA weight prefixes (contain lora_up or lora_down)
|
||||
if "lora_up" in key or "lora_down" in key:
|
||||
# Strip the lora_up/lora_down suffix to get the base prefix
|
||||
base = key
|
||||
for suffix in (".lora_up.weight", ".lora_down.weight",
|
||||
".lora_A.weight", ".lora_B.weight",
|
||||
".lora_up", ".lora_down"):
|
||||
if base.endswith(suffix):
|
||||
base = base[:-len(suffix)]
|
||||
break
|
||||
all_prefixes.add(base)
|
||||
|
||||
# Per-LoRA pairwise accumulators
|
||||
pairs = [(i, j) for i in range(n_loras) for j in range(i + 1, n_loras)]
|
||||
pair_accum = {(i, j): [0, 0, 0.0, 0.0, 0.0] for i, j in pairs}
|
||||
per_lora_key_count = [0] * n_loras
|
||||
|
||||
for lora_prefix in all_prefixes:
|
||||
# Compute diffs for each LoRA at this prefix
|
||||
diffs = []
|
||||
indices = []
|
||||
for idx, item in enumerate(active_loras):
|
||||
lora_info = self._get_lora_key_info(item["lora"], lora_prefix)
|
||||
if lora_info is None:
|
||||
continue
|
||||
mat_up, mat_down, alpha, mid = lora_info
|
||||
rank = mat_down.shape[0]
|
||||
|
||||
if use_gpu:
|
||||
mat_up = mat_up.to(compute_device)
|
||||
mat_down = mat_down.to(compute_device)
|
||||
if mid is not None:
|
||||
mid = mid.to(compute_device)
|
||||
|
||||
if mid is not None:
|
||||
final_shape = [mat_down.shape[1], mat_down.shape[0],
|
||||
mid.shape[2], mid.shape[3]]
|
||||
mat_down = (
|
||||
torch.mm(
|
||||
mat_down.transpose(0, 1).flatten(start_dim=1).float(),
|
||||
mid.transpose(0, 1).flatten(start_dim=1).float(),
|
||||
)
|
||||
.reshape(final_shape)
|
||||
.transpose(0, 1)
|
||||
)
|
||||
|
||||
diff = torch.mm(
|
||||
mat_up.flatten(start_dim=1).float(),
|
||||
mat_down.flatten(start_dim=1).float()
|
||||
)
|
||||
del mat_up, mat_down
|
||||
diff = diff * (alpha / rank)
|
||||
diffs.append(diff)
|
||||
indices.append(idx)
|
||||
per_lora_key_count[idx] += 1
|
||||
|
||||
# Pairwise conflicts at this prefix
|
||||
for a in range(len(diffs)):
|
||||
for b in range(a + 1, len(diffs)):
|
||||
i, j = indices[a], indices[b]
|
||||
ov, conf, dot, na_sq, nb_sq = self._sample_conflict(diffs[a], diffs[b], device=compute_device)
|
||||
pair_accum[(i, j)][0] += ov
|
||||
pair_accum[(i, j)][1] += conf
|
||||
pair_accum[(i, j)][2] += dot
|
||||
pair_accum[(i, j)][3] += na_sq
|
||||
pair_accum[(i, j)][4] += nb_sq
|
||||
|
||||
del diffs # free GPU memory
|
||||
|
||||
# --- Compute per-LoRA avg conflict ratio ---
|
||||
per_lora_conflict = [0.0] * n_loras
|
||||
per_lora_pair_count = [0] * n_loras
|
||||
pairwise_results = []
|
||||
total_overlap = 0
|
||||
total_conflict = 0
|
||||
|
||||
for i, j in pairs:
|
||||
ov, conf, dot, na_sq, nb_sq = pair_accum[(i, j)]
|
||||
ratio = conf / ov if ov > 0 else 0.0
|
||||
total_overlap += ov
|
||||
total_conflict += conf
|
||||
|
||||
# Cosine similarity
|
||||
cos_sim = None
|
||||
if na_sq > 0 and nb_sq > 0:
|
||||
cos_sim = dot / (math.sqrt(na_sq) * math.sqrt(nb_sq))
|
||||
|
||||
pairwise_results.append({
|
||||
"i": i, "j": j,
|
||||
"name_a": active_loras[i]["name"],
|
||||
"name_b": active_loras[j]["name"],
|
||||
"overlap": ov,
|
||||
"conflicts": conf,
|
||||
"ratio": ratio,
|
||||
"cosine_sim": cos_sim,
|
||||
})
|
||||
|
||||
per_lora_conflict[i] += ratio
|
||||
per_lora_conflict[j] += ratio
|
||||
per_lora_pair_count[i] += 1
|
||||
per_lora_pair_count[j] += 1
|
||||
|
||||
avg_conflict = total_conflict / total_overlap if total_overlap > 0 else 0.0
|
||||
|
||||
per_lora_avg_conflict = []
|
||||
for i in range(n_loras):
|
||||
if per_lora_pair_count[i] > 0:
|
||||
per_lora_avg_conflict.append(per_lora_conflict[i] / per_lora_pair_count[i])
|
||||
else:
|
||||
per_lora_avg_conflict.append(0.0)
|
||||
|
||||
# --- Resolve conflict modes ---
|
||||
resolved_modes = []
|
||||
for i in range(n_loras):
|
||||
cm = kwargs.get(f"conflict_mode_{i+1}", "auto")
|
||||
if cm != "auto":
|
||||
resolved_modes.append((cm, "manual"))
|
||||
else:
|
||||
# Auto-suggest based on avg pairwise conflict
|
||||
avg_c = per_lora_avg_conflict[i]
|
||||
if avg_c > 0.40:
|
||||
resolved_modes.append(("high_conflict", f"auto — high avg conflict {avg_c:.1%}"))
|
||||
elif avg_c > 0.15:
|
||||
resolved_modes.append(("low_conflict", f"auto — moderate avg conflict {avg_c:.1%}"))
|
||||
else:
|
||||
resolved_modes.append(("all", f"auto — low avg conflict {avg_c:.1%}"))
|
||||
|
||||
# --- Resolve merge strategy ---
|
||||
if merge_strategy == "auto":
|
||||
if avg_conflict > 0.25:
|
||||
resolved_strategy = "ties"
|
||||
strategy_reason = f"auto — avg conflict {avg_conflict:.1%} > 25%"
|
||||
else:
|
||||
resolved_strategy = "weighted_average"
|
||||
strategy_reason = f"auto — avg conflict {avg_conflict:.1%} <= 25%"
|
||||
else:
|
||||
resolved_strategy = merge_strategy
|
||||
strategy_reason = "manual"
|
||||
|
||||
# --- Build enriched output stack ---
|
||||
# Output in the same format as input, with conflict_mode baked in
|
||||
enriched = []
|
||||
first = lora_stack[0]
|
||||
if isinstance(first, dict):
|
||||
for i, item in enumerate(active_loras):
|
||||
enriched_item = dict(item)
|
||||
enriched_item["conflict_mode"] = resolved_modes[i][0]
|
||||
enriched.append(enriched_item)
|
||||
else:
|
||||
for i, item in enumerate(active_loras):
|
||||
enriched.append((
|
||||
item["name"],
|
||||
item["strength"],
|
||||
item["clip_strength"] if item["clip_strength"] is not None else item["strength"],
|
||||
resolved_modes[i][0],
|
||||
))
|
||||
|
||||
# --- Build analysis report ---
|
||||
report = self._build_conflict_report(
|
||||
active_loras, pairwise_results, per_lora_avg_conflict,
|
||||
per_lora_key_count, resolved_modes, resolved_strategy, strategy_reason
|
||||
)
|
||||
|
||||
return (enriched, report, resolved_strategy)
|
||||
```
|
||||
|
||||
**Step 2: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 3: Commit**
|
||||
```
|
||||
feat: implement LoRAConflictEditor analyze_and_enrich method
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Implement _build_conflict_report
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` (LoRAConflictEditor class)
|
||||
|
||||
**Step 1: Add the report builder**
|
||||
|
||||
Add to LoRAConflictEditor class:
|
||||
|
||||
```python
|
||||
@staticmethod
|
||||
def _build_conflict_report(active_loras, pairwise_results, per_lora_avg_conflict,
|
||||
per_lora_key_count, resolved_modes, resolved_strategy, strategy_reason):
|
||||
"""Format conflict analysis as a readable report."""
|
||||
n = len(active_loras)
|
||||
lines = []
|
||||
lines.append("=" * 46)
|
||||
lines.append("LORA CONFLICT EDITOR - ANALYSIS REPORT")
|
||||
lines.append("=" * 46)
|
||||
|
||||
# Stack overview
|
||||
lines.append("")
|
||||
lines.append(f"--- Stack ({n} LoRAs) ---")
|
||||
for i, item in enumerate(active_loras):
|
||||
lines.append(f" {i+1}. {item['name']} (strength: {item['strength']})")
|
||||
|
||||
# Pairwise conflicts
|
||||
if pairwise_results:
|
||||
lines.append("")
|
||||
lines.append("--- Pairwise Conflicts ---")
|
||||
for pc in pairwise_results:
|
||||
# Short names for readability
|
||||
name_a = os.path.splitext(os.path.basename(pc["name_a"]))[0]
|
||||
name_b = os.path.splitext(os.path.basename(pc["name_b"]))[0]
|
||||
lines.append(f" {name_a} \u00d7 {name_b}:")
|
||||
lines.append(f" Overlapping positions: {pc['overlap']:,}")
|
||||
lines.append(f" Sign conflicts: {pc['conflicts']:,} ({pc['ratio']:.1%})")
|
||||
if pc["cosine_sim"] is not None:
|
||||
lines.append(f" Cosine similarity: {pc['cosine_sim']:.3f}")
|
||||
|
||||
# Applied conflict modes
|
||||
lines.append("")
|
||||
lines.append("--- Applied Conflict Modes ---")
|
||||
for i, item in enumerate(active_loras):
|
||||
name = os.path.splitext(os.path.basename(item["name"]))[0]
|
||||
mode, reason = resolved_modes[i]
|
||||
lines.append(f" {i+1}. {name}: {mode} ({reason})")
|
||||
|
||||
# Merge strategy
|
||||
lines.append("")
|
||||
lines.append("--- Merge Strategy ---")
|
||||
lines.append(f" Applied: {resolved_strategy} ({strategy_reason})")
|
||||
|
||||
return "\n".join(lines)
|
||||
```
|
||||
|
||||
**Step 2: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 3: Commit**
|
||||
```
|
||||
feat: add conflict report builder for LoRAConflictEditor
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Add merge_strategy_override input to LoRAOptimizer
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:1422-1467` (LoRAOptimizer.INPUT_TYPES optional inputs)
|
||||
- Modify: `lora_optimizer.py:2249` (optimize_merge signature)
|
||||
- Modify: `lora_optimizer.py:2509-2510` (where _auto_select_params is called)
|
||||
- Modify: `lora_optimizer.py:1497-1509` (IS_CHANGED and _compute_cache_key)
|
||||
|
||||
**Step 1: Add optional input**
|
||||
|
||||
In `INPUT_TYPES` optional dict (after `sparsification_density`, line ~1466):
|
||||
|
||||
```python
|
||||
"merge_strategy_override": ("STRING", {
|
||||
"default": "",
|
||||
"forceInput": True,
|
||||
"tooltip": "Connect the merge_strategy output from a LoRA Conflict Editor to override the optimizer's auto-detected strategy."
|
||||
}),
|
||||
```
|
||||
|
||||
Note: `forceInput: True` hides the widget and only shows the input socket — the user can't type in it, only connect it.
|
||||
|
||||
**Step 2: Update optimize_merge signature**
|
||||
|
||||
Add `merge_strategy_override=""` parameter to the signature (line 2249).
|
||||
|
||||
**Step 3: Apply the override**
|
||||
|
||||
After `_auto_select_params` is called (around line 2509-2510), add:
|
||||
|
||||
```python
|
||||
# Apply merge strategy override from Conflict Editor
|
||||
if merge_strategy_override and merge_strategy_override in ("ties", "weighted_average", "weighted_sum"):
|
||||
mode = merge_strategy_override
|
||||
reasoning.append(f"Merge mode overridden to '{mode}' by Conflict Editor")
|
||||
```
|
||||
|
||||
For per_prefix mode, also apply in the per-prefix decision (around line 2622-2625): when override is set, use it instead of per-prefix auto-selection.
|
||||
|
||||
**Step 4: Update cache key and IS_CHANGED**
|
||||
|
||||
Add `merge_strategy_override` to `_compute_cache_key` hash string and to `IS_CHANGED` parameter list.
|
||||
|
||||
**Step 5: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 6: Commit**
|
||||
```
|
||||
feat: add merge_strategy_override input to LoRAOptimizer
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: Add JS widget visibility for conflict editor
|
||||
|
||||
**Files:**
|
||||
- Modify: `js/lora_stack_dynamic.js` (add extension for LoRAConflictEditor)
|
||||
|
||||
**Step 1: Add the extension**
|
||||
|
||||
Append to the JS file:
|
||||
|
||||
```javascript
|
||||
app.registerExtension({
|
||||
name: "LoRAOptimizer.LoRAConflictEditor",
|
||||
nodeCreated(node) {
|
||||
if (node.comfyClass !== "LoRAConflictEditor") return;
|
||||
|
||||
// After execution, hide conflict_mode slots beyond the actual LoRA count.
|
||||
// Since we don't know the count until execution, show all 10 initially.
|
||||
// The node's output report tells the user which slots are active.
|
||||
|
||||
// No dynamic hiding needed for now — all 10 slots are always visible.
|
||||
// Unused slots (beyond the stack size) are simply ignored.
|
||||
// Future: use PromptServer messages to hide extras after first run.
|
||||
},
|
||||
});
|
||||
```
|
||||
|
||||
**Step 2: Verify syntax**
|
||||
|
||||
Run: `node --check js/lora_stack_dynamic.js`
|
||||
Expected: no output (success)
|
||||
|
||||
**Step 3: Commit**
|
||||
```
|
||||
feat: add JS extension stub for LoRAConflictEditor
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 6: Integration test and final verification
|
||||
|
||||
**Files:**
|
||||
- Verify: `lora_optimizer.py` (full syntax check)
|
||||
- Verify: `js/lora_stack_dynamic.js` (syntax check)
|
||||
|
||||
**Step 1: Full syntax verification**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Run: `node --check js/lora_stack_dynamic.js`
|
||||
|
||||
**Step 2: Verify NODE_CLASS_MAPPINGS includes all nodes**
|
||||
|
||||
Run: `grep -A 10 "NODE_CLASS_MAPPINGS" lora_optimizer.py`
|
||||
|
||||
Expected: Should include `LoRAConflictEditor`.
|
||||
|
||||
**Step 3: Verify data flow**
|
||||
|
||||
Check that:
|
||||
- LoRAConflictEditor accepts LORA_STACK and outputs LORA_STACK (pass-through with enrichment)
|
||||
- Output LORA_STACK contains `conflict_mode` in each entry
|
||||
- `merge_strategy` STRING output can connect to optimizer's `merge_strategy_override` STRING input
|
||||
- Optimizer's `forceInput: True` means the widget is hidden unless connected
|
||||
|
||||
**Step 4: Commit**
|
||||
```
|
||||
chore: verify LoRAConflictEditor integration
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Verification Checklist
|
||||
|
||||
- [ ] LoRAConflictEditor appears in ComfyUI node list under "loaders/lora"
|
||||
- [ ] Node accepts LORA_STACK from any stacker (dict or tuple format)
|
||||
- [ ] Analysis report shows pairwise conflicts and applied modes
|
||||
- [ ] `auto` conflict modes produce sensible suggestions
|
||||
- [ ] Manual conflict modes (`all`/`low_conflict`/`high_conflict`) pass through unchanged and don't get overwritten on re-run
|
||||
- [ ] `merge_strategy` output connects to optimizer's `merge_strategy_override` input
|
||||
- [ ] Optimizer respects the override when connected
|
||||
- [ ] Node works with 1 LoRA (no pairwise analysis, all defaults)
|
||||
- [ ] Old workflows without the editor still work (backward compatible)
|
||||
@@ -1,88 +0,0 @@
|
||||
# Design: WanVideoLoRAOptimizer Node
|
||||
|
||||
## Problem
|
||||
|
||||
The WanVideoWrapper custom node has its own LoRA pipeline using `WANVIDEOMODEL` and `WANVIDLORA` types, separate from ComfyUI's standard `MODEL`/`CLIP` types. Users want to merge multiple WanVideo LoRAs using our optimizer's algorithms (TIES, DARE, auto-strength, SVD compression) but can't connect `WANVIDEOMODEL` to our `MODEL` input.
|
||||
|
||||
## Solution
|
||||
|
||||
A thin subclass of `LoRAOptimizer` that bridges the type gap. All merging logic is inherited — the subclass only overrides input/output types and wraps the return tuple.
|
||||
|
||||
### Why this works
|
||||
|
||||
1. **Merging algorithms are architecture-agnostic** — TIES, DARE/DELLA, SVD compression, auto-strength, and conflict detection all operate on generic tensors.
|
||||
2. **Wan key normalization already exists** — `_detect_architecture` returns `'wan'`, and `_normalize_keys_wan` handles LyCORIS, diffusers, Fun LoRA, finetrainer, and RS-LoRA formats.
|
||||
3. **Generic key mapping works** — `comfy.lora.model_lora_keys_unet()` creates `key_map[k[:-len(".weight")]] = k` for all model architectures, which matches standardized WanVideo LoRA key prefixes.
|
||||
4. **CLIP is already optional** — the parent class handles `clip=None` gracefully.
|
||||
5. **Patch format is compatible** — both `("diff", (tensor,))` tuples and `LoRAAdapter` objects work with WanVideoWrapper's `set_lora_params()` (which checks `hasattr(lora_obj, "weights")` and `tuple[0] == "diff"`).
|
||||
|
||||
## Implementation
|
||||
|
||||
### New class: `WanVideoLoRAOptimizer` (~30 lines)
|
||||
|
||||
Inherits from `LoRAOptimizer`. Overrides:
|
||||
|
||||
- **`INPUT_TYPES`**: Copies parent's inputs, replaces `MODEL` → `WANVIDEOMODEL`, removes `clip`, `clip_strength_multiplier`, `free_vram_between_passes`. Changes defaults: `normalize_keys="enabled"`, `cache_patches="disabled"`.
|
||||
- **`RETURN_TYPES`**: `("WANVIDEOMODEL", "STRING", "LORA_DATA")` — drops CLIP output.
|
||||
- **`optimize_merge()`**: Calls `super().optimize_merge(model, lora_stack, ..., clip=None, clip_strength_multiplier=0)`, remaps return tuple to skip CLIP.
|
||||
- **`CATEGORY`**: `"loaders/lora"` (same as parent).
|
||||
- **`DESCRIPTION`**: Updated for WanVideo context.
|
||||
|
||||
### Inputs
|
||||
|
||||
| Input | Type | Required | Default | Notes |
|
||||
|-------|------|----------|---------|-------|
|
||||
| model | WANVIDEOMODEL | Yes | — | From WanVideoModelLoader |
|
||||
| lora_stack | LORA_STACK | Yes | — | Same LoRA Stack nodes |
|
||||
| output_strength | FLOAT | Yes | 1.0 | Master strength |
|
||||
| auto_strength | enum | No | disabled | Same |
|
||||
| optimization_mode | enum | No | per_prefix | Same |
|
||||
| cache_patches | enum | No | **disabled** | Changed: video models are large |
|
||||
| compress_patches | enum | No | non_ties | Same |
|
||||
| svd_device | enum | No | gpu | Same |
|
||||
| normalize_keys | enum | No | **enabled** | Changed: WanVideo LoRAs vary |
|
||||
| sparsification | enum | No | disabled | Same |
|
||||
| sparsification_density | FLOAT | No | 0.7 | Same |
|
||||
| merge_strategy_override | STRING | No | "" | Same |
|
||||
|
||||
### Outputs
|
||||
|
||||
| Output | Type |
|
||||
|--------|------|
|
||||
| model | WANVIDEOMODEL |
|
||||
| analysis_report | STRING |
|
||||
| lora_data | LORA_DATA |
|
||||
|
||||
### Registration
|
||||
|
||||
```python
|
||||
NODE_CLASS_MAPPINGS["WanVideoLoRAOptimizer"] = WanVideoLoRAOptimizer
|
||||
NODE_DISPLAY_NAME_MAPPINGS["WanVideoLoRAOptimizer"] = "WanVideo LoRA Optimizer"
|
||||
```
|
||||
|
||||
### Workflow
|
||||
|
||||
```
|
||||
[WanVideoModelLoader] → WANVIDEOMODEL → [WanVideoLoRAOptimizer] → WANVIDEOMODEL → [WanVideoSampler]
|
||||
↑
|
||||
[LoRA Stack] ─────────────┘
|
||||
```
|
||||
|
||||
Chaining with individual (non-merged) LoRAs:
|
||||
```
|
||||
[WanVideoLoraSelect] → WANVIDLORA → [WanVideoModelLoader] → WANVIDEOMODEL → [WanVideoLoRAOptimizer] → WANVIDEOMODEL → [Sampler]
|
||||
↑
|
||||
[LoRA Stack] ───────────────┘
|
||||
```
|
||||
|
||||
## Files to modify
|
||||
|
||||
- `lora_optimizer.py` — add `WanVideoLoRAOptimizer` class (~30 lines), register in mappings
|
||||
|
||||
## Verification
|
||||
|
||||
1. Load WanVideo model (no LoRA) → optimizer with 2+ WanVideo LoRAs → sampler. Verify LoRA effect appears in generated video.
|
||||
2. Empty LoRA stack → optimizer returns unmodified WANVIDEOMODEL, no crash.
|
||||
3. Mixed trainer formats (e.g., one LyCORIS + one diffusers LoRA) with normalize_keys=enabled → both processed correctly.
|
||||
4. LORA_DATA output → SaveMergedLoRA → loadable via WanVideoLoraSelect.
|
||||
5. Chain: individual LoRAs via model loader + merged LoRAs via optimizer — both effects visible.
|
||||
@@ -1,153 +0,0 @@
|
||||
# WanVideoLoRAOptimizer Implementation Plan
|
||||
|
||||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||||
|
||||
**Goal:** Add a `WanVideoLoRAOptimizer` node that accepts `WANVIDEOMODEL` and outputs a patched `WANVIDEOMODEL` with merged LoRA patches, reusing all existing merging algorithms via inheritance.
|
||||
|
||||
**Architecture:** Thin subclass of `LoRAOptimizer` (~30 lines). Overrides `INPUT_TYPES` (replaces `MODEL` → `WANVIDEOMODEL`, removes CLIP inputs, changes defaults), `RETURN_TYPES` (drops CLIP output), and wraps `optimize_merge()` to pass `clip=None` and remap the return tuple.
|
||||
|
||||
**Tech Stack:** Python, ComfyUI node system, existing `LoRAOptimizer` class
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Add WanVideoLoRAOptimizer class
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:2938-2940` (insert after `LoRAOptimizer` class, before `SaveMergedLoRA`)
|
||||
|
||||
**Step 1: Add the class after line 2939 (the blank line between LoRAOptimizer and SaveMergedLoRA)**
|
||||
|
||||
Insert the following class:
|
||||
|
||||
```python
|
||||
class WanVideoLoRAOptimizer(LoRAOptimizer):
|
||||
"""
|
||||
WanVideo variant of the LoRA Optimizer. Accepts WANVIDEOMODEL instead of
|
||||
MODEL, skips CLIP, and applies merged LoRA patches in-memory.
|
||||
|
||||
All merging algorithms (TIES, DARE/DELLA, SVD compression, auto-strength,
|
||||
conflict analysis) are inherited from LoRAOptimizer. Wan LoRA key
|
||||
normalization (LyCORIS, diffusers, Fun LoRA, finetrainer, RS-LoRA) is
|
||||
already handled by the parent's _normalize_keys_wan.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
base = LoRAOptimizer.INPUT_TYPES()
|
||||
# Replace MODEL with WANVIDEOMODEL
|
||||
base["required"]["model"] = ("WANVIDEOMODEL", {
|
||||
"tooltip": "Your WanVideo model from WanVideoModelLoader."
|
||||
})
|
||||
# Remove CLIP-related inputs (WanVideo doesn't use CLIP)
|
||||
base["optional"].pop("clip", None)
|
||||
base["optional"].pop("clip_strength_multiplier", None)
|
||||
base["optional"].pop("free_vram_between_passes", None)
|
||||
# Change defaults for video models
|
||||
base["optional"]["cache_patches"] = (["disabled", "enabled"], {
|
||||
"default": "disabled",
|
||||
"tooltip": "Keep the merge result in memory so re-running the workflow is instant. Disabled by default for large video models to save RAM."
|
||||
})
|
||||
base["optional"]["normalize_keys"] = (["enabled", "disabled"], {
|
||||
"default": "enabled",
|
||||
"tooltip": "Normalizes LoRA keys from different training tools (LyCORIS, diffusers, finetrainer, etc.) to a common format. Enabled by default for WanVideo LoRAs."
|
||||
})
|
||||
return base
|
||||
|
||||
RETURN_TYPES = ("WANVIDEOMODEL", "STRING", "LORA_DATA")
|
||||
RETURN_NAMES = ("model", "analysis_report", "lora_data")
|
||||
CATEGORY = "loaders/lora"
|
||||
DESCRIPTION = (
|
||||
"WanVideo LoRA Optimizer — merges multiple WanVideo LoRAs using "
|
||||
"conflict-aware algorithms (TIES, DARE, auto-strength). "
|
||||
"Connect after WanVideoModelLoader, before WanVideoSampler."
|
||||
)
|
||||
|
||||
def optimize_merge(self, model, lora_stack, output_strength, **kwargs):
|
||||
kwargs.pop("clip", None)
|
||||
kwargs.pop("clip_strength_multiplier", None)
|
||||
kwargs.pop("free_vram_between_passes", None)
|
||||
result = super().optimize_merge(
|
||||
model, lora_stack, output_strength,
|
||||
clip=None, clip_strength_multiplier=0, **kwargs
|
||||
)
|
||||
# Parent returns (model, clip, report, lora_data) — drop clip
|
||||
return (result[0], result[2], result[3])
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, model, lora_stack, output_strength, **kwargs):
|
||||
kwargs.pop("clip", None)
|
||||
kwargs.pop("clip_strength_multiplier", None)
|
||||
kwargs.pop("free_vram_between_passes", None)
|
||||
return LoRAOptimizer.IS_CHANGED(
|
||||
model, lora_stack, output_strength,
|
||||
clip=None, clip_strength_multiplier=0, **kwargs
|
||||
)
|
||||
```
|
||||
|
||||
**Step 2: Verify the class is syntactically correct**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"` from the project directory.
|
||||
Expected: `OK`
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: add WanVideoLoRAOptimizer node"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: Register in NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:3513-3529` (the mappings dicts at end of file)
|
||||
|
||||
**Step 1: Add entries to both dicts**
|
||||
|
||||
In `NODE_CLASS_MAPPINGS` (after `MergedLoRAToHook` entry):
|
||||
```python
|
||||
"WanVideoLoRAOptimizer": WanVideoLoRAOptimizer,
|
||||
```
|
||||
|
||||
In `NODE_DISPLAY_NAME_MAPPINGS` (after `MergedLoRAToHook` entry):
|
||||
```python
|
||||
"WanVideoLoRAOptimizer": "WanVideo LoRA Optimizer",
|
||||
```
|
||||
|
||||
**Step 2: Verify ComfyUI can load the module**
|
||||
|
||||
Run: `cd /media/p5/Comfyui && python -c "import importlib.util; spec = importlib.util.spec_from_file_location('lora_optimizer', '/media/p5/ComfyUI-ZImage-LoRA-Merger/lora_optimizer.py'); mod = importlib.util.module_from_spec(spec); print('Import OK')"`
|
||||
Expected: `Import OK` (note: full exec may fail outside ComfyUI runtime, but import/parse should succeed)
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: register WanVideoLoRAOptimizer in node mappings"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Update README with WanVideoLoRAOptimizer documentation
|
||||
|
||||
**Files:**
|
||||
- Modify: `README.md`
|
||||
|
||||
**Step 1: Add a section for WanVideo LoRA Optimizer**
|
||||
|
||||
Add after the existing "Merged LoRA to Hook" section, following the same formatting pattern (heading, description, inputs/outputs table, collapsible details for workflow).
|
||||
|
||||
Key content:
|
||||
- Node name and purpose
|
||||
- Inputs table (WANVIDEOMODEL, LORA_STACK, same options as main optimizer minus CLIP)
|
||||
- Outputs table (WANVIDEOMODEL, analysis_report, lora_data)
|
||||
- Workflow diagram showing WanVideoModelLoader → WanVideoLoRAOptimizer → WanVideoSampler
|
||||
- Collapsible section with usage notes (chaining with individual LoRAs, normalize_keys default, etc.)
|
||||
|
||||
**Step 2: Commit**
|
||||
|
||||
```bash
|
||||
git add README.md
|
||||
git commit -m "docs: add WanVideoLoRAOptimizer to README"
|
||||
```
|
||||
@@ -1,186 +0,0 @@
|
||||
# LoRA AutoTuner & Merge Selector — Design
|
||||
|
||||
## Overview
|
||||
|
||||
Two new nodes that automatically sweep all merge parameters, score configurations using a hybrid heuristic+measurement approach, and output the best merge — with an optional selector node to pick alternatives.
|
||||
|
||||
## Nodes
|
||||
|
||||
### LoRAAutoTuner — "LoRA AutoTuner"
|
||||
|
||||
**Inherits**: `LoRAOptimizerBase`
|
||||
|
||||
**Inputs (required)**:
|
||||
- `model` — MODEL
|
||||
- `lora_stack` — LORA_STACK
|
||||
- `output_strength` — float (0.0–10.0, default 1.0)
|
||||
|
||||
**Inputs (optional)**:
|
||||
- `clip` — CLIP
|
||||
- `clip_strength_multiplier` — float (0.0–10.0, default 1.0)
|
||||
- `top_n` — int (1–10, default 3)
|
||||
|
||||
**Outputs**:
|
||||
- `MODEL` — merged model using #1 ranked config
|
||||
- `CLIP` — merged clip (if clip input provided)
|
||||
- `REPORT` — STRING, ranked table of top-N configs with scores
|
||||
- `TUNER_DATA` — custom type, top-N configs + analysis cache
|
||||
|
||||
**Behavior**:
|
||||
1. Run Pass 1 analysis once (existing `_analyze_prefix`)
|
||||
2. Score ~500–600 parameter combos via Phase 1 heuristic
|
||||
3. Merge top-N candidates via Phase 2, measure output quality
|
||||
4. Apply #1 config to model → output MODEL/CLIP
|
||||
5. Pack all top-N configs + scores + analysis cache into TUNER_DATA
|
||||
|
||||
### LoRAMergeSelector — "Merge Selector"
|
||||
|
||||
**Inherits**: `LoRAOptimizerBase`
|
||||
|
||||
**Inputs (required)**:
|
||||
- `model` — MODEL (original, unpatched)
|
||||
- `lora_stack` — LORA_STACK
|
||||
- `tuner_data` — TUNER_DATA
|
||||
- `selection` — int (1-based index into top-N)
|
||||
|
||||
**Inputs (optional)**:
|
||||
- `clip` — CLIP
|
||||
|
||||
**Outputs**:
|
||||
- `MODEL` — merged model using selected config
|
||||
- `CLIP` — merged clip
|
||||
- `REPORT` — STRING
|
||||
|
||||
**Behavior**:
|
||||
1. Validate `lora_hash` in tuner_data matches current lora_stack
|
||||
2. Read config at `tuner_data.top_n[selection - 1]`
|
||||
3. Run Pass 2 merge only (skip analysis, use cached data)
|
||||
4. Output merged model
|
||||
|
||||
## Workflow
|
||||
|
||||
```
|
||||
Simple: [Stack] -> [AutoTuner] -> [MODEL] -> [Sampler]
|
||||
|
||||
Override: [Stack] -> [AutoTuner] -> [TUNER_DATA] -> [Selector] -> [MODEL]
|
||||
| ^
|
||||
+-> [MODEL (unpatched)] --------+
|
||||
```
|
||||
|
||||
## Parameter Grid
|
||||
|
||||
| Parameter | Values | Count |
|
||||
|-----------|--------|-------|
|
||||
| merge_mode | weighted_average, slerp, consensus, ties | 4 |
|
||||
| sparsification | disabled, dare, della | 3 |
|
||||
| sparsification_density | 0.5, 0.7, 0.9 | 3 (when sparsification != disabled) |
|
||||
| dare_dampening | 0.0, 0.3, 0.6 | 3 (when sparsification == dare) |
|
||||
| merge_quality | standard, enhanced, maximum | 3 |
|
||||
| auto_strength | enabled, disabled | 2 |
|
||||
| optimization_mode | per_prefix, global | 2 |
|
||||
|
||||
Nonsensical combos pruned (density/dampening only when applicable). Estimated ~500–600 valid combos.
|
||||
|
||||
## Hybrid Scoring
|
||||
|
||||
### Phase 1 — Heuristic (no merge, milliseconds per combo)
|
||||
|
||||
Composite score from Pass 1 analysis metrics:
|
||||
|
||||
```
|
||||
score = w1 * conflict_score # lower sign conflicts = better
|
||||
+ w2 * alignment_score # cosine sim alignment with chosen mode
|
||||
+ w3 * energy_score # norm preservation prediction
|
||||
+ w4 * mode_fit_score # how well mode matches data profile
|
||||
```
|
||||
|
||||
Scoring intuitions:
|
||||
- **consensus** high when cos_sim > 0.5, conflict < 15%
|
||||
- **slerp/weighted_avg** high when LoRAs near-orthogonal
|
||||
- **ties** high when conflict > 25%
|
||||
- **sparsification** bonus for heavy-tailed magnitude distributions
|
||||
- **enhanced/maximum quality** bonus proportional to conflict level
|
||||
- **auto_strength** bonus when magnitude_ratio > 2×
|
||||
|
||||
### Phase 2 — Merge-and-Measure (top-N only, ~30–60s each)
|
||||
|
||||
Actual merge, then measure:
|
||||
- **Norm preservation ratio** — merged vs input energy
|
||||
- **Effective rank** — spectral rank of merged tensors
|
||||
- **Weight distribution stats** — kurtosis, sparsity
|
||||
|
||||
Phase 2 scores replace Phase 1 heuristics for final ranking.
|
||||
|
||||
## TUNER_DATA Format
|
||||
|
||||
```python
|
||||
{
|
||||
"version": 1,
|
||||
"top_n": [
|
||||
{
|
||||
"rank": 1,
|
||||
"score_heuristic": 0.87,
|
||||
"score_measured": 0.92,
|
||||
"config": {
|
||||
"merge_mode": "slerp",
|
||||
"sparsification": "dare",
|
||||
"sparsification_density": 0.7,
|
||||
"dare_dampening": 0.0,
|
||||
"merge_quality": "enhanced",
|
||||
"auto_strength": "enabled",
|
||||
"optimization_mode": "per_prefix",
|
||||
"adjusted_strengths": [0.8, 1.0, 0.6],
|
||||
},
|
||||
"metrics": {
|
||||
"norm_preservation": 0.73,
|
||||
"avg_conflict_ratio": 0.12,
|
||||
"avg_cosine_sim": 0.45,
|
||||
"effective_rank_mean": 28.3,
|
||||
},
|
||||
},
|
||||
# ... rank 2, 3, etc.
|
||||
],
|
||||
"analysis_cache": { ... },
|
||||
"lora_hash": "...",
|
||||
}
|
||||
```
|
||||
|
||||
## Report Format
|
||||
|
||||
```
|
||||
══════════════════════════════════════════
|
||||
LoRA AutoTuner Results
|
||||
══════════════════════════════════════════
|
||||
|
||||
Analysis Summary:
|
||||
LoRAs: 3 | Prefixes: 142 | Avg conflict: 18.2%
|
||||
Avg cosine similarity: 0.31 | Magnitude ratio: 2.4x
|
||||
|
||||
──────────────────────────────────────
|
||||
Top 3 Configurations
|
||||
──────────────────────────────────────
|
||||
|
||||
#1 ★ (applied to output) Score: 0.92
|
||||
Mode: slerp | Quality: enhanced
|
||||
Sparsification: dare (density=0.7, dampening=0.0)
|
||||
Auto-strength: enabled | Optimization: per_prefix
|
||||
Norm preservation: 0.73 | Effective rank: 28.3
|
||||
|
||||
#2 Score: 0.88
|
||||
Mode: consensus | Quality: standard
|
||||
Sparsification: disabled
|
||||
Auto-strength: enabled | Optimization: per_prefix
|
||||
Norm preservation: 0.81 | Effective rank: 25.1
|
||||
|
||||
#3 Score: 0.85
|
||||
Mode: ties (d=0.6) | Quality: enhanced
|
||||
Sparsification: della (density=0.5, dampening=0.3)
|
||||
Auto-strength: disabled | Optimization: global
|
||||
Norm preservation: 0.68 | Effective rank: 30.7
|
||||
|
||||
Suggested max output_strength: 1.37
|
||||
|
||||
To use a different config: connect TUNER_DATA
|
||||
to a Merge Selector node and set selection=N
|
||||
══════════════════════════════════════════
|
||||
```
|
||||
@@ -1,946 +0,0 @@
|
||||
# LoRA AutoTuner Implementation Plan
|
||||
|
||||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||||
|
||||
**Goal:** Add two new ComfyUI nodes — LoRAAutoTuner and LoRAMergeSelector — that sweep all merge parameters, score configurations via a hybrid heuristic+measurement approach, and output the best merge with the option to select alternatives.
|
||||
|
||||
**Architecture:** AutoTuner inherits from `LoRAOptimizer` (which inherits `_LoRAMergeBase`). It reuses Pass 1 analysis, adds a parameter grid generator, heuristic scorer, merge-and-measure scorer, and a two-pass workflow (explore → apply). MergeSelector also inherits `LoRAOptimizer` and takes cached `TUNER_DATA` to replay a specific config. Both nodes live in `lora_optimizer.py`.
|
||||
|
||||
**Tech Stack:** Python 3.10+, PyTorch, ComfyUI node API
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Add parameter grid generator
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` (add method to `_LoRAMergeBase` or as standalone function near line ~1928)
|
||||
|
||||
**Step 1: Write `_generate_param_grid` function**
|
||||
|
||||
Add before the `LoRAOptimizer` class definition (line 1930). This generates all valid parameter combinations, pruning nonsensical combos:
|
||||
|
||||
```python
|
||||
def _generate_param_grid():
|
||||
"""Generate all valid merge parameter combinations for AutoTuner sweep."""
|
||||
grid = []
|
||||
merge_modes = ["weighted_average", "slerp", "consensus", "ties"]
|
||||
sparsifications = ["disabled", "dare", "della"]
|
||||
densities = [0.5, 0.7, 0.9]
|
||||
dampenings = [0.0, 0.3, 0.6]
|
||||
qualities = ["standard", "enhanced", "maximum"]
|
||||
auto_strengths = ["enabled", "disabled"]
|
||||
opt_modes = ["per_prefix", "global"]
|
||||
|
||||
for mode in merge_modes:
|
||||
for spars in sparsifications:
|
||||
density_vals = densities if spars != "disabled" else [0.7]
|
||||
for density in density_vals:
|
||||
damp_vals = dampenings if spars == "dare" else [0.0]
|
||||
for dampening in damp_vals:
|
||||
for quality in qualities:
|
||||
for auto_str in auto_strengths:
|
||||
for opt_mode in opt_modes:
|
||||
grid.append({
|
||||
"merge_mode": mode,
|
||||
"sparsification": spars,
|
||||
"sparsification_density": density,
|
||||
"dare_dampening": dampening,
|
||||
"merge_quality": quality,
|
||||
"auto_strength": auto_str,
|
||||
"optimization_mode": opt_mode,
|
||||
})
|
||||
return grid
|
||||
```
|
||||
|
||||
**Step 2: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: OK
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat(autotuner): add parameter grid generator"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: Add heuristic scoring function
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` (add method after `_generate_param_grid`, before `LoRAOptimizer` class)
|
||||
|
||||
**Step 1: Write `_score_config_heuristic` function**
|
||||
|
||||
This scores a parameter config against Pass 1 analysis metrics WITHOUT merging. It predicts how well each config matches the data profile:
|
||||
|
||||
```python
|
||||
def _score_config_heuristic(config, avg_conflict_ratio, avg_cos_sim,
|
||||
magnitude_ratio, prefix_stats):
|
||||
"""
|
||||
Score a merge config against analysis metrics (no merge needed).
|
||||
Returns float score in [0, 1] where higher = better predicted quality.
|
||||
"""
|
||||
mode = config["merge_mode"]
|
||||
spars = config["sparsification"]
|
||||
density = config["sparsification_density"]
|
||||
quality = config["merge_quality"]
|
||||
auto_str = config["auto_strength"]
|
||||
opt_mode = config["optimization_mode"]
|
||||
|
||||
score = 0.0
|
||||
|
||||
# --- Mode fit score (0-0.4) ---
|
||||
# How well does the chosen mode match the data profile?
|
||||
if mode == "consensus":
|
||||
# Consensus excels with high similarity, low conflict
|
||||
if avg_cos_sim > 0.5 and avg_conflict_ratio < 0.15:
|
||||
score += 0.4
|
||||
elif avg_cos_sim > 0.3 and avg_conflict_ratio < 0.25:
|
||||
score += 0.25
|
||||
else:
|
||||
score += 0.05 # poor fit
|
||||
elif mode == "slerp":
|
||||
# SLERP excels with low-to-moderate conflict, any alignment
|
||||
if avg_conflict_ratio < 0.30:
|
||||
score += 0.35
|
||||
elif avg_conflict_ratio < 0.50:
|
||||
score += 0.20
|
||||
else:
|
||||
score += 0.10
|
||||
elif mode == "weighted_average":
|
||||
# Weighted average: decent baseline, best for orthogonal LoRAs
|
||||
if abs(avg_cos_sim) < 0.25 and avg_conflict_ratio < 0.60:
|
||||
score += 0.30
|
||||
elif avg_conflict_ratio < 0.40:
|
||||
score += 0.20
|
||||
else:
|
||||
score += 0.10
|
||||
elif mode == "ties":
|
||||
# TIES excels with high conflict
|
||||
if avg_conflict_ratio > 0.25:
|
||||
score += 0.35
|
||||
elif avg_conflict_ratio > 0.15:
|
||||
score += 0.20
|
||||
else:
|
||||
score += 0.10 # overkill for low conflict
|
||||
|
||||
# --- Sparsification fit (0-0.15) ---
|
||||
if spars != "disabled":
|
||||
# Sparsification helps when there's conflict to reduce
|
||||
conflict_benefit = min(avg_conflict_ratio / 0.5, 1.0) * 0.10
|
||||
score += conflict_benefit
|
||||
# Density: lower keeps less, higher preserves more
|
||||
# Moderate density (0.7) is generally safest
|
||||
density_penalty = abs(density - 0.7) * 0.05
|
||||
score += 0.05 - density_penalty
|
||||
else:
|
||||
# No sparsification: slight bonus for low-conflict stacks
|
||||
if avg_conflict_ratio < 0.15:
|
||||
score += 0.10
|
||||
|
||||
# --- Quality fit (0-0.15) ---
|
||||
if quality == "maximum":
|
||||
# Maximum quality always helps but diminishing returns for easy merges
|
||||
conflict_benefit = min(avg_conflict_ratio / 0.3, 1.0) * 0.10
|
||||
score += 0.05 + conflict_benefit
|
||||
elif quality == "enhanced":
|
||||
score += 0.08 + min(avg_conflict_ratio / 0.3, 1.0) * 0.07
|
||||
else:
|
||||
# Standard: slight bonus for very clean merges (no overhead needed)
|
||||
if avg_conflict_ratio < 0.10:
|
||||
score += 0.10
|
||||
else:
|
||||
score += 0.05
|
||||
|
||||
# --- Auto-strength fit (0-0.15) ---
|
||||
if auto_str == "enabled":
|
||||
# Auto-strength helps with unbalanced magnitudes
|
||||
if magnitude_ratio > 2.0:
|
||||
score += 0.15
|
||||
elif magnitude_ratio > 1.5:
|
||||
score += 0.10
|
||||
else:
|
||||
score += 0.05
|
||||
else:
|
||||
# Manual strength: slight bonus for balanced stacks
|
||||
if magnitude_ratio < 1.5:
|
||||
score += 0.10
|
||||
else:
|
||||
score += 0.03
|
||||
|
||||
# --- Optimization mode fit (0-0.15) ---
|
||||
if opt_mode == "per_prefix":
|
||||
# Per-prefix is generally better when prefixes vary
|
||||
if prefix_stats:
|
||||
conflict_ratios = [ps["conflict_ratio"] for ps in prefix_stats.values()
|
||||
if ps.get("n_loras", 0) > 1]
|
||||
if conflict_ratios:
|
||||
conflict_variance = max(conflict_ratios) - min(conflict_ratios)
|
||||
if conflict_variance > 0.2:
|
||||
score += 0.15 # high variance = per-prefix really helps
|
||||
else:
|
||||
score += 0.10
|
||||
else:
|
||||
score += 0.10
|
||||
else:
|
||||
score += 0.10
|
||||
else:
|
||||
# Global: simpler, slight bonus for uniform stacks
|
||||
score += 0.07
|
||||
|
||||
return score
|
||||
```
|
||||
|
||||
**Step 2: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: OK
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat(autotuner): add heuristic scoring function"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Add merge-and-measure scoring function
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` (add method after `_score_config_heuristic`)
|
||||
|
||||
**Step 1: Write `_score_merge_result` function**
|
||||
|
||||
This measures quality metrics on an actual merged result (patches dict). Called only for top-N candidates:
|
||||
|
||||
```python
|
||||
def _score_merge_result(model_patches, clip_patches):
|
||||
"""
|
||||
Score an actual merge result by measuring output quality metrics.
|
||||
Returns dict with individual metrics and composite score in [0, 1].
|
||||
"""
|
||||
import torch
|
||||
|
||||
norms = []
|
||||
effective_ranks = []
|
||||
sparsities = []
|
||||
|
||||
all_patches = list(model_patches.values()) + list(clip_patches.values())
|
||||
for patch in all_patches:
|
||||
# Extract the tensor from patch format
|
||||
if patch is None:
|
||||
continue
|
||||
if isinstance(patch, tuple) and len(patch) >= 2:
|
||||
# ("diff", (tensor,)) format
|
||||
tensor = patch[1][0] if isinstance(patch[1], tuple) else patch[1]
|
||||
elif hasattr(patch, 'diff'):
|
||||
# LoRAAdapter format — get effective diff
|
||||
tensor = None # low-rank, skip full analysis
|
||||
else:
|
||||
continue
|
||||
|
||||
if tensor is None:
|
||||
continue
|
||||
|
||||
t = tensor.float()
|
||||
|
||||
# Frobenius norm
|
||||
norms.append(t.norm().item())
|
||||
|
||||
# Effective rank via spectral analysis (sample for speed)
|
||||
if t.dim() == 2 and min(t.shape) > 1:
|
||||
try:
|
||||
# Use a small number of singular values for speed
|
||||
k = min(min(t.shape), 64)
|
||||
s = torch.linalg.svdvals(t)[:k]
|
||||
s_norm = s / (s.sum() + 1e-10)
|
||||
# Shannon entropy -> effective rank
|
||||
entropy = -(s_norm * (s_norm + 1e-10).log()).sum().item()
|
||||
eff_rank = min(math.exp(entropy), min(t.shape))
|
||||
effective_ranks.append(eff_rank)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Sparsity (fraction of near-zero weights)
|
||||
threshold = t.abs().max().item() * 0.01
|
||||
if threshold > 0:
|
||||
sparsity = (t.abs() < threshold).float().mean().item()
|
||||
sparsities.append(sparsity)
|
||||
|
||||
metrics = {}
|
||||
|
||||
# Norm stats
|
||||
if norms:
|
||||
metrics["norm_mean"] = sum(norms) / len(norms)
|
||||
metrics["norm_std"] = (sum((n - metrics["norm_mean"])**2 for n in norms)
|
||||
/ len(norms)) ** 0.5
|
||||
# Lower coefficient of variation = more uniform = better
|
||||
metrics["norm_cv"] = metrics["norm_std"] / (metrics["norm_mean"] + 1e-10)
|
||||
else:
|
||||
metrics["norm_mean"] = 0.0
|
||||
metrics["norm_cv"] = 1.0
|
||||
|
||||
# Effective rank (higher = more expressive)
|
||||
if effective_ranks:
|
||||
metrics["effective_rank_mean"] = sum(effective_ranks) / len(effective_ranks)
|
||||
else:
|
||||
metrics["effective_rank_mean"] = 0.0
|
||||
|
||||
# Sparsity (moderate is ideal — too sparse = lost info, too dense = noisy)
|
||||
if sparsities:
|
||||
avg_sparsity = sum(sparsities) / len(sparsities)
|
||||
metrics["sparsity_mean"] = avg_sparsity
|
||||
# Ideal sparsity around 0.3-0.5
|
||||
metrics["sparsity_fit"] = 1.0 - abs(avg_sparsity - 0.4) * 2.0
|
||||
metrics["sparsity_fit"] = max(0.0, metrics["sparsity_fit"])
|
||||
else:
|
||||
metrics["sparsity_mean"] = 0.0
|
||||
metrics["sparsity_fit"] = 0.5
|
||||
|
||||
# Composite score
|
||||
score = 0.0
|
||||
# Higher effective rank = more expressive (0-0.4)
|
||||
if metrics["effective_rank_mean"] > 0:
|
||||
# Normalize: rank 10-50 is typical range
|
||||
rank_score = min(metrics["effective_rank_mean"] / 40.0, 1.0)
|
||||
score += rank_score * 0.4
|
||||
|
||||
# Lower norm CV = more uniform patches (0-0.3)
|
||||
cv_score = max(0.0, 1.0 - metrics["norm_cv"])
|
||||
score += cv_score * 0.3
|
||||
|
||||
# Sparsity fitness (0-0.3)
|
||||
score += metrics["sparsity_fit"] * 0.3
|
||||
|
||||
metrics["composite_score"] = score
|
||||
return metrics
|
||||
```
|
||||
|
||||
**Step 2: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: OK
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat(autotuner): add merge-and-measure scoring function"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Add LoRAAutoTuner node class
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` (add new class after `LoRAOptimizerSimple`, before `WanVideoLoRAOptimizer` at line ~3691)
|
||||
|
||||
**Step 1: Write the LoRAAutoTuner class**
|
||||
|
||||
This is the main node. It inherits from `LoRAOptimizer` and overrides `optimize_merge` to add the sweep logic:
|
||||
|
||||
```python
|
||||
class LoRAAutoTuner(LoRAOptimizer):
|
||||
"""
|
||||
Automatic parameter sweep that finds the best merge configuration for
|
||||
a given LoRA stack. Runs Pass 1 analysis once, scores all parameter
|
||||
combinations via heuristic, then merges top-N candidates and measures
|
||||
output quality. Outputs the #1 result as MODEL/CLIP, plus a ranked
|
||||
report and TUNER_DATA for optional override via a Merge Selector node.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL", {
|
||||
"tooltip": "Base model to merge LoRAs into."
|
||||
}),
|
||||
"lora_stack": ("LORA_STACK", {
|
||||
"tooltip": "LoRA stack from a LoRA Stack node."
|
||||
}),
|
||||
"output_strength": ("FLOAT", {
|
||||
"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.05,
|
||||
"tooltip": "Master volume for the merged result."
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"clip": ("CLIP", {
|
||||
"tooltip": "Optional CLIP model for text-encoder LoRA keys."
|
||||
}),
|
||||
"clip_strength_multiplier": ("FLOAT", {
|
||||
"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.05,
|
||||
"tooltip": "Multiplier for CLIP LoRA strengths."
|
||||
}),
|
||||
"top_n": ("INT", {
|
||||
"default": 3, "min": 1, "max": 10, "step": 1,
|
||||
"tooltip": "Number of top configurations to evaluate via actual merge. Higher = slower but explores more options."
|
||||
}),
|
||||
"normalize_keys": (["disabled", "enabled"], {
|
||||
"default": "disabled",
|
||||
"tooltip": "Makes LoRAs from different training tools compatible."
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "STRING", "TUNER_DATA")
|
||||
RETURN_NAMES = ("model", "clip", "report", "tuner_data")
|
||||
FUNCTION = "auto_tune"
|
||||
CATEGORY = "LoRA Optimizer"
|
||||
DESCRIPTION = (
|
||||
"Automatically sweeps all merge parameters and finds the best "
|
||||
"configuration for your LoRA stack. Outputs the best merge directly. "
|
||||
"Connect TUNER_DATA to a Merge Selector node to try alternatives."
|
||||
)
|
||||
|
||||
def auto_tune(self, model, lora_stack, output_strength, clip=None,
|
||||
clip_strength_multiplier=1.0, top_n=3, normalize_keys="disabled"):
|
||||
import hashlib, json
|
||||
|
||||
# --- Normalize & validate stack ---
|
||||
active_loras = self._normalize_stack(lora_stack, normalize_keys=normalize_keys)
|
||||
if not active_loras:
|
||||
return (model, clip, "No active LoRAs in stack.", None)
|
||||
|
||||
# Compute lora_hash for cache validation
|
||||
hash_input = json.dumps([(l["name"], l["strength"]) for l in active_loras],
|
||||
sort_keys=True)
|
||||
lora_hash = hashlib.sha256(hash_input.encode()).hexdigest()[:16]
|
||||
|
||||
# --- Pass 1: Analysis (run once, reuse for all configs) ---
|
||||
# We need the raw analysis data. Run the parent's Pass 1 logic
|
||||
# by calling optimize_merge with a neutral config and extracting
|
||||
# the analysis. Instead, we replicate the analysis accumulation
|
||||
# here to avoid a full merge.
|
||||
|
||||
model_keys = self._get_model_keys(model)
|
||||
clip_keys = self._get_model_keys(clip) if clip is not None else set()
|
||||
|
||||
# Detect architecture
|
||||
first_lora_sd = self._load_lora(active_loras[0]["name"])
|
||||
detected_arch = self._detect_architecture(first_lora_sd)
|
||||
if normalize_keys == "enabled":
|
||||
for item in active_loras:
|
||||
lora_sd = self._load_lora(item["name"])
|
||||
self._normalize_keys(lora_sd, detected_arch)
|
||||
|
||||
# Collect all prefixes
|
||||
all_lora_prefixes = set()
|
||||
for item in active_loras:
|
||||
lora_sd = self._load_lora(item["name"])
|
||||
for key in lora_sd:
|
||||
if key.endswith(".lora_up.weight") or key.endswith(".lora_B.weight"):
|
||||
prefix = key.rsplit(".", 2)[0]
|
||||
all_lora_prefixes.add(prefix)
|
||||
all_lora_prefixes = sorted(all_lora_prefixes)
|
||||
|
||||
if not all_lora_prefixes:
|
||||
return (model, clip, "No LoRA prefixes found.", None)
|
||||
|
||||
# Determine compute device
|
||||
compute_device = self._get_compute_device()
|
||||
use_gpu = compute_device.type == "cuda"
|
||||
|
||||
# Run Pass 1 analysis (same as parent)
|
||||
all_key_targets = {}
|
||||
skipped_keys = 0
|
||||
per_lora_stats = [{
|
||||
"name": item["name"], "strength": item["strength"],
|
||||
"ranks": [], "key_count": 0, "l2_norms": [],
|
||||
} for item in active_loras]
|
||||
|
||||
pairs = [(i, j) for i in range(len(active_loras))
|
||||
for j in range(i + 1, len(active_loras))]
|
||||
pair_accum = {(i, j): [0, 0, 0.0, 0.0, 0.0] for i, j in pairs}
|
||||
all_magnitude_samples = []
|
||||
prefix_count = 0
|
||||
prefix_stats = {}
|
||||
|
||||
def _collect_analysis(result):
|
||||
nonlocal skipped_keys, prefix_count
|
||||
if result is None:
|
||||
return
|
||||
prefix, partial_stats, pair_conflicts, mag_samples, target_info, skips = result
|
||||
skipped_keys += skips
|
||||
if len(partial_stats) > 0:
|
||||
all_key_targets[prefix] = target_info
|
||||
prefix_count += 1
|
||||
for (idx, rank, l2) in partial_stats:
|
||||
per_lora_stats[idx]["ranks"].append(rank)
|
||||
per_lora_stats[idx]["key_count"] += 1
|
||||
per_lora_stats[idx]["l2_norms"].append(l2)
|
||||
for (i, j), (ov, conf, dot, na_sq, nb_sq) in pair_conflicts.items():
|
||||
pair_accum[(i, j)][0] += ov
|
||||
pair_accum[(i, j)][1] += conf
|
||||
pair_accum[(i, j)][2] += dot
|
||||
pair_accum[(i, j)][3] += na_sq
|
||||
pair_accum[(i, j)][4] += nb_sq
|
||||
all_magnitude_samples.extend(mag_samples)
|
||||
if len(partial_stats) > 0:
|
||||
n_contributing = len(partial_stats)
|
||||
pf_overlap = sum(ov for ov, conf, dot, na_sq, nb_sq in pair_conflicts.values())
|
||||
pf_conflict = sum(conf for ov, conf, dot, na_sq, nb_sq in pair_conflicts.values())
|
||||
pf_conflict_ratio = pf_conflict / pf_overlap if pf_overlap > 0 else 0.0
|
||||
pf_l2s = [l2 for _, _, l2 in partial_stats if l2 > 0]
|
||||
pf_mag_ratio = max(pf_l2s) / min(pf_l2s) if len(pf_l2s) >= 2 else 1.0
|
||||
pf_cos_sims = []
|
||||
for (ov, conf, dot, na_sq, nb_sq) in pair_conflicts.values():
|
||||
denom = (na_sq ** 0.5) * (nb_sq ** 0.5)
|
||||
if denom > 0:
|
||||
pf_cos_sims.append(dot / denom)
|
||||
avg_cs = sum(pf_cos_sims) / len(pf_cos_sims) if pf_cos_sims else 0.0
|
||||
prefix_stats[prefix] = {
|
||||
"n_loras": n_contributing,
|
||||
"conflict_ratio": pf_conflict_ratio,
|
||||
"magnitude_ratio": pf_mag_ratio,
|
||||
"magnitude_samples": list(mag_samples),
|
||||
"avg_cos_sim": avg_cs,
|
||||
}
|
||||
|
||||
logging.info(f"[LoRA AutoTuner] Pass 1: Analyzing {len(all_lora_prefixes)} prefixes...")
|
||||
t_start = time.time()
|
||||
if use_gpu:
|
||||
for lora_prefix in all_lora_prefixes:
|
||||
result = self._analyze_prefix(lora_prefix, active_loras,
|
||||
model_keys, clip_keys, model, clip, compute_device)
|
||||
_collect_analysis(result)
|
||||
else:
|
||||
max_workers = min(4, max(1, len(all_lora_prefixes)))
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
futures = {
|
||||
executor.submit(self._analyze_prefix, lora_prefix, active_loras,
|
||||
model_keys, clip_keys, model, clip, compute_device): lora_prefix
|
||||
for lora_prefix in all_lora_prefixes
|
||||
}
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
_collect_analysis(future.result())
|
||||
|
||||
if prefix_count == 0:
|
||||
return (model, clip, "No compatible LoRA keys found.", None)
|
||||
|
||||
t_analysis = time.time() - t_start
|
||||
logging.info(f"[LoRA AutoTuner] Analysis complete: {prefix_count} prefixes ({t_analysis:.1f}s)")
|
||||
|
||||
# Finalize global stats
|
||||
lora_stats = []
|
||||
l2_means = []
|
||||
for i, stat in enumerate(per_lora_stats):
|
||||
avg_rank = sum(stat["ranks"]) / len(stat["ranks"]) if stat["ranks"] else 0
|
||||
l2_mean = sum(stat["l2_norms"]) / len(stat["l2_norms"]) if stat["l2_norms"] else 0
|
||||
l2_means.append(l2_mean)
|
||||
lora_stats.append({
|
||||
"name": stat["name"], "strength": stat["strength"],
|
||||
"key_count": stat["key_count"], "avg_rank": avg_rank,
|
||||
"l2_mean": l2_mean,
|
||||
"conflict_mode": active_loras[i].get("conflict_mode", "all"),
|
||||
"key_filter": active_loras[i].get("key_filter", "all"),
|
||||
})
|
||||
|
||||
total_overlap = sum(pair_accum[p][0] for p in pairs)
|
||||
total_conflict = sum(pair_accum[p][1] for p in pairs)
|
||||
avg_conflict_ratio = total_conflict / total_overlap if total_overlap > 0 else 0
|
||||
|
||||
pairwise_similarities = {}
|
||||
for i, j in pairs:
|
||||
ov, conf, dot, na_sq, nb_sq = pair_accum[(i, j)]
|
||||
denom = math.sqrt(na_sq) * math.sqrt(nb_sq)
|
||||
pairwise_similarities[(i, j)] = dot / denom if denom > 0 else 0.0
|
||||
|
||||
valid_l2 = [m for m in l2_means if m > 0]
|
||||
magnitude_ratio = max(valid_l2) / min(valid_l2) if len(valid_l2) >= 2 else 1.0
|
||||
|
||||
avg_cos_sim = (sum(pairwise_similarities.values())
|
||||
/ len(pairwise_similarities)) if pairwise_similarities else 0.0
|
||||
|
||||
del all_magnitude_samples
|
||||
|
||||
# --- Phase 1: Score all parameter combos (heuristic, fast) ---
|
||||
logging.info("[LoRA AutoTuner] Phase 1: Scoring parameter grid...")
|
||||
grid = _generate_param_grid()
|
||||
scored = []
|
||||
for config in grid:
|
||||
h_score = _score_config_heuristic(
|
||||
config, avg_conflict_ratio, avg_cos_sim,
|
||||
magnitude_ratio, prefix_stats)
|
||||
scored.append((h_score, config))
|
||||
scored.sort(key=lambda x: x[0], reverse=True)
|
||||
logging.info(f"[LoRA AutoTuner] Scored {len(grid)} combos, top heuristic: {scored[0][0]:.3f}")
|
||||
|
||||
# --- Phase 2: Merge top-N and measure ---
|
||||
top_candidates = scored[:top_n]
|
||||
results = []
|
||||
|
||||
for rank_idx, (h_score, config) in enumerate(top_candidates):
|
||||
logging.info(f"[LoRA AutoTuner] Phase 2: Merging candidate #{rank_idx + 1} "
|
||||
f"({config['merge_mode']}, {config['merge_quality']})...")
|
||||
t_merge = time.time()
|
||||
|
||||
merged_model, merged_clip, report, lora_data = super().optimize_merge(
|
||||
model, lora_stack, output_strength,
|
||||
clip=clip,
|
||||
clip_strength_multiplier=clip_strength_multiplier,
|
||||
auto_strength=config["auto_strength"],
|
||||
optimization_mode=config["optimization_mode"],
|
||||
sparsification=config["sparsification"],
|
||||
sparsification_density=config["sparsification_density"],
|
||||
dare_dampening=config["dare_dampening"],
|
||||
merge_quality=config["merge_quality"],
|
||||
merge_strategy_override=config["merge_mode"],
|
||||
# Fixed defaults for sweep
|
||||
free_vram_between_passes="disabled",
|
||||
cache_patches="disabled",
|
||||
compress_patches="non_ties",
|
||||
svd_device="gpu",
|
||||
normalize_keys=normalize_keys,
|
||||
behavior_profile="v1.2",
|
||||
)
|
||||
|
||||
# Measure output quality
|
||||
m_patches = lora_data["model_patches"] if lora_data else {}
|
||||
c_patches = lora_data["clip_patches"] if lora_data else {}
|
||||
measured = _score_merge_result(m_patches, c_patches)
|
||||
|
||||
t_elapsed = time.time() - t_merge
|
||||
logging.info(f"[LoRA AutoTuner] Candidate #{rank_idx + 1}: "
|
||||
f"measured={measured['composite_score']:.3f} ({t_elapsed:.1f}s)")
|
||||
|
||||
# Extract auto-adjusted strengths from report if available
|
||||
adjusted_strengths = None
|
||||
if config["auto_strength"] == "enabled" and lora_data:
|
||||
adjusted_strengths = [ls["strength"] for ls in lora_stats]
|
||||
|
||||
results.append({
|
||||
"rank": rank_idx + 1,
|
||||
"score_heuristic": h_score,
|
||||
"score_measured": measured["composite_score"],
|
||||
"config": config,
|
||||
"metrics": {
|
||||
"norm_preservation": measured.get("norm_mean", 0.0),
|
||||
"effective_rank_mean": measured.get("effective_rank_mean", 0.0),
|
||||
"sparsity_mean": measured.get("sparsity_mean", 0.0),
|
||||
"norm_cv": measured.get("norm_cv", 0.0),
|
||||
},
|
||||
"merged_model": merged_model,
|
||||
"merged_clip": merged_clip,
|
||||
"lora_data": lora_data,
|
||||
})
|
||||
|
||||
# Sort by measured score (Phase 2 replaces Phase 1 ranking)
|
||||
results.sort(key=lambda x: x["score_measured"], reverse=True)
|
||||
for i, r in enumerate(results):
|
||||
r["rank"] = i + 1
|
||||
|
||||
# Best result
|
||||
best = results[0]
|
||||
|
||||
# Build TUNER_DATA (exclude model/clip objects — just configs and scores)
|
||||
tuner_data = {
|
||||
"version": 1,
|
||||
"lora_hash": lora_hash,
|
||||
"analysis_summary": {
|
||||
"n_loras": len(active_loras),
|
||||
"prefix_count": prefix_count,
|
||||
"avg_conflict_ratio": avg_conflict_ratio,
|
||||
"avg_cosine_sim": avg_cos_sim,
|
||||
"magnitude_ratio": magnitude_ratio,
|
||||
},
|
||||
"top_n": [{
|
||||
"rank": r["rank"],
|
||||
"score_heuristic": r["score_heuristic"],
|
||||
"score_measured": r["score_measured"],
|
||||
"config": r["config"],
|
||||
"metrics": r["metrics"],
|
||||
} for r in results],
|
||||
}
|
||||
|
||||
# Build report
|
||||
report = self._build_autotuner_report(
|
||||
results, tuner_data["analysis_summary"], output_strength)
|
||||
|
||||
return (best["merged_model"], best["merged_clip"], report, tuner_data)
|
||||
```
|
||||
|
||||
**Step 2: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: OK
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat(autotuner): add LoRAAutoTuner node class"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: Add AutoTuner report builder
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` (add `_build_autotuner_report` method to `LoRAAutoTuner` class)
|
||||
|
||||
**Step 1: Write the report builder method**
|
||||
|
||||
Add inside the `LoRAAutoTuner` class, after `auto_tune`:
|
||||
|
||||
```python
|
||||
def _build_autotuner_report(self, results, analysis_summary, output_strength):
|
||||
"""Build the ranked report for AutoTuner results."""
|
||||
lines = []
|
||||
lines.append("=" * 54)
|
||||
lines.append(" LoRA AutoTuner Results")
|
||||
lines.append("=" * 54)
|
||||
lines.append("")
|
||||
lines.append(" Analysis Summary:")
|
||||
s = analysis_summary
|
||||
lines.append(f" LoRAs: {s['n_loras']} | Prefixes: {s['prefix_count']} "
|
||||
f"| Avg conflict: {s['avg_conflict_ratio']:.1%}")
|
||||
lines.append(f" Avg cosine similarity: {s['avg_cosine_sim']:.2f} "
|
||||
f"| Magnitude ratio: {s['magnitude_ratio']:.1f}x")
|
||||
lines.append("")
|
||||
lines.append(" " + "-" * 38)
|
||||
lines.append(f" Top {len(results)} Configurations")
|
||||
lines.append(" " + "-" * 38)
|
||||
|
||||
for r in results:
|
||||
lines.append("")
|
||||
c = r["config"]
|
||||
m = r["metrics"]
|
||||
marker = " (applied to output)" if r["rank"] == 1 else ""
|
||||
star = " \u2605" if r["rank"] == 1 else ""
|
||||
lines.append(f" #{r['rank']}{star}{marker}"
|
||||
f" Score: {r['score_measured']:.2f}")
|
||||
lines.append(f" Mode: {c['merge_mode']} | Quality: {c['merge_quality']}")
|
||||
if c["sparsification"] != "disabled":
|
||||
spars_info = f"{c['sparsification']} (density={c['sparsification_density']}"
|
||||
if c["dare_dampening"] > 0:
|
||||
spars_info += f", dampening={c['dare_dampening']}"
|
||||
spars_info += ")"
|
||||
lines.append(f" Sparsification: {spars_info}")
|
||||
else:
|
||||
lines.append(f" Sparsification: disabled")
|
||||
lines.append(f" Auto-strength: {c['auto_strength']} "
|
||||
f"| Optimization: {c['optimization_mode']}")
|
||||
if m.get("effective_rank_mean", 0) > 0:
|
||||
lines.append(f" Effective rank: {m['effective_rank_mean']:.1f} "
|
||||
f"| Sparsity: {m.get('sparsity_mean', 0):.1%}")
|
||||
|
||||
lines.append("")
|
||||
lines.append(" To use a different config: connect TUNER_DATA")
|
||||
lines.append(" to a Merge Selector node and set selection=N")
|
||||
lines.append("=" * 54)
|
||||
return "\n".join(lines)
|
||||
```
|
||||
|
||||
**Step 2: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: OK
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat(autotuner): add report builder for AutoTuner"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 6: Add LoRAMergeSelector node class
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` (add new class after `LoRAAutoTuner`)
|
||||
|
||||
**Step 1: Write the LoRAMergeSelector class**
|
||||
|
||||
```python
|
||||
class LoRAMergeSelector(LoRAOptimizer):
|
||||
"""
|
||||
Applies a specific merge configuration from AutoTuner results.
|
||||
Connect TUNER_DATA from a LoRA AutoTuner node and set the selection
|
||||
index to choose which ranked configuration to apply.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL", {
|
||||
"tooltip": "Base model (same one connected to the AutoTuner)."
|
||||
}),
|
||||
"lora_stack": ("LORA_STACK", {
|
||||
"tooltip": "LoRA stack (same one connected to the AutoTuner)."
|
||||
}),
|
||||
"tuner_data": ("TUNER_DATA", {
|
||||
"tooltip": "Connect from the LoRA AutoTuner's tuner_data output."
|
||||
}),
|
||||
"selection": ("INT", {
|
||||
"default": 1, "min": 1, "max": 10, "step": 1,
|
||||
"tooltip": "Which ranked configuration to apply (1 = best, 2 = second best, etc.)."
|
||||
}),
|
||||
"output_strength": ("FLOAT", {
|
||||
"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.05,
|
||||
"tooltip": "Master volume for the merged result."
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"clip": ("CLIP", {
|
||||
"tooltip": "Optional CLIP model for text-encoder LoRA keys."
|
||||
}),
|
||||
"clip_strength_multiplier": ("FLOAT", {
|
||||
"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.05,
|
||||
"tooltip": "Multiplier for CLIP LoRA strengths."
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "STRING")
|
||||
RETURN_NAMES = ("model", "clip", "report")
|
||||
FUNCTION = "select_merge"
|
||||
CATEGORY = "LoRA Optimizer"
|
||||
DESCRIPTION = (
|
||||
"Applies a specific merge configuration from LoRA AutoTuner results. "
|
||||
"Set selection to choose which ranked configuration to use."
|
||||
)
|
||||
|
||||
def select_merge(self, model, lora_stack, tuner_data, selection,
|
||||
output_strength, clip=None, clip_strength_multiplier=1.0):
|
||||
import hashlib, json
|
||||
|
||||
if tuner_data is None or "top_n" not in tuner_data:
|
||||
return (model, clip, "Error: No valid TUNER_DATA provided.")
|
||||
|
||||
# Validate lora_hash
|
||||
active_loras = self._normalize_stack(lora_stack)
|
||||
hash_input = json.dumps([(l["name"], l["strength"]) for l in active_loras],
|
||||
sort_keys=True)
|
||||
current_hash = hashlib.sha256(hash_input.encode()).hexdigest()[:16]
|
||||
if current_hash != tuner_data.get("lora_hash", ""):
|
||||
logging.warning("[Merge Selector] LoRA stack has changed since AutoTuner ran. "
|
||||
"Results may not match. Re-run AutoTuner for accurate results.")
|
||||
|
||||
# Get selected config
|
||||
top_n = tuner_data["top_n"]
|
||||
if selection < 1 or selection > len(top_n):
|
||||
return (model, clip,
|
||||
f"Error: selection={selection} out of range (1-{len(top_n)}).")
|
||||
|
||||
entry = top_n[selection - 1]
|
||||
config = entry["config"]
|
||||
|
||||
logging.info(f"[Merge Selector] Applying config #{selection}: "
|
||||
f"{config['merge_mode']}, {config['merge_quality']}")
|
||||
|
||||
# Run merge with selected config
|
||||
merged_model, merged_clip, _report, _lora_data = super().optimize_merge(
|
||||
model, lora_stack, output_strength,
|
||||
clip=clip,
|
||||
clip_strength_multiplier=clip_strength_multiplier,
|
||||
auto_strength=config["auto_strength"],
|
||||
optimization_mode=config["optimization_mode"],
|
||||
sparsification=config["sparsification"],
|
||||
sparsification_density=config["sparsification_density"],
|
||||
dare_dampening=config["dare_dampening"],
|
||||
merge_quality=config["merge_quality"],
|
||||
merge_strategy_override=config["merge_mode"],
|
||||
free_vram_between_passes="disabled",
|
||||
cache_patches="enabled",
|
||||
compress_patches="non_ties",
|
||||
svd_device="gpu",
|
||||
normalize_keys="disabled",
|
||||
behavior_profile="v1.2",
|
||||
)
|
||||
|
||||
# Build report for this selection
|
||||
lines = []
|
||||
lines.append(f"Merge Selector — Applied config #{selection}")
|
||||
lines.append(f" Mode: {config['merge_mode']} | Quality: {config['merge_quality']}")
|
||||
if config["sparsification"] != "disabled":
|
||||
lines.append(f" Sparsification: {config['sparsification']} "
|
||||
f"(density={config['sparsification_density']})")
|
||||
lines.append(f" Auto-strength: {config['auto_strength']} "
|
||||
f"| Optimization: {config['optimization_mode']}")
|
||||
lines.append(f" Heuristic score: {entry['score_heuristic']:.3f} "
|
||||
f"| Measured score: {entry['score_measured']:.3f}")
|
||||
report = "\n".join(lines)
|
||||
|
||||
return (merged_model, merged_clip, report)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, model, lora_stack, tuner_data, selection,
|
||||
output_strength, clip=None, clip_strength_multiplier=1.0):
|
||||
return (id(tuner_data), selection, output_strength, clip_strength_multiplier)
|
||||
```
|
||||
|
||||
**Step 2: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: OK
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat(autotuner): add LoRAMergeSelector node class"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 7: Register nodes in mappings
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` (lines 4400-4422 — NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS)
|
||||
|
||||
**Step 1: Add both nodes to mappings**
|
||||
|
||||
Add to `NODE_CLASS_MAPPINGS` (after the `WanVideoLoRAOptimizer` entry):
|
||||
```python
|
||||
"LoRAAutoTuner": LoRAAutoTuner,
|
||||
"LoRAMergeSelector": LoRAMergeSelector,
|
||||
```
|
||||
|
||||
Add to `NODE_DISPLAY_NAME_MAPPINGS`:
|
||||
```python
|
||||
"LoRAAutoTuner": "LoRA AutoTuner",
|
||||
"LoRAMergeSelector": "Merge Selector",
|
||||
```
|
||||
|
||||
**Step 2: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: OK
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat(autotuner): register AutoTuner and MergeSelector nodes"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 8: Final verification and cleanup
|
||||
|
||||
**Step 1: Full syntax check**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: OK
|
||||
|
||||
**Step 2: Check imports**
|
||||
|
||||
Verify that `math`, `time`, `logging`, `concurrent.futures`, and `hashlib` are already imported at the top of `lora_optimizer.py`. If `hashlib` or `json` are missing from top-level imports, they're imported inline in the methods (via `import hashlib, json`), which is fine.
|
||||
|
||||
**Step 3: Smoke test import**
|
||||
|
||||
Run: `python -c "from lora_optimizer import NODE_CLASS_MAPPINGS; print(list(NODE_CLASS_MAPPINGS.keys()))"`
|
||||
Expected: List including `LoRAAutoTuner` and `LoRAMergeSelector`
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat(autotuner): final verification pass"
|
||||
```
|
||||
@@ -1,371 +0,0 @@
|
||||
# AutoTuner Diff Cache Implementation Plan
|
||||
|
||||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||||
|
||||
**Goal:** Cache raw LoRA diffs (A@B matmul results) across AutoTuner Phase 2 candidates to eliminate redundant computation.
|
||||
|
||||
**Architecture:** Add a `_diff_cache` parameter to `optimize_merge()` that `_merge_one_prefix()` checks before computing diffs. The AutoTuner initializes the cache before the Phase 2 loop and clears it after. Two modes: RAM (dict) and disk (torch.save + mmap).
|
||||
|
||||
**Tech Stack:** PyTorch, tempfile, shutil, torch.load(mmap=True)
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Add diff_cache_mode to AutoTuner INPUT_TYPES
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:4263-4267` (add input before `vram_budget`)
|
||||
|
||||
**Step 1: Add the combo input**
|
||||
|
||||
In `LoRAAutoTuner.INPUT_TYPES` at line 4263, insert before the `vram_budget` entry:
|
||||
|
||||
```python
|
||||
"diff_cache_mode": (["ram", "disk", "disabled"], {
|
||||
"default": "ram",
|
||||
"tooltip": "Cache raw LoRA diffs across candidates to avoid redundant computation. 'ram' is fastest (uses ~3-12GB RAM). 'disk' uses temp files with memory-mapping (slower but low RAM). 'disabled' recomputes diffs each time."
|
||||
}),
|
||||
```
|
||||
|
||||
**Step 2: Add parameter to auto_tune signature**
|
||||
|
||||
At line 4280, add `diff_cache_mode="ram"` to the `auto_tune` method signature:
|
||||
|
||||
```python
|
||||
def auto_tune(self, model, lora_stack, output_strength, clip=None,
|
||||
clip_strength_multiplier=1.0, top_n=3, normalize_keys="disabled",
|
||||
scoring_svd="disabled", scoring_device="gpu",
|
||||
architecture_preset="auto", record_dataset="disabled",
|
||||
vram_budget=0.0, diff_cache_mode="ram"):
|
||||
```
|
||||
|
||||
**Step 3: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat(autotuner): add diff_cache_mode input parameter"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: Add _diff_cache parameter to optimize_merge
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:3289` (optimize_merge signature)
|
||||
|
||||
**Step 1: Add _diff_cache to signature**
|
||||
|
||||
At line 3289, add `_diff_cache=None` after `_analysis_cache=None`:
|
||||
|
||||
```python
|
||||
def optimize_merge(self, model, lora_stack, output_strength, clip=None, clip_strength_multiplier=1.0, auto_strength="disabled", free_vram_between_passes="disabled", vram_budget=0.0, optimization_mode="per_prefix", cache_patches="enabled", compress_patches="non_ties", svd_device="gpu", normalize_keys="disabled", sparsification="disabled", sparsification_density=0.7, dare_dampening=0.0, merge_strategy_override="", merge_quality="standard", behavior_profile="v1.2", architecture_preset="auto", _analysis_cache=None, _diff_cache=None, _skip_report=False):
|
||||
```
|
||||
|
||||
**Step 2: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: add _diff_cache parameter to optimize_merge"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Add diff cache lookup/store in _merge_one_prefix
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:3789-3842` (diff computation loop in _merge_one_prefix)
|
||||
|
||||
**Step 1: Add cache lookup before diff computation**
|
||||
|
||||
The diff computation loop starts at line 3789. Inside the `for i, item in enumerate(active_loras):` loop, after the key_filter checks (line 3797) and before getting lora_info (line 3799), add the cache lookup. The full modified loop body should be:
|
||||
|
||||
After line 3797 (`continue`), add cache lookup:
|
||||
|
||||
```python
|
||||
# Check diff cache
|
||||
cache_key = (lora_prefix, i)
|
||||
if _diff_cache is not None and cache_key in _diff_cache:
|
||||
cached = _diff_cache[cache_key]
|
||||
if isinstance(cached, str):
|
||||
# Disk mode: load from file
|
||||
diff = torch.load(cached, map_location=compute_device if use_gpu else "cpu",
|
||||
weights_only=True, mmap=True).float()
|
||||
else:
|
||||
# RAM mode: tensor on CPU
|
||||
diff = cached.to(compute_device).float() if use_gpu else cached.float()
|
||||
|
||||
if storage_dtype is None:
|
||||
storage_dtype = diff.dtype
|
||||
|
||||
if is_clip_key and item["clip_strength"] is not None:
|
||||
eff_strength = item["clip_strength"]
|
||||
else:
|
||||
eff_strength = item["strength"]
|
||||
if scale_ratios:
|
||||
eff_strength *= scale_ratios.get(i, 1.0)
|
||||
|
||||
diffs_list.append((diff, eff_strength))
|
||||
diff_to_lora.append(i)
|
||||
continue
|
||||
```
|
||||
|
||||
This goes right after line 3797 and before line 3799 (`mat_up, mat_down, alpha, mid = lora_info`).
|
||||
|
||||
**Step 2: Add cache store after diff computation**
|
||||
|
||||
After line 3832 (`diff = diff * (alpha / rank)`), add:
|
||||
|
||||
```python
|
||||
# Store in diff cache
|
||||
if _diff_cache is not None and cache_key not in _diff_cache:
|
||||
if isinstance(_diff_cache, dict) and not hasattr(_diff_cache, '_cache_dir'):
|
||||
# RAM mode
|
||||
_diff_cache[cache_key] = diff.cpu()
|
||||
elif hasattr(_diff_cache, '_cache_dir'):
|
||||
# Disk mode
|
||||
cache_path = os.path.join(_diff_cache._cache_dir,
|
||||
f"{lora_prefix.replace('.', '_')}_{i}.pt")
|
||||
torch.save(diff.cpu(), cache_path)
|
||||
_diff_cache[cache_key] = cache_path
|
||||
```
|
||||
|
||||
Wait — this approach with `hasattr` is messy. Let's use a simple wrapper class instead.
|
||||
|
||||
**Revised approach: Use a DiffCache class**
|
||||
|
||||
Add before `_merge_one_prefix` definition (or at module level near top of file, after imports). Actually, add it as a standalone class near the top of the file (after the `LoRAAdapter` class, around line 100):
|
||||
|
||||
```python
|
||||
class _DiffCache:
|
||||
"""Cache for raw LoRA diffs across AutoTuner candidates."""
|
||||
|
||||
def __init__(self, mode="ram"):
|
||||
self.mode = mode
|
||||
self._store = {}
|
||||
self._cache_dir = None
|
||||
if mode == "disk":
|
||||
import tempfile
|
||||
self._cache_dir = tempfile.mkdtemp(prefix="lora_diff_cache_")
|
||||
|
||||
def get(self, key, device=None):
|
||||
if key not in self._store:
|
||||
return None
|
||||
val = self._store[key]
|
||||
if self.mode == "disk":
|
||||
return torch.load(val, map_location=device or "cpu",
|
||||
weights_only=True, mmap=True)
|
||||
else:
|
||||
return val.to(device) if device is not None else val
|
||||
|
||||
def put(self, key, tensor):
|
||||
if key in self._store:
|
||||
return
|
||||
if self.mode == "disk":
|
||||
prefix_str = key[0].replace(".", "_").replace("/", "_")
|
||||
path = os.path.join(self._cache_dir, f"{prefix_str}_{key[1]}.pt")
|
||||
torch.save(tensor.cpu(), path)
|
||||
self._store[key] = path
|
||||
else:
|
||||
self._store[key] = tensor.cpu()
|
||||
|
||||
def clear(self):
|
||||
self._store.clear()
|
||||
if self._cache_dir is not None:
|
||||
import shutil
|
||||
shutil.rmtree(self._cache_dir, ignore_errors=True)
|
||||
self._cache_dir = None
|
||||
|
||||
def __contains__(self, key):
|
||||
return key in self._store
|
||||
```
|
||||
|
||||
Then in `_merge_one_prefix`, the cache interaction is clean:
|
||||
|
||||
After the key_filter checks (~line 3797), before `mat_up, mat_down, alpha, mid = lora_info`:
|
||||
|
||||
```python
|
||||
# Check diff cache
|
||||
cache_key = (lora_prefix, i)
|
||||
if _diff_cache is not None and cache_key in _diff_cache:
|
||||
diff = _diff_cache.get(cache_key,
|
||||
device=compute_device if use_gpu else None).float()
|
||||
if storage_dtype is None:
|
||||
storage_dtype = diff.dtype
|
||||
|
||||
if is_clip_key and item["clip_strength"] is not None:
|
||||
eff_strength = item["clip_strength"]
|
||||
else:
|
||||
eff_strength = item["strength"]
|
||||
if scale_ratios:
|
||||
eff_strength *= scale_ratios.get(i, 1.0)
|
||||
|
||||
diffs_list.append((diff, eff_strength))
|
||||
diff_to_lora.append(i)
|
||||
continue
|
||||
```
|
||||
|
||||
After line 3832 (`diff = diff * (alpha / rank)`), add:
|
||||
|
||||
```python
|
||||
# Store in diff cache for subsequent candidates
|
||||
if _diff_cache is not None:
|
||||
_diff_cache.put(cache_key, diff)
|
||||
```
|
||||
|
||||
**Step 3: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: add _DiffCache class and cache lookup/store in merge"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Wire up the cache in AutoTuner Phase 2
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:4482-4573` (Phase 2 loop)
|
||||
|
||||
**Step 1: Initialize cache before Phase 2 loop**
|
||||
|
||||
At line 4482 (before `top_candidates = scored[:top_n]`), add:
|
||||
|
||||
```python
|
||||
# Initialize diff cache for Phase 2
|
||||
_diff_cache = None
|
||||
if diff_cache_mode != "disabled":
|
||||
_diff_cache = _DiffCache(mode=diff_cache_mode)
|
||||
logging.info(f"[LoRA AutoTuner] Diff cache enabled (mode={diff_cache_mode})")
|
||||
```
|
||||
|
||||
**Step 2: Pass cache to optimize_merge call**
|
||||
|
||||
At line 4499-4520 (the `super().optimize_merge()` call), add `_diff_cache=_diff_cache`:
|
||||
|
||||
```python
|
||||
merged_model, merged_clip, _report, lora_data = super().optimize_merge(
|
||||
model, lora_stack, output_strength,
|
||||
clip=clip,
|
||||
clip_strength_multiplier=clip_strength_multiplier,
|
||||
auto_strength=config["auto_strength"],
|
||||
optimization_mode=config["optimization_mode"],
|
||||
sparsification=config["sparsification"],
|
||||
sparsification_density=config["sparsification_density"],
|
||||
dare_dampening=config["dare_dampening"],
|
||||
merge_quality=config["merge_quality"],
|
||||
merge_strategy_override=config["merge_mode"],
|
||||
free_vram_between_passes="disabled",
|
||||
vram_budget=vram_budget,
|
||||
cache_patches="disabled",
|
||||
compress_patches="disabled",
|
||||
svd_device="gpu",
|
||||
normalize_keys=normalize_keys,
|
||||
behavior_profile="v1.2",
|
||||
architecture_preset=architecture_preset,
|
||||
_analysis_cache=_analysis_cache,
|
||||
_diff_cache=_diff_cache,
|
||||
_skip_report=True,
|
||||
)
|
||||
```
|
||||
|
||||
**Step 3: Clear cache after Phase 2**
|
||||
|
||||
At line 4567 (after the Phase 2 loop, where cleanup happens), add before `del all_magnitude_samples`:
|
||||
|
||||
```python
|
||||
if _diff_cache is not None:
|
||||
_diff_cache.clear()
|
||||
del _diff_cache
|
||||
```
|
||||
|
||||
**Step 4: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat(autotuner): wire diff cache into Phase 2 candidate loop"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: Handle storage_dtype properly for cached diffs
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` (in _merge_one_prefix, around the cache lookup)
|
||||
|
||||
**Step 1: Fix storage_dtype for cached path**
|
||||
|
||||
The `storage_dtype` variable (line 3787) tracks the native dtype before float32 upcast. In the cache lookup path, we need to get it from the LoRA matrices even though we skip the matmul. Update the cache-hit path:
|
||||
|
||||
Replace the `storage_dtype` line in the cache-hit block with:
|
||||
|
||||
```python
|
||||
if storage_dtype is None:
|
||||
lora_info = self._get_lora_key_info(item["lora"], lora_prefix)
|
||||
if lora_info is not None:
|
||||
storage_dtype = lora_info[0].dtype # mat_up.dtype
|
||||
```
|
||||
|
||||
**Step 2: Verify syntax**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "fix: resolve storage_dtype from LoRA matrices for cached diffs"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 6: Final verification and version bump
|
||||
|
||||
**Files:**
|
||||
- Modify: `pyproject.toml` (version bump)
|
||||
|
||||
**Step 1: Full syntax check**
|
||||
|
||||
Run: `python -c "import ast; ast.parse(open('lora_optimizer.py').read()); print('OK')"`
|
||||
Expected: `OK`
|
||||
|
||||
**Step 2: Verify the node loads in Python**
|
||||
|
||||
Run: `python -c "exec(open('lora_optimizer.py').read()); print('Loaded:', list(NODE_CLASS_MAPPINGS.keys()))"`
|
||||
Expected: Lists all node class names including LoRAAutoTuner
|
||||
|
||||
**Step 3: Bump version**
|
||||
|
||||
In `pyproject.toml`, change version to `1.2.4`.
|
||||
|
||||
**Step 4: Commit and push**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py pyproject.toml
|
||||
git commit -m "feat(autotuner): diff cache to eliminate redundant A@B across candidates
|
||||
|
||||
Caches raw LoRA diffs (A@B matmul results) during Phase 2 so subsequent
|
||||
candidates reuse them instead of recomputing. Two modes: 'ram' (default,
|
||||
fastest, ~3-12GB) and 'disk' (temp files with mmap, low RAM). Reduces
|
||||
Phase 2 wall time by 30-50% for top_n=3."
|
||||
git push origin main
|
||||
```
|
||||
@@ -1,73 +0,0 @@
|
||||
# LoRA Extract From Model — Design
|
||||
|
||||
## Context
|
||||
|
||||
Users sometimes work with finetuned checkpoints that already have a LoRA baked in (merged into the full model weights). When they try to add additional LoRAs via the merger, it creates issues: the merger's conflict analysis and auto-strength assume a clean base model, so scoring goes wrong and oversaturation occurs because the pre-merged LoRA is invisible to the pipeline.
|
||||
|
||||
The fix is to extract the baked-in LoRA as a proper low-rank file by subtracting the original base model weights, then inject it into the merge stack as a first-class entry.
|
||||
|
||||
## Approach
|
||||
|
||||
**SVD delta extraction with base model as prerequisite.**
|
||||
|
||||
Compute `delta = W_finetuned - W_base` per layer, SVD-decompose each delta to rank-r, and output a proper LoRA dict. The result is a virtual LoRA fully compatible with the existing merge pipeline — no special-casing needed downstream.
|
||||
|
||||
## New Node: `LoRAExtractFromModel`
|
||||
|
||||
**Inputs:**
|
||||
|
||||
| Name | Type | Default | Notes |
|
||||
|------|------|---------|-------|
|
||||
| `base_model` | MODEL | — | Clean base checkpoint |
|
||||
| `finetuned_model` | MODEL | — | Model with LoRA baked in |
|
||||
| `rank` | INT | 32 | Max rank for SVD truncation |
|
||||
| `rank_mode` | COMBO | `auto` | `auto` = retain % energy; `fixed` = always use `rank` |
|
||||
| `energy_threshold` | FLOAT | 0.99 | Fraction of delta energy to retain (auto mode only) |
|
||||
| `strength` | FLOAT | 1.0 | Scale applied to extracted LoRA in the stack |
|
||||
|
||||
**Outputs:**
|
||||
|
||||
| Name | Type | Purpose |
|
||||
|------|------|---------|
|
||||
| `LORA_STACK` | LORA_STACK | Feeds into LoRAOptimizerSimple / LoRAAutoTuner |
|
||||
| `LORA_DATA` | LORA_DATA | Feeds directly into SaveMergedLoRA (reuses existing node) |
|
||||
|
||||
## Extraction Pipeline
|
||||
|
||||
Per layer:
|
||||
|
||||
1. Fetch `W_base` and `W_finetuned` via `comfy.utils.get_attr()`
|
||||
2. Cast both to `float32`, compute `delta = W_finetuned - W_base`
|
||||
3. Skip layers where `delta.norm() < epsilon` (unaffected — avoids noise)
|
||||
4. Reshape delta to 2D if needed (conv layers: `[C_out, C_in * kH * kW]`)
|
||||
5. `U, S, Vh = torch.linalg.svd(delta, full_matrices=False)`
|
||||
6. Apply singular value floor: drop `S[i] < 1e-4 * S[0]` before energy calc (prevents noise inflation)
|
||||
7. Determine rank `r`:
|
||||
- `auto`: smallest `r` where `S[:r].pow(2).sum() / S.pow(2).sum() >= energy_threshold`
|
||||
- `fixed`: `min(rank, S.shape[0])`
|
||||
8. `lora_up = U[:, :r] * S[:r].sqrt()`; `lora_down = S[:r].sqrt().unsqueeze(1) * Vh[:r]`
|
||||
9. Set `alpha = r` (effective scale = 1.0; user controls via `strength`)
|
||||
|
||||
**Key mapping:** model weight keys → LoRA prefix format via `comfy.lora.model_lora_keys_unet()` / `model_lora_keys_clip()` (inverse lookup, already used elsewhere in the codebase).
|
||||
|
||||
**Memory:** process one layer at a time, free tensors immediately — never materialize the full delta dict at once. Critical for large models (Flux ~23GB × 2).
|
||||
|
||||
## Warnings & Validation
|
||||
|
||||
| Condition | Action |
|
||||
|-----------|--------|
|
||||
| Architectures differ between models | Raise error before processing |
|
||||
| `delta.norm()` low on >80% of layers | Warn: models may be identical or same base |
|
||||
| Significant deltas on nearly all layers | Warn: looks like a full finetune, not a LoRA merge — extraction may produce a noisy high-rank result |
|
||||
| Extracted rank consistently hitting `rank` cap | Warn: consider increasing `rank` or switching to `auto` mode |
|
||||
|
||||
## Integration
|
||||
|
||||
- `LORA_STACK` output is structurally identical to `LoRAStack` output — plugs into `LoRAOptimizerSimple`, `LoRAAutoTuner`, and `LoRAConflictEditor` with no changes
|
||||
- `LORA_DATA` output matches the merger's output format (`model_patches`, `clip_patches`, `key_map`, `output_strength`, `clip_strength`) — plugs directly into the existing `SaveMergedLoRA` node
|
||||
- Architecture detection reuses `_detect_architecture()` (line 683)
|
||||
- Node registered in `NODE_CLASS_MAPPINGS` in `lora_optimizer.py` (line 10476)
|
||||
|
||||
## Files to Modify
|
||||
|
||||
- `lora_optimizer.py` — add `LoRAExtractFromModel` class, register in `NODE_CLASS_MAPPINGS`
|
||||
@@ -1,496 +0,0 @@
|
||||
# LoRA Extract From Model — Implementation Plan
|
||||
|
||||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||||
|
||||
**Goal:** Add a `LoRAExtractFromModel` node that extracts a baked-in LoRA from a finetuned model by SVD-decomposing the delta against a clean base model, outputting both a `LORA_STACK` entry (feeds into the merger) and a `LORA_DATA` entry (feeds into `SaveMergedLoRA`).
|
||||
|
||||
**Architecture:** Subtract base model weights from finetuned weights per layer, SVD-decompose each delta to rank-r, and package the result in both formats. LORA_STACK uses the raw safetensors-style dict (same format `LoRAStack` already produces). LORA_DATA uses `LoRAAdapter` patches (same format the merger already produces).
|
||||
|
||||
**Tech Stack:** PyTorch (`torch.linalg.svd`), ComfyUI model APIs (`comfy.utils.get_attr`, `comfy.lora.model_lora_keys_unet/clip`), `LoRAAdapter` from `comfy.weight_adapter.lora` (already imported at line 24).
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Add unit tests for the SVD extraction helper
|
||||
|
||||
**Files:**
|
||||
- Modify: `tests/test_lora_optimizer.py`
|
||||
|
||||
**Step 1: Write the failing tests**
|
||||
|
||||
Add to the bottom of `tests/test_lora_optimizer.py`, after the existing test classes:
|
||||
|
||||
```python
|
||||
@unittest.skipIf(torch is None, "torch not available")
|
||||
class TestExtractLoRAFromDelta(unittest.TestCase):
|
||||
"""Tests for _extract_lora_svd() helper used by LoRAExtractFromModel."""
|
||||
|
||||
def _make_delta(self, rows, cols, rank):
|
||||
"""Create a synthetic rank-r delta matrix."""
|
||||
U = torch.randn(rows, rank)
|
||||
S = torch.rand(rank) + 0.1 # positive singular values
|
||||
V = torch.randn(rank, cols)
|
||||
return (U * S.unsqueeze(0)) @ V
|
||||
|
||||
def test_fixed_rank_output_shape(self):
|
||||
"""Fixed mode: output lora_up/down have correct shapes."""
|
||||
mod = _load_module()
|
||||
delta = self._make_delta(64, 32, 8)
|
||||
up, down, alpha = mod._extract_lora_svd(delta, rank=4, rank_mode="fixed", energy_threshold=0.99)
|
||||
self.assertEqual(up.shape, (64, 4))
|
||||
self.assertEqual(down.shape, (4, 32))
|
||||
self.assertEqual(alpha, 4.0)
|
||||
|
||||
def test_auto_rank_energy_retained(self):
|
||||
"""Auto mode: retained energy >= threshold."""
|
||||
mod = _load_module()
|
||||
delta = self._make_delta(64, 32, 16)
|
||||
up, down, alpha = mod._extract_lora_svd(delta, rank=32, rank_mode="auto", energy_threshold=0.95)
|
||||
# Reconstruct and check energy
|
||||
reconstructed = up @ down
|
||||
original_energy = (delta ** 2).sum().item()
|
||||
reconstructed_energy = (reconstructed ** 2).sum().item()
|
||||
self.assertGreaterEqual(reconstructed_energy / original_energy, 0.90) # allow small numerical error
|
||||
|
||||
def test_near_zero_delta_returns_none(self):
|
||||
"""Near-zero delta (unaffected layer) returns None."""
|
||||
mod = _load_module()
|
||||
delta = torch.zeros(64, 32)
|
||||
result = mod._extract_lora_svd(delta, rank=4, rank_mode="fixed", energy_threshold=0.99)
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_singular_value_floor_applied(self):
|
||||
"""Noise-only singular values below floor are excluded even in fixed mode."""
|
||||
mod = _load_module()
|
||||
# Create a rank-2 signal with tiny noise
|
||||
signal = self._make_delta(32, 32, 2)
|
||||
noise = torch.randn(32, 32) * 1e-6
|
||||
delta = signal + noise
|
||||
up, down, alpha = mod._extract_lora_svd(delta, rank=16, rank_mode="fixed", energy_threshold=0.99)
|
||||
# Rank should be clamped to signal rank (2), not 16, due to floor
|
||||
self.assertLessEqual(alpha, 4.0) # floor should cut most of the 16 slots
|
||||
```
|
||||
|
||||
Also add a `_load_module()` helper near the top of the test file (after `_install_stubs()`):
|
||||
|
||||
```python
|
||||
def _load_module():
|
||||
"""Load lora_optimizer module with stubs in place."""
|
||||
_install_stubs()
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"lora_optimizer",
|
||||
os.path.join(os.path.dirname(__file__), "..", "lora_optimizer.py"),
|
||||
)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod)
|
||||
return mod
|
||||
```
|
||||
|
||||
**Step 2: Run tests to verify they fail**
|
||||
|
||||
```bash
|
||||
cd /media/p5/ComfyUI-ZImage-LoRA-Merger
|
||||
python -m pytest tests/test_lora_optimizer.py::TestExtractLoRAFromDelta -v
|
||||
```
|
||||
|
||||
Expected: `ERROR` — `module '_load_module' not found` or `AttributeError: module has no attribute '_extract_lora_svd'`
|
||||
|
||||
**Step 3: Commit the failing tests**
|
||||
|
||||
```bash
|
||||
git add tests/test_lora_optimizer.py
|
||||
git commit -m "test: add failing tests for _extract_lora_svd helper"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: Implement `_extract_lora_svd()` module-level helper
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` — add after `_compress_to_lowrank` (search for it, around line 2800)
|
||||
|
||||
**Step 1: Find the insertion point**
|
||||
|
||||
Search for `def _compress_to_lowrank` in `lora_optimizer.py` and insert the new function immediately after its closing line.
|
||||
|
||||
**Step 2: Add the helper**
|
||||
|
||||
```python
|
||||
def _extract_lora_svd(delta: torch.Tensor, rank: int, rank_mode: str, energy_threshold: float):
|
||||
"""
|
||||
SVD-decompose a weight delta into (lora_up, lora_down, alpha).
|
||||
|
||||
Returns None if the delta is near-zero (layer unaffected by the LoRA).
|
||||
Returns (lora_up, lora_down, alpha) otherwise.
|
||||
|
||||
lora_up shape: (rows, r)
|
||||
lora_down shape: (r, cols)
|
||||
alpha = float(r) — effective scale = 1.0 when loaded at strength 1.0
|
||||
"""
|
||||
NEAR_ZERO_NORM = 1e-8
|
||||
SV_FLOOR_RATIO = 1e-4 # drop singular values below floor_ratio * S[0]
|
||||
|
||||
if delta.norm().item() < NEAR_ZERO_NORM:
|
||||
return None
|
||||
|
||||
# Ensure 2D: conv layers arrive as [C_out, C_in, kH, kW]
|
||||
original_shape = delta.shape
|
||||
if delta.ndim > 2:
|
||||
delta = delta.reshape(delta.shape[0], -1)
|
||||
|
||||
delta = delta.float()
|
||||
|
||||
U, S, Vh = torch.linalg.svd(delta, full_matrices=False)
|
||||
|
||||
# Apply singular value floor to exclude noise
|
||||
sv_floor = SV_FLOOR_RATIO * S[0].item()
|
||||
valid = S > sv_floor
|
||||
if not valid.any():
|
||||
return None
|
||||
U, S, Vh = U[:, valid], S[valid], Vh[valid, :]
|
||||
|
||||
# Determine rank r
|
||||
if rank_mode == "auto":
|
||||
energy_cumsum = S.pow(2).cumsum(0) / S.pow(2).sum()
|
||||
r = int((energy_cumsum < energy_threshold).sum().item()) + 1
|
||||
r = min(r, S.shape[0])
|
||||
else:
|
||||
r = min(rank, S.shape[0])
|
||||
|
||||
sqrt_s = S[:r].sqrt()
|
||||
lora_up = U[:, :r] * sqrt_s.unsqueeze(0) # (rows, r)
|
||||
lora_down = sqrt_s.unsqueeze(1) * Vh[:r, :] # (r, cols)
|
||||
|
||||
return lora_up, lora_down, float(r)
|
||||
```
|
||||
|
||||
**Step 3: Run the tests**
|
||||
|
||||
```bash
|
||||
python -m pytest tests/test_lora_optimizer.py::TestExtractLoRAFromDelta -v
|
||||
```
|
||||
|
||||
Expected: all 4 tests PASS.
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: add _extract_lora_svd helper for LoRA extraction"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Implement `LoRAExtractFromModel` node class
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` — add the class before `NODE_CLASS_MAPPINGS` (near line 10476)
|
||||
|
||||
**Step 1: Find the insertion point**
|
||||
|
||||
Search for `NODE_CLASS_MAPPINGS` in `lora_optimizer.py`. Insert the new class in the block of node classes just before it.
|
||||
|
||||
**Step 2: Add the node class**
|
||||
|
||||
```python
|
||||
class LoRAExtractFromModel:
|
||||
"""
|
||||
Extracts a baked-in LoRA from a finetuned model by subtracting the clean
|
||||
base model weights and SVD-decomposing the per-layer delta.
|
||||
|
||||
Requires the original base model as a reference.
|
||||
Outputs:
|
||||
- LORA_STACK: feeds directly into LoRAOptimizerSimple / LoRAAutoTuner
|
||||
- LORA_DATA: feeds directly into SaveMergedLoRA
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"base_model": ("MODEL", {
|
||||
"tooltip": "The original clean base model (e.g. the Flux or SDXL base checkpoint "
|
||||
"that was used as the starting point for finetuning)."
|
||||
}),
|
||||
"finetuned_model": ("MODEL", {
|
||||
"tooltip": "The finetuned model with a LoRA already baked into its weights."
|
||||
}),
|
||||
"rank": ("INT", {
|
||||
"default": 32, "min": 1, "max": 512, "step": 1,
|
||||
"tooltip": "Maximum rank for SVD decomposition. Used directly in 'fixed' mode; "
|
||||
"acts as an upper bound in 'auto' mode."
|
||||
}),
|
||||
"rank_mode": (["auto", "fixed"], {
|
||||
"default": "auto",
|
||||
"tooltip": "'auto': choose rank to retain the given energy fraction (recommended). "
|
||||
"'fixed': always use exactly the specified rank."
|
||||
}),
|
||||
"energy_threshold": ("FLOAT", {
|
||||
"default": 0.99, "min": 0.5, "max": 1.0, "step": 0.01,
|
||||
"tooltip": "Fraction of delta energy to retain when rank_mode='auto'. "
|
||||
"0.99 = retain 99% of the signal. Higher = more accurate, higher rank."
|
||||
}),
|
||||
"strength": ("FLOAT", {
|
||||
"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05,
|
||||
"tooltip": "Strength to assign this extracted LoRA in the output stack."
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LORA_STACK", "LORA_DATA")
|
||||
RETURN_NAMES = ("lora_stack", "lora_data")
|
||||
FUNCTION = "extract"
|
||||
CATEGORY = "LoRA Optimizer"
|
||||
DESCRIPTION = (
|
||||
"Extracts a LoRA that was baked into a finetuned model by comparing it against "
|
||||
"the original base model. Requires the base model as a reference. "
|
||||
"Connect lora_stack to LoRAOptimizerSimple to merge additional LoRAs on top, "
|
||||
"or connect lora_data to SaveMergedLoRA to save the extracted LoRA to disk."
|
||||
)
|
||||
|
||||
def extract(self, base_model, finetuned_model, rank, rank_mode, energy_threshold, strength):
|
||||
import comfy.lora
|
||||
|
||||
base_sd = base_model.model.state_dict()
|
||||
fine_sd = finetuned_model.model.state_dict()
|
||||
|
||||
# Build inverse key map: model_weight_key → lora_prefix
|
||||
# comfy.lora.model_lora_keys_unet returns {lora_prefix: model_key}
|
||||
lora_to_model_unet = {}
|
||||
lora_to_model_clip = {}
|
||||
try:
|
||||
lora_to_model_unet = comfy.lora.model_lora_keys_unet(base_model.model, {})
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
lora_to_model_clip = comfy.lora.model_lora_keys_clip(base_model.model, {})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
model_to_lora = {} # model_key → (lora_prefix, is_clip)
|
||||
for lora_prefix, model_key in lora_to_model_unet.items():
|
||||
model_to_lora[model_key] = (lora_prefix, False)
|
||||
for lora_prefix, model_key in lora_to_model_clip.items():
|
||||
model_to_lora[model_key] = (lora_prefix, True)
|
||||
|
||||
# Validate that both models share the same keys
|
||||
base_keys = set(base_sd.keys())
|
||||
fine_keys = set(fine_sd.keys())
|
||||
if base_keys != fine_keys:
|
||||
missing = base_keys - fine_keys
|
||||
extra = fine_keys - base_keys
|
||||
logging.warning(
|
||||
f"[LoRAExtract] Key mismatch between base and finetuned models. "
|
||||
f"Missing in finetuned: {len(missing)}, extra in finetuned: {len(extra)}. "
|
||||
f"Proceeding with shared keys only."
|
||||
)
|
||||
|
||||
# Counters for diagnostics
|
||||
n_processed = 0
|
||||
n_skipped_zero = 0
|
||||
n_skipped_no_map = 0
|
||||
n_rank_capped = 0
|
||||
|
||||
# Output containers
|
||||
raw_lora_dict = {} # for LORA_STACK: {prefix.lora_up.weight: tensor, ...}
|
||||
model_patches = {} # for LORA_DATA
|
||||
clip_patches = {}
|
||||
key_map = {} # target_key → {"canonical_prefix": lora_prefix}
|
||||
|
||||
shared_keys = base_keys & fine_keys
|
||||
|
||||
for model_key in sorted(shared_keys):
|
||||
if model_key not in model_to_lora:
|
||||
n_skipped_no_map += 1
|
||||
continue
|
||||
|
||||
W_base = base_sd[model_key]
|
||||
W_fine = fine_sd[model_key]
|
||||
|
||||
if W_base.shape != W_fine.shape:
|
||||
logging.warning(f"[LoRAExtract] Shape mismatch for {model_key}, skipping.")
|
||||
continue
|
||||
|
||||
# Skip non-float layers (embeddings, norms, etc.)
|
||||
if not W_base.is_floating_point():
|
||||
continue
|
||||
|
||||
delta = W_fine.float() - W_base.float()
|
||||
result = _extract_lora_svd(delta, rank=rank, rank_mode=rank_mode, energy_threshold=energy_threshold)
|
||||
|
||||
if result is None:
|
||||
n_skipped_zero += 1
|
||||
continue
|
||||
|
||||
lora_up, lora_down, alpha = result
|
||||
lora_prefix, is_clip = model_to_lora[model_key]
|
||||
|
||||
if alpha >= rank and rank_mode == "fixed":
|
||||
n_rank_capped += 1
|
||||
|
||||
# Build raw LoRA dict entry (for LORA_STACK)
|
||||
raw_lora_dict[f"{lora_prefix}.lora_up.weight"] = lora_up
|
||||
raw_lora_dict[f"{lora_prefix}.lora_down.weight"] = lora_down
|
||||
raw_lora_dict[f"{lora_prefix}.alpha"] = torch.tensor(alpha)
|
||||
|
||||
# Build patch entry (for LORA_DATA)
|
||||
patch = LoRAAdapter(set(), (lora_up, lora_down, alpha, None, None, None))
|
||||
if is_clip:
|
||||
clip_patches[model_key] = patch
|
||||
else:
|
||||
model_patches[model_key] = patch
|
||||
|
||||
key_map[model_key] = {"canonical_prefix": lora_prefix}
|
||||
n_processed += 1
|
||||
|
||||
logging.info(
|
||||
f"[LoRAExtract] Extracted {n_processed} layers "
|
||||
f"({len(model_patches)} model, {len(clip_patches)} CLIP). "
|
||||
f"Skipped: {n_skipped_zero} near-zero, {n_skipped_no_map} unmapped. "
|
||||
f"Rank-capped layers: {n_rank_capped}."
|
||||
)
|
||||
|
||||
if n_processed == 0:
|
||||
logging.warning("[LoRAExtract] No layers extracted — models may be identical or incompatible.")
|
||||
|
||||
# Warn if this looks like a full finetune rather than a LoRA merge
|
||||
total_float_layers = sum(
|
||||
1 for k in shared_keys
|
||||
if k in base_sd and base_sd[k].is_floating_point()
|
||||
)
|
||||
if total_float_layers > 0 and n_processed > total_float_layers * 0.5:
|
||||
logging.warning(
|
||||
f"[LoRAExtract] {n_processed}/{total_float_layers} layers have significant deltas. "
|
||||
f"This may be a full finetune rather than a LoRA merge — "
|
||||
f"extraction will produce a high-rank approximation."
|
||||
)
|
||||
|
||||
# LORA_STACK output
|
||||
lora_stack = [{
|
||||
"name": "<extracted from finetuned_model>",
|
||||
"lora": raw_lora_dict,
|
||||
"strength": strength,
|
||||
"conflict_mode": "all",
|
||||
"key_filter": "all",
|
||||
"metadata": {},
|
||||
}]
|
||||
|
||||
# LORA_DATA output
|
||||
lora_data = {
|
||||
"model_patches": model_patches,
|
||||
"clip_patches": clip_patches,
|
||||
"key_map": key_map,
|
||||
"output_strength": strength,
|
||||
"clip_strength": strength,
|
||||
"suggested_max_strength": None,
|
||||
"sum_rank": rank,
|
||||
"merge_metadata": {
|
||||
"source_loras": [{"name": "<extracted>", "strength": strength}],
|
||||
"mode": "extract",
|
||||
"architecture": "unknown",
|
||||
},
|
||||
}
|
||||
|
||||
return (lora_stack, lora_data)
|
||||
```
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: add LoRAExtractFromModel node class"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Register the node in NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py` — lines around 10476
|
||||
|
||||
**Step 1: Find NODE_CLASS_MAPPINGS**
|
||||
|
||||
Search for `NODE_CLASS_MAPPINGS = {` in `lora_optimizer.py`.
|
||||
|
||||
**Step 2: Add the entry**
|
||||
|
||||
Add to `NODE_CLASS_MAPPINGS`:
|
||||
```python
|
||||
"LoRAExtractFromModel": LoRAExtractFromModel,
|
||||
```
|
||||
|
||||
Add to `NODE_DISPLAY_NAME_MAPPINGS` (search for it immediately below):
|
||||
```python
|
||||
"LoRAExtractFromModel": "LoRA Extract From Model",
|
||||
```
|
||||
|
||||
**Step 3: Run the existing tests to verify no regressions**
|
||||
|
||||
```bash
|
||||
python -m pytest tests/test_lora_optimizer.py -v
|
||||
```
|
||||
|
||||
Expected: all existing tests PASS + the 4 new extraction tests PASS.
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "feat: register LoRAExtractFromModel in node mappings"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: End-to-end verification
|
||||
|
||||
No automated test is possible here (requires real ComfyUI models), but verify the module loads cleanly:
|
||||
|
||||
**Step 1: Verify module imports without error**
|
||||
|
||||
```bash
|
||||
cd /media/p5/ComfyUI-ZImage-LoRA-Merger
|
||||
python -c "
|
||||
import sys, types, tempfile, os
|
||||
tmpdir = tempfile.gettempdir()
|
||||
|
||||
# Minimal stubs
|
||||
fp = types.ModuleType('folder_paths')
|
||||
fp.models_dir = tmpdir
|
||||
fp.get_temp_directory = lambda: tmpdir
|
||||
fp.get_user_directory = lambda: tmpdir
|
||||
fp.get_folder_paths = lambda _: [tmpdir]
|
||||
fp.get_filename_list = lambda _: []
|
||||
fp.get_full_path_or_raise = lambda _, n: n
|
||||
sys.modules['folder_paths'] = fp
|
||||
|
||||
import importlib.util
|
||||
spec = importlib.util.spec_from_file_location('lora_optimizer', 'lora_optimizer.py')
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
# Expect ImportError on comfy — that's fine, means the file parsed correctly
|
||||
try:
|
||||
spec.loader.exec_module(mod)
|
||||
except Exception as e:
|
||||
if 'comfy' in str(e) or 'safetensors' in str(e):
|
||||
print('OK: module parsed, expected import error on ComfyUI deps')
|
||||
else:
|
||||
raise
|
||||
"
|
||||
```
|
||||
|
||||
Expected output: `OK: module parsed, expected import error on ComfyUI deps`
|
||||
|
||||
**Step 2: Verify node appears in mappings**
|
||||
|
||||
```bash
|
||||
python -c "
|
||||
# Same stubs as above, then:
|
||||
# print('LoRAExtractFromModel' in mod.NODE_CLASS_MAPPINGS)
|
||||
"
|
||||
```
|
||||
|
||||
Expected: `True`
|
||||
|
||||
**Step 3: Final commit if any fixups needed, then tag**
|
||||
|
||||
```bash
|
||||
git log --oneline -5
|
||||
```
|
||||
@@ -1,100 +0,0 @@
|
||||
# AutoTuner Analysis Cache Design
|
||||
|
||||
**Date:** 2026-04-02
|
||||
**Status:** Approved
|
||||
|
||||
## Problem
|
||||
|
||||
The AutoTuner analysis pass (computing LoRA diffs, pairwise conflict sampling, magnitude statistics) is the most expensive part of a tuning run. Currently the results cache (`lora_hash`) includes LoRA strengths in its key, so changing any strength invalidates the cache entirely and forces a full re-analysis — even though conflict ratios, cosine similarities, and SVD directions are scale-invariant and do not change with strength.
|
||||
|
||||
This makes strength exploration slow: each strength sweep runs the full GPU-intensive analysis from scratch.
|
||||
|
||||
## Goal
|
||||
|
||||
Cache the scale-invariant analysis results separately, keyed only by LoRA file identity (not strengths or settings), so that strength changes reuse the analysis pass and only recompute the cheap magnitude-dependent synthesis.
|
||||
|
||||
## Design
|
||||
|
||||
### Cache Key
|
||||
|
||||
`names_only_hash`: 16-char SHA256 of the sorted list of `(name, mtime, file_size)` tuples for all LoRAs in the stack. No strengths, no settings.
|
||||
|
||||
- Fast: single `os.stat()` call per file, no content read
|
||||
- Reliable: covers 99%+ of real changes (file replacement changes mtime or size)
|
||||
- Invalidates automatically on: file replacement, LoRA set changes, algo version bumps
|
||||
|
||||
### File Format
|
||||
|
||||
Stored at `AUTOTUNER_MEMORY_DIR/<names_only_hash>.analysis.json`.
|
||||
|
||||
```json
|
||||
{
|
||||
"analysis_version": 1,
|
||||
"algo_version": "<AUTOTUNER_ALGO_VERSION>",
|
||||
"created_at": "2026-04-02T...",
|
||||
"source_loras": [
|
||||
{"name": "lora_a.safetensors", "mtime": 1234567890.0, "size": 102400}
|
||||
],
|
||||
"per_prefix": {
|
||||
"<lora_prefix>": {
|
||||
"pair_conflicts": {
|
||||
"0,1": [n_overlap, n_conflict, dot, norm_a_sq, norm_b_sq]
|
||||
},
|
||||
"diff_norms": {"0": 1.234, "1": 0.987},
|
||||
"ranks": {"0": 16, "1": 32},
|
||||
"target_key": "...",
|
||||
"is_clip": false
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
`pair_conflicts` stores raw scale-invariant components (not derived ratios) so strength rescaling can be correctly applied at synthesis time. `diff_norms` stores `diff.norm()` without strength applied.
|
||||
|
||||
### Read/Write Flow
|
||||
|
||||
In `auto_tune`, before the analysis pass:
|
||||
|
||||
1. Compute `names_only_hash` from `(name, mtime, size)` per LoRA file
|
||||
2. Attempt to load `<names_only_hash>.analysis.json`
|
||||
3. **Cache hit** (and `algo_version` matches):
|
||||
- Reconstruct `pair_conflicts` directly (scale-invariant, no change needed)
|
||||
- Reconstruct magnitude stats by scaling `diff_norms` by `abs(strength)` per LoRA
|
||||
- Check for sign flips vs cached data: if any strength changed sign, treat as miss and run full analysis (conservative, sign flips are rare)
|
||||
- Skip `_analyze_prefix` entirely
|
||||
4. **Cache miss**: run analysis as normal, then save `.analysis.json` atomically (`tmp` → `os.replace`)
|
||||
|
||||
The existing results cache (`memory_lora_hash + settings_hash`) is unchanged — it sits on top and continues to work as before.
|
||||
|
||||
### Dataset Persistence
|
||||
|
||||
`_save_tuner_dataset_entry` gets a `raw_analysis` field added with the same `per_prefix` structure. This ensures analysis data survives even if `.analysis.json` files are manually cleared, and feeds any future cross-run reuse from the dataset.
|
||||
|
||||
### Invalidation Rules
|
||||
|
||||
| Trigger | Effect |
|
||||
|---|---|
|
||||
| File `mtime` or `size` changes | `names_only_hash` changes → miss |
|
||||
| LoRA added/removed from stack | Different set → miss |
|
||||
| `algo_version` bump | File rejected → miss |
|
||||
| Strength sign flip | Conservative full-analysis fallback |
|
||||
| `memory_mode = "clear_and_run"` | Also deletes `.analysis.json` for current set |
|
||||
| Corrupted/unparseable file | Caught, treated as miss |
|
||||
|
||||
### What Does Not Change
|
||||
|
||||
- No new UI settings or toggles — fully transparent to the user
|
||||
- `memory_mode` semantics unchanged
|
||||
- Existing `.memory.json` file format unchanged
|
||||
- `lora_hash` and `settings_hash` computation unchanged
|
||||
|
||||
## Expected Benefit
|
||||
|
||||
For a strength sweep of N candidates over the same LoRA set:
|
||||
|
||||
- **Before**: O(N × full_analysis_cost)
|
||||
- **After**: O(analysis_cost + N × synthesis_cost)
|
||||
|
||||
Since `_analyze_prefix` (diff matmuls + GPU conflict sampling) dominates cost, synthesis (magnitude rescaling + heuristic scoring) is cheap arithmetic. Speedup scales with N.
|
||||
|
||||
Cross-run benefit: returning to the same LoRA set in a future ComfyUI session skips the full analysis pass on first run too.
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,580 +0,0 @@
|
||||
# AutoTuner Speed Optimizations
|
||||
|
||||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||||
|
||||
**Goal:** Reduce autotuner tuning loop time by 15-30% through 6 targeted micro-optimizations.
|
||||
|
||||
**Architecture:** All changes are in `lora_optimizer.py`. Each optimization is independent and touches a different function/section. No new files, no API changes, no behavior changes — only faster execution of the same logic.
|
||||
|
||||
**Tech Stack:** Python, PyTorch, existing test harness in `tests/test_lora_optimizer.py`
|
||||
|
||||
---
|
||||
|
||||
## Task 1: Lighten GC + remove `torch.cuda.empty_cache()` in Phase 2 loop
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:9270-9273`
|
||||
|
||||
**Context:** Inside the Phase 2 per-candidate loop, every iteration calls `gc.collect()` (full heap scan, 100-500ms on large processes) and `torch.cuda.empty_cache()` (forces CUDA sync). With `top_n=3`, that's 3 forced full GC cycles + 3 CUDA syncs.
|
||||
|
||||
**Why we can't remove gc.collect() entirely:** `merged_model` is a ComfyUI `ModelPatcher` wrapping a PyTorch `nn.Module`, which has circular references (parent→child→parent via `_modules`). Python's refcounting alone won't free these — gc is needed to break cycles. Removing it risks accumulating ~10-40GB of unreleased model clones until the post-loop `gc.collect()` at line 9354.
|
||||
|
||||
**Safe approach:** Replace `gc.collect()` (full scan of all 3 generations) with `gc.collect(0)` (youngest generation only — ~10x faster, catches most freshly-created cycles). Remove `torch.cuda.empty_cache()` — the CUDA allocator handles block reuse internally, and the forced sync is pure overhead. The post-loop cleanup at line 9354-9356 still does a full `gc.collect()` + `empty_cache()`.
|
||||
|
||||
**Step 1: Make the code change**
|
||||
|
||||
Current code (lines 9270-9273):
|
||||
```python
|
||||
del m_patches, c_patches # Drop patch-dict references so tensors can free
|
||||
gc.collect()
|
||||
if use_gpu:
|
||||
torch.cuda.empty_cache()
|
||||
```
|
||||
|
||||
New code:
|
||||
```python
|
||||
del m_patches, c_patches # Drop patch-dict references so tensors can free
|
||||
gc.collect(0) # gen-0 only: ~10x faster, catches fresh cycles from ModelPatcher
|
||||
```
|
||||
|
||||
**Step 2: Verify existing tests still pass**
|
||||
|
||||
Run: `cd /media/p5/ComfyUI-ZImage-LoRA-Merger && python -m pytest tests/test_lora_optimizer.py -v`
|
||||
Expected: All tests PASS
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "perf: lighten GC in Phase 2 loop — gen-0 only, drop empty_cache
|
||||
|
||||
Replace full gc.collect() with gc.collect(0) (youngest generation only,
|
||||
~10x faster). Remove torch.cuda.empty_cache() — CUDA allocator handles
|
||||
block reuse, forced sync is pure overhead. Post-loop cleanup at line
|
||||
9354-9356 still does full gc.collect() + empty_cache()."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 2: Replace `torch.randperm` with `torch.randint` in `_sample_pair_metrics`
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:2107-2113`
|
||||
- Test: `tests/test_lora_optimizer.py`
|
||||
|
||||
**Context:** `_sample_pair_metrics` downsamples to 100k elements when vectors are large (line 2108). Currently uses `torch.randperm(n)[:100000]` which generates a full permutation of n elements (O(n) time and memory) then takes the first 100k. For a 10M-element vector, this allocates a 10M-element int64 tensor just to pick 100k indices. `torch.randint(0, n, (100000,))` is O(k) — only allocates the 100k indices needed. The analysis path at line 4632 already uses `randint` correctly.
|
||||
|
||||
Note: `randint` samples with replacement (may produce duplicate indices), but for 100k out of millions the collision rate is <1% and doesn't affect conflict statistics meaningfully.
|
||||
|
||||
**Step 1: Write the test**
|
||||
|
||||
```python
|
||||
def test_sample_pair_metrics_downsamples_large_vectors(self):
|
||||
"""Pair metrics should work correctly with large vectors that trigger downsampling."""
|
||||
optimizer = lora_optimizer.LoRAOptimizer()
|
||||
# Create vectors larger than 100k to trigger downsampling path
|
||||
a = torch.randn(200000)
|
||||
b = torch.randn(200000)
|
||||
result = optimizer._sample_pair_metrics(a, b)
|
||||
# Should still produce valid metrics
|
||||
self.assertIn("overlap", result)
|
||||
self.assertIn("conflict", result)
|
||||
self.assertIn("dot", result)
|
||||
self.assertGreater(result["overlap"], 0)
|
||||
# Verify determinism: same inputs should give same results
|
||||
result2 = optimizer._sample_pair_metrics(a, b)
|
||||
self.assertEqual(result["overlap"], result2["overlap"])
|
||||
self.assertEqual(result["conflict"], result2["conflict"])
|
||||
self.assertAlmostEqual(result["dot"], result2["dot"], places=4)
|
||||
```
|
||||
|
||||
**Step 2: Run test to verify it passes (baseline)**
|
||||
|
||||
Run: `python -m pytest tests/test_lora_optimizer.py::LoRAOptimizerTests::test_sample_pair_metrics_downsamples_large_vectors -v`
|
||||
Expected: PASS
|
||||
|
||||
**Step 3: Make the code change**
|
||||
|
||||
In `lora_optimizer.py` at lines 2107-2113, replace `randperm` with `randint`:
|
||||
|
||||
Current:
|
||||
```python
|
||||
n = flat_a.numel()
|
||||
if n > 100000:
|
||||
target_device = flat_a.device
|
||||
g = torch.Generator(device=target_device).manual_seed(42)
|
||||
indices = torch.randperm(n, device=target_device, generator=g)[:100000]
|
||||
flat_a = flat_a[indices]
|
||||
flat_b = flat_b[indices]
|
||||
```
|
||||
|
||||
New:
|
||||
```python
|
||||
n = flat_a.numel()
|
||||
if n > 100000:
|
||||
target_device = flat_a.device
|
||||
g = torch.Generator(device=target_device).manual_seed(42)
|
||||
indices = torch.randint(0, n, (100000,), device=target_device, generator=g)
|
||||
flat_a = flat_a[indices]
|
||||
flat_b = flat_b[indices]
|
||||
```
|
||||
|
||||
**Step 4: Run tests to verify everything passes**
|
||||
|
||||
Run: `python -m pytest tests/test_lora_optimizer.py -v`
|
||||
Expected: All tests PASS
|
||||
|
||||
**Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py tests/test_lora_optimizer.py
|
||||
git commit -m "perf: use randint instead of randperm in _sample_pair_metrics
|
||||
|
||||
randperm(n)[:100k] allocates O(n) memory for full permutation then
|
||||
discards most of it. randint is O(k) — only allocates the 100k
|
||||
indices needed. For 10M-element vectors, avoids ~80MB temp allocation."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 3: Transfer tensors before float conversion in `_score_merge_result`
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:3696-3700`
|
||||
- Test: `tests/test_lora_optimizer.py`
|
||||
|
||||
**Context:** In the LoRAAdapter scoring branch of `_score_merge_result`, the code converts to float32 THEN transfers to score_device. Moving fp32 data is 2x the bandwidth of moving fp16/bf16 data. The fix is to transfer first (in native dtype), then convert to float32 on the target device. The float conversion still happens before the gram matrix computation, so numerical accuracy is preserved.
|
||||
|
||||
**Step 1: Write the test**
|
||||
|
||||
```python
|
||||
def test_score_merge_result_lora_adapter_on_device(self):
|
||||
"""_score_merge_result should handle LoRAAdapter patches correctly."""
|
||||
LoRAAdapter = lora_optimizer.LoRAAdapter
|
||||
up = torch.randn(4, 8)
|
||||
down = torch.randn(4, 16)
|
||||
alpha = 4.0
|
||||
adapter = LoRAAdapter(
|
||||
loaded_keys=set(),
|
||||
weights=(up, down, alpha, None, None, None),
|
||||
)
|
||||
patches = {("key1",): adapter}
|
||||
result = lora_optimizer._score_merge_result(patches, {}, compute_svd=False)
|
||||
self.assertIn("norm_mean", result)
|
||||
self.assertGreater(result["norm_mean"], 0)
|
||||
self.assertIn("composite_score", result)
|
||||
```
|
||||
|
||||
**Step 2: Run test to verify it passes (baseline)**
|
||||
|
||||
Run: `python -m pytest tests/test_lora_optimizer.py -k test_score_merge_result_lora_adapter -v`
|
||||
Expected: PASS
|
||||
|
||||
**Step 3: Make the code change**
|
||||
|
||||
In `lora_optimizer.py` at lines 3696-3700, reorder to transfer then convert:
|
||||
|
||||
Current:
|
||||
```python
|
||||
up_flat = mat_up.flatten(start_dim=1).float()
|
||||
down_flat = mat_down.flatten(start_dim=1).float()
|
||||
if score_device is not None:
|
||||
up_flat = up_flat.to(score_device)
|
||||
down_flat = down_flat.to(score_device)
|
||||
```
|
||||
|
||||
New:
|
||||
```python
|
||||
up_flat = mat_up.flatten(start_dim=1)
|
||||
down_flat = mat_down.flatten(start_dim=1)
|
||||
if score_device is not None:
|
||||
up_flat = up_flat.to(score_device)
|
||||
down_flat = down_flat.to(score_device)
|
||||
up_flat = up_flat.float()
|
||||
down_flat = down_flat.float()
|
||||
```
|
||||
|
||||
**Step 4: Run tests**
|
||||
|
||||
Run: `python -m pytest tests/test_lora_optimizer.py -v`
|
||||
Expected: All tests PASS
|
||||
|
||||
**Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py tests/test_lora_optimizer.py
|
||||
git commit -m "perf: transfer LoRA factors before float conversion in scoring
|
||||
|
||||
Moving fp16/bf16 tensors to score_device then converting to fp32
|
||||
halves the PCIe transfer bandwidth vs converting first. For ~240
|
||||
LoRAAdapter patches per candidate, saves meaningful transfer time."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 4: Gate psutil memory logging behind first/last iteration
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:9274-9283`
|
||||
|
||||
**Context:** Every Phase 2 iteration imports psutil and reads `/proc/self/status` for memory diagnostics. This is debug-only information. Gate it to only run on the first and last iteration of the loop.
|
||||
|
||||
**Step 1: Make the code change**
|
||||
|
||||
Current (lines 9274-9283):
|
||||
```python
|
||||
# Log memory usage to help diagnose leaks on large models
|
||||
try:
|
||||
import psutil
|
||||
proc = psutil.Process()
|
||||
rss_gb = proc.memory_info().rss / (1024**3)
|
||||
dc_mb = _diff_cache.size_mb() if _diff_cache else 0
|
||||
logging.info(f"[LoRA AutoTuner] Memory: process={rss_gb:.1f}GB"
|
||||
f"{f', diff_cache={dc_mb:.0f}MB' if dc_mb > 0 else ''}")
|
||||
except ImportError:
|
||||
pass
|
||||
```
|
||||
|
||||
New:
|
||||
```python
|
||||
# Log memory usage on first/last iteration to diagnose leaks
|
||||
if rank_idx == 0 or rank_idx == len(top_candidates) - 1:
|
||||
try:
|
||||
import psutil
|
||||
proc = psutil.Process()
|
||||
rss_gb = proc.memory_info().rss / (1024**3)
|
||||
dc_mb = _diff_cache.size_mb() if _diff_cache else 0
|
||||
logging.info(f"[LoRA AutoTuner] Memory: process={rss_gb:.1f}GB"
|
||||
f"{f', diff_cache={dc_mb:.0f}MB' if dc_mb > 0 else ''}")
|
||||
except ImportError:
|
||||
pass
|
||||
```
|
||||
|
||||
**Step 2: Verify tests pass**
|
||||
|
||||
Run: `python -m pytest tests/test_lora_optimizer.py -v`
|
||||
Expected: All PASS
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py
|
||||
git commit -m "perf: gate psutil memory logging to first/last Phase 2 iteration
|
||||
|
||||
Memory diagnostics on every iteration adds overhead without value.
|
||||
First and last iteration still show whether memory leaks occur."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 5: Cache LoRA format detection per lora_dict
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:684` (`_LoRAMergeBase.__init__`)
|
||||
- Modify: `lora_optimizer.py:1786-1821` (`_get_lora_key_info`)
|
||||
- Test: `tests/test_lora_optimizer.py`
|
||||
|
||||
**Context:** `_get_lora_key_info` (defined on `_LoRAMergeBase`, line 1786) tries 4 LoRA key formats (8 dict lookups) per call. It's called from 5 sites: `_prepare_group_diffs` (line 1672), `_process_prefix` (lines 4343, 4472), `_build_exact_linear_patch` (line 5400), and `analyze_and_enrich` (line 11302). A LoRA file uses one format consistently for all keys. Caching the detected format index avoids up to 6 wasted dict lookups per subsequent call.
|
||||
|
||||
**Important:** Initialize `_lora_format_cache` in `_LoRAMergeBase.__init__` (line 684), NOT `LoRAOptimizer.__init__`, since the method lives on `_LoRAMergeBase`. Stale cache entries are safe — if the cached format doesn't match a prefix, it falls through to the full cascade.
|
||||
|
||||
**Step 1: Write the failing test**
|
||||
|
||||
```python
|
||||
def test_lora_format_cache_avoids_repeated_detection(self):
|
||||
"""After detecting a LoRA's format once, subsequent prefixes should reuse it."""
|
||||
optimizer = lora_optimizer.LoRAOptimizer()
|
||||
# Create a LoRA with diffusers2 format (lora_B/lora_A) — the 3rd format tried
|
||||
lora_dict = {
|
||||
"unet.a.lora_B.weight": torch.tensor([[1.0]], dtype=torch.float32),
|
||||
"unet.a.lora_A.weight": torch.tensor([[1.0]], dtype=torch.float32),
|
||||
"unet.b.lora_B.weight": torch.tensor([[2.0]], dtype=torch.float32),
|
||||
"unet.b.lora_A.weight": torch.tensor([[1.0]], dtype=torch.float32),
|
||||
}
|
||||
# First call should detect format and cache it
|
||||
result1 = optimizer._get_lora_key_info(lora_dict, "unet.a")
|
||||
self.assertIsNotNone(result1)
|
||||
# Verify cache was populated
|
||||
self.assertIn(id(lora_dict), optimizer._lora_format_cache)
|
||||
# Second call with different prefix should use cache
|
||||
result2 = optimizer._get_lora_key_info(lora_dict, "unet.b")
|
||||
self.assertIsNotNone(result2)
|
||||
```
|
||||
|
||||
**Step 2: Run test to verify it fails**
|
||||
|
||||
Run: `python -m pytest tests/test_lora_optimizer.py -k test_lora_format_cache -v`
|
||||
Expected: FAIL — `_lora_format_cache` doesn't exist yet
|
||||
|
||||
**Step 3: Implement format caching**
|
||||
|
||||
Add `_lora_format_cache` to `_LoRAMergeBase.__init__` (line 684):
|
||||
|
||||
```python
|
||||
def __init__(self):
|
||||
self.loaded_loras = {}
|
||||
self._lora_format_cache = {} # id(lora_dict) -> format_index (0-3)
|
||||
```
|
||||
|
||||
Modify `_get_lora_key_info` (lines 1786-1821). To avoid duplicating the extraction body, extract it into a local helper:
|
||||
|
||||
```python
|
||||
def _get_lora_key_info(self, lora_dict, key_prefix):
|
||||
formats = [
|
||||
("{}.lora_up.weight", "{}.lora_down.weight"), # regular
|
||||
("{}_lora.up.weight", "{}_lora.down.weight"), # diffusers
|
||||
("{}.lora_B.weight", "{}.lora_A.weight"), # diffusers2
|
||||
("{}.lora.up.weight", "{}.lora.down.weight"), # diffusers3
|
||||
]
|
||||
|
||||
def _extract(up_key, down_key):
|
||||
mat_up = lora_dict[up_key]
|
||||
mat_down = lora_dict[down_key]
|
||||
alpha_key = "{}.alpha".format(key_prefix)
|
||||
alpha = lora_dict.get(alpha_key, None)
|
||||
if alpha is not None:
|
||||
alpha = alpha.item()
|
||||
else:
|
||||
alpha = mat_down.shape[0]
|
||||
mid_key = "{}.lora_mid.weight".format(key_prefix)
|
||||
mid = lora_dict.get(mid_key, None)
|
||||
return (mat_up, mat_down, alpha, mid)
|
||||
|
||||
# Try cached format first
|
||||
dict_id = id(lora_dict)
|
||||
cached_fmt = self._lora_format_cache.get(dict_id)
|
||||
if cached_fmt is not None:
|
||||
up_key = formats[cached_fmt][0].format(key_prefix)
|
||||
down_key = formats[cached_fmt][1].format(key_prefix)
|
||||
if up_key in lora_dict and down_key in lora_dict:
|
||||
return _extract(up_key, down_key)
|
||||
|
||||
for fmt_idx, (up_fmt, down_fmt) in enumerate(formats):
|
||||
up_key = up_fmt.format(key_prefix)
|
||||
down_key = down_fmt.format(key_prefix)
|
||||
if up_key in lora_dict and down_key in lora_dict:
|
||||
self._lora_format_cache[dict_id] = fmt_idx
|
||||
return _extract(up_key, down_key)
|
||||
|
||||
return None
|
||||
```
|
||||
|
||||
**Step 4: Run tests**
|
||||
|
||||
Run: `python -m pytest tests/test_lora_optimizer.py -v`
|
||||
Expected: All PASS
|
||||
|
||||
**Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py tests/test_lora_optimizer.py
|
||||
git commit -m "perf: cache LoRA format detection per lora_dict
|
||||
|
||||
A LoRA file uses one key format consistently. Caching the winning
|
||||
format index (0-3) after first detection avoids up to 6 wasted
|
||||
dict lookups per prefix on subsequent calls. Cache miss on stale
|
||||
entries is safe — falls through to full cascade."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 6: Skip single-LoRA patches in per-candidate scoring
|
||||
|
||||
**Files:**
|
||||
- Modify: `lora_optimizer.py:9119-9138` (Phase 2 scoring call site)
|
||||
- Modify: `lora_optimizer.py:3617-3850` (`_score_merge_result` — add `_baseline` and `_return_raw` params)
|
||||
- Test: `tests/test_lora_optimizer.py`
|
||||
|
||||
**Context:** `_score_merge_result` iterates ALL patches (model + clip), including single-LoRA prefixes. For single-LoRA prefixes, the merge path returns early at line 2868 (`if len(diffs_with_weights) == 1: return diff * weight`) — no TIES/DARE/merge strategy is applied. The patches are produced identically across candidates.
|
||||
|
||||
**Auto-strength caveat:** Different candidates may have different `auto_strength` settings ("enabled" vs "disabled"), which applies a uniform scalar (`model_auto_scale`) to ALL patches. However, the scoring metrics used in composite_score (norm_cv, effective_rank, sparsity) are all scale-invariant — the scalar cancels out. The one non-scale-invariant metric, `norm_energy_sq`, is used in `energy_preservation` (lines 9160-9206), but that's computed separately in the Phase 2 loop from `branch_energy` data — not from `_score_merge_result`'s energy_sq. So the cached single-LoRA contribution to the composite score is valid regardless of auto_strength.
|
||||
|
||||
**Approach:** Add `_baseline` (pre-computed raw lists) and `_return_raw` (return raw lists) parameters to `_score_merge_result`. On the first candidate, score single-LoRA patches separately with `_return_raw=True` to capture their raw norms/sparsities/ranks. On all candidates, score only multi-LoRA patches with `_baseline` set to the cached single-LoRA lists. The accumulators get the cached lists prepended, so composite metrics (mean, CV, etc.) are computed correctly from the full combined data.
|
||||
|
||||
**Step 1: Write the test**
|
||||
|
||||
```python
|
||||
def test_score_merge_result_baseline_matches_full(self):
|
||||
"""Scoring multi-LoRA patches with single-LoRA baseline should match full scoring."""
|
||||
LoRAAdapter = lora_optimizer.LoRAAdapter
|
||||
# Create 10 patches: 5 "single-LoRA" and 5 "multi-LoRA"
|
||||
single_patches = {}
|
||||
multi_patches = {}
|
||||
all_patches = {}
|
||||
for i in range(10):
|
||||
up = torch.randn(4, 8)
|
||||
down = torch.randn(4, 16)
|
||||
adapter = LoRAAdapter(set(), (up, down, 4.0, None, None, None))
|
||||
key = f"key{i}"
|
||||
all_patches[key] = adapter
|
||||
if i < 5:
|
||||
single_patches[key] = adapter
|
||||
else:
|
||||
multi_patches[key] = adapter
|
||||
|
||||
# Score everything together
|
||||
full = lora_optimizer._score_merge_result(all_patches, {}, compute_svd=False)
|
||||
|
||||
# Score single-LoRA first with _return_raw, then multi with _baseline
|
||||
sl = lora_optimizer._score_merge_result(
|
||||
single_patches, {}, compute_svd=False, _return_raw=True)
|
||||
baseline = sl["_raw"]
|
||||
combined = lora_optimizer._score_merge_result(
|
||||
multi_patches, {}, compute_svd=False, _baseline=baseline)
|
||||
|
||||
# Composite score should match
|
||||
self.assertAlmostEqual(full["composite_score"], combined["composite_score"], places=6)
|
||||
self.assertAlmostEqual(full["norm_mean"], combined["norm_mean"], places=6)
|
||||
self.assertAlmostEqual(full["norm_cv"], combined["norm_cv"], places=6)
|
||||
self.assertAlmostEqual(full["sparsity_mean"], combined["sparsity_mean"], places=6)
|
||||
self.assertAlmostEqual(full["norm_energy_sq"], combined["norm_energy_sq"], places=4)
|
||||
```
|
||||
|
||||
**Step 2: Run test to verify it fails**
|
||||
|
||||
Run: `python -m pytest tests/test_lora_optimizer.py -k test_score_merge_result_baseline -v`
|
||||
Expected: FAIL — `_return_raw` and `_baseline` parameters don't exist yet
|
||||
|
||||
**Step 3: Implement `_baseline` and `_return_raw` in `_score_merge_result`**
|
||||
|
||||
Modify the function signature (line 3617):
|
||||
```python
|
||||
def _score_merge_result(model_patches, clip_patches, compute_svd=True,
|
||||
score_device=None, arch_preset=None, lora_svd=False,
|
||||
_baseline=None, _return_raw=False):
|
||||
```
|
||||
|
||||
Replace the accumulator initialization (lines 3630-3634):
|
||||
|
||||
Current:
|
||||
```python
|
||||
norms = []
|
||||
importance_values = []
|
||||
effective_ranks = []
|
||||
sparsities = []
|
||||
```
|
||||
|
||||
New:
|
||||
```python
|
||||
if _baseline is not None:
|
||||
norms = list(_baseline["norms"])
|
||||
importance_values = list(_baseline["importance_values"])
|
||||
effective_ranks = list(_baseline["effective_ranks"])
|
||||
sparsities = list(_baseline["sparsities"])
|
||||
else:
|
||||
norms = []
|
||||
importance_values = []
|
||||
effective_ranks = []
|
||||
sparsities = []
|
||||
```
|
||||
|
||||
Before the `return metrics` at the end (line ~3849), add:
|
||||
```python
|
||||
if _return_raw:
|
||||
metrics["_raw"] = {
|
||||
"norms": norms,
|
||||
"importance_values": importance_values,
|
||||
"effective_ranks": effective_ranks,
|
||||
"sparsities": sparsities,
|
||||
}
|
||||
```
|
||||
|
||||
**Step 4: Run test to verify _baseline/_return_raw works**
|
||||
|
||||
Run: `python -m pytest tests/test_lora_optimizer.py -k test_score_merge_result_baseline -v`
|
||||
Expected: PASS
|
||||
|
||||
**Step 5: Wire up single-LoRA caching in Phase 2 loop**
|
||||
|
||||
Before the Phase 2 loop (around line 9070), build the set of single-LoRA target keys:
|
||||
|
||||
```python
|
||||
# Pre-identify single-LoRA target keys for scoring cache
|
||||
single_lora_keys = set()
|
||||
for pfx, info in all_key_targets.items():
|
||||
if prefix_stats.get(pfx, {}).get("n_loras", 0) <= 1:
|
||||
single_lora_keys.add(info[0]) # info is (target_key, is_clip)
|
||||
_cached_sl_baseline = None
|
||||
```
|
||||
|
||||
Inside the loop, replace the scoring block (lines 9121-9138) with:
|
||||
|
||||
```python
|
||||
m_patches = lora_data["model_patches"] if lora_data else {}
|
||||
c_patches = lora_data["clip_patches"] if lora_data else {}
|
||||
is_ortho_score = (
|
||||
abs(avg_cos_sim) < tuner_arch_preset["orthogonal_cos_sim_max"]
|
||||
and avg_subspace_overlap < 0.35
|
||||
)
|
||||
compute_svd = scoring_svd in ("merge_quality", "full") and not is_ortho_score
|
||||
compute_lora_svd = scoring_svd in ("lora_rank", "full")
|
||||
score_dev = torch.device("cuda") if scoring_device == "gpu" and torch.cuda.is_available() else None
|
||||
t_score = time.time()
|
||||
score_arch = tuner_arch_preset if scoring_formula == "v2" else None
|
||||
|
||||
if single_lora_keys:
|
||||
# Filter patches into single-LoRA and multi-LoRA sets
|
||||
def _target_key(k):
|
||||
return k[0] if isinstance(k, tuple) else k
|
||||
|
||||
m_multi = {k: v for k, v in m_patches.items()
|
||||
if _target_key(k) not in single_lora_keys}
|
||||
c_multi = {k: v for k, v in c_patches.items()
|
||||
if _target_key(k) not in single_lora_keys}
|
||||
|
||||
if _cached_sl_baseline is None:
|
||||
# First candidate: score single-LoRA patches, cache raw lists
|
||||
m_single = {k: v for k, v in m_patches.items()
|
||||
if _target_key(k) in single_lora_keys}
|
||||
c_single = {k: v for k, v in c_patches.items()
|
||||
if _target_key(k) in single_lora_keys}
|
||||
sl_measured = _score_merge_result(
|
||||
m_single, c_single, compute_svd=compute_svd,
|
||||
score_device=score_dev, arch_preset=score_arch,
|
||||
lora_svd=compute_lora_svd, _return_raw=True)
|
||||
_cached_sl_baseline = sl_measured.get("_raw")
|
||||
del m_single, c_single, sl_measured
|
||||
|
||||
# Score multi-LoRA patches with single-LoRA baseline merged in
|
||||
measured = _score_merge_result(
|
||||
m_multi, c_multi, compute_svd=compute_svd,
|
||||
score_device=score_dev, arch_preset=score_arch,
|
||||
lora_svd=compute_lora_svd,
|
||||
_baseline=_cached_sl_baseline)
|
||||
else:
|
||||
measured = _score_merge_result(
|
||||
m_patches, c_patches, compute_svd=compute_svd,
|
||||
score_device=score_dev, arch_preset=score_arch,
|
||||
lora_svd=compute_lora_svd)
|
||||
```
|
||||
|
||||
**Step 6: Run all tests**
|
||||
|
||||
Run: `python -m pytest tests/test_lora_optimizer.py -v`
|
||||
Expected: All PASS
|
||||
|
||||
**Step 7: Commit**
|
||||
|
||||
```bash
|
||||
git add lora_optimizer.py tests/test_lora_optimizer.py
|
||||
git commit -m "perf: cache single-LoRA patch scoring across Phase 2 candidates
|
||||
|
||||
Single-LoRA patches bypass all merge strategies (early return at line
|
||||
2868). Score them once on the first candidate and merge as baseline
|
||||
for subsequent candidates. All composite_score components (norm_cv,
|
||||
effective_rank, sparsity) are scale-invariant, so auto_strength
|
||||
differences between candidates don't affect correctness."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
| Task | Optimization | Est. Impact |
|
||||
|------|-------------|-------------|
|
||||
| 1 | gc.collect(0) + drop empty_cache in loop | 5-15% Phase 2 |
|
||||
| 2 | randperm -> randint | 1-3% Phase 1 |
|
||||
| 3 | Transfer before float conversion | 2-5% scoring |
|
||||
| 4 | Gate psutil logging | Negligible |
|
||||
| 5 | Cache LoRA format detection | 1-3% Phase 1 |
|
||||
| 6 | Skip single-LoRA patch re-scoring | 10-30% scoring |
|
||||
Reference in New Issue
Block a user