Finegrain/refiners
1
0
Fork 0
mirror of https://github.com/finegrain-ai/refiners.git synced 2024-11-23 06:38:45 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

allow passing inclusions and exlusions to SDLoraManager

Browse source
This commit is contained in:
Pierre Chapuis 2024-03-05 15:14:21 +01:00
parent cce2a98fa6
commit ed8ec26e63
1 changed files with 29 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
src/refiners/foundationals/latent_diffusion/lora.py
View file

@ -18,6 +18,10 @@ class SDLoraManager:
def __init__(
self,
target: LatentDiffusionModel,
unet_inclusions: list[str] | None = None,
unet_exclusions: list[str] | None = None,
text_encoder_inclusions: list[str] | None = None,
text_encoder_exclusions: list[str] | None = None,
) -> None:
"""Initialize the LoRA manager.
@ -25,6 +29,10 @@ class SDLoraManager:
target: The target model to manage the LoRAs for.
"""
self.target = target
self.unet_inclusions = unet_inclusions
self.unet_exclusions = unet_exclusions
self.text_encoder_inclusions = text_encoder_inclusions
self.text_encoder_exclusions = text_encoder_exclusions
@property
def unet(self) -> fl.Chain:
@ -114,7 +122,12 @@ class SDLoraManager:
(keys are the names of the LoRAs, values are the LoRAs to add to the text encoder)
"""
text_encoder_loras = {key: loras[key] for key in loras.keys() if "text" in key}
auto_attach_loras(text_encoder_loras, self.clip_text_encoder)
auto_attach_loras(
text_encoder_loras,
self.clip_text_encoder,
exclude=self.text_encoder_exclusions,
include=self.text_encoder_inclusions,
)
def add_loras_to_unet(self, loras: dict[str, Lora[Any]], /) -> None:
"""Add multiple LoRAs to the U-Net.
@ -124,10 +137,21 @@ class SDLoraManager:
(keys are the names of the LoRAs, values are the LoRAs to add to the U-Net)
"""
unet_loras = {key: loras[key] for key in loras.keys() if "unet" in key}
exclude = [
block for s, block in self.unet_exclusions.items() if all([s not in key for key in unet_loras.keys()])
if self.unet_exclusions is None:
auto_exclusions = {
"time": "TimestepEncoder",
"res": "ResidualBlock",
"downsample": "Downsample",
"upsample": "Upsample",
}
exclusions = [
block for s, block in auto_exclusions.items() if all([s not in key for key in unet_loras.keys()])
]
auto_attach_loras(unet_loras, self.unet, exclude=exclude)
else:
exclusions = self.unet_exclusions
auto_attach_loras(unet_loras, self.unet, exclude=exclusions, include=self.unet_inclusions)
def remove_loras(self, *names: str) -> None:
"""Remove multiple LoRAs from the target.
@ -206,15 +230,6 @@ class SDLoraManager:
"""List of all the LoraAdapters managed by the SDLoraManager."""
return list(self.unet.layers(LoraAdapter)) + list(self.clip_text_encoder.layers(LoraAdapter))
@property
def unet_exclusions(self) -> dict[str, str]:
return {
"time": "TimestepEncoder",
"res": "ResidualBlock",
"downsample": "Downsample",
"upsample": "Upsample",
}
@property
def scales(self) -> dict[str, float]:
"""The scales of all the LoRAs managed by the SDLoraManager."""