feat: refactor conversion module, add test for svd correctness

Author: simpletrontdip

invokeai/backend/patches/layers/utils.py

@@ -46,6 +46,7 @@ def swap_shift_scale_for_linear_weight(weight: torch.Tensor) -> torch.Tensor:
 def decomposite_weight_matric_with_rank(
     delta: torch.Tensor,
     rank: int,
+    epsilon: float = 1e-8,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Decompose given matrix with a specified rank."""
     U, S, V = torch.svd(delta)
@@ -55,50 +56,9 @@ def decomposite_weight_matric_with_rank(
     S_r = S[:rank]
     V_r = V[:, :rank]
 
-    S_sqrt = torch.sqrt(S_r)
+    S_sqrt = torch.sqrt(S_r + epsilon) # regularization
 
     up = torch.matmul(U_r, torch.diag(S_sqrt))
     down = torch.matmul(torch.diag(S_sqrt), V_r.T)
 
     return up, down
-
-
-def approximate_flux_adaLN_lora_layer_from_diffusers_state_dict(state_dict: Dict[str, torch.Tensor]) -> LoRALayer:
-    '''Approximate given diffusers AdaLN loRA layer in our Flux model'''
-
-    if not "lora_up.weight" in state_dict:
-        raise ValueError(f"Unsupported lora format: {state_dict.keys()}, missing lora_up")
-    
-    if not "lora_down.weight" in state_dict:
-        raise ValueError(f"Unsupported lora format: {state_dict.keys()}, missing lora_down")
-    
-    up = state_dict.pop('lora_up.weight')
-    down = state_dict.pop('lora_down.weight')
-
-    dtype = up.dtype
-    device = up.device
-    up_shape = up.shape
-    down_shape = down.shape
-    
-    # desired low rank
-    rank = up_shape[1]
-
-    # up scaling for more precise
-    up.double()
-    down.double()
-    weight  = up.reshape(up.shape[0], -1) @ down.reshape(down.shape[0], -1)
-
-    # swap to our linear format
-    swapped = swap_shift_scale_for_linear_weight(weight)
-
-    _up, _down = decomposite_weight_matric_with_rank(swapped, rank)
-
-    assert(_up.shape == up_shape)
-    assert(_down.shape == down_shape)
-
-    # down scaling to original dtype, device
-    state_dict['lora_up.weight'] = _up.to(dtype).to(device=device)
-    state_dict['lora_down.weight'] = _down.to(dtype).to(device=device)
-
-    return LoRALayer.from_state_dict_values(state_dict)
-

invokeai/backend/patches/lora_conversions/flux_diffusers_lora_conversion_utils.py

@@ -2,9 +2,10 @@
 
 import torch
 
+from invokeai.backend.patches.layers.lora_layer import LoRALayer
 from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
 from invokeai.backend.patches.layers.merged_layer_patch import MergedLayerPatch, Range
-from invokeai.backend.patches.layers.utils import any_lora_layer_from_state_dict, approximate_flux_adaLN_lora_layer_from_diffusers_state_dict
+from invokeai.backend.patches.layers.utils import any_lora_layer_from_state_dict, swap_shift_scale_for_linear_weight, decomposite_weight_matric_with_rank
 from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX
 from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 
@@ -29,6 +30,49 @@ def is_state_dict_likely_in_flux_diffusers_format(state_dict: Dict[str, torch.Te
 
     return all_keys_in_peft_format and all_expected_keys_present
 
+def approximate_flux_adaLN_lora_layer_from_diffusers_state_dict(state_dict: Dict[str, torch.Tensor]) -> LoRALayer:
+    '''Approximate given diffusers AdaLN loRA layer in our Flux model'''
+
+    if not "lora_up.weight" in state_dict:
+        raise ValueError(f"Unsupported lora format: {state_dict.keys()}, missing lora_up")
+    
+    if not "lora_down.weight" in state_dict:
+        raise ValueError(f"Unsupported lora format: {state_dict.keys()}, missing lora_down")
+    
+    up = state_dict.pop('lora_up.weight')
+    down = state_dict.pop('lora_down.weight')
+
+    # layer-patcher upcast things to f32, 
+    # we want to maintain a better precison for this one
+    dtype = torch.float32
+
+    device = up.device
+    up_shape = up.shape
+    down_shape = down.shape
+    
+    # desired low rank
+    rank = up_shape[1]
+
+    # up scaling for more precise
+    up = up.to(torch.float32)
+    down = down.to(torch.float32)
+
+    weight  = up.reshape(up_shape[0], -1) @ down.reshape(down_shape[0], -1)
+
+    # swap to our linear format
+    swapped = swap_shift_scale_for_linear_weight(weight)
+
+    _up, _down = decomposite_weight_matric_with_rank(swapped, rank)
+
+    assert(_up.shape == up_shape)
+    assert(_down.shape == down_shape)
+
+    # down scaling to original dtype, device
+    state_dict['lora_up.weight'] = _up.to(dtype).to(device=device)
+    state_dict['lora_down.weight'] = _down.to(dtype).to(device=device)
+
+    return LoRALayer.from_state_dict_values(state_dict)
+
 
 def lora_model_from_flux_diffusers_state_dict(
     state_dict: Dict[str, torch.Tensor], alpha: float | None

tests/backend/patches/lora_conversions/test_flux_diffusers_lora_conversion_utils.py

@@ -1,9 +1,12 @@
 import pytest
 import torch
 
+
+from invokeai.backend.patches.layers.utils import swap_shift_scale_for_linear_weight
 from invokeai.backend.patches.lora_conversions.flux_diffusers_lora_conversion_utils import (
     is_state_dict_likely_in_flux_diffusers_format,
     lora_model_from_flux_diffusers_state_dict,
+    approximate_flux_adaLN_lora_layer_from_diffusers_state_dict,
 )
 from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX
 from tests.backend.patches.lora_conversions.lora_state_dicts.flux_dora_onetrainer_format import (
@@ -78,3 +81,55 @@ def test_lora_model_from_flux_diffusers_state_dict_extra_keys_error():
     # Check that an error is raised.
     with pytest.raises(AssertionError):
         lora_model_from_flux_diffusers_state_dict(state_dict, alpha=8.0)
+
+
+@pytest.mark.parametrize("layer_sd_keys",[
+    {}, # no keys
+    {'lora_A.weight': [1024, 8], 'lora_B.weight': [8, 512]}, # wrong keys
+    {'lora_up.weight': [1024, 8],}, # missing key
+    {'lora_down.weight': [8, 512],}, # missing key
+])
+def test_approximate_adaLN_from_state_dict_should_only_accept_vanilla_LoRA_format(layer_sd_keys: dict[str, list[int]]):
+    """Should only accept the valid state dict"""
+    layer_state_dict = keys_to_mock_state_dict(layer_sd_keys)
+
+    with pytest.raises(ValueError):
+        approximate_flux_adaLN_lora_layer_from_diffusers_state_dict(layer_state_dict)
+
+
+@pytest.mark.parametrize("dtype, rtol", [
+   (torch.float32, 1e-4),
+   (torch.half, 1e-3),
+])
+def test_approximate_adaLN_from_state_dict_should_work(dtype: torch.dtype, rtol: float, rate: float = 0.99):
+    """Test that we should approximate good enough adaLN layer from diffusers state dict.
+    This should tolorance some kind of errorness respect to input dtype"""
+    input_dim = 1024
+    output_dim = 512
+    rank = 8  # Low rank
+    total = input_dim * output_dim
+
+    up = torch.randn(input_dim, rank, dtype=dtype)
+    down = torch.randn(rank, output_dim, dtype=dtype)
+
+    layer_state_dict = {
+        'lora_up.weight': up,
+        'lora_down.weight': down
+    }
+
+    # XXX Layer patcher cast things to f32
+    original = up.float() @ down.float()
+    swapped = swap_shift_scale_for_linear_weight(original)
+
+    layer = approximate_flux_adaLN_lora_layer_from_diffusers_state_dict(layer_state_dict)
+    weight = layer.get_weight(original).float()
+
+    print(weight.dtype, swapped.dtype, layer.up.dtype)
+
+    close_count = torch.isclose(weight, swapped, rtol=rtol).sum().item()
+    close_rate = close_count / total
+
+    assert close_rate > rate
+
+
+