diff --git a/__init__.py b/__init__.py index fc2a2ed..25da4c4 100644 --- a/__init__.py +++ b/__init__.py @@ -42,6 +42,7 @@ from .nodes.text_encoder_extract import ( TextEncoderDoRAExtractKnee, TextEncoderDoRAExtractFrobenius, ) from .nodes.lora_resize import ( + LoRANormalizeAlpha, LoRAResizeFixed, LoRAResizeRatio, LoRAResizeFrobenius, LoRAResizeCumulative, LoRAMergeToModel @@ -98,6 +99,7 @@ class ModelUtilsExtension(ComfyExtension): TextEncoderDoRAExtractRatio, TextEncoderDoRAExtractQuantile, TextEncoderDoRAExtractKnee, TextEncoderDoRAExtractFrobenius, # LoRA Resize + LoRANormalizeAlpha, LoRAResizeFixed, LoRAResizeRatio, LoRAResizeFrobenius, LoRAResizeCumulative, # LoRA Multi-Merge diff --git a/nodes/lora_resize.py b/nodes/lora_resize.py index cbf3620..d1f4c53 100644 --- a/nodes/lora_resize.py +++ b/nodes/lora_resize.py @@ -19,7 +19,7 @@ from .device_utils import estimate_model_size, prepare_for_large_operation, clea from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned from .uel_io import atomic_uel_writer, stream_work_units from .artifact_paths import canonical_model_artifact_path -from .lora_alpha import normalize_lora_pair +from .lora_alpha import lora_alpha_scale, normalize_lora_pair from .quantization_guard import inspect_low_bit_input, layer_has_low_bit, write_preserved_tensor from typing import Optional, Dict, Tuple, List @@ -633,6 +633,116 @@ def _rank_summary(rank_counts: Counter) -> str: # Main Resize Function # ============================================================================= +def normalize_lora_alpha_file( + lora_path: str, + output_filename: str, + verbose: bool = True, +) -> str: + """Materialize LoRA alpha scaling and remove alpha tensors.""" + handler = MemoryEfficientSafeOpen(lora_path, low_memory=True) + try: + low_bit_keys = inspect_low_bit_input( + handler, f"LoRA ({lora_path})", "LoRA Alpha Normalize" + ) + metadata = (handler.metadata() or {}).copy() + all_keys = handler.keys() + pairs, _ = parse_lora_layers(all_keys) + + work_units = [] + normalized_up_keys = set() + consumed_alpha_keys = set() + for block_name, block_keys in pairs.items(): + alpha_key = block_keys.get("alpha") + if alpha_key is None: + continue + down_key = block_keys.get("down") + up_key = block_keys.get("up") + if down_key is None or up_key is None: + raise ValueError( + f"Alpha tensor for '{block_name}' has no complete LoRA pair." + ) + if down_key in low_bit_keys or up_key in low_bit_keys: + raise ValueError( + f"Cannot alpha-normalize low-bit LoRA factors for '{block_name}'." + ) + down_shape = handler.get_shape(down_key) + up_shape = handler.get_shape(up_key) + if len(down_shape) < 2 or len(up_shape) < 2: + raise ValueError( + f"LoRA factors for '{block_name}' must have at least two dimensions." + ) + rank = int(down_shape[0]) + if rank <= 0 or int(up_shape[1]) != rank: + raise ValueError(f"LoRA factor rank mismatch for '{block_name}'.") + work_units.append((block_name, {0: [up_key, alpha_key]})) + normalized_up_keys.add(up_key) + consumed_alpha_keys.add(alpha_key) + + all_alpha_keys = {key for key in all_keys if key.endswith(".alpha")} + orphan_alpha_keys = sorted(all_alpha_keys - consumed_alpha_keys) + if orphan_alpha_keys: + raise ValueError( + f"Alpha tensor has no complete LoRA pair: {orphan_alpha_keys[0]}" + ) + if not work_units: + raise ValueError("Input LoRA contains no alpha tensors.") + + metadata["alpha_normalized"] = "true" + metadata["alpha_normalization"] = ( + "up factor multiplied by alpha divided by rank; alpha tensors removed" + ) + output_path, _ = canonical_model_artifact_path("loras", output_filename) + + with atomic_uel_writer(output_path, metadata) as writer: + for key in all_keys: + if key in all_alpha_keys or key in normalized_up_keys: + continue + write_preserved_tensor(writer, key, handler, force_raw=True) + + with torch.no_grad(): + iterator = stream_work_units( + {0: handler}, work_units, pin_memory=False + ) + try: + for block_name, loaded in iterator: + block_keys = pairs[block_name] + up_key = block_keys["up"] + alpha_key = block_keys["alpha"] + rank = int(handler.get_shape(block_keys["down"])[0]) + up = loaded[(0, up_key)] + alpha = loaded[(0, alpha_key)] + if alpha.numel() != 1 or not torch.isfinite(alpha).all().item(): + raise ValueError( + f"LoRA alpha for '{block_name}' must be a finite scalar." + ) + scale = lora_alpha_scale(alpha, rank, layer=block_name) + output = up + try: + if scale != 1.0: + output = ( + up.to(dtype=torch.float32) + .mul_(scale) + .to(dtype=up.dtype) + .contiguous() + ) + writer.write_batch([(up_key, output.cpu().contiguous())]) + finally: + del output + del up + del alpha + finally: + iterator.close() + + if verbose: + print( + f"[LoRA Alpha Normalize] Removed {len(work_units)} alpha tensors" + ) + print(f"[LoRA Alpha Normalize] Saved to {output_path}") + return output_path + finally: + handler.__exit__(None, None, None) + cleanup_after_operation() + def resize_lora_file( lora_path: str, new_rank: int, @@ -897,6 +1007,39 @@ def resize_lora_file( # Node Definitions # ============================================================================= +class LoRANormalizeAlpha(io.ComfyNode): + """Materialize LoRA alpha scaling into its up factors.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LoRANormalizeAlpha", + display_name="LoRA Normalize Alpha", + category="ModelUtils/LoRA/Utilities", + description="Apply each alpha divided by rank scale to its LoRA up factor and save an alpha-free LoRA.", + inputs=[ + io.Combo.Input( + "lora_name", + options=folder_paths.get_filename_list("loras"), + tooltip="LoRA whose per-layer alpha scaling will be materialized.", + ), + io.String.Input( + "output_filename", + default="normalized_lora", + tooltip="Output filename without extension, written under ComfyUI's LoRA directory.", + ), + ], + outputs=[io.AnyType.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, lora_name, output_filename) -> io.NodeOutput: + lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) + normalize_lora_alpha_file(lora_path, output_filename) + _, output_name = canonical_model_artifact_path("loras", output_filename) + return io.NodeOutput(output_name) + class LoRAResizeFixed(io.ComfyNode): """Resize LoRA to a fixed rank.""" diff --git a/tests/test_lora_resize_factor_space.py b/tests/test_lora_resize_factor_space.py index 1d3fd90..be11372 100644 --- a/tests/test_lora_resize_factor_space.py +++ b/tests/test_lora_resize_factor_space.py @@ -189,6 +189,91 @@ def test_resize_file_uses_only_async_uel_and_preserves_per_layer_scale( assert "diffusion_model.r3.alpha" not in output_tensors +def test_normalize_alpha_uses_bounded_uel_and_removes_alpha( + monkeypatch, tmp_path, resize +): + source = tmp_path / "alpha_input.safetensors" + down = torch.eye(2, dtype=torch.float16) + up = torch.tensor([[2.0, 4.0], [6.0, 8.0]], dtype=torch.float16) + _write_uel( + source, + { + "diffusion_model.block.lora_A.weight": down, + "diffusion_model.block.lora_B.weight": up, + "diffusion_model.block.alpha": torch.tensor(1.0), + "other.weight": torch.tensor([3.0]), + }, + metadata={"source": "test"}, + ) + monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path)) + monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None) + + original_async = resize.MemoryEfficientSafeOpen.async_stream + original_mark = resize.MemoryEfficientSafeOpen.mark_processed + async_calls = [] + marked = [] + + def track_async(self, keys, **kwargs): + async_calls.append((tuple(keys), kwargs, self.low_memory)) + yield from original_async(self, keys, **kwargs) + + def track_mark(self, key): + marked.append(key) + return original_mark(self, key) + + def reject_sync_get(*args, **kwargs): + raise AssertionError("alpha normalization used synchronous get_tensor") + + monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "async_stream", track_async) + monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "mark_processed", track_mark) + monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "get_tensor", reject_sync_get) + + output = resize.normalize_lora_alpha_file( + str(source), "alpha_normalized", verbose=False + ) + + assert len(async_calls) == 1 + streamed_keys, options, low_memory = async_calls[0] + assert streamed_keys == ( + "diffusion_model.block.lora_B.weight", + "diffusion_model.block.alpha", + ) + assert options == { + "batch_size": 1, + "prefetch_batches": 1, + "pin_memory": False, + } + assert low_memory is True + assert Counter(marked) == Counter(streamed_keys) + + output_tensors = _read_uel(output, original_async, original_mark) + assert "diffusion_model.block.alpha" not in output_tensors + assert torch.equal(output_tensors["diffusion_model.block.lora_A.weight"], down) + assert torch.equal( + output_tensors["diffusion_model.block.lora_B.weight"], up * 0.5 + ) + assert torch.equal(output_tensors["other.weight"], torch.tensor([3.0])) + with MemoryEfficientSafeOpen(output, low_memory=True) as loader: + assert loader.metadata()["alpha_normalized"] == "true" + + +def test_normalize_alpha_rejects_alpha_free_input(monkeypatch, tmp_path, resize): + source = tmp_path / "no_alpha.safetensors" + _write_uel( + source, + { + "diffusion_model.block.lora_A.weight": torch.eye(2), + "diffusion_model.block.lora_B.weight": torch.eye(2), + }, + ) + monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None) + + with pytest.raises(ValueError, match="contains no alpha tensors"): + resize.normalize_lora_alpha_file( + str(source), "unused", verbose=False + ) + + def test_resize_does_not_invent_alpha_for_peft_pair(monkeypatch, tmp_path, resize): source = tmp_path / "peft_no_alpha.safetensors" _write_uel(