Fix RenormCFG node for batches

This commit is contained in:
rockerBOO 2025-03-17 01:03:13 -04:00
parent 55a1b09ddc
commit 8a50599bd8
No known key found for this signature in database
GPG Key ID: 0D4EAF00DCABC97B

View File

@ -33,12 +33,13 @@ class RenormCFG:
ori_pos_norm = torch.linalg.vector_norm(cond_eps
, dim=tuple(range(1, len(cond_eps.shape))), keepdim=True
)
max_new_norm = ori_pos_norm * float(renorm_cfg)
new_pos_norm = torch.linalg.vector_norm(
max_new_norms = ori_pos_norm * float(renorm_cfg)
new_pos_norms = torch.linalg.vector_norm(
half_eps, dim=tuple(range(1, len(half_eps.shape))), keepdim=True
)
if new_pos_norm >= max_new_norm:
half_eps = half_eps * (max_new_norm / new_pos_norm)
for max_new_norm, new_pos_norm, denoise in zip(max_new_norms, new_pos_norms, half_eps):
if new_pos_norm >= max_new_norm:
denoise = denoise * (max_new_norm / new_pos_norm)
else:
cond_eps, uncond_eps = cond_denoised[:, :in_channels], uncond_denoised[:, :in_channels]
cond_rest, _ = cond_denoised[:, in_channels:], uncond_denoised[:, in_channels:]