chore: ruff fix

Author: simpletrontdip

invokeai/backend/patches/layers/utils.py

@@ -35,14 +35,14 @@ def any_lora_layer_from_state_dict(state_dict: Dict[str, torch.Tensor]) -> BaseL
         raise ValueError(f"Unsupported lora format: {state_dict.keys()}")
 
 
-
 def swap_shift_scale_for_linear_weight(weight: torch.Tensor) -> torch.Tensor:
     """Swap shift/scale for given linear layer back and forth"""
     # In SD3 and Flux implementation of AdaLayerNormContinuous, it split linear projection output into shift, scale;
     # while in diffusers it split into scale, shift. This will flip them around
-    chunk1, chunk2 = weight.chunk(2, dim=0) 
+    chunk1, chunk2 = weight.chunk(2, dim=0)
     return torch.cat([chunk2, chunk1], dim=0)
 
+
 def decomposite_weight_matric_with_rank(
     delta: torch.Tensor,
     rank: int,
@@ -56,7 +56,7 @@ def decomposite_weight_matric_with_rank(
     S_r = S[:rank]
     V_r = V[:, :rank]
 
-    S_sqrt = torch.sqrt(S_r + epsilon) # regularization
+    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)

invokeai/backend/patches/lora_conversions/flux_diffusers_lora_conversion_utils.py

@@ -2,10 +2,14 @@
 
 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.lora_layer import LoRALayer
 from invokeai.backend.patches.layers.merged_layer_patch import MergedLayerPatch, Range
-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.layers.utils import (
+    any_lora_layer_from_state_dict,
+    decomposite_weight_matric_with_rank,
+    swap_shift_scale_for_linear_weight,
+)
 from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX
 from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 
@@ -30,46 +34,47 @@ 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'''
+    """Approximate given diffusers AdaLN loRA layer in our Flux model"""
 
-    if not "lora_up.weight" in state_dict:
+    if "lora_up.weight" not in state_dict:
         raise ValueError(f"Unsupported lora format: {state_dict.keys()}, missing lora_up")
-    
-    if not "lora_down.weight" in state_dict:
+
+    if "lora_down.weight" not 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, 
+    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)
+    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)
+    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)
+    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)
 
@@ -131,7 +136,7 @@ def add_adaLN_lora_layer_if_present(src_key: str, dst_key: str) -> None:
             src_layer_dict = grouped_state_dict.pop(src_key)
             values = get_lora_layer_values(src_layer_dict)
             layers[dst_key] = approximate_flux_adaLN_lora_layer_from_diffusers_state_dict(values)
-    
+
     def add_qkv_lora_layer_if_present(
         src_keys: list[str],
         src_weight_shapes: list[tuple[int, int]],
@@ -274,8 +279,8 @@ def add_qkv_lora_layer_if_present(
     # Final layer.
     add_lora_layer_if_present("proj_out", "final_layer.linear")
     add_adaLN_lora_layer_if_present(
-        'norm_out.linear',
-        'final_layer.adaLN_modulation.1',
+        "norm_out.linear",
+        "final_layer.adaLN_modulation.1",
     )
 
     # Assert that all keys were processed.

tests/backend/patches/layers/test_layer_utils.py

@@ -1,6 +1,9 @@
 import torch
 
-from invokeai.backend.patches.layers.utils import decomposite_weight_matric_with_rank, swap_shift_scale_for_linear_weight
+from invokeai.backend.patches.layers.utils import (
+    decomposite_weight_matric_with_rank,
+    swap_shift_scale_for_linear_weight,
+)
 
 
 def test_swap_shift_scale_for_linear_weight():
@@ -9,38 +12,37 @@ def test_swap_shift_scale_for_linear_weight():
     expected = torch.Tensor([2, 1])
 
     swapped = swap_shift_scale_for_linear_weight(original)
-    assert(torch.allclose(expected, swapped))
+    assert torch.allclose(expected, swapped)
 
-    size= (3, 4)
+    size = (3, 4)
     first = torch.randn(size)
     second = torch.randn(size)
 
     original = torch.concat([first, second])
     expected = torch.concat([second, first])
 
     swapped = swap_shift_scale_for_linear_weight(original)
-    assert(torch.allclose(expected, swapped))
+    assert torch.allclose(expected, swapped)
 
     # call this twice will reconstruct the original
     reconstructed = swap_shift_scale_for_linear_weight(swapped)
-    assert(torch.allclose(reconstructed, original))
+    assert torch.allclose(reconstructed, original)
+
 
 def test_decomposite_weight_matric_with_rank():
     """Test that decompsition of given matrix into 2 low rank matrices work"""
     input_dim = 1024
     output_dim = 1024
     rank = 8  # Low rank
 
-
     A = torch.randn(input_dim, rank).double()
     B = torch.randn(rank, output_dim).double()
     W0 = A @ B
 
     C, D = decomposite_weight_matric_with_rank(W0, rank)
     R = C @ D
 
-    assert(C.shape == A.shape)
-    assert(D.shape == B.shape)
+    assert C.shape == A.shape
+    assert D.shape == B.shape
 
     assert torch.allclose(W0, R)
-