Skip to content

Commit

Permalink
fix exclusions for Downsample and Upsample
Browse files Browse the repository at this point in the history
(also simplify list comprehension for exclusion list)
  • Loading branch information
catwell committed Jan 29, 2024
1 parent 83c95fc commit deb5e93
Showing 1 changed file with 3 additions and 6 deletions.
9 changes: 3 additions & 6 deletions src/refiners/foundationals/latent_diffusion/lora.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,11 +67,8 @@ def add_loras_to_text_encoder(self, loras: dict[str, Lora], /) -> None:

def add_loras_to_unet(self, loras: dict[str, Lora], /) -> None:
unet_loras = {key: loras[key] for key in loras.keys() if "unet" in key}
exclude: list[str] = []
exclude = [
self.unet_exclusions[exclusion]
for exclusion in self.unet_exclusions
if all([exclusion not in key for key in unet_loras.keys()])
block for s, block in self.unet_exclusions.items() if all([s not in key for key in unet_loras.keys()])
]
SDLoraManager.auto_attach(unet_loras, self.unet, exclude=exclude)

Expand Down Expand Up @@ -121,8 +118,8 @@ def unet_exclusions(self) -> dict[str, str]:
return {
"time": "TimestepEncoder",
"res": "ResidualBlock",
"downsample": "DownsampleBlock",
"upsample": "UpsampleBlock",
"downsample": "Downsample",
"upsample": "Upsample",
}

@property
Expand Down

0 comments on commit deb5e93

Please sign in to comment.