Spaces:
Running on Zero
Running on Zero
File size: 2,069 Bytes
65635a9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | """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
@torch.no_grad()
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)
|