Adaptive Projected Guidance

Author: hlky

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

@@ -78,6 +78,47 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
     return noise_cfg
 
 
+class MomentumBuffer:
+    def __init__(self, momentum: float):
+        self.momentum = momentum
+        self.running_average = 0
+
+    def update(self, update_value: torch.Tensor):
+        new_average = self.momentum * self.running_average
+        self.running_average = update_value + new_average
+
+
+def normalized_guidance(
+    pred_cond: torch.Tensor,
+    pred_uncond: torch.Tensor,
+    guidance_scale: float,
+    momentum_buffer: MomentumBuffer = None,
+    eta: float = 1.0,
+    norm_threshold: float = 0.0,
+):
+    """
+    Based on the findings of [Eliminating Oversaturation and Artifacts of High Guidance Scales
+    in Diffusion Models](https://arxiv.org/pdf/2410.02416)
+    """
+    diff = pred_cond - pred_uncond
+    if momentum_buffer is not None:
+        momentum_buffer.update(diff)
+        diff = momentum_buffer.running_average
+    if norm_threshold > 0:
+        ones = torch.ones_like(diff)
+        diff_norm = diff.norm(p=2, dim=[-1, -2, -3], keepdim=True)
+        scale_factor = torch.minimum(ones, norm_threshold / diff_norm)
+        diff = diff * scale_factor
+    v0, v1 = diff.double(), pred_cond.double()
+    v1 = torch.nn.functional.normalize(v1, dim=[-1, -2, -3])
+    v0_parallel = (v0 * v1).sum(dim=[-1, -2, -3], keepdim=True) * v1
+    v0_orthogonal = v0 - v0_parallel
+    diff_parallel, diff_orthogonal = v0_parallel.to(diff.dtype), v0_orthogonal.to(diff.dtype)
+    normalized_update = diff_orthogonal + eta * diff_parallel
+    pred_guided = pred_cond + (guidance_scale - 1) * normalized_update
+    return pred_guided
+
+
 def retrieve_timesteps(
     scheduler,
     num_inference_steps: Optional[int] = None,
@@ -730,6 +771,14 @@ def guidance_scale(self):
     def guidance_rescale(self):
         return self._guidance_rescale
 
+    @property
+    def adaptive_projected_guidance(self):
+        return self._adaptive_projected_guidance
+
+    @property
+    def adaptive_projected_guidance_momentum(self):
+        return self._adaptive_projected_guidance_momentum
+
     @property
     def clip_skip(self):
         return self._clip_skip
@@ -777,6 +826,8 @@ def __call__(
         return_dict: bool = True,
         cross_attention_kwargs: Optional[Dict[str, Any]] = None,
         guidance_rescale: float = 0.0,
+        adaptive_projected_guidance: Optional[bool] = False,
+        adaptive_projected_guidance_momentum: Optional[float] = -0.75,
         clip_skip: Optional[int] = None,
         callback_on_step_end: Optional[
             Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
@@ -847,6 +898,11 @@ def __call__(
                 Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are
                 Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when
                 using zero terminal SNR.
+            adaptive_projected_guidance (`bool`, *optional*, defaults to `False`):
+                Use adaptive projected guidance from [Eliminating Oversaturation and Artifacts of High Guidance Scales
+                in Diffusion Models](https://arxiv.org/pdf/2410.02416)
+            adaptive_projected_guidance_momentum (`float`, *optional*, defaults to `-0.75`):
+                Momentum value to use with adaptive projected guidance. Use `None` to disable momentum.
             clip_skip (`int`, *optional*):
                 Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
                 the output of the pre-final layer will be used for computing the prompt embeddings.
@@ -910,6 +966,8 @@ def __call__(
 
         self._guidance_scale = guidance_scale
         self._guidance_rescale = guidance_rescale
+        self._adaptive_projected_guidance = adaptive_projected_guidance
+        self._adaptive_projected_guidance_momentum = adaptive_projected_guidance_momentum
         self._clip_skip = clip_skip
         self._cross_attention_kwargs = cross_attention_kwargs
         self._interrupt = False
@@ -992,6 +1050,11 @@ def __call__(
                 guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
             ).to(device=device, dtype=latents.dtype)
 
+        if adaptive_projected_guidance and adaptive_projected_guidance_momentum is not None:
+            momentum_buffer = MomentumBuffer(adaptive_projected_guidance_momentum)
+        else:
+            momentum_buffer = None
+
         # 7. Denoising loop
         num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
         self._num_timesteps = len(timesteps)
@@ -1018,7 +1081,12 @@ def __call__(
                 # perform guidance
                 if self.do_classifier_free_guidance:
                     noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+                    if adaptive_projected_guidance:
+                        noise_pred = normalized_guidance(
+                            noise_pred_text, noise_pred_uncond, self.guidance_scale, momentum_buffer, eta
+                        )
+                    else:
+                        noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
 
                 if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
                     # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf