| 90 | } |
| 91 | |
| 92 | std::vector<float> adg_guidance( |
| 93 | const std::vector<float> & latents, |
| 94 | const std::vector<float> & pred_cond, |
| 95 | const std::vector<float> & pred_uncond, |
| 96 | float sigma, |
| 97 | float guidance_scale, |
| 98 | int64_t frames, |
| 99 | int64_t channels, |
| 100 | float angle_clip) { |
| 101 | constexpr float kEps = 1.0e-8F; |
| 102 | require_vector_shape(pred_cond.size(), frames, channels, "ADG branch shape mismatch"); |
| 103 | require_same_size(latents.size(), pred_cond.size(), "ADG latent/condition size mismatch"); |
| 104 | require_same_size(pred_cond.size(), pred_uncond.size(), "ADG branch size mismatch"); |
| 105 | if (std::abs(sigma) <= kEps) { |
| 106 | throw std::runtime_error("ADG sigma must be nonzero"); |
| 107 | } |
| 108 | float weight = guidance_scale - 1.0F; |
| 109 | weight = weight * (weight > 0.0F ? 1.0F : 0.0F) + 1.0e-3F; |
| 110 | |
| 111 | std::vector<float> out(pred_cond.size(), 0.0F); |
| 112 | for (int64_t frame = 0; frame < frames; ++frame) { |
| 113 | const size_t offset = static_cast<size_t>(frame * channels); |
| 114 | float text_norm_sq = 0.0F; |
| 115 | float uncond_norm_sq = 0.0F; |
| 116 | float text_uncond_dot = 0.0F; |
| 117 | float diff_uncond_dot = 0.0F; |
| 118 | float uncond_norm_square_for_projection = 0.0F; |
| 119 | for (int64_t channel = 0; channel < channels; ++channel) { |
| 120 | const size_t index = offset + static_cast<size_t>(channel); |
| 121 | const float latent_hat_text = latents[index] - sigma * pred_cond[index]; |
| 122 | const float latent_hat_uncond = latents[index] - sigma * pred_uncond[index]; |
| 123 | const float latent_diff = latent_hat_text - latent_hat_uncond; |
| 124 | text_norm_sq += latent_hat_text * latent_hat_text; |
| 125 | uncond_norm_sq += latent_hat_uncond * latent_hat_uncond; |
| 126 | text_uncond_dot += latent_hat_text * latent_hat_uncond; |
| 127 | diff_uncond_dot += latent_diff * latent_hat_uncond; |
| 128 | uncond_norm_square_for_projection += latent_hat_uncond * latent_hat_uncond; |
| 129 | } |
| 130 | const float norm_product = std::sqrt(text_norm_sq) * std::sqrt(uncond_norm_sq); |
| 131 | float cos_theta = text_uncond_dot / std::max(norm_product, kEps); |
| 132 | cos_theta = std::clamp(cos_theta, -1.0F + 1.0e-6F, 1.0F - 1.0e-6F); |
| 133 | const float theta = std::acos(cos_theta); |
| 134 | const float theta_new = std::clamp(weight * theta, -angle_clip, angle_clip); |
| 135 | const float sin_theta = std::sin(theta); |
| 136 | const float sin_scale = sin_theta > 1.0e-3F ? std::sin(theta_new) / sin_theta : weight; |
| 137 | const float projection_scale = diff_uncond_dot / (uncond_norm_square_for_projection + kEps); |
| 138 | for (int64_t channel = 0; channel < channels; ++channel) { |
| 139 | const size_t index = offset + static_cast<size_t>(channel); |
| 140 | const float latent_hat_text = latents[index] - sigma * pred_cond[index]; |
| 141 | const float latent_hat_uncond = latents[index] - sigma * pred_uncond[index]; |
| 142 | const float latent_diff = latent_hat_text - latent_hat_uncond; |
| 143 | const float projection = projection_scale * latent_hat_uncond; |
| 144 | const float perpendicular = latent_diff - projection; |
| 145 | const float latent_new = std::cos(theta_new) * latent_hat_text + perpendicular * sin_scale; |
| 146 | out[index] = (latents[index] - latent_new) / sigma; |
| 147 | } |
| 148 | } |
| 149 | return out; |