diff --git a/ai_diffusion/api.py b/ai_diffusion/api.py index 74351107a4..b9de298e39 100644 --- a/ai_diffusion/api.py +++ b/ai_diffusion/api.py @@ -58,6 +58,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 3abbcb3ab7..bb4f3743ff 100644 --- a/ai_diffusion/comfy_workflow.py +++ b/ai_diffusion/comfy_workflow.py @@ -483,6 +483,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 523fddbcee..622595d5fd 100644 --- a/ai_diffusion/style.py +++ b/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 diff --git a/ai_diffusion/ui/style.py b/ai_diffusion/ui/style.py index 5b3754c3ff..f792e975f3 100644 --- a/ai_diffusion/ui/style.py +++ b/ai_diffusion/ui/style.py @@ -467,6 +467,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)) diff --git a/ai_diffusion/workflow.py b/ai_diffusion/workflow.py index 614532fc76..c2140345e5 100644 --- a/ai_diffusion/workflow.py +++ b/ai_diffusion/workflow.py @@ -122,6 +122,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