[diffusion] chore: ensure CFG Zero Star numerical stability for Helios model (#20091)

Signed-off-by: Lancer <maruixiang6688@gmail.com>
This commit is contained in:
Lancer
2026-03-08 14:25:14 +08:00
committed by GitHub
parent 72f6dfcc31
commit a73369c39f

View File

@@ -29,6 +29,8 @@ logger = init_logger(__name__)
def optimized_scale(positive_flat, negative_flat):
"""CFG Zero Star: compute optimal guidance scale."""
positive_flat = positive_flat.float()
negative_flat = negative_flat.float()
dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
return dot_product / squared_norm