update

Author: a-r-r-o-w

src/diffusers/guiders/adaptive_projected_guidance.py

@@ -39,7 +39,8 @@ class AdaptiveProjectedGuidance(GuidanceMixin):
             Flawed](https://huggingface.co/papers/2305.08891).
         use_original_formulation (`bool`, defaults to `False`):
             Whether to use the original formulation of classifier-free guidance as proposed in the paper. By default,
-            we use the diffusers-native implementation that has been in the codebase for a long time.
+            we use the diffusers-native implementation that has been in the codebase for a long time. See
+            [~guiders.classifier_free_guidance.ClassifierFreeGuidance] for more details.
     """
 
     _input_predictions = ["pred_cond", "pred_uncond"]

src/diffusers/guiders/classifier_free_guidance.py

@@ -31,8 +31,9 @@ class ClassifierFreeGuidance(GuidanceMixin):
     The original paper proposes scaling and shifting the conditional distribution based on the difference between
     conditional and unconditional predictions. [x_pred = x_cond + scale * (x_cond - x_uncond)]
 
-    Diffusers implemented the scaling and shifting on the unconditional prediction instead, which is equivalent to what
-    the original paper proposed in theory. [x_pred = x_uncond + scale * (x_cond - x_uncond)]
+    Diffusers implemented the scaling and shifting on the unconditional prediction instead based on the [Imagen
+    paper](https://huggingface.co/papers/2205.11487), which is equivalent to what the original paper proposed in
+    theory. [x_pred = x_uncond + scale * (x_cond - x_uncond)]
 
     The intution behind the original formulation can be thought of as moving the conditional distribution estimates
     further away from the unconditional distribution estimates, while the diffusers-native implementation can be
@@ -53,7 +54,8 @@ class ClassifierFreeGuidance(GuidanceMixin):
             Flawed](https://huggingface.co/papers/2305.08891).
         use_original_formulation (`bool`, defaults to `False`):
             Whether to use the original formulation of classifier-free guidance as proposed in the paper. By default,
-            we use the diffusers-native implementation that has been in the codebase for a long time.
+            we use the diffusers-native implementation that has been in the codebase for a long time. See
+            [~guiders.classifier_free_guidance.ClassifierFreeGuidance] for more details.
     """
 
     _input_predictions = ["pred_cond", "pred_uncond"]

src/diffusers/guiders/perturbed_attention_guidance.py

@@ -109,26 +109,26 @@ def prepare_outputs(self, pred: torch.Tensor) -> None:
                 key = "pred_perturbed"
         self._preds[key] = pred
 
-        # Prepare denoiser for perturbed attention prediction if needed
-        if not self._is_pag_enabled():
-            return
-        should_register_pag = (self._is_cfg_enabled() and self._num_outputs_prepared == 2) or (
-            not self._is_cfg_enabled() and self._num_outputs_prepared == 1
-        )
-        if should_register_pag:
-            self._is_pag_batch = True
-            self._original_processors = _replace_attention_processors(
-                self._denoiser,
-                self.pag_applied_layers,
-                skip_context_attention=self.skip_context_attention,
-                metadata_name="perturbed_attention_guidance_processor_cls",
-            )
-        elif self._is_pag_batch:
-            # Restore the original attention processors
+        # Restore the original attention processors if previously replaced
+        if self._is_pag_batch:
             _replace_attention_processors(self._denoiser, processors=self._original_processors)
             self._is_pag_batch = False
             self._original_processors = None
 
+        # Prepare denoiser for perturbed attention prediction if needed
+        if self._is_pag_enabled():
+            should_register_pag = (self._is_cfg_enabled() and self._num_outputs_prepared == 2) or (
+                not self._is_cfg_enabled() and self._num_outputs_prepared == 1
+            )
+            if should_register_pag:
+                self._is_pag_batch = True
+                self._original_processors = _replace_attention_processors(
+                    self._denoiser,
+                    self.pag_applied_layers,
+                    skip_context_attention=self.skip_context_attention,
+                    metadata_name="perturbed_attention_guidance_processor_cls",
+                )
+
     def cleanup_models(self, denoiser: torch.nn.Module):
         self._denoiser = None