mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
feat(nodes): add freeu support
Add support for FreeU. See: - https://huggingface.co/docs/diffusers/main/en/using-diffusers/freeu - https://github.com/ChenyangSi/FreeU Implementation: - `ModelPatcher.apply_freeu()` handles the enabling freeu (which is very simple with diffusers). - `FreeUConfig` model added to hold the hyperparameters. - `freeu_config` added as optional sub-field on `UNetField`. - `FreeUInvocation` added, works like LoRA - chain it to add the FreeU config to the UNet - No support for model-dependent presets, this will be a future workflow editor enhancement Closes #4845
This commit is contained in:
@ -12,6 +12,8 @@ from diffusers.models import UNet2DConditionModel
|
||||
from safetensors.torch import load_file
|
||||
from transformers import CLIPTextModel, CLIPTokenizer
|
||||
|
||||
from invokeai.app.invocations.shared import FreeUConfig
|
||||
|
||||
from .models.lora import LoRAModel
|
||||
|
||||
"""
|
||||
@ -231,6 +233,25 @@ class ModelPatcher:
|
||||
while len(skipped_layers) > 0:
|
||||
text_encoder.text_model.encoder.layers.append(skipped_layers.pop())
|
||||
|
||||
@classmethod
|
||||
@contextmanager
|
||||
def apply_freeu(
|
||||
cls,
|
||||
unet: UNet2DConditionModel,
|
||||
freeu_config: Optional[FreeUConfig] = None,
|
||||
):
|
||||
did_apply_freeu = False
|
||||
try:
|
||||
if freeu_config is not None:
|
||||
unet.enable_freeu(b1=freeu_config.b1, b2=freeu_config.b2, s1=freeu_config.s1, s2=freeu_config.s2)
|
||||
did_apply_freeu = True
|
||||
|
||||
yield
|
||||
|
||||
finally:
|
||||
if did_apply_freeu:
|
||||
unet.disable_freeu()
|
||||
|
||||
|
||||
class TextualInversionModel:
|
||||
embedding: torch.Tensor # [n, 768]|[n, 1280]
|
||||
|
Reference in New Issue
Block a user