268 lines
11 KiB
Python
268 lines
11 KiB
Python
import os
|
|
import math
|
|
import torch
|
|
from collections import OrderedDict
|
|
|
|
from scepter.modules.utils.config import dict_to_yaml, Config
|
|
from scepter.modules.utils.distribute import we
|
|
from scepter.modules.utils.file_system import FS
|
|
from scepter.modules.model.registry import DIFFUSIONS, NOISE_SCHEDULERS, DIFFUSION_SAMPLERS
|
|
from tqdm import trange
|
|
|
|
@DIFFUSIONS.register_class()
|
|
class BaseDiffusion(object):
|
|
para_dict = {
|
|
"NOISE_SCHEDULER": {},
|
|
"SAMPLER_SCHEDULER": {},
|
|
"MIN_SNR_GAMMA": {
|
|
"value": None,
|
|
"description": "The minimum SNR gamma value for the loss function."
|
|
},
|
|
"PREDICTION_TYPE": {
|
|
"value": "eps",
|
|
"description": "The type of prediction to use for the loss function."
|
|
}
|
|
}
|
|
def __init__(self, cfg, logger=None):
|
|
super(BaseDiffusion, self).__init__()
|
|
self.logger = logger
|
|
self.cfg = cfg
|
|
self.init_params()
|
|
|
|
def init_params(self):
|
|
self.min_snr_gamma = self.cfg.get("MIN_SNR_GAMMA", None)
|
|
self.prediction_type = self.cfg.get("PREDICTION_TYPE", "eps")
|
|
self.noise_scheduler = NOISE_SCHEDULERS.build(self.cfg.NOISE_SCHEDULER, logger=self.logger)
|
|
self.sampler_scheduler = NOISE_SCHEDULERS.build(self.cfg.get("SAMPLER_SCHEDULER", self.cfg.NOISE_SCHEDULER),
|
|
logger=self.logger)
|
|
self.num_timesteps = self.noise_scheduler.num_timesteps
|
|
if self.cfg.have("WORK_DIR") and we.rank == 0:
|
|
schedule_visualization = os.path.join(self.cfg.WORK_DIR, "noise_schedule.png")
|
|
with FS.put_to(schedule_visualization) as local_path:
|
|
self.noise_scheduler.plot_noise_sampling_map(local_path)
|
|
schedule_visualization = os.path.join(self.cfg.WORK_DIR, "sampler_schedule.png")
|
|
with FS.put_to(schedule_visualization) as local_path:
|
|
self.sampler_scheduler.plot_noise_sampling_map(local_path)
|
|
|
|
|
|
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, **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 - timestamp.item()) / steps) ** 5.0)) / 2)
|
|
else:
|
|
guidance_scale = guide_scale
|
|
y_out = model(x=x_t, t=t, **model_kwargs[0])
|
|
u_out = model(x=x_t, t=t, **model_kwargs[1])
|
|
out = u_out + guidance_scale * (y_out - u_out)
|
|
if guide_rescale is not None and guide_rescale > 0.0:
|
|
ratio = (
|
|
y_out.flatten(1).std(dim=1) /
|
|
(out.flatten(1).std(dim=1) + 1e-12)).view((-1, ) + (1, ) *
|
|
(y_out.ndim - 1))
|
|
out *= guide_rescale * ratio + (1 - guide_rescale) * 1.0
|
|
|
|
if self.prediction_type == 'x0':
|
|
x0 = out
|
|
elif self.prediction_type == 'eps':
|
|
x0 = (x_t - sigma * out) / alpha
|
|
elif self.prediction_type == 'v':
|
|
x0 = alpha * x_t - sigma * out
|
|
else:
|
|
raise NotImplementedError(
|
|
f'prediction_type {self.prediction_type} not implemented')
|
|
|
|
# print("torch.sum(y_out):", torch.sum(y_out), "torch.sum(u_out):", torch.sum(u_out), "torch.sum(out):",
|
|
# torch.sum(out), "torch.sum(x0):", torch.sum(x0), "sigmas", sigma, "alphas", alpha)
|
|
return x0
|
|
|
|
sampler_ins = self.get_sampler(sampler)
|
|
|
|
# this is ignored for schnell
|
|
sampler_output = sampler_ins.preprare_sampler(
|
|
noise,
|
|
steps=steps,
|
|
prediction_type=self.prediction_type,
|
|
scheduler_ins=self.sampler_scheduler,
|
|
callback_fn=callback_fn
|
|
)
|
|
|
|
for _ in trange(steps, disable=not show_progress):
|
|
trange.desc = sampler_output.msg
|
|
sampler_output = sampler_ins.step(sampler_output)
|
|
if return_intermediate == 'x_0':
|
|
intermediates.append(sampler_output.x_0)
|
|
elif return_intermediate == 'x_t':
|
|
intermediates.append(sampler_output.x_t)
|
|
if intermediate_callback is not None:
|
|
intermediate_callback(intermediates[-1])
|
|
return (sampler_output.x_0, intermediates) if return_intermediate is not None else sampler_output.x_0
|
|
|
|
|
|
def loss(self, x_0, model, model_kwargs={}, reduction='mean', noise=None, **kwargs):
|
|
# use noise scheduler to add noise
|
|
if noise is None:
|
|
noise = torch.randn_like(x_0)
|
|
schedule_output = self.noise_scheduler.add_noise(x_0, noise)
|
|
x_t, t, sigma, alpha = schedule_output.x_t, schedule_output.t, schedule_output.sigma, schedule_output.alpha
|
|
out = model(x=x_t, t=t, **model_kwargs)
|
|
|
|
# mse loss
|
|
target = {
|
|
'eps': noise,
|
|
'x0': x_0,
|
|
'v': alpha * noise - sigma * x_0
|
|
}[self.prediction_type]
|
|
|
|
loss = (out - target).pow(2)
|
|
if reduction == 'mean':
|
|
loss = loss.flatten(1).mean(dim=1)
|
|
|
|
if self.min_snr_gamma is not None:
|
|
alphas = self.noise_scheduler.alphas.to(x_0.device)[t]
|
|
sigmas = self.noise_scheduler.sigmas.pow(2).to(x_0.device)[t]
|
|
snrs = (alphas / sigmas).clamp(min=1e-20)
|
|
min_snrs = snrs.clamp(max=self.min_snr_gamma)
|
|
weights = min_snrs / snrs
|
|
else:
|
|
weights = 1
|
|
|
|
loss = loss * weights
|
|
return loss
|
|
|
|
def get_sampler(self, sampler):
|
|
if isinstance(sampler, str):
|
|
if sampler not in DIFFUSION_SAMPLERS.class_map:
|
|
if self.logger is not None:
|
|
self.logger.info(f"{sampler} not in the defined samplers list {DIFFUSION_SAMPLERS.class_map.keys()}")
|
|
else:
|
|
print(f"{sampler} not in the defined samplers list {DIFFUSION_SAMPLERS.class_map.keys()}")
|
|
return None
|
|
sampler_cfg = Config(cfg_dict={"NAME": sampler}, load=False)
|
|
sampler_ins = DIFFUSION_SAMPLERS.build(sampler_cfg, logger=self.logger)
|
|
elif isinstance(sampler, (Config, dict, OrderedDict)):
|
|
if isinstance(sampler, (dict, OrderedDict)):
|
|
sampler = Config(cfg_dict={k.upper():v for k, v in dict(sampler).items()}, load=False)
|
|
sampler_ins = DIFFUSION_SAMPLERS.build(sampler, logger=self.logger)
|
|
else:
|
|
raise NotImplementedError
|
|
return sampler_ins
|
|
|
|
def __repr__(self) -> str:
|
|
return f'{self.__class__.__name__}' + ' ' + super().__repr__()
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
return dict_to_yaml('DIFFUSIONS',
|
|
__class__.__name__,
|
|
BaseDiffusion.para_dict,
|
|
set_name=True)
|
|
|
|
@DIFFUSIONS.register_class()
|
|
class DiffusionFluxRF(BaseDiffusion):
|
|
para_dict = {
|
|
"PREDICTION_TYPE": {
|
|
"value": "raw",
|
|
"description": "The type of prediction to use for the loss function."
|
|
}
|
|
}
|
|
para_dict.update(BaseDiffusion.para_dict)
|
|
def __init__(self, cfg, logger=None):
|
|
super(DiffusionFluxRF, self).__init__(cfg, logger=logger)
|
|
self.prediction_type = self.cfg.get("PREDICTION_TYPE", "raw")
|
|
|
|
def loss(self, x_0, model, model_kwargs={}, reduction='mean', noise=None, **kwargs):
|
|
if noise is None:
|
|
noise = torch.randn_like(x_0)
|
|
schedule_output = self.noise_scheduler.add_noise(x_0, noise)
|
|
x_t, t, sigma = schedule_output.x_t, schedule_output.t, schedule_output.sigma
|
|
out = model(x=x_t, t=sigma, **model_kwargs)
|
|
# raw
|
|
if self.prediction_type == "raw":
|
|
target = noise - x_0
|
|
out = out
|
|
elif self.prediction_type == "sigma_scaled":
|
|
target = x_0
|
|
out = out * (-sigma) + x_t
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
loss = (target - out) ** 2
|
|
if reduction == 'mean':
|
|
loss = loss.flatten(1).mean(dim=1)
|
|
|
|
if self.min_snr_gamma is not None:
|
|
alphas = self.noise_scheduler.alphas.to(x_0.device)[t]
|
|
sigmas = self.noise_scheduler.sigmas.pow(2).to(x_0.device)[t]
|
|
snrs = (alphas / sigmas).clamp(min=1e-20)
|
|
min_snrs = snrs.clamp(max=self.min_snr_gamma)
|
|
weights = min_snrs / snrs
|
|
else:
|
|
weights = 1
|
|
|
|
loss = loss * weights
|
|
return loss
|
|
|
|
@torch.no_grad()
|
|
def sample(self,
|
|
noise,
|
|
model,
|
|
model_kwargs={},
|
|
steps=20,
|
|
sampler = None,
|
|
show_progress=False,
|
|
return_intermediate=None,
|
|
intermediate_callback=None,
|
|
**kwargs):
|
|
# sanity check
|
|
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):
|
|
sigma = torch.full((x_t.shape[0],), sigma, dtype=x_t.dtype, device=x_t.device)
|
|
x_0 = model(x = x_t, t=sigma, **model_kwargs)
|
|
return x_0
|
|
|
|
sampler_ins = self.get_sampler(sampler)
|
|
|
|
# this is ignored for schnell
|
|
sampler_output = sampler_ins.preprare_sampler(
|
|
noise,
|
|
steps = steps,
|
|
prediction_type=self.prediction_type,
|
|
scheduler_ins = self.sampler_scheduler,
|
|
callback_fn=callback_fn
|
|
)
|
|
|
|
for _ in trange(steps, disable=not show_progress):
|
|
trange.desc = sampler_output.msg
|
|
sampler_output = sampler_ins.step(sampler_output)
|
|
if return_intermediate == 'x_0':
|
|
intermediates.append(sampler_output.x_0)
|
|
elif return_intermediate == 'x_t':
|
|
intermediates.append(sampler_output.x_t)
|
|
if intermediate_callback is not None:
|
|
intermediate_callback(intermediates[-1])
|
|
return (sampler_output.x_0, intermediates) if return_intermediate is not None else sampler_output.x_t
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
return dict_to_yaml('DIFFUSIONS',
|
|
__class__.__name__,
|
|
DiffusionFluxRF.para_dict,
|
|
set_name=True) |