Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,856 Bytes
9e426da |
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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
import math
from src.diffusion.base.sampling import *
from src.diffusion.base.scheduling import *
from src.diffusion.pre_integral import *
from typing import Callable, List, Tuple
def ode_step_fn(x, v, dt, s, w):
return x + v * dt
def t2snr(t):
if isinstance(t, torch.Tensor):
return (t.clip(min=1e-8)/(1-t + 1e-8))
if isinstance(t, List) or isinstance(t, Tuple):
return [t2snr(t) for t in t]
t = max(t, 1e-8)
return (t/(1-t + 1e-8))
def t2logsnr(t):
if isinstance(t, torch.Tensor):
return torch.log(t.clip(min=1e-3)/(1-t + 1e-3))
if isinstance(t, List) or isinstance(t, Tuple):
return [t2logsnr(t) for t in t]
t = max(t, 1e-3)
return math.log(t/(1-t + 1e-3))
def t2isnr(t):
return 1/t2snr(t)
def nop(t):
return t
def shift_respace_fn(t, shift=3.0):
return t / (t + (1 - t) * shift)
import logging
logger = logging.getLogger(__name__)
class AdamLMSampler(BaseSampler):
def __init__(
self,
order: int = 2,
timeshift: float = 1.0,
state_refresh_rate: int = 1,
lms_transform_fn: Callable = nop,
w_scheduler: BaseScheduler = None,
step_fn: Callable = ode_step_fn,
*args,
**kwargs
):
super().__init__(*args, **kwargs)
self.step_fn = step_fn
self.w_scheduler = w_scheduler
self.state_refresh_rate = state_refresh_rate
assert self.scheduler is not None
assert self.w_scheduler is not None or self.step_fn in [ode_step_fn, ]
self.order = order
self.lms_transform_fn = lms_transform_fn
timesteps = torch.linspace(0.0, 1 - self.last_step, self.num_steps)
timesteps = torch.cat([timesteps, torch.tensor([1.0])], dim=0)
self.timesteps = shift_respace_fn(timesteps, timeshift)
self.timedeltas = timesteps[1:] - self.timesteps[:-1]
self._reparameterize_coeffs()
def _reparameterize_coeffs(self):
solver_coeffs = [[] for _ in range(self.num_steps)]
for i in range(0, self.num_steps):
pre_vs = [1.0, ]*(i+1)
pre_ts = self.lms_transform_fn(self.timesteps[:i+1])
int_t_start = self.lms_transform_fn(self.timesteps[i])
int_t_end = self.lms_transform_fn(self.timesteps[i+1])
order_annealing = self.order #self.num_steps - i
order = min(self.order, i + 1, order_annealing)
_, coeffs = lagrange_preint(order, pre_vs, pre_ts, int_t_start, int_t_end)
solver_coeffs[i] = coeffs
self.solver_coeffs = solver_coeffs
def _impl_sampling(self, net, noise, condition, uncondition):
"""
sampling process of Euler sampler
-
"""
batch_size = noise.shape[0]
cfg_condition = torch.cat([uncondition, condition], dim=0)
x = x0 = noise
state = None
pred_trajectory = []
t_cur = torch.zeros([batch_size,]).to(noise.device, noise.dtype)
timedeltas = self.timedeltas
solver_coeffs = self.solver_coeffs
for i in range(self.num_steps):
cfg_x = torch.cat([x, x], dim=0)
cfg_t = t_cur.repeat(2)
if i % self.state_refresh_rate == 0:
state = None
out, state = net(cfg_x, cfg_t, cfg_condition, state)
out = self.guidance_fn(out, self.guidances[i])
pred_trajectory.append(out)
out = torch.zeros_like(out)
order = len(self.solver_coeffs[i])
for j in range(order):
out += solver_coeffs[i][j] * pred_trajectory[-order:][j]
v = out
dt = timedeltas[i]
x0 = self.step_fn(x, v, 1-t_cur[0], s=0, w=0)
x = self.step_fn(x, v, dt, s=0, w=0)
t_cur += dt
return x |