| """The paper's DDPM sampler (model.sample_ddpm) with two additions for the demo. |
| |
| - Guidance, as the old listener had it: a second pass with the encoded grid zeroed (its |
| ``force_uncond``), eps = eps_uncond + w * (eps_cond - eps_uncond). Following Kynkäänniemi et al. |
| (NeurIPS 2024), it can be limited to an interval of noise levels; w = 1 is the paper's sampler. |
| - Frozen frames: latent frames given in ``known`` are replaced at every step by the known latent |
| noised to that step's level (RePaint-style; "frozen overlap" continuation), so a chunk can be |
| sampled as the continuation of the previous one. |
| |
| With w = 1 and nothing known this is exactly model.sample_ddpm. |
| """ |
| from __future__ import annotations |
|
|
| from typing import Any, Mapping |
|
|
| import torch |
|
|
| from model import _prepare_batch_tensors, apply_seq_mask, clip_x0, resolve_inference_geometry |
|
|
|
|
| @torch.no_grad() |
| def sample( |
| *, |
| model: Any, |
| diffusion: Any, |
| batch: Mapping[str, Any], |
| device: torch.device, |
| sample_seed: int, |
| x0_clip_bound: float, |
| target_token_rate_hz: float, |
| inference_num_beats: int = 4, |
| inference_geometry: Mapping[str, Any] | None = None, |
| guidance: float = 1.0, |
| guidance_steps: tuple[int, int] | None = None, |
| known: torch.Tensor | None = None, |
| known_mask: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| """Returns the normalized latent [B,T,D]. ``guidance_steps`` = (lo, hi) applies guidance only for |
| lo <= step <= hi (steps count down from num_steps - 1). ``known``/``known_mask``: [B,T,D] / [B,T].""" |
| prepared = _prepare_batch_tensors(batch, device, require_target=False, require_timing=False) |
| geometry = inference_geometry or resolve_inference_geometry( |
| prepared, |
| use_bpm_inference_geometry=True, |
| inference_num_beats=int(inference_num_beats), |
| target_token_rate_hz=float(target_token_rate_hz), |
| ) |
| target_mask = geometry["target_valid_mask_bt"] |
| cond, cond_mask = model.encode_conditioning( |
| grid=prepared["grid"], |
| grid_ids=prepared["grid_ids"], |
| grid_times_sec=geometry["grid_times_sec"], |
| token_times_sec=geometry["token_times_sec"], |
| target_valid_mask_bt=target_mask, |
| grid_valid_mask_bt=prepared["grid_valid_mask"], |
| ) |
| uncond = torch.zeros_like(cond) |
| gen = torch.Generator(device=device).manual_seed(int(sample_seed)) |
| x = apply_seq_mask( |
| torch.randn(target_mask.shape[0], target_mask.shape[1], int(model.cfg.x_dim), device=device, generator=gen), |
| target_mask, |
| ) |
| keep = known_mask[..., None].to(device) if known is not None else None |
|
|
| def eps_at(x_t: torch.Tensor, t: torch.Tensor, c: torch.Tensor) -> torch.Tensor: |
| return model( |
| x_t=x_t, |
| t=t, |
| target_valid_mask_bt=target_mask, |
| token_times_sec=geometry["token_times_sec"], |
| cond_btd=c, |
| cond_valid_mask_bt=cond_mask, |
| ) |
|
|
| for step in reversed(range(diffusion.num_steps)): |
| t = torch.full((x.shape[0],), step, device=device, dtype=torch.long) |
| if keep is not None: |
| x = torch.where(keep, diffusion.q_sample(known.to(device), t, torch.randn(x.shape, device=device, generator=gen)), x) |
| eps = eps_at(x, t, cond) |
| guided = guidance != 1.0 and (guidance_steps is None or guidance_steps[0] <= step <= guidance_steps[1]) |
| if guided: |
| eps_u = eps_at(x, t, uncond) |
| eps = eps_u + float(guidance) * (eps - eps_u) |
| x0_hat = apply_seq_mask(clip_x0(diffusion.predict_x0_from_eps(x, t, eps), x0_clip_bound), target_mask) |
| mean = diffusion.posterior_mean_from_x0(x, t, x0_hat) |
| if step > 0: |
| z = torch.randn(x.shape, dtype=x.dtype, device=device, generator=gen) |
| x = mean + torch.sqrt(diffusion.posterior_variance[t].view(-1, 1, 1)) * z |
| else: |
| x = mean |
| x = apply_seq_mask(x, target_mask) |
| if keep is not None: |
| x = torch.where(keep, known.to(device), x) |
| return x |
|
|
