| 337 | |
| 338 | |
| 339 | def adaptive_projected_guidance( |
| 340 | diff: torch.Tensor, # [B, C, T, H, W] |
| 341 | pred_cond: torch.Tensor, # [B, C, T, H, W] |
| 342 | momentum_buffer: MomentumBuffer = None, |
| 343 | eta: float = 0.0, |
| 344 | norm_threshold: float = 55, |
| 345 | ): |
| 346 | if momentum_buffer is not None: |
| 347 | momentum_buffer.update(diff) |
| 348 | diff = momentum_buffer.running_average |
| 349 | if norm_threshold > 0: |
| 350 | ones = torch.ones_like(diff) |
| 351 | diff_norm = diff.norm(p=2, dim=[-1, -2, -3, -4], keepdim=True) |
| 352 | print(f"diff_norm: {diff_norm}") |
| 353 | scale_factor = torch.minimum(ones, norm_threshold / diff_norm) |
| 354 | diff = diff * scale_factor |
| 355 | diff_parallel, diff_orthogonal = project(diff, pred_cond) |
| 356 | normalized_update = diff_orthogonal + eta * diff_parallel |
| 357 | return normalized_update |
| 358 | |
| 359 | |
| 360 | |