Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(nodes): add freeu support (#4846)

Browse Source 
### What type of PR is this? (check all applicable)

- [ ] Refactor
- [x] Feature
- [ ] Bug Fix
- [ ] Optimization
- [ ] Documentation Update
- [ ] Community Node Submission


## Have you discussed this change with the InvokeAI team?
- [ ] Yes
- [x] No, because:

      
## Have you updated all relevant documentation?
- [ ] Yes
- [x] No


## Description

**Note: FreeU is not in the current release of diffusers. Looks like it
will be in release 0.22. This PR needs to wait until that is released.**

[feat(nodes): add freeu
support](15b33ad501)

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 https://github.com/invoke-ai/InvokeAI/issues/4845

## Related Tickets & Documents

<!--
For pull requests that relate or close an issue, please include them
below. 

For example having the text: "closes #1234" would connect the current
pull
request to issue 1234.  And when we merge the pull request, Github will
automatically close the issue.
-->

- Closes #4845 

## QA Instructions, Screenshots, Recordings

You'll need to install diffusers from their github repo before testing
this:
`pip install git+https://github.com/huggingface/diffusers`

1. Create a graph like this:

![image](https://github.com/invoke-ai/InvokeAI/assets/4822129/af17719b-b001-4534-8c4e-883484fd7465)
2. Get a free lunch!

No FreeU:

![image](https://github.com/invoke-ai/InvokeAI/assets/4822129/279d1a69-1577-4c31-ab82-ebf67f65920d)
With FreeU:

![image](https://github.com/invoke-ai/InvokeAI/assets/4822129/c332c778-0b87-4215-8a36-d4822e06f4de)

No FreeU:

![image](https://github.com/invoke-ai/InvokeAI/assets/4822129/ebec097b-ad54-4295-b734-33656738a2cf)
With FreeU:

![image](https://github.com/invoke-ai/InvokeAI/assets/4822129/3423140d-c9ce-4697-9993-d2bb0d0f5634)

No FreeU:

![image](https://github.com/invoke-ai/InvokeAI/assets/4822129/7cb0e39d-aa87-4a48-a3af-b9f47a866814)
With FreeU:

![image](https://github.com/invoke-ai/InvokeAI/assets/4822129/9113d2fe-5bd3-474f-8f33-82cdeb7abf82)
This commit is contained in:
Kent Keirsey 2023-11-06 13:58:32 -08:00 committed by GitHub
commit 935e4632c2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 92 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
invokeai/app/invocations/baseinvocation.py
View File

@ -92,6 +92,10 @@ class FieldDescriptions:
inclusive_low = "The inclusive low value"
exclusive_high = "The exclusive high value"
decimal_places = "The number of decimal places to round to"
freeu_s1 = 'Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to mitigate the "oversmoothing effect" in the enhanced denoising process.'
freeu_s2 = 'Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to mitigate the "oversmoothing effect" in the enhanced denoising process.'
freeu_b1 = "Scaling factor for stage 1 to amplify the contributions of backbone features."
freeu_b2 = "Scaling factor for stage 2 to amplify the contributions of backbone features."
class Input(str, Enum):

2
invokeai/app/invocations/latent.py
View File

@ -710,6 +710,8 @@ class DenoiseLatentsInvocation(BaseInvocation):
)
with (
ExitStack() as exit_stack,
ModelPatcher.apply_lora_unet(unet_info.context.model, _lora_loader()),
ModelPatcher.apply_freeu(unet_info.context.model, self.unet.freeu_config),
set_seamless(unet_info.context.model, self.unet.seamless_axes),
unet_info as unet,
# Apply the LoRA after unet has been moved to its target device for faster patching.

62
invokeai/app/invocations/model.py
View File

@ -3,6 +3,8 @@ from typing import List, Optional
from pydantic import BaseModel, ConfigDict, Field
from invokeai.app.invocations.shared import FreeUConfig
from ...backend.model_management import BaseModelType, ModelType, SubModelType
from .baseinvocation import (
BaseInvocation,
@ -36,6 +38,7 @@ class UNetField(BaseModel):
scheduler: ModelInfo = Field(description="Info to load scheduler submodel")
loras: List[LoraInfo] = Field(description="Loras to apply on model loading")
seamless_axes: List[str] = Field(default_factory=list, description='Axes("x" and "y") to which apply seamless')
freeu_config: Optional[FreeUConfig] = Field(default=None, description="FreeU configuration")
class ClipField(BaseModel):
@ -51,15 +54,34 @@ class VaeField(BaseModel):
seamless_axes: List[str] = Field(default_factory=list, description='Axes("x" and "y") to which apply seamless')
@invocation_output("model_loader_output")
class ModelLoaderOutput(BaseInvocationOutput):
"""Model loader output"""
@invocation_output("unet_output")
class UNetOutput(BaseInvocationOutput):
"""Base class for invocations that output a UNet field"""
unet: UNetField = OutputField(description=FieldDescriptions.unet, title="UNet")
clip: ClipField = OutputField(description=FieldDescriptions.clip, title="CLIP")
@invocation_output("vae_output")
class VAEOutput(BaseInvocationOutput):
"""Base class for invocations that output a VAE field"""
vae: VaeField = OutputField(description=FieldDescriptions.vae, title="VAE")
@invocation_output("clip_output")
class CLIPOutput(BaseInvocationOutput):
"""Base class for invocations that output a CLIP field"""
clip: ClipField = OutputField(description=FieldDescriptions.clip, title="CLIP")
@invocation_output("model_loader_output")
class ModelLoaderOutput(UNetOutput, CLIPOutput, VAEOutput):
"""Model loader output"""
pass
class MainModelField(BaseModel):
"""Main model field"""
@ -366,13 +388,6 @@ class VAEModelField(BaseModel):
model_config = ConfigDict(protected_namespaces=())
@invocation_output("vae_loader_output")
class VaeLoaderOutput(BaseInvocationOutput):
"""VAE output"""
vae: VaeField = OutputField(description=FieldDescriptions.vae, title="VAE")
@invocation("vae_loader", title="VAE", tags=["vae", "model"], category="model", version="1.0.0")
class VaeLoaderInvocation(BaseInvocation):
"""Loads a VAE model, outputting a VaeLoaderOutput"""
@ -384,7 +399,7 @@ class VaeLoaderInvocation(BaseInvocation):
title="VAE",
)
def invoke(self, context: InvocationContext) -> VaeLoaderOutput:
def invoke(self, context: InvocationContext) -> VAEOutput:
base_model = self.vae_model.base_model
model_name = self.vae_model.model_name
model_type = ModelType.Vae
@ -395,7 +410,7 @@ class VaeLoaderInvocation(BaseInvocation):
model_type=model_type,
):
raise Exception(f"Unkown vae name: {model_name}!")
return VaeLoaderOutput(
return VAEOutput(
vae=VaeField(
vae=ModelInfo(
model_name=model_name,
@ -457,3 +472,24 @@ class SeamlessModeInvocation(BaseInvocation):
vae.seamless_axes = seamless_axes_list
return SeamlessModeOutput(unet=unet, vae=vae)
@invocation("freeu", title="FreeU", tags=["freeu"], category="unet", version="1.0.0")
class FreeUInvocation(BaseInvocation):
"""
Applies FreeU to the UNet. Suggested values (b1/b2/s1/s2):
SD1.5: 1.2/1.4/0.9/0.2,
SD2: 1.1/1.2/0.9/0.2,
SDXL: 1.1/1.2/0.6/0.4,
"""
unet: UNetField = InputField(description=FieldDescriptions.unet, input=Input.Connection, title="UNet")
b1: float = InputField(default=1.2, ge=-1, le=3, description=FieldDescriptions.freeu_b1)
b2: float = InputField(default=1.4, ge=-1, le=3, description=FieldDescriptions.freeu_b2)
s1: float = InputField(default=0.9, ge=-1, le=3, description=FieldDescriptions.freeu_s1)
s2: float = InputField(default=0.2, ge=-1, le=3, description=FieldDescriptions.freeu_s2)
def invoke(self, context: InvocationContext) -> UNetOutput:
self.unet.freeu_config = FreeUConfig(s1=self.s1, s2=self.s2, b1=self.b1, b2=self.b2)
return UNetOutput(unet=self.unet)

16
invokeai/app/invocations/shared.py Normal file
View File

@ -0,0 +1,16 @@
from pydantic import BaseModel, Field
from invokeai.app.invocations.baseinvocation import FieldDescriptions
class FreeUConfig(BaseModel):
"""
Configuration for the FreeU hyperparameters.
- https://huggingface.co/docs/diffusers/main/en/using-diffusers/freeu
- https://github.com/ChenyangSi/FreeU
"""
s1: float = Field(ge=-1, le=3, description=FieldDescriptions.freeu_s1)
s2: float = Field(ge=-1, le=3, description=FieldDescriptions.freeu_s2)
b1: float = Field(ge=-1, le=3, description=FieldDescriptions.freeu_b1)
b2: float = Field(ge=-1, le=3, description=FieldDescriptions.freeu_b2)

21
invokeai/backend/model_management/lora.py
View File

@ -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
"""
@ -240,6 +242,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]