80 lines
3.0 KiB
Python
80 lines
3.0 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# This code is borrowed and modified from Guided Diffusion Model,
|
|
# made publicly available under MIT license
|
|
# at https://github.com/IDEA-CCNL/Fengshenbang-LM/tree/main/fengshen/examples/disco_project
|
|
|
|
import numpy as np
|
|
import torch as th
|
|
|
|
from .gaussian_diffusion import GaussianDiffusion
|
|
|
|
|
|
class SpacedDiffusion(GaussianDiffusion):
|
|
"""
|
|
A diffusion process which can skip steps in a base diffusion process.
|
|
|
|
:param use_timesteps: a collection (sequence or set) of timesteps from the
|
|
original diffusion process to retain.
|
|
:param kwargs: the kwargs to create the base diffusion process.
|
|
"""
|
|
|
|
def __init__(self, use_timesteps, **kwargs):
|
|
self.use_timesteps = set(use_timesteps)
|
|
self.timestep_map = []
|
|
self.original_num_steps = len(kwargs['betas'])
|
|
|
|
base_diffusion = GaussianDiffusion(**kwargs) # pylint: disable=missing-kwoa
|
|
last_alpha_cumprod = 1.0
|
|
new_betas = []
|
|
for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):
|
|
if i in self.use_timesteps:
|
|
new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
|
|
last_alpha_cumprod = alpha_cumprod
|
|
self.timestep_map.append(i)
|
|
kwargs['betas'] = np.array(new_betas)
|
|
super().__init__(**kwargs)
|
|
|
|
def p_mean_variance(self, model, *args, **kwargs): # pylint: disable=signature-differs
|
|
return super().p_mean_variance(
|
|
self._wrap_model(model), *args, **kwargs)
|
|
|
|
def training_losses(self, model, *args, **kwargs): # pylint: disable=signature-differs
|
|
return super().training_losses(
|
|
self._wrap_model(model), *args, **kwargs)
|
|
|
|
def condition_mean(self, cond_fn, *args, **kwargs):
|
|
return super().condition_mean(
|
|
self._wrap_model(cond_fn), *args, **kwargs)
|
|
|
|
def condition_score(self, cond_fn, *args, **kwargs):
|
|
return super().condition_score(
|
|
self._wrap_model(cond_fn), *args, **kwargs)
|
|
|
|
def _wrap_model(self, model):
|
|
if isinstance(model, _WrappedModel):
|
|
return model
|
|
return _WrappedModel(model, self.timestep_map, self.rescale_timesteps,
|
|
self.original_num_steps)
|
|
|
|
def _scale_timesteps(self, t):
|
|
# Scaling is done by the wrapped model.
|
|
return t
|
|
|
|
|
|
class _WrappedModel:
|
|
|
|
def __init__(self, model, timestep_map, rescale_timesteps,
|
|
original_num_steps):
|
|
self.model = model
|
|
self.timestep_map = timestep_map
|
|
self.rescale_timesteps = rescale_timesteps
|
|
self.original_num_steps = original_num_steps
|
|
|
|
def __call__(self, x, ts, **kwargs):
|
|
map_tensor = th.tensor(
|
|
self.timestep_map, device=ts.device, dtype=ts.dtype)
|
|
new_ts = map_tensor[ts]
|
|
if self.rescale_timesteps:
|
|
new_ts = new_ts.float() * (1000.0 / self.original_num_steps)
|
|
return self.model(x, new_ts, **kwargs)
|