Spaces:
Running on Zero
Running on Zero
| """CFG + DDIM sampling for the DM baseline and PaTh, shared by evaluate.py and app.py.""" | |
| import torch | |
| from . import common | |
| from .diffusion import cfg as guide | |
| from .thinker import think | |
| class Sampler: | |
| """Samples x0 for a batch of (cond, mask). | |
| PaTh options (efficient variant, Sec. 3.2 / App. N): | |
| carry keep the Thinker's (y, z) across denoising steps instead of resetting it every step; | |
| think_frac call the Thinker only for the first fraction of the steps, then reuse its last residuals. | |
| """ | |
| def __init__(self, sched, dev, dm=None, thinker=None, painter=None, cfg_scale=2.0, n_sup=16, carry=False, | |
| think_frac=1.0): | |
| self.sched, self.dev, self.dm, self.thinker, self.painter = sched, dev, dm, thinker, painter | |
| self.cfg_scale, self.n_sup, self.carry, self.think_frac = cfg_scale, n_sup, carry, think_frac | |
| def __call__(self, cond, mask, noise, steps=20, callback=None): | |
| B = len(cond) | |
| # conditional and unconditional (empty puzzle) halves in one batch for classifier-free guidance | |
| cond = torch.cat([cond, torch.full_like(cond, -1.0)]) | |
| mask = torch.cat([mask, torch.zeros_like(mask)]) | |
| state = {"carry": None, "residuals": None} | |
| def model_fn(x, t, i): | |
| x2, t2 = torch.cat([x, x]), torch.cat([t, t]) | |
| with common.autocast(self.dev): | |
| if self.dm is not None: | |
| out = self.dm(torch.cat([x2, cond, mask], 1), t2) | |
| else: | |
| if state["residuals"] is None or i < self.think_frac * steps: | |
| state["residuals"], carry = think(self.thinker, x2, t2, cond, mask, | |
| state["carry"] if self.carry else None, self.n_sup) | |
| state["carry"] = carry | |
| out = self.painter(x2, t2, state["residuals"]) | |
| return guide(out[:B], out[B:], self.cfg_scale).float() | |
| return self.sched.ddim(model_fn, noise, steps, callback=callback) | |