akshatvishuandHongsheng Liu
2fa4c7a4a9
Signed-off-by: akshatvishu <akshatnayak197@gmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
69 lines
2.0 KiB
Python
69 lines
2.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import torch
|
|
|
|
from .dit import get_epss_timesteps
|
|
|
|
|
|
def apply_sway_sampling(t, sway_sampling_coef):
|
|
if sway_sampling_coef is None:
|
|
return t
|
|
return t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
|
|
|
|
|
|
def build_timesteps(steps, device, dtype, use_epss=True, sway_sampling_coef=-1.0):
|
|
if steps <= 0:
|
|
raise ValueError(f"steps must be positive, got {steps}")
|
|
if use_epss:
|
|
t = get_epss_timesteps(steps, device=device, dtype=dtype)
|
|
else:
|
|
t = torch.linspace(0, 1, steps + 1, device=device, dtype=dtype)
|
|
return apply_sway_sampling(t, sway_sampling_coef)
|
|
|
|
|
|
class Solver:
|
|
def __init__(self, func, y0, sigma=0.25, temperature=1.5) -> None:
|
|
self.func = func
|
|
self.y0 = y0
|
|
self.sigma = sigma
|
|
self.temperature = temperature
|
|
|
|
def integrate(self, t):
|
|
solution = torch.empty(len(t), *self.y0.shape, dtype=self.y0.dtype, device=self.y0.device)
|
|
solution[0] = self.y0
|
|
|
|
j = 1
|
|
y0 = self.y0
|
|
for t0, t1 in zip(t[:-1], t[1:]):
|
|
dt = t1 - t0
|
|
f0 = self.func(t0, y0)
|
|
dy = dt * f0
|
|
y1 = y0 + dy
|
|
|
|
while j < len(t) and t1 >= t[j]:
|
|
solution[j] = self._linear_interp(t0, t1, y0, y1, t[j])
|
|
j += 1
|
|
|
|
noise = torch.randn_like(y0)
|
|
shift = self.sigma * (self.temperature**0.5) * (abs(dt) ** 0.5) * noise
|
|
y0 = y1 + shift
|
|
|
|
return solution
|
|
|
|
def _linear_interp(self, t0, t1, y0, y1, t):
|
|
if t == t0:
|
|
return y0
|
|
if t == t1:
|
|
return y1
|
|
slope = (t - t0) / (t1 - t0)
|
|
return y0 + slope * (y1 - y0)
|
|
|
|
|
|
def integrate_cfm_steps(fn, y0, t, sde_args, sde_rnd, steps):
|
|
for step in range(steps):
|
|
dt = t[step + 1] - t[step]
|
|
y0 = y0 + fn(t[step], y0) * dt
|
|
y0 = y0 + sde_args[1] * (sde_args[2] ** 0.5) * (dt.abs() ** 0.5) * sde_rnd[step]
|
|
return y0
|