refact project

apg_guidance.py

@@ -52,3 +52,35 @@ def apg_forward(
 
 def cfg_forward(cond_output, uncond_output, cfg_strength):
     return uncond_output + cfg_strength * (cond_output - uncond_output)
+
+
+
+def optimized_scale(positive_flat, negative_flat):
+
+    # Calculate dot production
+    dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
+
+    # Squared norm of uncondition
+    squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8
+
+    # st_star = v_cond^T * v_uncond / ||v_uncond||^2
+    st_star = dot_product / squared_norm
+    
+    return st_star
+
+
+def cfg_zero_star(noise_pred_with_cond, noise_pred_uncond, guidance_scale, i, zero_steps=1, use_zero_init=True):
+    bsz = noise_pred_with_cond.shape[0]
+    print("debug noise_pred_with_cond", noise_pred_with_cond.shape)
+    positive_flat = noise_pred_with_cond.view(bsz, -1)
+    negative_flat = noise_pred_uncond.view(bsz, -1)
+    print(f"debug {positive_flat.shape=} {negative_flat.shape=}")
+    alpha = optimized_scale(positive_flat, negative_flat)
+    print(f"debug {alpha.shape=}")
+    alpha = alpha.view(bsz, 1, 1, 1)
+    if (i <= zero_steps) and use_zero_init:
+        noise_pred = noise_pred_with_cond * 0.
+    else:
+        noise_pred = noise_pred_uncond * alpha + guidance_scale * (noise_pred_with_cond - noise_pred_uncond * alpha)
+    print(f"debug {noise_pred.shape=}")
+    return noise_pred