Add Perturbed-Attention Guidance toggle. See https://ku-cvlab.github.…

…io/Perturbed-Attention-Guidance/

ai_diffusion/api.py

@@ -63,6 +63,7 @@ class CheckpointInput:
     rescale_cfg: float = 0.7
     self_attention_guidance: bool = False
     dynamic_caching: bool = False
+    perturbed_attention_guidance: bool = False
 
 
 @dataclass

ai_diffusion/comfy_workflow.py

@@ -711,6 +711,9 @@ def apply_ip_adapter_face(
     def apply_self_attention_guidance(self, model: Output):
         return self.add("SelfAttentionGuidance", 1, model=model, scale=0.5, blur_sigma=2.0)
 
+    def apply_perturbed_attention_guidance(self, model: Output):
+        return self.add("PerturbedAttentionGuidance", 1, model=model)
+
     def inpaint_preprocessor(self, image: Output, mask: Output):
         return self.add("InpaintPreprocessor", 1, image=image, mask=mask)
 

ai_diffusion/style.py

@@ -84,6 +84,12 @@ class StyleSettings:
         _("Pay more attention to difficult parts of the image. Can improve fine details."),
     )
 
+    perturbed_attention_guidance = Setting(
+        "Enable PAG / Perturbed-Attention Guidance",
+        False,
+        'Deliberately introduce errors in "difficult" parts to steer away from. Can improve coherence.',
+    )
+
     preferred_resolution = Setting(
         _("Preferred Resolution"), 0, _("Image resolution the checkpoint was trained on")
     )
@@ -121,6 +127,7 @@ class Style:
     v_prediction_zsnr: bool = StyleSettings.v_prediction_zsnr.default
     rescale_cfg: float = StyleSettings.rescale_cfg.default
     self_attention_guidance: bool = StyleSettings.self_attention_guidance.default
+    perturbed_attention_guidance: bool = StyleSettings.perturbed_attention_guidance.default
     preferred_resolution: int = StyleSettings.preferred_resolution.default
     sampler: str = StyleSettings.sampler.default
     sampler_steps: int = StyleSettings.sampler_steps.default
@@ -199,6 +206,7 @@ def get_models(self, available_checkpoints: Iterable[str]):
             rescale_cfg=self.rescale_cfg,
             loras=[LoraInput.from_dict(l) for l in self.loras if l.get("enabled", True)],
             self_attention_guidance=self.self_attention_guidance,
+            perturbed_attention_guidance=self.perturbed_attention_guidance,
         )
         return result
 

ai_diffusion/ui/style.py

@@ -662,13 +662,19 @@ def add(name: str, widget: SettingWidget):
             SwitchSetting(StyleSettings.self_attention_guidance, parent=self),
         )
 
+        self._pag = add(
+            "perturbed_attention_guidance",
+            SwitchSetting(StyleSettings.perturbed_attention_guidance, parent=self),
+        )
+
         self._checkpoint_advanced_widgets = [
             self._arch_select,
             self._vae,
             self._clip_skip,
             self._resolution_spin,
             self._zsnr,
             self._sag,
+            self._pag,
         ]
         for widget in self._checkpoint_advanced_widgets:
             widget.indent = 1
@@ -867,6 +873,7 @@ def _enable_checkpoint_advanced(self):
         self._clip_skip.enabled = arch.supports_clip_skip and self.current_style.clip_skip > 0
         self._zsnr.enabled = arch.supports_attention_guidance
         self._sag.enabled = arch.supports_attention_guidance
+        self._pag.enabled = arch.supports_attention_guidance
 
     def _read_style(self, style: Style):
         with self._write_guard:

ai_diffusion/workflow.py

@@ -153,6 +153,9 @@ def load_checkpoint_with_lora(w: ComfyWorkflow, checkpoint: CheckpointInput, mod
     if arch.supports_attention_guidance and checkpoint.self_attention_guidance:
         model = w.apply_self_attention_guidance(model)
 
+    if checkpoint.perturbed_attention_guidance:
+        model = w.apply_perturbed_attention_guidance(model)
+
     return model, clip, vae