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

…io/Perturbed-Attention-Guidance/

ai_diffusion/api.py

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

@@ -346,6 +346,9 @@ def apply_ip_adapter_face(
     def apply_self_attention_guidance(self, model: Output):
         return self.add("SelfAttentionGuidance", 1, model=model)
 
+    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

@@ -75,6 +75,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"
     )
@@ -111,6 +117,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
@@ -180,6 +187,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

@@ -484,6 +484,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),
+            )
+        )
         self._toggle_checkpoint_advanced(False)
 
         add("loras", LoraList(StyleSettings.loras, self))

ai_diffusion/workflow.py

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