update
This commit is contained in:
@@ -46,19 +46,23 @@ class BaseDiffusion(object):
|
||||
|
||||
|
||||
def sample(self, noise, model, model_kwargs={}, steps=20, sampler=None, use_dynamic_cfg=False, guide_scale=None, guide_rescale=None,
|
||||
show_progress=False, return_intermediate=None, intermediate_callback=None):
|
||||
show_progress=False, return_intermediate=None, intermediate_callback=None, **kwargs):
|
||||
assert isinstance(steps, (int, torch.LongTensor))
|
||||
assert return_intermediate in (None, 'x0', 'xt')
|
||||
assert isinstance(sampler, (str, dict, Config))
|
||||
intermediates = []
|
||||
|
||||
def callback_fn(x_t, t, sigma=None, alpha=None):
|
||||
timestamp = t
|
||||
t = t.repeat(len(x_t)).round().long().to(x_t.device)
|
||||
sigma = sigma.repeat(len(x_t), *([1] * (len(sigma.shape) - 1)))
|
||||
alpha = alpha.repeat(len(x_t), *([1] * (len(alpha.shape) - 1)))
|
||||
|
||||
if guide_scale is None or guide_scale == 1.0:
|
||||
out = model(x=x_t, t=t, **model_kwargs)
|
||||
else:
|
||||
if use_dynamic_cfg:
|
||||
guidance_scale = 1 + guide_scale * ((1 - math.cos(math.pi * ((steps - t.item()) / steps) ** 5.0)) / 2)
|
||||
guidance_scale = 1 + guide_scale * ((1 - math.cos(math.pi * ((steps - timestamp.item()) / steps) ** 5.0)) / 2)
|
||||
else:
|
||||
guidance_scale = guide_scale
|
||||
y_out = model(x=x_t, t=t, **model_kwargs[0])
|
||||
|
||||
@@ -145,12 +145,12 @@ class DDIMSampler(BaseDiffusionSampler):
|
||||
return output
|
||||
|
||||
def step(self, sampler_output):
|
||||
step = sampler_output.step
|
||||
x_t = sampler_output.x_t
|
||||
step = sampler_output.step
|
||||
t = sampler_output.ts[step]
|
||||
sigmas_vp = sampler_output.sigmas_vp.to(x_t.device)
|
||||
alpha_init = _i(sampler_output.alphas_init, step, x_t)
|
||||
sigma_init = _i(sampler_output.sigmas_init, step, x_t)
|
||||
alpha_init = _i(sampler_output.alphas_init, step, x_t[:1])
|
||||
sigma_init = _i(sampler_output.sigmas_init, step, x_t[:1])
|
||||
|
||||
x = sampler_output.callback_fn(x_t, t, sigma_init, alpha_init)
|
||||
noise_factor = self.eta * (sigmas_vp[step + 1] ** 2 / sigmas_vp[step] ** 2 *
|
||||
|
||||
Reference in New Issue
Block a user