Skip to content

Commit c88fee6

Browse files
tidy(mm): patcher types and import paths
1 parent 20a0231 commit c88fee6

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

invokeai/backend/model_patcher.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,10 @@
55

66
import pickle
77
from contextlib import contextmanager
8-
from typing import Any, Iterator, List, Optional, Tuple, Type, Union
8+
from typing import Any, Generator, Iterator, List, Optional, Tuple, Type, Union
99

1010
import torch
11-
from diffusers import UNet2DConditionModel
11+
from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
1212
from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
1313

1414
from invokeai.app.shared.models import FreeUConfig
@@ -146,7 +146,7 @@ def apply_clip_skip(
146146
cls,
147147
text_encoder: Union[CLIPTextModel, CLIPTextModelWithProjection],
148148
clip_skip: int,
149-
) -> None:
149+
) -> Generator[None, Any, Any]:
150150
skipped_layers = []
151151
try:
152152
for _i in range(clip_skip):
@@ -164,7 +164,7 @@ def apply_freeu(
164164
cls,
165165
unet: UNet2DConditionModel,
166166
freeu_config: Optional[FreeUConfig] = None,
167-
) -> None:
167+
) -> Generator[None, Any, Any]:
168168
did_apply_freeu = False
169169
try:
170170
assert hasattr(unet, "enable_freeu") # mypy doesn't pick up this attribute?

0 commit comments

Comments
 (0)