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

…io/Perturbed-Attention-Guidance/

ai_diffusion/api.py

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

@@ -80,6 +80,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")
     )
@@ -116,6 +122,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
@@ -185,6 +192,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

@@ -472,6 +472,12 @@ def add(name: str, widget: SettingWidget):
                 SwitchSetting(StyleSettings.self_attention_guidance, parent=self),
             )
         )
+        self._checkpoint_advanced_widgets.append(
+            add(
+                "perturbed_attention_guidance",
+                SwitchSetting(StyleSettings.perturbed_attention_guidance, parent=self),
+            )
+        )
         for widget in self._checkpoint_advanced_widgets:
             widget.indent = 1
         self._toggle_checkpoint_advanced(False)

ai_diffusion/workflow.py

@@ -124,6 +124,9 @@ def load_checkpoint_with_lora(w: ComfyWorkflow, checkpoint: CheckpointInput, mod
     if 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