Suggested changes

Co-Authored-By: Ryan Dick <[email protected]>

Author: StAlKeR7779

invokeai/backend/stable_diffusion/diffusion_backend.py

@@ -100,8 +100,8 @@ def apply_cfg(ctx: DenoiseContext) -> torch.Tensor:
         if isinstance(guidance_scale, list):
             guidance_scale = guidance_scale[ctx.step_index]
 
-        # Note: Although logically it same, it seams that precision errors differs.
-        # This sometimes results in slightly different output.
+        # Note: Although this `torch.lerp(...)` line is logically equivalent to the current CFG line, it seems to result
+        # in slightly different outputs. It is suspected that this is caused by small precision differences.
         # return torch.lerp(ctx.negative_noise_pred, ctx.positive_noise_pred, guidance_scale)
         return ctx.negative_noise_pred + guidance_scale * (ctx.positive_noise_pred - ctx.negative_noise_pred)
 

invokeai/backend/stable_diffusion/extensions/base.py

@@ -2,7 +2,7 @@
 
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Callable, Dict, List
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional
 
 import torch
 from diffusers import UNet2DConditionModel
@@ -56,5 +56,5 @@ def patch_extension(self, context: DenoiseContext):
         yield None
 
     @contextmanager
-    def patch_unet(self, state_dict: Dict[str, torch.Tensor], unet: UNet2DConditionModel):
+    def patch_unet(self, unet: UNet2DConditionModel, cached_weights: Optional[Dict[str, torch.Tensor]] = None):
         yield None

invokeai/backend/stable_diffusion/extensions/freeu.py

@@ -15,28 +15,21 @@
 class FreeUExt(ExtensionBase):
     def __init__(
         self,
-        freeu_config: Optional[FreeUConfig],
+        freeu_config: FreeUConfig,
     ):
         super().__init__()
-        self.freeu_config = freeu_config
+        self._freeu_config = freeu_config
 
     @contextmanager
     def patch_unet(self, unet: UNet2DConditionModel, cached_weights: Optional[Dict[str, torch.Tensor]] = None):
-        did_apply_freeu = False
-        try:
-            assert hasattr(unet, "enable_freeu")  # mypy doesn't pick up this attribute?
-            if self.freeu_config is not None:
-                unet.enable_freeu(
-                    b1=self.freeu_config.b1,
-                    b2=self.freeu_config.b2,
-                    s1=self.freeu_config.s1,
-                    s2=self.freeu_config.s2,
-                )
-                did_apply_freeu = True
+        unet.enable_freeu(
+            b1=self._freeu_config.b1,
+            b2=self._freeu_config.b2,
+            s1=self._freeu_config.s1,
+            s2=self._freeu_config.s2,
+        )
 
+        try:
             yield
-
         finally:
-            assert hasattr(unet, "disable_freeu")  # mypy doesn't pick up this attribute?
-            if did_apply_freeu:
-                unet.disable_freeu()
+            unet.disable_freeu()