Mirror apg_guidance.py from ACE-Step/acestep-v15-base

checkpoints/acestep-v15-base/apg_guidance.py

@@ -0,0 +1,220 @@
+import torch
+import torch.nn.functional as F
+
+
+class MomentumBuffer:
+
+    def __init__(self, momentum: float = -0.75):
+        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 project(
+    v0: torch.Tensor,  # [B, C, T]
+    v1: torch.Tensor,  # [B, C, T]
+    dims=[-1],
+):
+    dtype = v0.dtype
+    device_type = v0.device.type
+    if device_type == "mps":
+        v0, v1 = v0.cpu(), v1.cpu()
+
+    v0, v1 = v0.double(), v1.double()
+    v1 = torch.nn.functional.normalize(v1, dim=dims)
+    v0_parallel = (v0 * v1).sum(dim=dims, keepdim=True) * v1
+    v0_orthogonal = v0 - v0_parallel
+    return v0_parallel.to(dtype).to(device_type), v0_orthogonal.to(dtype).to(device_type)
+
+
+def apg_forward(
+    pred_cond: torch.Tensor,  # [B, C, T]
+    pred_uncond: torch.Tensor,  # [B, C, T]
+    guidance_scale: float,
+    momentum_buffer: MomentumBuffer = None,
+    eta: float = 0.0,
+    norm_threshold: float = 2.5,
+    dims=[-1],
+):
+    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=dims, keepdim=True)
+        scale_factor = torch.minimum(ones, norm_threshold / diff_norm)
+        diff = diff * scale_factor
+
+    diff_parallel, diff_orthogonal = project(diff, pred_cond, dims)
+    normalized_update = diff_orthogonal + eta * diff_parallel
+    pred_guided = pred_cond + (guidance_scale - 1) * normalized_update
+    return pred_guided
+
+
+def cfg_forward(cond_output, uncond_output, cfg_strength):
+    return uncond_output + cfg_strength * (cond_output - uncond_output)
+
+
+def call_cos_tensor(tensor1, tensor2):
+    """
+    Calculate cosine similarity between two normalized tensors.
+    
+    Args:
+        tensor1: First tensor [B, ...]
+        tensor2: Second tensor [B, ...]
+    
+    Returns:
+        Cosine similarity value [B, 1]
+    """
+    tensor1 = tensor1 / torch.linalg.norm(tensor1, dim=1, keepdim=True)
+    tensor2 = tensor2 / torch.linalg.norm(tensor2, dim=1, keepdim=True)
+    cosvalue = torch.sum(tensor1 * tensor2, dim=1, keepdim=True)
+    return cosvalue
+
+
+def compute_perpendicular_component(latent_diff, latent_hat_uncond):
+    """
+    Decompose latent_diff into parallel and perpendicular components relative to latent_hat_uncond.
+    
+    Args:
+        latent_diff: Difference tensor [B, C, ...]
+        latent_hat_uncond: Unconditional prediction tensor [B, C, ...]
+    
+    Returns:
+        projection: Parallel component
+        perpendicular_component: Perpendicular component
+    """
+    n, t, c = latent_diff.shape
+    latent_diff = latent_diff.view(n * t, c).float()
+    latent_hat_uncond = latent_hat_uncond.view(n * t, c).float()
+
+    if latent_diff.size() != latent_hat_uncond.size():
+        raise ValueError("latent_diff and latent_hat_uncond must have the same shape [n, d].")
+
+    dot_product = torch.sum(latent_diff * latent_hat_uncond, dim=1, keepdim=True)  # [n, 1]
+    norm_square = torch.sum(latent_hat_uncond * latent_hat_uncond, dim=1, keepdim=True)  # [n, 1]
+    projection = (dot_product / (norm_square + 1e-8)) * latent_hat_uncond
+    perpendicular_component = latent_diff - projection
+
+    return projection.view(n, t, c), perpendicular_component.reshape(n, t, c)
+
+
+def adg_forward(
+    latents: torch.Tensor,
+    noise_pred_cond: torch.Tensor,
+    noise_pred_uncond: torch.Tensor,
+    sigma: torch.Tensor,
+    guidance_scale: float,
+    angle_clip: float = 3.14 / 6,  # pi/6 by default
+    apply_norm: bool = False,
+    apply_clip: bool = True,
+):
+    """
+    ADG (Angle-based Dynamic Guidance) forward pass for Flow Matching.
+    
+    In flow matching (including SD3), sigma represents the current timestep t_curr.
+    The predictions are velocity fields v(x_t, t).
+    
+    Args:
+        latents: Current state x_t [N, T, d] where d=64
+        noise_pred_cond: Conditional velocity prediction v_cond [N, T, d]
+        noise_pred_uncond: Unconditional velocity prediction v_uncond [N, T, d]
+        sigma: Current timestep t_curr (not t_prev!)
+        guidance_scale: Guidance strength
+        angle_clip: Maximum angle for clipping (default: pi/6)
+        apply_norm: Whether to normalize the result (ADG_w_norm variant)
+        apply_clip: Whether to clip the angle (ADG_wo_clip when False)
+    
+    Returns:
+        Guided velocity prediction [N, T, d]
+    """
+    # Get batch size
+    n = noise_pred_cond.shape[0]
+    noise_pred_text = noise_pred_cond
+    n, t, c = noise_pred_text.shape
+
+    # Ensure sigma/t has the right shape for broadcasting [N, 1, 1]
+    if isinstance(sigma, (int, float)):
+        sigma = torch.tensor(sigma, device=latents.device, dtype=latents.dtype)
+        sigma = sigma.view(1, 1, 1).expand(n, 1, 1)
+    elif torch.is_tensor(sigma):
+        if sigma.numel() == 1:
+            sigma = sigma.view(1, 1, 1).expand(n, 1, 1)
+        elif sigma.numel() == n:
+            sigma = sigma.view(n, 1, 1)
+        else:
+            raise ValueError(f"sigma has incompatible shape. Expected scalar or size {n}, got {sigma.shape}")
+    else:
+        raise TypeError(f"sigma must be a number or tensor, got {type(sigma)}")
+
+    # Adjust guidance weight
+    weight = guidance_scale - 1
+    weight = weight * (weight > 0) + 1e-3
+
+    latent_hat_text = latents - sigma * noise_pred_text
+    latent_hat_uncond = latents - sigma * noise_pred_uncond
+    latent_diff = latent_hat_text - latent_hat_uncond
+
+    # Calculate angle between conditional and unconditional predicted data
+    latent_theta = torch.acos(
+        call_cos_tensor(latent_hat_text.view(-1, c).to(float),
+                        latent_hat_uncond.reshape(-1, c).contiguous().to(float)))
+    latent_theta_new = torch.clip(weight * latent_theta, -angle_clip, angle_clip) if apply_clip else weight * latent_theta
+    proj, perp = compute_perpendicular_component(latent_diff, latent_hat_uncond)
+    latent_v_new = torch.cos(latent_theta_new) * latent_hat_text
+
+    latent_p_new = perp * torch.sin(latent_theta_new) / torch.sin(latent_theta) * (
+        torch.sin(latent_theta) > 1e-3) + perp * weight * (torch.sin(latent_theta) <= 1e-3)
+    latent_new = latent_v_new + latent_p_new
+    if apply_norm:
+        latent_new = latent_new * torch.linalg.norm(latent_hat_text, dim=1, keepdim=True) / torch.linalg.norm(
+            latent_new, dim=1, keepdim=True)
+
+    noise_pred = (latents - latent_new) / sigma
+    noise_pred = noise_pred.reshape(n, t, c).to(latents.dtype)
+    return noise_pred
+
+
+def adg_w_norm_forward(
+    latents: torch.Tensor,
+    noise_pred_cond: torch.Tensor,
+    noise_pred_uncond: torch.Tensor,
+    sigma: float,
+    guidance_scale: float,
+    angle_clip: float = 3.14 / 3,
+):
+    """
+    ADG with normalization - preserves the magnitude of latent predictions.
+    
+    This variant normalizes the final latent to maintain the same norm as the
+    conditional prediction, which can help preserve image quality.
+    """
+    return adg_forward(latents,
+                       noise_pred_cond,
+                       noise_pred_uncond,
+                       sigma,
+                       guidance_scale,
+                       angle_clip=angle_clip,
+                       apply_norm=True,
+                       apply_clip=True)
+
+
+def adg_wo_clip_forward(
+    latents: torch.Tensor,
+    noise_pred_cond: torch.Tensor,
+    noise_pred_uncond: torch.Tensor,
+    sigma: float,
+    guidance_scale: float,
+):
+    """
+    ADG without angle clipping - allows unbounded angle adjustments.
+    
+    This variant doesn't clip the angle, which may result in more aggressive
+    guidance but could be less stable.
+    """
+    return adg_forward(latents, noise_pred_cond, noise_pred_uncond, sigma, guidance_scale, apply_norm=False, apply_clip=False)