Skip to content

Commit a05642c

Browse files
committed
Rename ModelPatcher -> LayerPatcher to avoid conflicts with another ModelPatcher definition.
1 parent d0de265 commit a05642c

File tree

9 files changed

+33
-33
lines changed

9 files changed

+33
-33
lines changed

invokeai/app/invocations/compel.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
from invokeai.app.util.ti_utils import generate_ti_list
2222
from invokeai.backend.model_patcher import ModelPatcher
2323
from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
24-
from invokeai.backend.patches.model_patcher import ModelPatcher
24+
from invokeai.backend.patches.model_patcher import LayerPatcher
2525
from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
2626
BasicConditioningInfo,
2727
ConditioningFieldData,
@@ -82,7 +82,7 @@ def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]:
8282
# apply all patches while the model is on the target device
8383
text_encoder_info.model_on_device() as (cached_weights, text_encoder),
8484
tokenizer_info as tokenizer,
85-
ModelPatcher.apply_model_patches(
85+
LayerPatcher.apply_model_patches(
8686
model=text_encoder,
8787
patches=_lora_loader(),
8888
prefix="lora_te_",
@@ -179,7 +179,7 @@ def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]:
179179
# apply all patches while the model is on the target device
180180
text_encoder_info.model_on_device() as (cached_weights, text_encoder),
181181
tokenizer_info as tokenizer,
182-
ModelPatcher.apply_model_patches(
182+
LayerPatcher.apply_model_patches(
183183
text_encoder,
184184
patches=_lora_loader(),
185185
prefix=lora_prefix,

invokeai/app/invocations/denoise_latents.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@
4040
from invokeai.backend.model_manager import BaseModelType, ModelVariantType
4141
from invokeai.backend.model_patcher import ModelPatcher
4242
from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
43-
from invokeai.backend.patches.model_patcher import ModelPatcher
43+
from invokeai.backend.patches.model_patcher import LayerPatcher
4444
from invokeai.backend.stable_diffusion import PipelineIntermediateState
4545
from invokeai.backend.stable_diffusion.denoise_context import DenoiseContext, DenoiseInputs
4646
from invokeai.backend.stable_diffusion.diffusers_pipeline import (
@@ -1003,7 +1003,7 @@ def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]:
10031003
ModelPatcher.apply_freeu(unet, self.unet.freeu_config),
10041004
SeamlessExt.static_patch_model(unet, self.unet.seamless_axes), # FIXME
10051005
# Apply the LoRA after unet has been moved to its target device for faster patching.
1006-
ModelPatcher.apply_model_patches(
1006+
LayerPatcher.apply_model_patches(
10071007
model=unet,
10081008
patches=_lora_loader(),
10091009
prefix="lora_unet_",

invokeai/app/invocations/flux_denoise.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@
5050
from invokeai.backend.model_manager.config import ModelFormat
5151
from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX
5252
from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
53-
from invokeai.backend.patches.model_patcher import ModelPatcher
53+
from invokeai.backend.patches.model_patcher import LayerPatcher
5454
from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
5555
from invokeai.backend.stable_diffusion.diffusion.conditioning_data import FLUXConditioningInfo
5656
from invokeai.backend.util.devices import TorchDevice
@@ -306,7 +306,7 @@ def _run_diffusion(
306306
if config.format in [ModelFormat.Checkpoint]:
307307
# The model is non-quantized, so we can apply the LoRA weights directly into the model.
308308
exit_stack.enter_context(
309-
ModelPatcher.apply_model_patches(
309+
LayerPatcher.apply_model_patches(
310310
model=transformer,
311311
patches=self._lora_iterator(context),
312312
prefix=FLUX_LORA_TRANSFORMER_PREFIX,
@@ -321,7 +321,7 @@ def _run_diffusion(
321321
# The model is quantized, so apply the LoRA weights as sidecar layers. This results in slower inference,
322322
# than directly patching the weights, but is agnostic to the quantization format.
323323
exit_stack.enter_context(
324-
ModelPatcher.apply_model_sidecar_patches(
324+
LayerPatcher.apply_model_sidecar_patches(
325325
model=transformer,
326326
patches=self._lora_iterator(context),
327327
prefix=FLUX_LORA_TRANSFORMER_PREFIX,

invokeai/app/invocations/flux_text_encoder.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
from invokeai.backend.model_manager.config import ModelFormat
2121
from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_CLIP_PREFIX
2222
from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
23-
from invokeai.backend.patches.model_patcher import ModelPatcher
23+
from invokeai.backend.patches.model_patcher import LayerPatcher
2424
from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData, FLUXConditioningInfo
2525

2626

@@ -111,7 +111,7 @@ def _clip_encode(self, context: InvocationContext) -> torch.Tensor:
111111
if clip_text_encoder_config.format in [ModelFormat.Diffusers]:
112112
# The model is non-quantized, so we can apply the LoRA weights directly into the model.
113113
exit_stack.enter_context(
114-
ModelPatcher.apply_model_patches(
114+
LayerPatcher.apply_model_patches(
115115
model=clip_text_encoder,
116116
patches=self._clip_lora_iterator(context),
117117
prefix=FLUX_LORA_CLIP_PREFIX,

invokeai/app/invocations/sd3_text_encoder.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
from invokeai.backend.model_manager.config import ModelFormat
2020
from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_CLIP_PREFIX
2121
from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
22-
from invokeai.backend.patches.model_patcher import ModelPatcher
22+
from invokeai.backend.patches.model_patcher import LayerPatcher
2323
from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData, SD3ConditioningInfo
2424

2525
# The SD3 T5 Max Sequence Length set based on the default in diffusers.
@@ -150,7 +150,7 @@ def _clip_encode(
150150
if clip_text_encoder_config.format in [ModelFormat.Diffusers]:
151151
# The model is non-quantized, so we can apply the LoRA weights directly into the model.
152152
exit_stack.enter_context(
153-
ModelPatcher.apply_model_patches(
153+
LayerPatcher.apply_model_patches(
154154
model=clip_text_encoder,
155155
patches=self._clip_lora_iterator(context, clip_model),
156156
prefix=FLUX_LORA_CLIP_PREFIX,

invokeai/app/invocations/tiled_multi_diffusion_denoise_latents.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
from invokeai.app.invocations.primitives import LatentsOutput
2424
from invokeai.app.services.shared.invocation_context import InvocationContext
2525
from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
26-
from invokeai.backend.patches.model_patcher import ModelPatcher
26+
from invokeai.backend.patches.model_patcher import LayerPatcher
2727
from invokeai.backend.stable_diffusion.diffusers_pipeline import ControlNetData, PipelineIntermediateState
2828
from invokeai.backend.stable_diffusion.multi_diffusion_pipeline import (
2929
MultiDiffusionPipeline,
@@ -207,7 +207,7 @@ def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]:
207207
with (
208208
ExitStack() as exit_stack,
209209
unet_info as unet,
210-
ModelPatcher.apply_model_patches(model=unet, patches=_lora_loader(), prefix="lora_unet_"),
210+
LayerPatcher.apply_model_patches(model=unet, patches=_lora_loader(), prefix="lora_unet_"),
211211
):
212212
assert isinstance(unet, UNet2DConditionModel)
213213
latents = latents.to(device=unet.device, dtype=unet.dtype)

invokeai/backend/patches/model_patcher.py

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
from invokeai.backend.util.original_weights_storage import OriginalWeightsStorage
1414

1515

16-
class ModelPatcher:
16+
class LayerPatcher:
1717
@staticmethod
1818
@torch.no_grad()
1919
@contextmanager
@@ -37,7 +37,7 @@ def apply_model_patches(
3737
original_weights = OriginalWeightsStorage(cached_weights)
3838
try:
3939
for patch, patch_weight in patches:
40-
ModelPatcher.apply_model_patch(
40+
LayerPatcher.apply_model_patch(
4141
model=model,
4242
prefix=prefix,
4343
patch=patch,
@@ -85,11 +85,11 @@ def apply_model_patch(
8585
if not layer_key.startswith(prefix):
8686
continue
8787

88-
module_key, module = ModelPatcher._get_submodule(
88+
module_key, module = LayerPatcher._get_submodule(
8989
model, layer_key[prefix_len:], layer_key_is_flattened=layer_keys_are_flattened
9090
)
9191

92-
ModelPatcher._apply_model_layer_patch(
92+
LayerPatcher._apply_model_layer_patch(
9393
module_to_patch=module,
9494
module_to_patch_key=module_key,
9595
patch=layer,
@@ -169,7 +169,7 @@ def apply_model_sidecar_patches(
169169
original_modules: dict[str, torch.nn.Module] = {}
170170
try:
171171
for patch, patch_weight in patches:
172-
ModelPatcher._apply_model_sidecar_patch(
172+
LayerPatcher._apply_model_sidecar_patch(
173173
model=model,
174174
prefix=prefix,
175175
patch=patch,
@@ -182,9 +182,9 @@ def apply_model_sidecar_patches(
182182
# Restore original modules.
183183
# Note: This logic assumes no nested modules in original_modules.
184184
for module_key, orig_module in original_modules.items():
185-
module_parent_key, module_name = ModelPatcher._split_parent_key(module_key)
185+
module_parent_key, module_name = LayerPatcher._split_parent_key(module_key)
186186
parent_module = model.get_submodule(module_parent_key)
187-
ModelPatcher._set_submodule(parent_module, module_name, orig_module)
187+
LayerPatcher._set_submodule(parent_module, module_name, orig_module)
188188

189189
@staticmethod
190190
def _apply_model_sidecar_patch(
@@ -212,11 +212,11 @@ def _apply_model_sidecar_patch(
212212
if not layer_key.startswith(prefix):
213213
continue
214214

215-
module_key, module = ModelPatcher._get_submodule(
215+
module_key, module = LayerPatcher._get_submodule(
216216
model, layer_key[prefix_len:], layer_key_is_flattened=layer_keys_are_flattened
217217
)
218218

219-
ModelPatcher._apply_model_layer_wrapper_patch(
219+
LayerPatcher._apply_model_layer_wrapper_patch(
220220
model=model,
221221
module_to_patch=module,
222222
module_to_patch_key=module_key,
@@ -242,9 +242,9 @@ def _apply_model_layer_wrapper_patch(
242242
if not isinstance(module_to_patch, BaseSidecarWrapper):
243243
wrapped_module = wrap_module_with_sidecar_wrapper(orig_module=module_to_patch)
244244
original_modules[module_to_patch_key] = module_to_patch
245-
module_parent_key, module_name = ModelPatcher._split_parent_key(module_to_patch_key)
245+
module_parent_key, module_name = LayerPatcher._split_parent_key(module_to_patch_key)
246246
module_parent = model.get_submodule(module_parent_key)
247-
ModelPatcher._set_submodule(module_parent, module_name, wrapped_module)
247+
LayerPatcher._set_submodule(module_parent, module_name, wrapped_module)
248248
else:
249249
assert module_to_patch_key in original_modules
250250
wrapped_module = module_to_patch

invokeai/backend/stable_diffusion/extensions/lora.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
from diffusers import UNet2DConditionModel
77

88
from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
9-
from invokeai.backend.patches.model_patcher import ModelPatcher
9+
from invokeai.backend.patches.model_patcher import LayerPatcher
1010
from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase
1111

1212
if TYPE_CHECKING:
@@ -31,7 +31,7 @@ def __init__(
3131
def patch_unet(self, unet: UNet2DConditionModel, original_weights: OriginalWeightsStorage):
3232
lora_model = self._node_context.models.load(self._model_id).model
3333
assert isinstance(lora_model, ModelPatchRaw)
34-
ModelPatcher.apply_model_patch(
34+
LayerPatcher.apply_model_patch(
3535
model=unet,
3636
prefix="lora_unet_",
3737
patch=lora_model,

tests/backend/patches/test_lora_patcher.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
from invokeai.backend.patches.layers.lora_layer import LoRALayer
55
from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
6-
from invokeai.backend.patches.model_patcher import ModelPatcher
6+
from invokeai.backend.patches.model_patcher import LayerPatcher
77

88

99
class DummyModule(torch.nn.Module):
@@ -53,7 +53,7 @@ def test_apply_lora_patches(device: str, num_layers: int):
5353
orig_linear_weight = model.linear_layer_1.weight.data.detach().clone()
5454
expected_patched_linear_weight = orig_linear_weight + (lora_rank * lora_weight * num_layers)
5555

56-
with ModelPatcher.apply_model_patches(model=model, patches=lora_models, prefix=""):
56+
with LayerPatcher.apply_model_patches(model=model, patches=lora_models, prefix=""):
5757
# After patching, all LoRA layer weights should have been moved back to the cpu.
5858
for lora, _ in lora_models:
5959
assert lora.layers["linear_layer_1"].up.device.type == "cpu"
@@ -93,7 +93,7 @@ def test_apply_lora_patches_change_device():
9393

9494
orig_linear_weight = model.linear_layer_1.weight.data.detach().clone()
9595

96-
with ModelPatcher.apply_model_patches(model=model, patches=[(lora, 0.5)], prefix=""):
96+
with LayerPatcher.apply_model_patches(model=model, patches=[(lora, 0.5)], prefix=""):
9797
# After patching, all LoRA layer weights should have been moved back to the cpu.
9898
assert lora_layers["linear_layer_1"].up.device.type == "cpu"
9999
assert lora_layers["linear_layer_1"].down.device.type == "cpu"
@@ -146,7 +146,7 @@ def test_apply_lora_sidecar_patches(device: str, num_layers: int):
146146
output_before_patch = model(input)
147147

148148
# Patch the model and run inference during the patch.
149-
with ModelPatcher.apply_model_sidecar_patches(model=model, patches=lora_models, prefix="", dtype=dtype):
149+
with LayerPatcher.apply_model_sidecar_patches(model=model, patches=lora_models, prefix="", dtype=dtype):
150150
output_during_patch = model(input)
151151

152152
# Run inference after unpatching.
@@ -186,10 +186,10 @@ def test_apply_lora_sidecar_patches_matches_apply_lora_patches(num_layers: int):
186186

187187
input = torch.randn(1, linear_in_features, device="cpu", dtype=dtype)
188188

189-
with ModelPatcher.apply_model_patches(model=model, patches=lora_models, prefix=""):
189+
with LayerPatcher.apply_model_patches(model=model, patches=lora_models, prefix=""):
190190
output_lora_patches = model(input)
191191

192-
with ModelPatcher.apply_model_sidecar_patches(model=model, patches=lora_models, prefix="", dtype=dtype):
192+
with LayerPatcher.apply_model_sidecar_patches(model=model, patches=lora_models, prefix="", dtype=dtype):
193193
output_lora_sidecar_patches = model(input)
194194

195195
# Note: We set atol=1e-5 because the test failed occasionally with the default atol=1e-8. Slight numerical

0 commit comments

Comments
 (0)