From 4e2145c6c4290afbc468cbfd21486c7df99653aa Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Tue, 23 Sep 2025 18:26:38 +1000
Subject: [PATCH] tidy(mm): patcher types and import paths

---
 invokeai/backend/model_patcher.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/invokeai/backend/model_patcher.py b/invokeai/backend/model_patcher.py
index a1d8bbed0a..04f9949560 100644
--- a/invokeai/backend/model_patcher.py
+++ b/invokeai/backend/model_patcher.py
@@ -5,10 +5,10 @@ from __future__ import annotations
 
 import pickle
 from contextlib import contextmanager
-from typing import Any, Iterator, List, Optional, Tuple, Type, Union
+from typing import Any, Generator, Iterator, List, Optional, Tuple, Type, Union
 
 import torch
-from diffusers import UNet2DConditionModel
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
 from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
 
 from invokeai.app.shared.models import FreeUConfig
@@ -146,7 +146,7 @@ class ModelPatcher:
         cls,
         text_encoder: Union[CLIPTextModel, CLIPTextModelWithProjection],
         clip_skip: int,
-    ) -> None:
+    ) -> Generator[None, Any, Any]:
         skipped_layers = []
         try:
             for _i in range(clip_skip):
@@ -164,7 +164,7 @@ class ModelPatcher:
         cls,
         unet: UNet2DConditionModel,
         freeu_config: Optional[FreeUConfig] = None,
-    ) -> None:
+    ) -> Generator[None, Any, Any]:
         did_apply_freeu = False
         try:
             assert hasattr(unet, "enable_freeu")  # mypy doesn't pick up this attribute?