guysss/ACE-Step

1import torch2 3 4class MomentumBuffer:5    def __init__(self, momentum: float = -0.75):6        self.momentum = momentum7        self.running_average = 08 9    def update(self, update_value: torch.Tensor):10        new_average = self.momentum * self.running_average11        self.running_average = update_value + new_average12 13 14def project(15    v0: torch.Tensor,  # [B, C, H, W]16    v1: torch.Tensor,  # [B, C, H, W]17    dims=[-1, -2],18):19    dtype = v0.dtype20    device_type = v0.device.type21    if device_type == "mps":22        v0, v1 = v0.cpu(), v1.cpu()23 24    v0, v1 = v0.double(), v1.double()25    v1 = torch.nn.functional.normalize(v1, dim=dims)26    v0_parallel = (v0 * v1).sum(dim=dims, keepdim=True) * v127    v0_orthogonal = v0 - v0_parallel28    return v0_parallel.to(dtype).to(device_type), v0_orthogonal.to(dtype).to(device_type)29 30 31def apg_forward(32    pred_cond: torch.Tensor,  # [B, C, H, W]33    pred_uncond: torch.Tensor,  # [B, C, H, W]34    guidance_scale: float,35    momentum_buffer: MomentumBuffer = None,36    eta: float = 0.0,37    norm_threshold: float = 2.5,38    dims=[-1, -2],39):40    diff = pred_cond - pred_uncond41    if momentum_buffer is not None:42        momentum_buffer.update(diff)43        diff = momentum_buffer.running_average44 45    if norm_threshold > 0:46        ones = torch.ones_like(diff)47        diff_norm = diff.norm(p=2, dim=dims, keepdim=True)48        scale_factor = torch.minimum(ones, norm_threshold / diff_norm)49        diff = diff * scale_factor50 51    diff_parallel, diff_orthogonal = project(diff, pred_cond, dims)52    normalized_update = diff_orthogonal + eta * diff_parallel53    pred_guided = pred_cond + (guidance_scale - 1) * normalized_update54    return pred_guided55 56 57def cfg_forward(cond_output, uncond_output, cfg_strength):58    return uncond_output + cfg_strength * (cond_output - uncond_output)59 60 61def cfg_double_condition_forward(62    cond_output,63    uncond_output,64    only_text_cond_output,65    guidance_scale_text,66    guidance_scale_lyric,67):68    return (1 - guidance_scale_text) * uncond_output + (guidance_scale_text - guidance_scale_lyric) * only_text_cond_output + guidance_scale_lyric * cond_output 69 70 71def optimized_scale(positive_flat, negative_flat):72 73    # Calculate dot production74    dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)75 76    # Squared norm of uncondition77    squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-878 79    # st_star = v_cond^T * v_uncond / ||v_uncond||^280    st_star = dot_product / squared_norm81    82    return st_star83 84 85def cfg_zero_star(noise_pred_with_cond, noise_pred_uncond, guidance_scale, i, zero_steps=1, use_zero_init=True):86    bsz = noise_pred_with_cond.shape[0]87    positive_flat = noise_pred_with_cond.view(bsz, -1)88    negative_flat = noise_pred_uncond.view(bsz, -1)89    alpha = optimized_scale(positive_flat, negative_flat)90    alpha = alpha.view(bsz, 1, 1, 1)91    if (i <= zero_steps) and use_zero_init:92        noise_pred = noise_pred_with_cond * 0.93    else:94        noise_pred = noise_pred_uncond * alpha + guidance_scale * (noise_pred_with_cond - noise_pred_uncond * alpha)95    return noise_pred96