import ast
from safetensors import safe_open
import torch
from dataclasses import dataclass
from typing import Optional, Union, List
def update_args_from_yaml(group, args, parser):
for key, value in group.items():
if isinstance(value, dict):
update_args_from_yaml(value, args, parser)
else:
if value == 'None' or value == 'null':
value = None
else:
arg_type = next((action.type for action in parser._actions if action.dest == key), str)
if arg_type is ast.literal_eval:
pass
elif arg_type is not None and not isinstance(value, arg_type):
try:
value = arg_type(value)
except ValueError as e:
raise ValueError(f"Cannot convert {key} to {arg_type}: {e}")
setattr(args, key, value)
def safe_load(model_path):
assert "safetensors" in model_path
state_dict = {}
with safe_open(model_path, framework="pt", device="cpu") as f:
for k in f.keys():
state_dict[k] = f.get_tensor(k)
return state_dict
@dataclass
class DDIMSchedulerStepOutput:
prev_sample: torch.Tensor # x_{t-1}
pred_original_sample: Optional[torch.Tensor] = None # x0
@dataclass
class DDIMSchedulerConversionOutput:
pred_epsilon: torch.Tensor
pred_original_sample: torch.Tensor
pred_velocity: torch.Tensor
class DDIMScheduler:
prediction_types = ["epsilon", "sample", "v_prediction"]
def __init__(
self,
num_train_timesteps: int,
num_inference_timesteps: int,
betas: torch.Tensor,
set_alpha_to_one: bool = True,
set_inference_timesteps_from_pure_noise: bool = True,
inference_timesteps: Union[str, List[int]] = "trailing",
device: Optional[Union[str, torch.device]] = None,
dtype: torch.dtype = torch.float32,
skip_step:bool = False,
original_inference_step: int=20,
steps_offset: int=0,
):
assert num_train_timesteps > 0
assert num_train_timesteps >= num_inference_timesteps
assert num_train_timesteps == betas.size(0)
assert betas.ndim == 1
# self.user_name = user_name
# self.run_time = Recorder.format_time()
# self.task_name = 'AutoAIGC_%s' % str(self.run_time)
self.module_name = 'AutoAIGC'
self.config_list = {"num_train_timesteps": num_train_timesteps,
"num_inference_timesteps": num_inference_timesteps,
"betas": betas,
"set_alpha_to_one": set_alpha_to_one,
"set_inference_timesteps_from_pure_noise": set_inference_timesteps_from_pure_noise,
"inference_timesteps": inference_timesteps}
self.module_info = str(self.config_list)
# self.upload_logger(user_name=user_name)
device = device or betas.device
self.num_train_timesteps = num_train_timesteps
self.num_inference_steps = num_inference_timesteps
self.steps_offset = steps_offset
self.betas = betas # .to(device=device, dtype=dtype)
self.alphas = 1.0 - self.betas
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
self.final_alpha_cumprod = torch.tensor(1.0, device=device, dtype=dtype) if set_alpha_to_one else self.alphas_cumprod[0]
if isinstance(inference_timesteps, torch.Tensor):
assert len(inference_timesteps) == num_inference_timesteps
self.timesteps = inference_timesteps.cpu().numpy().tolist()
elif set_inference_timesteps_from_pure_noise:
if inference_timesteps == "trailing":
# [999, 949, 899, 849, 799, 749, 699, 649, 599, 549, 499, 449, 399, 349, 299, 249, 199, 149, 99, 49]
if skip_step: # ?
original_timesteps = torch.arange(num_train_timesteps - 1, -1, -num_train_timesteps / original_inference_step, device=device).round().int().tolist()
skipping_step = len(original_timesteps) // num_inference_timesteps
self.timesteps = original_timesteps[::skipping_step][:num_inference_timesteps]
else: # [999, 899, 799, 699, 599, 499, 399, 299, 199, 99]
self.timesteps = torch.arange(num_train_timesteps - 1, -1, -num_train_timesteps / num_inference_timesteps, device=device).round().int().tolist()
elif inference_timesteps == "linspace":
# Fixed DDIM timestep. Make sure the timestep starts from 999.
# Example 20 steps:
# [999, 946, 894, 841, 789, 736, 684, 631, 578, 526, 473, 421, 368, 315, 263, 210, 158, 105, 53, 0]
# [999, 888, 777, 666, 555, 444, 333, 222, 111, 0]
self.timesteps = torch.linspace(0, num_train_timesteps - 1, num_inference_timesteps, device=device).round().int().flip(0).tolist()
elif inference_timesteps == "leading":
step_ratio = num_train_timesteps // num_inference_timesteps
# # creates integer timesteps by multiplying by ratio
# # casting to int to avoid issues when num_inference_step is power of 3
self.timesteps = torch.arange(0, num_inference_timesteps).mul(step_ratio).round().flip(dims=[0]) #.clone().long()
# self.timesteps += self.steps_offset
# Original SD and DDIM paper may have a bug: <https://github.com/huggingface/diffusers/issues/2585>
# The inference timestep does not start from 999.
# Example 20 steps:
# [950, 900, 850, 800, 750, 700, 650, 600, 550, 500, 450, 400, 350, 300, 250, 200, 150, 100, 50, 0]
# [ 900, 800, 700, 600, 500, 400, 300, 200, 100, 0]
# self.timesteps = torch.arange(0, num_train_timesteps, num_train_timesteps // num_inference_timesteps, device=self.device, dtype=torch.int).flip(0)
# self.timesteps = list(reversed(range(0, num_train_timesteps, num_train_timesteps // num_inference_timesteps)))
else:
raise NotImplementedError
elif inference_timesteps == "leading":
# Original SD and DDIM paper may have a bug: <https://github.com/huggingface/diffusers/issues/2585>
# The inference timestep does not start from 999.
# Example 20 steps:
# [950, 900, 850, 800, 750, 700, 650, 600, 550, 500, 450, 400, 350, 300, 250, 200, 150, 100, 50, 0]
# [ 900, 800, 700, 600, 500, 400, 300, 200, 100, 0]
# self.timesteps = torch.arange(0, num_train_timesteps, num_train_timesteps // num_inference_timesteps, device=self.device, dtype=torch.int).flip(0)
self.timesteps = list(reversed(range(0, num_train_timesteps, num_train_timesteps // num_inference_timesteps)))
else:
self.timesteps = list(reversed(range(0, num_train_timesteps, num_train_timesteps // num_inference_timesteps)))
# raise NotImplementedError
self.to(device=device)
def to(self, device):
self.betas = self.betas.to(device)
self.alphas_cumprod = self.alphas_cumprod.to(device)
self.final_alpha_cumprod = self.final_alpha_cumprod.to(device)
# self.timesteps = self.timesteps.to(device)
return self
def step(
self,
model_output: torch.Tensor,
model_output_type: str,
timestep: Union[torch.Tensor, int],
sample: torch.Tensor,
eta: float = 0.0,
clip_sample: bool = False,
dynamic_threshold: Optional[float] = None,
variance_noise: Optional[torch.Tensor] = None,
) -> DDIMSchedulerStepOutput:
# 1. get previous step value (t-1)
if isinstance(timestep, int):
# 1. get previous step value (t-1)
idx = self.timesteps.index(timestep)
prev_timestep = self.timesteps[idx + 1] if idx < self.num_inference_steps - 1 else None
# 2. compute alphas, betas
alpha_prod_t = self.alphas_cumprod[timestep]
alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep is not None else self.final_alpha_cumprod
beta_prod_t = 1 - alpha_prod_t
beta_prod_t_prev = 1 - alpha_prod_t_prev
else:
timesteps = torch.tensor(self.timesteps).to(timestep.device)
idx = timestep.reshape(-1, 1).eq(timesteps.reshape(1, -1)).nonzero()[:, 1] # 找到 timestep 在 timesteps 中的索引 idx
# 根据idx找到idx+1对应的timesteps元素，也就是下一个时间步。如果idx+1超出了timesteps的长度，它会被限制在self.num_inference_steps - 1
prev_timestep = timesteps[idx.add(1).clamp_max(self.num_inference_steps - 1)]
assert (prev_timestep is not None)
# 2. compute alphas, betas
alpha_prod_t = self.alphas_cumprod[timestep]
alpha_prod_t_prev = self.alphas_cumprod[prev_timestep]
alpha_prod_t_prev = torch.where(prev_timestep < 0, self.final_alpha_cumprod, alpha_prod_t_prev)
beta_prod_t = 1 - alpha_prod_t
beta_prod_t_prev = 1 - alpha_prod_t_prev
bs = timestep.size(0)
alpha_prod_t = alpha_prod_t.view(bs, 1, 1, 1)
alpha_prod_t_prev = alpha_prod_t_prev.view(bs, 1, 1, 1)
beta_prod_t = beta_prod_t.view(bs, 1, 1, 1)
beta_prod_t_prev = beta_prod_t_prev.view(bs, 1, 1, 1)
# # 2. compute alphas, betas
# alpha_prod_t = self.alphas_cumprod[timestep]
# alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep is not None else self.final_alpha_cumprod
# beta_prod_t = 1 - alpha_prod_t
# beta_prod_t_prev = 1 - alpha_prod_t_prev
# rcfg
self.stock_alpha_prod_t_prev = alpha_prod_t_prev
self.stock_beta_prod_t_prev = beta_prod_t_prev
# rcfg
self.stock_alpha_prod_t_prev = alpha_prod_t_prev
self.stock_beta_prod_t_prev = beta_prod_t_prev
# 3. compute predicted original sample from predicted noise also called
model_output_conversion = self.convert_output(model_output, model_output_type, sample, timestep)
pred_original_sample = model_output_conversion.pred_original_sample
pred_epsilon = model_output_conversion.pred_epsilon
# 4. Clip or threshold "predicted x_0"
if clip_sample:
pred_original_sample = torch.clamp(pred_original_sample, -1, 1)
pred_epsilon = self.convert_output(pred_original_sample, "sample", sample, timestep).pred_epsilon
if dynamic_threshold is not None:
# Dynamic thresholding in https://arxiv.org/abs/2205.11487
dynamic_max_val = pred_original_sample \
.flatten(1) \
.abs() \
.float() \
.quantile(dynamic_threshold, dim=1) \
.type_as(pred_original_sample) \
.clamp_min(1) \
.view(-1, *([1] * (pred_original_sample.ndim - 1)))
pred_original_sample = pred_original_sample.clamp(-dynamic_max_val, dynamic_max_val) / dynamic_max_val
pred_epsilon = self.convert_output(pred_original_sample, "sample", sample, timestep).pred_epsilon
# 5. compute variance: "sigma_t(η)" -> see formula (16) from https://arxiv.org/pdf/2010.02502.pdf
# σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
std_dev_t = eta * variance ** (0.5)
# 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * pred_epsilon
# 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
# 8. add "random noise" if needed.
if eta > 0:
if variance_noise is None:
variance_noise = torch.randn_like(model_output)
prev_sample = prev_sample + std_dev_t * variance_noise
return DDIMSchedulerStepOutput(
prev_sample=prev_sample, # x_{t-1}
pred_original_sample=pred_original_sample # x0
)
def add_noise(
self,
original_samples: torch.Tensor,
noise: torch.Tensor,
timesteps: Union[torch.Tensor, int],
replace_noise=True
) -> torch.Tensor:
alpha_prod_t = self.alphas_cumprod[timesteps].reshape(-1, *([1] * (original_samples.ndim - 1)))
if replace_noise:
indices = (timesteps == 999).nonzero()
if indices.numel() > 0:
alpha_prod_t[indices] = 0
return alpha_prod_t ** (0.5) * original_samples + (1 - alpha_prod_t) ** (0.5) * noise
def add_noise_lcm(
self,
original_samples: torch.Tensor,
noise: torch.Tensor,
timestep: Union[torch.Tensor, int],
) -> torch.Tensor:
if isinstance(timestep, int):
# 1. get previous step value (t-1)
idx = self.timesteps.index(timestep)
prev_timestep = self.timesteps[idx + 1] if idx < self.num_inference_steps - 1 else None
# 2. compute alphas, betas
alpha_prod_t = self.alphas_cumprod[timestep]
alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep is not None else self.final_alpha_cumprod
beta_prod_t = 1 - alpha_prod_t
beta_prod_t_prev = 1 - alpha_prod_t_prev
else:
timesteps = torch.tensor(self.timesteps).to(timestep.device)
idx = timestep.reshape(-1, 1).eq(timesteps.reshape(1, -1)).nonzero()[:, 1] # 找到 timestep 在 timesteps 中的索引 idx
prev_timestep = timesteps[idx.add(1).clamp_max(self.num_inference_steps - 1)]
assert (prev_timestep is not None)
# 2. compute alphas, betas
alpha_prod_t = self.alphas_cumprod[timestep]
alpha_prod_t_prev = self.alphas_cumprod[prev_timestep]
alpha_prod_t_prev = torch.where(prev_timestep < 0, self.final_alpha_cumprod, alpha_prod_t_prev)
beta_prod_t = 1 - alpha_prod_t
beta_prod_t_prev = 1 - alpha_prod_t_prev
bs = timestep.size(0)
alpha_prod_t = alpha_prod_t.view(bs, 1, 1, 1)
alpha_prod_t_prev = alpha_prod_t_prev.view(bs, 1, 1, 1)
beta_prod_t = beta_prod_t.view(bs, 1, 1, 1)
beta_prod_t_prev = beta_prod_t_prev.view(bs, 1, 1, 1)
alpha_prod_t_prev = alpha_prod_t_prev.reshape(-1, *([1] * (original_samples.ndim - 1)))
return alpha_prod_t_prev ** (0.5) * original_samples + (1 - alpha_prod_t_prev) ** (0.5) * noise
def convert_output(
self,
model_output: torch.Tensor,
model_output_type: str,
sample: torch.Tensor,
timesteps: Union[torch.Tensor, int]
) -> DDIMSchedulerConversionOutput:
assert model_output_type in self.prediction_types
alpha_prod_t = self.alphas_cumprod[timesteps].reshape(-1, *([1] * (sample.ndim - 1)))
beta_prod_t = 1 - alpha_prod_t
if model_output_type == "epsilon":
pred_epsilon = model_output
pred_original_sample = (sample - beta_prod_t ** (0.5) * pred_epsilon) / alpha_prod_t ** (0.5)
pred_velocity = alpha_prod_t ** (0.5) * pred_epsilon - (1 - alpha_prod_t) ** (0.5) * pred_original_sample
elif model_output_type == "sample":
pred_original_sample = model_output
pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)
pred_velocity = alpha_prod_t ** (0.5) * pred_epsilon - (1 - alpha_prod_t) ** (0.5) * pred_original_sample
elif model_output_type == "v_prediction":
pred_velocity = model_output
pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
else:
raise ValueError("Unknown prediction type")
return DDIMSchedulerConversionOutput(
pred_epsilon=pred_epsilon,
pred_original_sample=pred_original_sample,
pred_velocity=pred_velocity)
def get_velocity(
self,
sample: torch.Tensor,
noise: torch.Tensor,
timesteps: torch.Tensor
) -> torch.FloatTensor:
alpha_prod_t = self.alphas_cumprod[timesteps].reshape(-1, *([1] * (sample.ndim - 1)))
return alpha_prod_t ** (0.5) * noise - (1 - alpha_prod_t) ** (0.5) * sample