| """Equal-grid Multi-Trace objective and 1..K sampling from Untitled0.ipynb.""" |
| import torch |
| from trace_imf import TraceIMF |
|
|
|
|
| class MultiTraceIMF(TraceIMF): |
| def __init__(self, max_nfe=5, jvp_precision="bf16", channels_last=True, **kwargs): |
| super().__init__(**kwargs) |
| if max_nfe < 1 or jvp_precision not in ("fp32", "bf16"): |
| raise ValueError("Invalid max_nfe or JVP precision") |
| self.max_nfe = max_nfe |
| self.jvp_precision = jvp_precision |
| self.channels_last = channels_last |
|
|
| def draw_interval(self, device, generator=None, nfe=None): |
| |
| k = int(torch.randint(1, self.max_nfe + 1, (), device=device, generator=generator)) if nfe is None else nfe |
| if not 1 <= k <= self.max_nfe: |
| raise ValueError("NFE outside the trained range") |
| j = int(torch.randint(k, (), device=device, generator=generator)) |
| return k, j |
|
|
| def loss(self, model, x, labels=None, *, interval=None, nfe=None, generator=None, |
| t=None, r=None, noise=None, derivative_model=None): |
| if t is None: |
| k, j = interval or self.draw_interval(x.device, generator, nfe) |
| r = torch.full((len(x),), j / k, device=x.device) |
| t = r + torch.rand(len(x), device=x.device, generator=generator) / k |
| elif r is None: |
| raise ValueError("Explicit Multi-Trace t requires explicit r") |
| return super().loss(model, x, labels, derivative_model=derivative_model, |
| t=t, r=r, noise=noise, generator=generator) |
|
|
| @torch.no_grad() |
| def sample(self, model, n_samples=None, labels=None, *, device=None, noise=None, nfe=1): |
| if not 1 <= nfe <= self.max_nfe: |
| raise ValueError("NFE outside the trained range") |
| device = device or next(model.parameters()).device |
| if labels is None: |
| raise ValueError("Class-conditioned sampling requires labels") |
| labels = labels.to(device) |
| z = (torch.randn(len(labels), self.channels, self.image_size, self.image_size, device=device) |
| if noise is None else noise.to(device).clone()) |
| if len(z) != len(labels): |
| raise ValueError("Noise and labels must have the same batch size") |
| if self.channels_last: |
| z = z.contiguous(memory_format=torch.channels_last) |
| was_training = model.training |
| model.eval() |
| try: |
| for index in range(nfe, 0, -1): |
| t = torch.full((len(z),), index / nfe, device=device) |
| r = torch.full_like(t, (index - 1) / nfe) |
| z = z - model(z, t, r, y=labels, use_flash_attention=False) / nfe |
| return (z.clamp(-1, 1) + 1) * 0.5 |
| finally: |
| model.train(was_training) |
|
|
