diff --git a/ai_diffusion/api.py b/ai_diffusion/api.py index c099d5f3a7..62b55db78c 100644 --- a/ai_diffusion/api.py +++ b/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 diff --git a/ai_diffusion/comfy_workflow.py b/ai_diffusion/comfy_workflow.py index b985a1457b..95a67c8d8c 100644 --- a/ai_diffusion/comfy_workflow.py +++ b/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) diff --git a/ai_diffusion/style.py b/ai_diffusion/style.py index 4cff5ebc01..9f4092c57a 100644 --- a/ai_diffusion/style.py +++ b/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 diff --git a/ai_diffusion/ui/style.py b/ai_diffusion/ui/style.py index e7e96abb0e..8a684f34c3 100644 --- a/ai_diffusion/ui/style.py +++ b/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) diff --git a/ai_diffusion/workflow.py b/ai_diffusion/workflow.py index 2716b98ab5..711b2ea8aa 100644 --- a/ai_diffusion/workflow.py +++ b/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