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

…io/Perturbed-Attention-Guidance/

ai_diffusion/api.py

@@ -61,6 +61,7 @@ class CheckpointInput:
     clip_skip: int = 0
     v_prediction_zsnr: bool = False
     self_attention_guidance: bool = False
+    perturbed_attention_guidance: bool = False
 
 
 @dataclass

ai_diffusion/comfy_workflow.py

@@ -510,6 +510,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

@@ -81,6 +81,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")
     )
@@ -117,6 +123,7 @@ class Style:
     clip_skip: int = StyleSettings.clip_skip.default
     v_prediction_zsnr: bool = StyleSettings.v_prediction_zsnr.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
@@ -188,6 +195,7 @@ def get_models(self):
             v_prediction_zsnr=self.v_prediction_zsnr,
             loras=[LoraInput.from_dict(l) for l in self.loras],
             self_attention_guidance=self.self_attention_guidance,
+            perturbed_attention_guidance=self.perturbed_attention_guidance,
         )
         return result
 

ai_diffusion/ui/style.py

@@ -617,6 +617,12 @@ def add(name: str, widget: SettingWidget):
         )
         self._checkpoint_advanced_widgets.append(self._sag)
 
+        self._pag = add(
+            "perturbed_attention_guidance",
+            SwitchSetting(StyleSettings.perturbed_attention_guidance, parent=self),
+        )
+        self._checkpoint_advanced_widgets.append(self._pag)
+
         for widget in self._checkpoint_advanced_widgets:
             widget.indent = 1
         self._toggle_checkpoint_advanced(False)
@@ -762,6 +768,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

@@ -126,6 +126,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