标签：None CogVideoX self torch --- 源码 return sigma def
CogVideo & CogVideoX 微调代码源码解析（十）
.\cogvideo-finetune\sat\sgm\modules\diffusionmodules\sampling.py
# 部分移植自 https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/sampling.py
"""
# 从类型提示模块导入字典和联合类型
from typing import Dict, Union
# 导入 PyTorch 库
import torch
# 从 OmegaConf 导入列表配置和 OmegaConf 类
from omegaconf import ListConfig, OmegaConf
# 从 tqdm 导入进度条功能
from tqdm import tqdm
# 从自定义模块导入必要的函数
from ...modules.diffusionmodules.sampling_utils import (
get_ancestral_step, # 获取祖先步骤的函数
linear_multistep_coeff, # 线性多步骤系数的函数
to_d, # 转换到 d 的函数
to_neg_log_sigma, # 转换为负对数 sigma 的函数
to_sigma, # 转换为 sigma 的函数
)
# 从 util 模块导入辅助函数
from ...util import append_dims, default, instantiate_from_config
from ...util import SeededNoise # 导入带种子的噪声生成器
# 从 guiders 模块导入动态 CFG 类
from .guiders import DynamicCFG
# 默认引导器配置
DEFAULT_GUIDER = {"target": "sgm.modules.diffusionmodules.guiders.IdentityGuider"}
# 定义基础扩散采样器类
class BaseDiffusionSampler:
# 初始化方法，设置相关配置
def __init__(
self,
discretization_config: Union[Dict, ListConfig, OmegaConf], # 离散化配置
num_steps: Union[int, None] = None, # 步骤数
guider_config: Union[Dict, ListConfig, OmegaConf, None] = None, # 引导器配置
verbose: bool = False, # 是否详细输出
device: str = "cuda", # 使用的设备
):
self.num_steps = num_steps # 保存步骤数
# 根据配置实例化离散化对象
self.discretization = instantiate_from_config(discretization_config)
# 根据配置实例化引导器对象，使用默认值如果未提供
self.guider = instantiate_from_config(
default(
guider_config,
DEFAULT_GUIDER,
)
)
self.verbose = verbose # 设置详细输出标志
self.device = device # 设置设备
# 准备采样循环的方法
def prepare_sampling_loop(self, x, cond, uc=None, num_steps=None):
# 计算 sigma 值
sigmas = self.discretization(self.num_steps if num_steps is None else num_steps, device=self.device)
# 默认情况下使用条件输入
uc = default(uc, cond)
# 对输入进行缩放
x *= torch.sqrt(1.0 + sigmas[0] ** 2.0)
num_sigmas = len(sigmas) # 获取 sigma 数量
# 创建与输入样本数相同的全 1 张量
s_in = x.new_ones([x.shape[0]]).float()
# 返回准备好的参数
return x, s_in, sigmas, num_sigmas, cond, uc
# 去噪声的方法
def denoise(self, x, denoiser, sigma, cond, uc):
# 准备输入并进行去噪声处理
denoised = denoiser(*self.guider.prepare_inputs(x, sigma, cond, uc))
# 使用引导器进一步处理去噪声结果
denoised = self.guider(denoised, sigma)
# 返回去噪声后的结果
return denoised
# 获取 sigma 生成器的方法
def get_sigma_gen(self, num_sigmas):
sigma_generator = range(num_sigmas - 1) # 创建 sigma 生成器范围
if self.verbose: # 如果启用详细输出
print("#" * 30, " Sampling setting ", "#" * 30) # 输出分隔符
print(f"Sampler: {self.__class__.__name__}") # 输出采样器类名
print(f"Discretization: {self.discretization.__class__.__name__}") # 输出离散化类名
print(f"Guider: {self.guider.__class__.__name__}") # 输出引导器类名
# 包装 sigma 生成器为 tqdm 对象，以便显示进度条
sigma_generator = tqdm(
sigma_generator,
total=num_sigmas,
desc=f"Sampling with {self.__class__.__name__} for {num_sigmas} steps",
)
# 返回 sigma 生成器
return sigma_generator
# 定义单步扩散采样器类，继承自基础扩散采样器
class SingleStepDiffusionSampler(BaseDiffusionSampler):
# 定义采样步骤的抽象方法
def sampler_step(self, sigma, next_sigma, denoiser, x, cond, uc, *args, **kwargs):
raise NotImplementedError # 抛出未实现错误
# 欧拉步骤的方法
def euler_step(self, x, d, dt):
# 计算下一个状态
return x + dt * d
# 定义 EDM 采样器类，继承自单步扩散采样器
class EDMSampler(SingleStepDiffusionSampler):
# 初始化方法，设置相关参数
def __init__(self, s_churn=0.0, s_tmin=0.0, s_tmax=float("inf"), s_noise=1.0, *args, **kwargs):
super().__init__(*args, **kwargs) # 调用父类初始化
self.s_churn = s_churn # 设置 churn 参数
self.s_tmin = s_tmin # 设置最小时间
self.s_tmax = s_tmax # 设置最大时间
self.s_noise = s_noise # 设置噪声参数
# 定义sampler_step函数，用于执行采样步骤
def sampler_step(self, sigma, next_sigma, denoiser, x, cond, uc=None, gamma=0.0):
# 计算sigma_hat，用于计算噪声
sigma_hat = sigma * (gamma + 1.0)
# 如果gamma大于0，生成服从标准正态分布的随机数eps，并将其乘以s_noise，再加到x上
if gamma > 0:
eps = torch.randn_like(x) * self.s_noise
x = x + eps * append_dims(sigma_hat**2 - sigma**2, x.ndim) ** 0.5
# 使用denoiser对x进行去噪，得到denoised
denoised = self.denoise(x, denoiser, sigma_hat, cond, uc)
# 计算d，用于后续步骤
d = to_d(x, sigma_hat, denoised)
# 计算dt，用于后续步骤
dt = append_dims(next_sigma - sigma_hat, x.ndim)
# 使用欧拉步骤更新x
euler_step = self.euler_step(x, d, dt)
# 使用可能的修正步骤更新x
x = self.possible_correction_step(euler_step, x, d, dt, next_sigma, denoiser, cond, uc)
# 返回更新后的x
return x
# 定义__call__函数，用于执行整个采样过程
def __call__(self, denoiser, x, cond, uc=None, num_steps=None):
# 准备采样循环所需的变量
x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop(x, cond, uc, num_steps)
# 遍历sigmas，执行采样步骤
for i in self.get_sigma_gen(num_sigmas):
# 计算gamma，用于控制噪声的大小
gamma = (
min(self.s_churn / (num_sigmas - 1), 2**0.5 - 1) if self.s_tmin <= sigmas[i] <= self.s_tmax else 0.0
)
# 执行采样步骤
x = self.sampler_step(
s_in * sigmas[i],
s_in * sigmas[i + 1],
denoiser,
x,
cond,
uc,
gamma,
)
# 返回最终的x
return x
# 定义 DDIMSampler 类，继承自 SingleStepDiffusionSampler 类
class DDIMSampler(SingleStepDiffusionSampler):
# 初始化方法，接受 s_noise 参数，默认值为 0.1
def __init__(self, s_noise=0.1, *args, **kwargs):
# 调用父类的初始化方法
super().__init__(*args, **kwargs)
# 设置实例属性 s_noise 为传入的 s_noise 参数值
self.s_noise = s_noise
# sampler_step 方法，接受 sigma、next_sigma、denoiser、x、cond、uc 和 s_noise 参数
def sampler_step(self, sigma, next_sigma, denoiser, x, cond, uc=None, s_noise=0.0):
# 使用 denoiser 对 x 进行去噪，得到 denoised
denoised = self.denoise(x, denoiser, sigma, cond, uc)
# 计算 d，使用 to_d 函数
d = to_d(x, sigma, denoised)
# 计算 dt
dt = append_dims(next_sigma * (1 - s_noise**2) ** 0.5 - sigma, x.ndim)
# 计算 euler_step
euler_step = x + dt * d + s_noise * append_dims(next_sigma, x.ndim) * torch.randn_like(x)
# 调用 possible_correction_step 方法，得到 x
x = self.possible_correction_step(euler_step, x, d, dt, next_sigma, denoiser, cond, uc)
# 返回 x
return x
# __call__ 方法，接受 denoiser、x、cond、uc 和 num_steps 参数
def __call__(self, denoiser, x, cond, uc=None, num_steps=None):
# 调用 prepare_sampling_loop 方法，得到 x、s_in、sigmas、num_sigmas、cond 和 uc
x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop(x, cond, uc, num_steps)
# 遍历 sigmas
for i in self.get_sigma_gen(num_sigmas):
# 调用 sampler_step 方法，得到 x
x = self.sampler_step(
s_in * sigmas[i],
s_in * sigmas[i + 1],
denoiser,
x,
cond,
uc,
self.s_noise,
)
# 返回 x
return x
# 定义 AncestralSampler 类，继承自 SingleStepDiffusionSampler 类
class AncestralSampler(SingleStepDiffusionSampler):
# 初始化方法，接受 eta 和 s_noise 参数，默认值分别为 1.0 和 1.0
def __init__(self, eta=1.0, s_noise=1.0, *args, **kwargs):
# 调用父类的初始化方法
super().__init__(*args, **kwargs)
# 设置实例属性 eta 为传入的 eta 参数值
self.eta = eta
# 设置实例属性 s_noise 为传入的 s_noise 参数值
self.s_noise = s_noise
# 设置实例属性 noise_sampler 为一个 lambda 函数，用于生成噪声
self.noise_sampler = lambda x: torch.randn_like(x)
# ancestral_euler_step 方法，接受 x、denoised、sigma 和 sigma_down 参数
def ancestral_euler_step(self, x, denoised, sigma, sigma_down):
# 计算 d，使用 to_d 函数
d = to_d(x, sigma, denoised)
# 计算 dt
dt = append_dims(sigma_down - sigma, x.ndim)
# 调用 euler_step 方法，得到结果
return self.euler_step(x, d, dt)
# ancestral_step 方法，接受 x、sigma、next_sigma 和 sigma_up 参数
def ancestral_step(self, x, sigma, next_sigma, sigma_up):
# 根据条件进行赋值操作
x = torch.where(
append_dims(next_sigma, x.ndim) > 0.0,
x + self.noise_sampler(x) * self.s_noise * append_dims(sigma_up, x.ndim),
x,
)
# 返回结果
return x
# __call__ 方法，接受 denoiser、x、cond、uc 和 num_steps 参数
def __call__(self, denoiser, x, cond, uc=None, num_steps=None):
# 调用 prepare_sampling_loop 方法，得到 x、s_in、sigmas、num_sigmas、cond 和 uc
x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop(x, cond, uc, num_steps)
# 遍历 sigmas
for i in self.get_sigma_gen(num_sigmas):
# 调用 sampler_step 方法，得到 x
x = self.sampler_step(
s_in * sigmas[i],
s_in * sigmas[i + 1],
denoiser,
x,
cond,
uc,
)
# 返回 x
return x
# 定义 LinearMultistepSampler 类，继承自 BaseDiffusionSampler 类
class LinearMultistepSampler(BaseDiffusionSampler):
# 初始化方法，接受 order 参数，默认值为 4
def __init__(
self,
order=4,
*args,
**kwargs,
):
# 调用父类的初始化方法
super().__init__(*args, **kwargs)
# 设置实例属性 order 为传入的 order 参数值
self.order = order
# 定义可调用方法，用于执行去噪操作
def __call__(self, denoiser, x, cond, uc=None, num_steps=None, **kwargs):
# 准备采样循环，返回处理后的输入、sigma、以及条件信息
x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop(x, cond, uc, num_steps)
# 初始化去噪结果列表
ds = []
# 将 sigma 从计算图中分离并转移到 CPU，然后转换为 NumPy 数组
sigmas_cpu = sigmas.detach().cpu().numpy()
# 遍历生成的 sigma 值
for i in self.get_sigma_gen(num_sigmas):
# 计算当前 sigma
sigma = s_in * sigmas[i]
# 使用去噪器处理输入数据，获取去噪结果
denoised = denoiser(*self.guider.prepare_inputs(x, sigma, cond, uc), **kwargs)
# 进一步处理去噪结果
denoised = self.guider(denoised, sigma)
# 将当前输入、sigma 和去噪结果转换为目标格式
d = to_d(x, sigma, denoised)
# 将去噪结果添加到结果列表
ds.append(d)
# 如果结果列表超过预设顺序，则移除最旧的结果
if len(ds) > self.order:
ds.pop(0)
# 计算当前顺序
cur_order = min(i + 1, self.order)
# 计算线性多步法的系数
coeffs = [linear_multistep_coeff(cur_order, sigmas_cpu, i, j) for j in range(cur_order)]
# 更新输入 x，使用加权和的方式结合去噪结果
x = x + sum(coeff * d for coeff, d in zip(coeffs, reversed(ds)))
# 返回最终的去噪结果
return x
# 定义一个基于 Euler 方法的 EDM 采样器类，继承自 EDMSampler
class EulerEDMSampler(EDMSampler):
# 可能的修正步骤，接受多个参数
def possible_correction_step(self, euler_step, x, d, dt, next_sigma, denoiser, cond, uc):
# 直接返回 euler_step，没有进行任何修正
return euler_step
# 定义一个基于 Heun 方法的 EDM 采样器类，继承自 EDMSampler
class HeunEDMSampler(EDMSampler):
# 可能的修正步骤，接受多个参数
def possible_correction_step(self, euler_step, x, d, dt, next_sigma, denoiser, cond, uc):
# 检查 next_sigma 的总和是否小于 1e-14
if torch.sum(next_sigma) < 1e-14:
# 如果所有噪声水平为 0，则返回 euler_step，避免网络评估
return euler_step
else:
# 使用 denoiser 去噪 euler_step
denoised = self.denoise(euler_step, denoiser, next_sigma, cond, uc)
# 将 euler_step 和去噪结果转换为新的 d 值
d_new = to_d(euler_step, next_sigma, denoised)
# 计算 d 的新值，取 d 和 d_new 的平均
d_prime = (d + d_new) / 2.0
# 如果噪声水平不为 0，则应用修正
x = torch.where(append_dims(next_sigma, x.ndim) > 0.0, x + d_prime * dt, euler_step)
# 返回修正后的 x
return x
# 定义一个基于 Euler 的祖先采样器类，继承自 AncestralSampler
class EulerAncestralSampler(AncestralSampler):
# 进行采样步骤，接受多个参数
def sampler_step(self, sigma, next_sigma, denoiser, x, cond, uc):
# 获取祖先步骤的上下界
sigma_down, sigma_up = get_ancestral_step(sigma, next_sigma, eta=self.eta)
# 去噪 x
denoised = self.denoise(x, denoiser, sigma, cond, uc)
# 进行 Euler 采样步骤
x = self.ancestral_euler_step(x, denoised, sigma, sigma_down)
# 进行祖先步骤
x = self.ancestral_step(x, sigma, next_sigma, sigma_up)
# 返回最终的 x
return x
# 定义一个基于 DPM++ 的祖先采样器类，继承自 AncestralSampler
class DPMPP2SAncestralSampler(AncestralSampler):
# 获取变量，接受 sigma 和 sigma_down
def get_variables(self, sigma, sigma_down):
# 将 sigma 和 sigma_down 转换为负对数
t, t_next = [to_neg_log_sigma(s) for s in (sigma, sigma_down)]
# 计算时间差 h
h = t_next - t
# 计算 s 值
s = t + 0.5 * h
# 返回 h, s, t, t_next
return h, s, t, t_next
# 计算多重值，接受 h, s, t, t_next
def get_mult(self, h, s, t, t_next):
# 计算多个乘数
mult1 = to_sigma(s) / to_sigma(t)
mult2 = (-0.5 * h).expm1()
mult3 = to_sigma(t_next) / to_sigma(t)
mult4 = (-h).expm1()
# 返回所有计算出的乘数
return mult1, mult2, mult3, mult4
# 进行采样步骤，接受多个参数
def sampler_step(self, sigma, next_sigma, denoiser, x, cond, uc=None, **kwargs):
# 获取祖先步骤的上下界
sigma_down, sigma_up = get_ancestral_step(sigma, next_sigma, eta=self.eta)
# 去噪 x
denoised = self.denoise(x, denoiser, sigma, cond, uc)
# 进行 Euler 采样步骤
x_euler = self.ancestral_euler_step(x, denoised, sigma, sigma_down)
# 检查 sigma_down 的总和是否小于 1e-14
if torch.sum(sigma_down) < 1e-14:
# 如果所有噪声水平为 0，则返回 x_euler，避免网络评估
x = x_euler
else:
# 获取变量
h, s, t, t_next = self.get_variables(sigma, sigma_down)
# 计算多重值并调整维度
mult = [append_dims(mult, x.ndim) for mult in self.get_mult(h, s, t, t_next)]
# 计算新的 x2
x2 = mult[0] * x - mult[1] * denoised
# 对 x2 进行去噪
denoised2 = self.denoise(x2, denoiser, to_sigma(s), cond, uc)
# 计算最终的 x_dpmpp2s
x_dpmpp2s = mult[2] * x - mult[3] * denoised2
# 如果噪声水平不为 0，则应用修正
x = torch.where(append_dims(sigma_down, x.ndim) > 0.0, x_dpmpp2s, x_euler)
# 进行祖先步骤
x = self.ancestral_step(x, sigma, next_sigma, sigma_up)
# 返回最终的 x
return x
# 定义一个基于 DPM++ 的采样器类，继承自 BaseDiffusionSampler
class DPMPP2MSampler(BaseDiffusionSampler):
# 定义一个获取变量的函数，接受当前和下一个噪声级别，以及可选的上一个噪声级别
def get_variables(self, sigma, next_sigma, previous_sigma=None):
# 将 sigma 和 next_sigma 转换为负对数形式
t, t_next = [to_neg_log_sigma(s) for s in (sigma, next_sigma)]
# 计算两个时间点之间的差值
h = t_next - t
# 如果上一个噪声级别存在
if previous_sigma is not None:
# 计算当前时间与上一个时间的差值
h_last = t - to_neg_log_sigma(previous_sigma)
# 计算当前和上一个时间差的比值
r = h_last / h
# 返回差值 h、比值 r、当前和下一个时间
return h, r, t, t_next
else:
# 返回差值 h 和当前、下一个时间，但不返回比值 r
return h, None, t, t_next
# 定义一个获取乘数的函数，接受多个参数
def get_mult(self, h, r, t, t_next, previous_sigma):
# 计算当前和下一个时间的 sigma 乘数
mult1 = to_sigma(t_next) / to_sigma(t)
# 计算 h 的负值的指数减一
mult2 = (-h).expm1()
# 如果上一个噪声级别存在
if previous_sigma is not None:
# 计算与 r 相关的乘数
mult3 = 1 + 1 / (2 * r)
mult4 = 1 / (2 * r)
# 返回所有乘数
return mult1, mult2, mult3, mult4
else:
# 返回前两个乘数
return mult1, mult2
# 定义采样步骤函数，接受多个参数
def sampler_step(
self,
old_denoised,
previous_sigma,
sigma,
next_sigma,
denoiser,
x,
cond,
uc=None,
):
# 使用去噪器对输入 x 进行去噪
denoised = self.denoise(x, denoiser, sigma, cond, uc)
# 获取变量 h、r、t、t_next
h, r, t, t_next = self.get_variables(sigma, next_sigma, previous_sigma)
# 获取乘数，并将维度调整以匹配 x 的维度
mult = [append_dims(mult, x.ndim) for mult in self.get_mult(h, r, t, t_next, previous_sigma)]
# 计算标准化后的 x
x_standard = mult[0] * x - mult[1] * denoised
# 如果没有旧的去噪结果或下一个噪声级别接近零
if old_denoised is None or torch.sum(next_sigma) < 1e-14:
# 保存网络评估，如果所有噪声级别为 0 或为第一步
return x_standard, denoised
else:
# 计算去噪后的旧结果
denoised_d = mult[2] * denoised - mult[3] * old_denoised
# 计算高级 x
x_advanced = mult[0] * x - mult[1] * denoised_d
# 如果噪声级别不为 0 且不是第一步，则应用修正
x = torch.where(append_dims(next_sigma, x.ndim) > 0.0, x_advanced, x_standard)
# 返回最终的 x 和去噪结果
return x, denoised
# 定义调用函数，接受去噪器、输入 x、条件和其他可选参数
def __call__(self, denoiser, x, cond, uc=None, num_steps=None, **kwargs):
# 准备采样循环的参数，包括对输入 x、条件、噪声级别等的处理
x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop(x, cond, uc, num_steps)
# 初始化旧去噪结果
old_denoised = None
# 遍历生成的噪声级别
for i in self.get_sigma_gen(num_sigmas):
# 执行采样步骤，更新 x 和旧去噪结果
x, old_denoised = self.sampler_step(
old_denoised,
None if i == 0 else s_in * sigmas[i - 1],
s_in * sigmas[i],
s_in * sigmas[i + 1],
denoiser,
x,
cond,
uc=uc,
)
# 返回最终的 x
return x
# 定义 SDEDPMPP2MSampler 类，继承自 BaseDiffusionSampler
class SDEDPMPP2MSampler(BaseDiffusionSampler):
# 获取变量 h、r 和时间参数 t、t_next
def get_variables(self, sigma, next_sigma, previous_sigma=None):
# 将 sigma 和 next_sigma 转换为负对数形式
t, t_next = [to_neg_log_sigma(s) for s in (sigma, next_sigma)]
# 计算 h 为 t_next 和 t 的差值
h = t_next - t
# 如果 previous_sigma 不为 None
if previous_sigma is not None:
# 计算上一个 sigma 的负对数值
h_last = t - to_neg_log_sigma(previous_sigma)
# 计算 r 为 h_last 和 h 的比值
r = h_last / h
# 返回 h、r、t 和 t_next
return h, r, t, t_next
else:
# 返回 h 和 None（无 r），以及 t 和 t_next
return h, None, t, t_next
# 计算乘数值
def get_mult(self, h, r, t, t_next, previous_sigma):
# 计算 mult1 为 t_next 和 t 的 sigma 比值乘以 h 的负指数
mult1 = to_sigma(t_next) / to_sigma(t) * (-h).exp()
# 计算 mult2 为 (-2*h) 的 expm1 值
mult2 = (-2 * h).expm1()
# 如果 previous_sigma 不为 None
if previous_sigma is not None:
# 计算 mult3 为 1 + 1/(2*r)
mult3 = 1 + 1 / (2 * r)
# 计算 mult4 为 1/(2*r)
mult4 = 1 / (2 * r)
# 返回 mult1、mult2、mult3 和 mult4
return mult1, mult2, mult3, mult4
else:
# 返回 mult1 和 mult2
return mult1, mult2
# 执行采样步骤
def sampler_step(
self,
old_denoised,
previous_sigma,
sigma,
next_sigma,
denoiser,
x,
cond,
uc=None,
):
# 使用 denoiser 对 x 进行去噪处理
denoised = self.denoise(x, denoiser, sigma, cond, uc)
# 获取 h、r、t 和 t_next
h, r, t, t_next = self.get_variables(sigma, next_sigma, previous_sigma)
# 计算乘数，并调整维度
mult = [append_dims(mult, x.ndim) for mult in self.get_mult(h, r, t, t_next, previous_sigma)]
# 计算噪声乘数并调整维度
mult_noise = append_dims(next_sigma * (1 - (-2 * h).exp()) ** 0.5, x.ndim)
# 计算标准化后的 x
x_standard = mult[0] * x - mult[1] * denoised + mult_noise * torch.randn_like(x)
# 如果 old_denoised 为 None 或 next_sigma 的和小于 1e-14
if old_denoised is None or torch.sum(next_sigma) < 1e-14:
# 返回标准化后的 x 和去噪后的结果
return x_standard, denoised
else:
# 计算去噪的差异
denoised_d = mult[2] * denoised - mult[3] * old_denoised
# 计算高级 x
x_advanced = mult[0] * x - mult[1] * denoised_d + mult_noise * torch.randn_like(x)
# 如果噪声水平不为 0 且不是第一步，应用修正
x = torch.where(append_dims(next_sigma, x.ndim) > 0.0, x_advanced, x_standard)
# 返回最终的 x 和去噪后的结果
return x, denoised
# 调用采样器，执行采样循环
def __call__(self, denoiser, x, cond, uc=None, num_steps=None, scale=None, **kwargs):
# 准备采样循环，初始化输入和 sigma
x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop(x, cond, uc, num_steps)
# 初始化 old_denoised 为 None
old_denoised = None
# 遍历 sigma 生成器
for i in self.get_sigma_gen(num_sigmas):
# 执行采样步骤并更新 x 和 old_denoised
x, old_denoised = self.sampler_step(
old_denoised,
None if i == 0 else s_in * sigmas[i - 1],
s_in * sigmas[i],
s_in * sigmas[i + 1],
denoiser,
x,
cond,
uc=uc,
)
# 返回最终的 x
return x
# 定义 SdeditEDMSampler 类，继承自 EulerEDMSampler
class SdeditEDMSampler(EulerEDMSampler):
# 初始化函数，设置编辑比例
def __init__(self, edit_ratio=0.5, *args, **kwargs):
# 调用父类的初始化函数
super().__init__(*args, **kwargs)
# 设置编辑比例
self.edit_ratio = edit_ratio
# 定义一个可调用的方法，接受多个参数进行图像去噪
def __call__(self, denoiser, image, randn, cond, uc=None, num_steps=None, edit_ratio=None):
# 克隆 randn，创建 randn_unit 用于后续计算
randn_unit = randn.clone()
# 准备采样循环，处理 randn、条件、未条件和步骤数
randn, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop(randn, cond, uc, num_steps)
# 如果未指定 num_steps，则使用对象的默认步骤数
if num_steps is None:
num_steps = self.num_steps
# 如果未指定 edit_ratio，则使用对象的默认编辑比例
if edit_ratio is None:
edit_ratio = self.edit_ratio
# 初始化 x 为 None，用于后续存储结果
x = None
# 遍历 sigma 生成器，获取每个 sigma 的值
for i in self.get_sigma_gen(num_sigmas):
# 如果当前步骤比例小于 edit_ratio，则跳过此次循环
if i / num_steps < edit_ratio:
continue
# 如果 x 为 None，则初始化 x 为图像与噪声的组合
if x is None:
x = image + randn_unit * append_dims(s_in * sigmas[i], len(randn_unit.shape))
# 计算 gamma 值，依据条件限制调整
gamma = (
min(self.s_churn / (num_sigmas - 1), 2**0.5 - 1) if self.s_tmin <= sigmas[i] <= self.s_tmax else 0.0
)
# 进行一次采样步骤，更新 x 的值
x = self.sampler_step(
s_in * sigmas[i], # 当前 sigma 的输入
s_in * sigmas[i + 1], # 下一个 sigma 的输入
denoiser, # 去噪器
x, # 当前图像
cond, # 条件信息
uc, # 未条件信息
gamma, # gamma 值
)
# 返回最终处理后的图像
return x
# 定义一个名为 VideoDDIMSampler 的类，继承自 BaseDiffusionSampler
class VideoDDIMSampler(BaseDiffusionSampler):
# 初始化函数，接受固定帧数和 sdedit 标志，及其他参数
def __init__(self, fixed_frames=0, sdedit=False, **kwargs):
# 调用父类的初始化函数
super().__init__(**kwargs)
# 设置固定帧数
self.fixed_frames = fixed_frames
# 设置 sdedit 标志
self.sdedit = sdedit
# 准备采样循环，接受输入数据和条件
def prepare_sampling_loop(self, x, cond, uc=None, num_steps=None):
# 进行离散化，计算 alpha 的平方根累积乘积和时间步
alpha_cumprod_sqrt, timesteps = self.discretization(
self.num_steps if num_steps is None else num_steps, # 使用给定的步数或默认步数
device=self.device, # 指定设备
return_idx=True, # 返回索引
do_append_zero=False, # 不追加零
)
# 在 alpha_cumprod_sqrt 末尾添加一个值为 1 的新张量
alpha_cumprod_sqrt = torch.cat([alpha_cumprod_sqrt, alpha_cumprod_sqrt.new_ones([1])])
# 创建一个新的时间步张量，并在开头添加一个值为 -1 的零张量
timesteps = torch.cat([torch.tensor(list(timesteps)).new_zeros([1]) - 1, torch.tensor(list(timesteps))])
# 如果 uc 为空，使用 cond 作为默认值
uc = default(uc, cond)
# 计算 alpha_cumprod_sqrt 的元素数量
num_sigmas = len(alpha_cumprod_sqrt)
# 创建一个新的张量 s_in，初始值为 1
s_in = x.new_ones([x.shape[0]])
# 返回多个变量
return x, s_in, alpha_cumprod_sqrt, num_sigmas, cond, uc, timesteps
# 去噪函数，接受多个输入参数
def denoise(self, x, denoiser, alpha_cumprod_sqrt, cond, uc, timestep=None, idx=None, scale=None, scale_emb=None):
# 初始化额外模型输入的字典
additional_model_inputs = {}
# 检查 scale 是否为张量且不为 1
if isinstance(scale, torch.Tensor) == False and scale == 1:
# 为额外模型输入添加当前时间步的索引
additional_model_inputs["idx"] = x.new_ones([x.shape[0]]) * timestep
# 如果 scale_emb 不为 None，添加到额外输入
if scale_emb is not None:
additional_model_inputs["scale_emb"] = scale_emb
# 调用去噪器进行去噪，并转换为 float32 类型
denoised = denoiser(x, alpha_cumprod_sqrt, cond, **additional_model_inputs).to(torch.float32)
else:
# 创建一个新的索引张量，包含当前时间步的重复值
additional_model_inputs["idx"] = torch.cat([x.new_ones([x.shape[0]]) * timestep] * 2)
# 调用去噪器进行去噪，准备输入并转换为 float32 类型
denoised = denoiser(
*self.guider.prepare_inputs(x, alpha_cumprod_sqrt, cond, uc), **additional_model_inputs
).to(torch.float32)
# 如果 guider 是 DynamicCFG 的实例，进行动态调整
if isinstance(self.guider, DynamicCFG):
denoised = self.guider(
denoised, (1 - alpha_cumprod_sqrt**2) ** 0.5, step_index=self.num_steps - timestep, scale=scale
)
else:
# 否则，进行普通的调整
denoised = self.guider(denoised, (1 - alpha_cumprod_sqrt**2) ** 0.5, scale=scale)
# 返回去噪后的结果
return denoised
# 采样步骤函数，接受多个输入参数
def sampler_step(
self,
alpha_cumprod_sqrt,
next_alpha_cumprod_sqrt,
denoiser,
x,
cond,
uc=None,
idx=None,
timestep=None,
scale=None,
scale_emb=None,
):
# 调用 denoise 方法获取去噪结果
denoised = self.denoise(
x, denoiser, alpha_cumprod_sqrt, cond, uc, timestep, idx, scale=scale, scale_emb=scale_emb
).to(torch.float32)
# 计算 a_t 和 b_t 值
a_t = ((1 - next_alpha_cumprod_sqrt**2) / (1 - alpha_cumprod_sqrt**2)) ** 0.5
b_t = next_alpha_cumprod_sqrt - alpha_cumprod_sqrt * a_t
# 更新 x 的值，通过加权当前值和去噪后的值
x = append_dims(a_t, x.ndim) * x + append_dims(b_t, x.ndim) * denoised
# 返回更新后的 x
return x
# 定义一个可调用的方法，用于处理去噪过程
def __call__(self, denoiser, x, cond, uc=None, num_steps=None, scale=None, scale_emb=None):
# 准备采样循环的输入，返回处理后的数据和相关参数
x, s_in, alpha_cumprod_sqrt, num_sigmas, cond, uc, timesteps = self.prepare_sampling_loop(
x, cond, uc, num_steps
)
# 遍历生成的 sigma 值
for i in self.get_sigma_gen(num_sigmas):
# 执行采样步骤，更新输入数据
x = self.sampler_step(
s_in * alpha_cumprod_sqrt[i], # 当前 sigma 的缩放输入
s_in * alpha_cumprod_sqrt[i + 1], # 下一个 sigma 的缩放输入
denoiser, # 去噪器对象
x, # 当前输入
cond, # 条件输入
uc, # 可选的额外条件
idx=self.num_steps - i, # 当前步骤索引
timestep=timesteps[-(i + 1)], # 当前时间步
scale=scale, # 缩放因子
scale_emb=scale_emb, # 嵌入的缩放因子
)
# 返回处理后的结果
return x
# 定义 VPSDEDPMPP2M 采样器类，继承自 VideoDDIMSampler
class VPSDEDPMPP2MSampler(VideoDDIMSampler):
# 获取变量，计算多个参数
def get_variables(self, alpha_cumprod_sqrt, next_alpha_cumprod_sqrt, previous_alpha_cumprod_sqrt=None):
# 计算 alpha 的累积乘积
alpha_cumprod = alpha_cumprod_sqrt**2
# 计算 lamb 的对数值
lamb = ((alpha_cumprod / (1 - alpha_cumprod)) ** 0.5).log()
# 计算下一个 alpha 的累积乘积
next_alpha_cumprod = next_alpha_cumprod_sqrt**2
# 计算下一个 lamb 的对数值
lamb_next = ((next_alpha_cumprod / (1 - next_alpha_cumprod)) ** 0.5).log()
# 计算 h 值
h = lamb_next - lamb
# 如果存在前一个 alpha 的累积乘积
if previous_alpha_cumprod_sqrt is not None:
# 计算前一个 alpha 的累积乘积
previous_alpha_cumprod = previous_alpha_cumprod_sqrt**2
# 计算前一个 lamb 的对数值
lamb_previous = ((previous_alpha_cumprod / (1 - previous_alpha_cumprod)) ** 0.5).log()
# 计算 h_last 值
h_last = lamb - lamb_previous
# 计算 r 值
r = h_last / h
# 返回 h、r、lamb 和 lamb_next
return h, r, lamb, lamb_next
else:
# 返回 h、None、lamb 和 lamb_next
return h, None, lamb, lamb_next
# 计算乘数
def get_mult(self, h, r, alpha_cumprod_sqrt, next_alpha_cumprod_sqrt, previous_alpha_cumprod_sqrt):
# 计算第一个乘数
mult1 = ((1 - next_alpha_cumprod_sqrt**2) / (1 - alpha_cumprod_sqrt**2)) ** 0.5 * (-h).exp()
# 计算第二个乘数
mult2 = (-2 * h).expm1() * next_alpha_cumprod_sqrt
# 如果存在前一个 alpha 的累积乘积
if previous_alpha_cumprod_sqrt is not None:
# 计算第三个乘数
mult3 = 1 + 1 / (2 * r)
# 计算第四个乘数
mult4 = 1 / (2 * r)
# 返回所有乘数
return mult1, mult2, mult3, mult4
else:
# 返回前两个乘数
return mult1, mult2
# 执行采样步骤
def sampler_step(
self,
old_denoised,
previous_alpha_cumprod_sqrt,
alpha_cumprod_sqrt,
next_alpha_cumprod_sqrt,
denoiser,
x,
cond,
uc=None,
idx=None,
timestep=None,
scale=None,
scale_emb=None,
):
# 使用去噪器处理输入，得到去噪后的结果
denoised = self.denoise(
x, denoiser, alpha_cumprod_sqrt, cond, uc, timestep, idx, scale=scale, scale_emb=scale_emb
).to(torch.float32)
# 如果索引为 1，返回去噪结果
if idx == 1:
return denoised, denoised
# 获取相关变量
h, r, lamb, lamb_next = self.get_variables(
alpha_cumprod_sqrt, next_alpha_cumprod_sqrt, previous_alpha_cumprod_sqrt
)
# 获取乘数
mult = [
append_dims(mult, x.ndim)
for mult in self.get_mult(h, r, alpha_cumprod_sqrt, next_alpha_cumprod_sqrt, previous_alpha_cumprod_sqrt)
]
# 计算噪声乘数
mult_noise = append_dims((1 - next_alpha_cumprod_sqrt**2) ** 0.5 * (1 - (-2 * h).exp()) ** 0.5, x.ndim)
# 计算标准化 x
x_standard = mult[0] * x - mult[1] * denoised + mult_noise * torch.randn_like(x)
# 如果 old_denoised 为 None 或者下一个 alpha 的累积乘积小于阈值
if old_denoised is None or torch.sum(next_alpha_cumprod_sqrt) < 1e-14:
# 返回标准化的 x 和去噪后的结果
return x_standard, denoised
else:
# 计算去噪后的差异
denoised_d = mult[2] * denoised - mult[3] * old_denoised
# 计算高级 x
x_advanced = mult[0] * x - mult[1] * denoised_d + mult_noise * torch.randn_like(x)
# 更新 x
x = x_advanced
# 返回最终的 x 和去噪结果
return x, denoised
# 定义可调用方法，接收去噪器、输入数据、条件、上采样因子及其它参数
def __call__(self, denoiser, x, cond, uc=None, num_steps=None, scale=None, scale_emb=None):
# 准备采样循环所需的输入及参数
x, s_in, alpha_cumprod_sqrt, num_sigmas, cond, uc, timesteps = self.prepare_sampling_loop(
x, cond, uc, num_steps
)
# 如果固定帧数大于0，提取前固定帧数的图像
if self.fixed_frames > 0:
prefix_frames = x[:, : self.fixed_frames]
# 初始化去噪后的图像为 None
old_denoised = None
# 遍历生成的 sigma 值
for i in self.get_sigma_gen(num_sigmas):
# 如果固定帧数大于0，进行处理
if self.fixed_frames > 0:
# 如果启用 SD 编辑模式
if self.sdedit:
# 生成与前缀帧同形状的随机噪声
rd = torch.randn_like(prefix_frames)
# 计算带噪声的前缀帧
noised_prefix_frames = alpha_cumprod_sqrt[i] * prefix_frames + rd * append_dims(
s_in * (1 - alpha_cumprod_sqrt[i] ** 2) ** 0.5, len(prefix_frames.shape)
)
# 将带噪声的前缀帧与剩余帧连接
x = torch.cat([noised_prefix_frames, x[:, self.fixed_frames :]], dim=1)
else:
# 直接将前缀帧与剩余帧连接
x = torch.cat([prefix_frames, x[:, self.fixed_frames :]], dim=1)
# 执行去噪步骤
x, old_denoised = self.sampler_step(
old_denoised,
None if i == 0 else s_in * alpha_cumprod_sqrt[i - 1],
s_in * alpha_cumprod_sqrt[i],
s_in * alpha_cumprod_sqrt[i + 1],
denoiser,
x,
cond,
uc=uc,
idx=self.num_steps - i,
timestep=timesteps[-(i + 1)],
scale=scale,
scale_emb=scale_emb,
)
# 如果固定帧数大于0，重构最终输出
if self.fixed_frames > 0:
x = torch.cat([prefix_frames, x[:, self.fixed_frames :]], dim=1)
# 返回最终的去噪结果
return x
# 定义 VPODEDPMPP2MSampler 类，继承自 VideoDDIMSampler
class VPODEDPMPP2MSampler(VideoDDIMSampler):
# 获取变量，计算当前和下一个 alpha 的平方根
def get_variables(self, alpha_cumprod_sqrt, next_alpha_cumprod_sqrt, previous_alpha_cumprod_sqrt=None):
# 计算 alpha 的平方
alpha_cumprod = alpha_cumprod_sqrt**2
# 计算 lambda 值并取其对数
lamb = ((alpha_cumprod / (1 - alpha_cumprod)) ** 0.5).log()
# 计算下一个 alpha 的平方
next_alpha_cumprod = next_alpha_cumprod_sqrt**2
# 计算下一个 lambda 值并取其对数
lamb_next = ((next_alpha_cumprod / (1 - next_alpha_cumprod)) ** 0.5).log()
# 计算 h 值
h = lamb_next - lamb
# 如果提供了上一个 alpha 的平方根
if previous_alpha_cumprod_sqrt is not None:
# 计算上一个 alpha 的平方
previous_alpha_cumprod = previous_alpha_cumprod_sqrt**2
# 计算上一个 lambda 值并取其对数
lamb_previous = ((previous_alpha_cumprod / (1 - previous_alpha_cumprod)) ** 0.5).log()
# 计算上一个 h 值
h_last = lamb - lamb_previous
# 计算 r 值
r = h_last / h
# 返回 h, r, lamb, lamb_next
return h, r, lamb, lamb_next
else:
# 如果没有上一个 alpha，返回 h 和其他计算值
return h, None, lamb, lamb_next
# 获取乘数
def get_mult(self, h, r, alpha_cumprod_sqrt, next_alpha_cumprod_sqrt, previous_alpha_cumprod_sqrt):
# 计算第一个乘数
mult1 = ((1 - next_alpha_cumprod_sqrt**2) / (1 - alpha_cumprod_sqrt**2)) ** 0.5
# 计算第二个乘数
mult2 = (-h).expm1() * next_alpha_cumprod_sqrt
# 如果提供了上一个 alpha 的平方根
if previous_alpha_cumprod_sqrt is not None:
# 计算第三个乘数
mult3 = 1 + 1 / (2 * r)
# 计算第四个乘数
mult4 = 1 / (2 * r)
# 返回所有乘数
return mult1, mult2, mult3, mult4
else:
# 返回前两个乘数
return mult1, mult2
# 采样步骤
def sampler_step(
self,
old_denoised,
previous_alpha_cumprod_sqrt,
alpha_cumprod_sqrt,
next_alpha_cumprod_sqrt,
denoiser,
x,
cond,
uc=None,
idx=None,
timestep=None,
):
# 使用去噪器对输入 x 进行去噪处理
denoised = self.denoise(x, denoiser, alpha_cumprod_sqrt, cond, uc, timestep, idx).to(torch.float32)
# 如果索引为 1，返回去噪结果
if idx == 1:
return denoised, denoised
# 获取变量
h, r, lamb, lamb_next = self.get_variables(
alpha_cumprod_sqrt, next_alpha_cumprod_sqrt, previous_alpha_cumprod_sqrt
)
# 获取乘数并调整维度
mult = [
append_dims(mult, x.ndim)
for mult in self.get_mult(h, r, alpha_cumprod_sqrt, next_alpha_cumprod_sqrt, previous_alpha_cumprod_sqrt)
]
# 计算标准化的 x
x_standard = mult[0] * x - mult[1] * denoised
# 如果没有旧的去噪结果或下一个 alpha 的平方根总和接近 0
if old_denoised is None or torch.sum(next_alpha_cumprod_sqrt) < 1e-14:
# 返回标准化的 x 和去噪结果
return x_standard, denoised
else:
# 计算去噪后的结果
denoised_d = mult[2] * denoised - mult[3] * old_denoised
# 计算高级的 x
x_advanced = mult[0] * x - mult[1] * denoised_d
# 更新 x
x = x_advanced
# 返回最终的 x 和去噪结果
return x, denoised
# 定义可调用对象，接受去噪器、输入数据、条件及其他参数
def __call__(self, denoiser, x, cond, uc=None, num_steps=None, scale=None, **kwargs):
# 准备采样循环所需的输入数据和参数
x, s_in, alpha_cumprod_sqrt, num_sigmas, cond, uc, timesteps = self.prepare_sampling_loop(
x, cond, uc, num_steps
)
# 初始化旧去噪结果为 None
old_denoised = None
# 遍历生成的 sigma 值，进行采样步骤
for i in self.get_sigma_gen(num_sigmas):
# 执行单步采样，并更新当前输入和旧去噪结果
x, old_denoised = self.sampler_step(
old_denoised,
None if i == 0 else s_in * alpha_cumprod_sqrt[i - 1], # 第一步不使用旧去噪
s_in * alpha_cumprod_sqrt[i], # 当前 sigma 的值
s_in * alpha_cumprod_sqrt[i + 1], # 下一个 sigma 的值
denoiser, # 去噪器
x, # 当前输入数据
cond, # 条件输入
uc=uc, # 额外条件（可选）
idx=self.num_steps - i, # 当前步骤索引
timestep=timesteps[-(i + 1)], # 当前时间步
)
# 返回最终生成的输入数据
return x
.\cogvideo-finetune\sat\sgm\modules\diffusionmodules\sampling_utils.py
# 导入 PyTorch 库
import torch
# 从 scipy 导入积分功能
from scipy import integrate
# 从上层模块导入 append_dims 函数
from ...util import append_dims
# 从 einops 导入 rearrange 函数，用于重排张量
from einops import rearrange
# 定义无动态阈值类
class NoDynamicThresholding:
# 定义调用方法，接受无条件输入、条件输入和缩放因子
def __call__(self, uncond, cond, scale):
# 如果 scale 是张量，则调整其维度以匹配条件的维度
scale = append_dims(scale, cond.ndim) if isinstance(scale, torch.Tensor) else scale
# 返回无条件输入和缩放后的条件与无条件差值之和
return uncond + scale * (cond - uncond)
# 定义静态阈值类
class StaticThresholding:
# 定义调用方法，接受无条件输入、条件输入和缩放因子
def __call__(self, uncond, cond, scale):
# 计算无条件输入和缩放后的条件与无条件差值之和
result = uncond + scale * (cond - uncond)
# 将结果限制在 -1.0 到 1.0 之间
result = torch.clamp(result, min=-1.0, max=1.0)
# 返回处理后的结果
return result
# 定义动态阈值函数
def dynamic_threshold(x, p=0.95):
# 获取输入张量的维度
N, T, C, H, W = x.shape
# 将张量重排为适合计算的形状
x = rearrange(x, "n t c h w -> n c (t h w)")
# 计算给定分位数的左侧和右侧阈值
l, r = x.quantile(q=torch.tensor([1 - p, p], device=x.device), dim=-1, keepdim=True)
# 计算阈值的最大值
s = torch.maximum(-l, r)
# 创建阈值掩码，用于过滤
threshold_mask = (s > 1).expand(-1, -1, H * W * T)
# 如果阈值掩码中有任何 True 值，则应用阈值处理
if threshold_mask.any():
x = torch.where(threshold_mask, x.clamp(min=-1 * s, max=s), x)
# 恢复张量的原始形状
x = rearrange(x, "n c (t h w) -> n t c h w", t=T, h=H, w=W)
# 返回处理后的张量
return x
# 定义第二种动态阈值处理函数
def dynamic_thresholding2(x0):
p = 0.995 # 参考论文“Imagen”中的超参数
# 保存输入张量的原始数据类型
origin_dtype = x0.dtype
# 将输入张量转换为浮点数
x0 = x0.to(torch.float32)
# 计算输入张量绝对值的分位数
s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
# 将分位数与 1 进行比较并调整维度
s = append_dims(torch.maximum(s, torch.ones_like(s).to(s.device)), x0.dim())
# 限制输入张量的值在 -s 和 s 之间
x0 = torch.clamp(x0, -s, s) # / s
# 返回转换为原始数据类型的张量
return x0.to(origin_dtype)
# 定义潜在动态阈值处理函数
def latent_dynamic_thresholding(x0):
p = 0.9995 # 参考论文中的超参数
# 保存输入张量的原始数据类型
origin_dtype = x0.dtype
# 将输入张量转换为浮点数
x0 = x0.to(torch.float32)
# 计算输入张量绝对值的分位数
s = torch.quantile(torch.abs(x0), p, dim=2)
# 调整分位数的维度以匹配输入张量
s = append_dims(s, x0.dim())
# 限制输入张量的值在 -s 和 s 之间，并进行归一化处理
x0 = torch.clamp(x0, -s, s) / s
# 返回转换为原始数据类型的张量
return x0.to(origin_dtype)
# 定义第三种动态阈值处理函数
def dynamic_thresholding3(x0):
p = 0.995 # 参考论文“Imagen”中的超参数
# 保存输入张量的原始数据类型
origin_dtype = x0.dtype
# 将输入张量转换为浮点数
x0 = x0.to(torch.float32)
# 计算输入张量绝对值的分位数
s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
# 将分位数与 1 进行比较并调整维度
s = append_dims(torch.maximum(s, torch.ones_like(s).to(s.device)), x0.dim())
# 限制输入张量的值在 -s 和 s 之间
x0 = torch.clamp(x0, -s, s) # / s
# 返回转换为原始数据类型的张量
return x0.to(origin_dtype)
# 定义动态阈值类
class DynamicThresholding:
# 定义调用方法，接受无条件输入、条件输入和缩放因子
def __call__(self, uncond, cond, scale):
# 计算无条件输入的均值和标准差
mean = uncond.mean()
std = uncond.std()
# 计算无条件输入和缩放后的条件与无条件差值之和
result = uncond + scale * (cond - uncond)
# 计算结果的均值和标准差
result_mean, result_std = result.mean(), result.std()
# 标准化结果，使其具有相同的标准差
result = (result - result_mean) / result_std * std
# result = dynamic_thresholding3(result) # 可选的进一步处理
# 返回处理后的结果
return result
# 定义动态阈值版本1类
class DynamicThresholdingV1:
# 初始化时接收缩放因子
def __init__(self, scale_factor):
self.scale_factor = scale_factor
# 定义一个函数，接受三个参数 uncond, cond, scale
def __call__(self, uncond, cond, scale):
# 计算结果，根据公式 uncond + scale * (cond - uncond)
result = uncond + scale * (cond - uncond)
# 对结果进行反缩放，除以缩放因子
unscaled_result = result / self.scale_factor
# 获取结果的形状信息，分别为 Batch size, Time steps, Channels, Height, Width
B, T, C, H, W = unscaled_result.shape
# 将结果重新排列成 "b t c h w" 的形式
flattened = rearrange(unscaled_result, "b t c h w -> b c (t h w)")
# 计算每个通道的均值，并在第二维度上增加一个维度
means = flattened.mean(dim=2).unsqueeze(2)
# 对结果进行重新中心化，减去均值
recentered = flattened - means
# 计算每个通道的绝对值的最大值
magnitudes = recentered.abs().max()
# 对结果进行归一化，除以最大值
normalized = recentered / magnitudes
# 对结果进行动态阈值处理
thresholded = latent_dynamic_thresholding(normalized)
# 对结果进行反归一化，乘以最大值
denormalized = thresholded * magnitudes
# 对结果进行重新中心化，加上均值
uncentered = denormalized + means
# 将结果重新排列成 "b c (t h w)" 的形式
unflattened = rearrange(uncentered, "b c (t h w) -> b t c h w", t=T, h=H, w=W)
# 对结果进行缩放，乘以缩放因子
scaled_result = unflattened * self.scale_factor
# 返回缩放后的结果
return scaled_result
# 定义一个动态阈值处理类
class DynamicThresholdingV2:
# 定义类的调用方法，接受无条件值、条件值和缩放比例作为参数
def __call__(self, uncond, cond, scale):
# 获取无条件值的形状信息
B, T, C, H, W = uncond.shape
# 计算条件值和无条件值的差值
diff = cond - uncond
# 计算最小目标值
mim_target = uncond + diff * 4.0
# 计算配置目标值
cfg_target = uncond + diff * 8.0
# 将最小目标值展平为二维数组
mim_flattened = rearrange(mim_target, "b t c h w -> b c (t h w)")
# 将配置目标值展平为二维数组
cfg_flattened = rearrange(cfg_target, "b t c h w -> b c (t h w)")
# 计算最小目标值的均值
mim_means = mim_flattened.mean(dim=2).unsqueeze(2)
# 计算配置目标值的均值
cfg_means = cfg_flattened.mean(dim=2).unsqueeze(2)
# 计算最小目标值的中心化值
mim_centered = mim_flattened - mim_means
# 计算配置目标值的中心化值
cfg_centered = cfg_flattened - cfg_means
# 计算最小目标值的标准差
mim_scaleref = mim_centered.std(dim=2).unsqueeze(2)
# 计算配置目标值的标准差
cfg_scaleref = cfg_centered.std(dim=2).unsqueeze(2)
# 对配置目标值进行重新归一化
cfg_renormalized = cfg_centered / cfg_scaleref * mim_scaleref
# 将结果还原为原始形状
result = cfg_renormalized + cfg_means
unflattened = rearrange(result, "b c (t h w) -> b t c h w", t=T, h=H, w=W)
return unflattened
# 定义一个线性多步系数函数
def linear_multistep_coeff(order, t, i, j, epsrel=1e-4):
# 如果阶数大于i，则抛出异常
if order - 1 > i:
raise ValueError(f"Order {order} too high for step {i}")
# 定义一个内部函数，接受tau作为参数
def fn(tau):
prod = 1.0
# 遍历阶数
for k in range(order):
# 如果j等于k，则跳过
if j == k:
continue
# 计算乘积
prod *= (tau - t[i - k]) / (t[i - j] - t[i - k])
return prod
# 对内部函数进行积分
return integrate.quad(fn, t[i], t[i + 1], epsrel=epsrel)[0]
# 定义一个获取祖先步长的函数
def get_ancestral_step(sigma_from, sigma_to, eta=1.0):
# 如果eta为假，则返回sigma_to和0.0
if not eta:
return sigma_to, 0.0
# 计算上行步长
sigma_up = torch.minimum(
sigma_to,
eta * (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5,
)
# 计算下行步长
sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5
return sigma_down, sigma_up
# 定义一个转换为d的函数
def to_d(x, sigma, denoised):
return (x - denoised) / append_dims(sigma, x.ndim)
# 定义一个转换为负对数sigma的函数
def to_neg_log_sigma(sigma):
return sigma.log().neg()
# 定义一个转换为sigma的函数
def to_sigma(neg_log_sigma):
return neg_log_sigma.neg().exp()
.\cogvideo-finetune\sat\sgm\modules\diffusionmodules\sigma_sampling.py
import torch
import torch.distributed
from sat import mpu
from ...util import default, instantiate_from_config
class EDMSampling:
def __init__(self, p_mean=-1.2, p_std=1.2):
# 初始化函数，设置默认的 p_mean 和 p_std 值
self.p_mean = p_mean
self.p_std = p_std
def __call__(self, n_samples, rand=None):
# 调用函数，生成 n_samples 个样本
# 通过随机数生成 log_sigma，其均值为 p_mean，标准差为 p_std
log_sigma = self.p_mean + self.p_std * default(rand, torch.randn((n_samples,)))
# 返回 log_sigma 的指数形式
return log_sigma.exp()
class DiscreteSampling:
def __init__(self, discretization_config, num_idx, do_append_zero=False, flip=True, uniform_sampling=False):
# 初始化函数，设置离散化配置、索引数量、是否追加零、是否翻转和是否均匀采样
self.num_idx = num_idx
# 实例化离散化配置，得到 sigmas
self.sigmas = instantiate_from_config(discretization_config)(num_idx, do_append_zero=do_append_zero, flip=flip)
# 获取数据并行的世界大小
world_size = mpu.get_data_parallel_world_size()
self.uniform_sampling = uniform_sampling
if self.uniform_sampling:
# 如果进行均匀采样
i = 1
while True:
if world_size % i != 0 or num_idx % (world_size // i) != 0:
i += 1
else:
# 计算组数
self.group_num = world_size // i
break
assert self.group_num > 0
assert world_size % self.group_num == 0
# 计算每组的宽度，即每组中的排名数量
self.group_width = world_size // self.group_num # the number of rank in one group
# 计算每个组的 sigma 区间
self.sigma_interval = self.num_idx // self.group_num
def idx_to_sigma(self, idx):
# 根据索引获取 sigma
return self.sigmas[idx]
def __call__(self, n_samples, rand=None, return_idx=False):
# 调用函数，生成 n_samples 个样本
if self.uniform_sampling:
# 如果进行均匀采样
rank = mpu.get_data_parallel_rank()
# 计算组索引
group_index = rank // self.group_width
# 生成索引，范围为 group_index * sigma_interval 到 (group_index + 1) * sigma_interval
idx = default(
rand,
torch.randint(
group_index * self.sigma_interval, (group_index + 1) * self.sigma_interval, (n_samples,)
),
)
else:
# 如果不进行均匀采样
# 生成索引，范围为 0 到 num_idx
idx = default(
rand,
torch.randint(0, self.num_idx, (n_samples,)),
)
if return_idx:
# 如果返回索引
# 返回索引对应的 sigma 和索引
return self.idx_to_sigma(idx), idx
else:
# 如果不返回索引
# 返回索引对应的 sigma
return self.idx_to_sigma(idx)
class PartialDiscreteSampling:
def __init__(self, discretization_config, total_num_idx, partial_num_idx, do_append_zero=False, flip=True):
# 初始化函数，设置离散化配置、总索引数量、部分索引数量、是否追加零和是否翻转
self.total_num_idx = total_num_idx
self.partial_num_idx = partial_num_idx
# 实例化离散化配置，得到 sigmas
self.sigmas = instantiate_from_config(discretization_config)(
total_num_idx, do_append_zero=do_append_zero, flip=flip
)
def idx_to_sigma(self, idx):
# 根据索引获取 sigma
return self.sigmas[idx]
def __call__(self, n_samples, rand=None):
# 调用函数，生成 n_samples 个样本
# 生成索引，范围为 0 到 partial_num_idx
idx = default(
rand,
torch.randint(0, self.partial_num_idx, (n_samples,)),
)
# 返回索引对应的 sigma
return self.idx_to_sigma(idx)
.\cogvideo-finetune\sat\sgm\modules\diffusionmodules\util.py
# 该模块源自多个开源项目，感谢它们的贡献
"""
adopted from
https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
and
https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
and
https://github.com/openai/guided-diffusion/blob/0ba878e517b276c45d1195eb29f6f5f72659a05b/guided_diffusion/nn.py
thanks!
"""
# 导入数学库
import math
# 从 typing 导入可选类型
from typing import Optional
# 导入 PyTorch 库
import torch
import torch.nn as nn
# 从 einops 导入重排和重复函数
from einops import rearrange, repeat
# 创建 beta 调度的函数
def make_beta_schedule(
schedule,
n_timestep,
linear_start=1e-4,
linear_end=2e-2,
):
# 如果调度方式为线性，生成 beta 值的序列
if schedule == "linear":
betas = torch.linspace(linear_start**0.5, linear_end**0.5, n_timestep, dtype=torch.float64) ** 2
# 返回 numpy 格式的 beta 值
return betas.numpy()
# 从张量中提取特定的值并调整形状
def extract_into_tensor(a, t, x_shape):
# 获取 t 的第一个维度的大小
b, *_ = t.shape
# 根据索引 t 从 a 中提取数据
out = a.gather(-1, t)
# 返回调整形状后的输出
return out.reshape(b, *((1,) * (len(x_shape) - 1)))
# 混合检查点函数，用于在不缓存中间激活的情况下评估函数
def mixed_checkpoint(func, inputs: dict, params, flag):
"""
在不缓存中间激活的情况下评估函数，减少内存消耗，但会增加反向传播的计算量。
该实现允许非张量输入。
:param func: 要评估的函数。
:param inputs: 传递给 `func` 的参数字典。
:param params: func 依赖但不作为参数的参数序列。
:param flag: 如果为 False，禁用梯度检查点。
"""
# 如果启用标志，处理张量输入
if flag:
# 获取所有张量类型的输入键
tensor_keys = [key for key in inputs if isinstance(inputs[key], torch.Tensor)]
# 获取所有张量类型的输入值
tensor_inputs = [inputs[key] for key in inputs if isinstance(inputs[key], torch.Tensor)]
# 获取所有非张量类型的输入键
non_tensor_keys = [key for key in inputs if not isinstance(inputs[key], torch.Tensor)]
# 获取所有非张量类型的输入值
non_tensor_inputs = [inputs[key] for key in inputs if not isinstance(inputs[key], torch.Tensor)]
# 构建参数元组
args = tuple(tensor_inputs) + tuple(non_tensor_inputs) + tuple(params)
# 应用混合检查点函数
return MixedCheckpointFunction.apply(
func,
len(tensor_inputs),
len(non_tensor_inputs),
tensor_keys,
non_tensor_keys,
*args,
)
else:
# 如果禁用标志，直接调用函数
return func(**inputs)
# 定义混合检查点函数的类
class MixedCheckpointFunction(torch.autograd.Function):
@staticmethod
# 定义前向传播方法
def forward(
ctx,
run_function,
length_tensors,
length_non_tensors,
tensor_keys,
non_tensor_keys,
*args,
):
# 将长度张量赋值给上下文对象的属性
ctx.end_tensors = length_tensors
# 将长度非张量赋值给上下文对象的属性
ctx.end_non_tensors = length_tensors + length_non_tensors
# 创建包含 GPU 自动混合精度参数的字典
ctx.gpu_autocast_kwargs = {
# 检查自动混合精度是否启用
"enabled": torch.is_autocast_enabled(),
# 获取自动混合精度的 GPU 数据类型
"dtype": torch.get_autocast_gpu_dtype(),
# 检查自动混合精度缓存是否启用
"cache_enabled": torch.is_autocast_cache_enabled(),
}
# 断言张量键和非张量键的长度与预期相符
assert len(tensor_keys) == length_tensors and len(non_tensor_keys) == length_non_tensors
# 创建输入张量字典，将键与对应的值配对
ctx.input_tensors = {key: val for (key, val) in zip(tensor_keys, list(args[: ctx.end_tensors]))}
# 创建输入非张量字典，将键与对应的值配对
ctx.input_non_tensors = {
key: val for (key, val) in zip(non_tensor_keys, list(args[ctx.end_tensors : ctx.end_non_tensors]))
}
# 将运行函数赋值给上下文对象的属性
ctx.run_function = run_function
# 将输入参数赋值为剩余的参数
ctx.input_params = list(args[ctx.end_non_tensors :])
# 在无梯度计算上下文中执行
with torch.no_grad():
# 调用运行函数，并传入输入张量和非张量
output_tensors = ctx.run_function(**ctx.input_tensors, **ctx.input_non_tensors)
# 返回输出张量
return output_tensors
@staticmethod
def backward(ctx, *output_grads):
# 创建额外参数的字典（已注释掉）
# additional_args = {key: ctx.input_tensors[key] for key in ctx.input_tensors if not isinstance(ctx.input_tensors[key],torch.Tensor)}
# 将输入张量设为不跟踪梯度并要求梯度
ctx.input_tensors = {key: ctx.input_tensors[key].detach().requires_grad_(True) for key in ctx.input_tensors}
# 在启用梯度计算和混合精度上下文中执行
with torch.enable_grad(), torch.cuda.amp.autocast(**ctx.gpu_autocast_kwargs):
# 创建输入张量的浅拷贝以避免修改原张量
shallow_copies = {key: ctx.input_tensors[key].view_as(ctx.input_tensors[key]) for key in ctx.input_tensors}
# shallow_copies.update(additional_args) # 更新额外参数（已注释掉）
# 调用运行函数，并传入浅拷贝和非张量
output_tensors = ctx.run_function(**shallow_copies, **ctx.input_non_tensors)
# 计算输出张量相对于输入张量和输入参数的梯度
input_grads = torch.autograd.grad(
output_tensors,
list(ctx.input_tensors.values()) + ctx.input_params,
output_grads,
allow_unused=True,
)
# 删除输入张量
del ctx.input_tensors
# 删除输入参数
del ctx.input_params
# 删除输出张量
del output_tensors
# 返回梯度和占位符
return (
(None, None, None, None, None)
+ input_grads[: ctx.end_tensors]
+ (None,) * (ctx.end_non_tensors - ctx.end_tensors)
+ input_grads[ctx.end_tensors :]
)
# 定义一个检查点函数，允许不缓存中间激活以降低内存使用
def checkpoint(func, inputs, params, flag):
# 文档字符串，解释函数的用途和参数
"""
Evaluate a function without caching intermediate activations, allowing for
reduced memory at the expense of extra compute in the backward pass.
:param func: the function to evaluate.
:param inputs: the argument sequence to pass to `func`.
:param params: a sequence of parameters `func` depends on but does not
explicitly take as arguments.
:param flag: if False, disable gradient checkpointing.
"""
# 如果标志为真，执行检查点功能
if flag:
# 将输入和参数组合为一个元组
args = tuple(inputs) + tuple(params)
# 应用检查点功能并返回结果
return CheckpointFunction.apply(func, len(inputs), *args)
else:
# 直接调用函数并返回结果
return func(*inputs)
# 定义检查点功能类，继承自torch的自动梯度功能
class CheckpointFunction(torch.autograd.Function):
# 静态方法，用于前向传播
@staticmethod
def forward(ctx, run_function, length, *args):
# 将运行的函数和输入张量保存到上下文中
ctx.run_function = run_function
ctx.input_tensors = list(args[:length])
ctx.input_params = list(args[length:])
# 保存GPU自动混合精度的相关设置
ctx.gpu_autocast_kwargs = {
"enabled": torch.is_autocast_enabled(),
"dtype": torch.get_autocast_gpu_dtype(),
"cache_enabled": torch.is_autocast_cache_enabled(),
}
# 在无梯度计算模式下执行函数
with torch.no_grad():
output_tensors = ctx.run_function(*ctx.input_tensors)
# 返回输出张量
return output_tensors
# 静态方法，用于反向传播
@staticmethod
def backward(ctx, *output_grads):
# 将输入张量标记为需要梯度计算
ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
# 启用梯度计算和混合精度模式
with torch.enable_grad(), torch.cuda.amp.autocast(**ctx.gpu_autocast_kwargs):
# 创建输入张量的浅拷贝，以避免存储修改问题
shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
# 执行反向传播，获取输出张量
output_tensors = ctx.run_function(*shallow_copies)
# 计算输入梯度
input_grads = torch.autograd.grad(
output_tensors,
ctx.input_tensors + ctx.input_params,
output_grads,
allow_unused=True,
)
# 清理上下文中的临时变量
del ctx.input_tensors
del ctx.input_params
del output_tensors
# 返回None和输入梯度
return (None, None) + input_grads
# 定义时间步嵌入函数，创建正弦时间步嵌入
def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False, dtype=torch.float32):
# 文档字符串，解释函数的用途和参数
"""
Create sinusoidal timestep embeddings.
:param timesteps: a 1-D Tensor of N indices, one per batch element.
These may be fractional.
:param dim: the dimension of the output.
:param max_period: controls the minimum frequency of the embeddings.
:return: an [N x dim] Tensor of positional embeddings.
"""
# 检查是否只重复使用嵌入
if not repeat_only:
# 计算维度的一半
half = dim // 2
# 计算频率值，使用指数衰减函数生成频率序列
freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
device=timesteps.device
)
# 计算时间步长与频率的乘积，形成角度参数
args = timesteps[:, None].float() * freqs[None]
# 计算余弦和正弦值并在最后一维合并
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
# 如果维度为奇数，添加一列全零以保持维度一致
if dim % 2:
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
else:
# 如果只重复，直接生成与时间步长相同的嵌入
embedding = repeat(timesteps, "b -> b d", d=dim)
# 将嵌入转换为指定的数据类型并返回
return embedding.to(dtype)
# 定义一个将模块参数归零的函数，并返回该模块
def zero_module(module):
"""
将模块的参数置为零，并返回该模块。
"""
# 遍历模块的所有参数
for p in module.parameters():
# 将参数分离并置为零
p.detach().zero_()
# 返回处理后的模块
return module
# 定义一个对模块参数进行缩放的函数，并返回该模块
def scale_module(module, scale):
"""
将模块的参数进行缩放，并返回该模块。
"""
# 遍历模块的所有参数
for p in module.parameters():
# 将参数分离并进行缩放
p.detach().mul_(scale)
# 返回处理后的模块
return module
# 定义一个计算非批次维度均值的函数
def mean_flat(tensor):
"""
对所有非批次维度进行求均值。
"""
# 计算并返回张量的均值，忽略第一维（批次维度）
return tensor.mean(dim=list(range(1, len(tensor.shape))))
# 定义一个创建标准归一化层的函数
def normalization(channels):
"""
创建一个标准归一化层。
:param channels: 输入通道的数量。
:return: 一个用于归一化的 nn.Module。
"""
# 返回一个具有指定输入通道数的 GroupNorm32 对象
return GroupNorm32(32, channels)
# 定义一个 SiLU 激活函数类，支持 PyTorch 1.5
class SiLU(nn.Module):
# 定义前向传播方法
def forward(self, x):
# 返回 x 与其 sigmoid 值的乘积
return x * torch.sigmoid(x)
# 定义一个扩展自 nn.GroupNorm 的 GroupNorm32 类
class GroupNorm32(nn.GroupNorm):
# 定义前向传播方法
def forward(self, x):
# 调用父类的 forward 方法，并返回与输入相同数据类型的结果
return super().forward(x).type(x.dtype)
# 定义一个创建卷积模块的函数，支持 1D、2D 和 3D 卷积
def conv_nd(dims, *args, **kwargs):
"""
创建一个 1D、2D 或 3D 卷积模块。
"""
# 判断维度是否为 1
if dims == 1:
# 返回 1D 卷积模块
return nn.Conv1d(*args, **kwargs)
# 判断维度是否为 2
elif dims == 2:
# 返回 2D 卷积模块
return nn.Conv2d(*args, **kwargs)
# 判断维度是否为 3
elif dims == 3:
# 返回 3D 卷积模块
return nn.Conv3d(*args, **kwargs)
# 抛出不支持的维度错误
raise ValueError(f"unsupported dimensions: {dims}")
# 定义一个创建线性模块的函数
def linear(*args, **kwargs):
"""
创建一个线性模块。
"""
# 返回一个线性模块
return nn.Linear(*args, **kwargs)
# 定义一个创建平均池化模块的函数，支持 1D、2D 和 3D 池化
def avg_pool_nd(dims, *args, **kwargs):
"""
创建一个 1D、2D 或 3D 平均池化模块。
"""
# 判断维度是否为 1
if dims == 1:
# 返回 1D 平均池化模块
return nn.AvgPool1d(*args, **kwargs)
# 判断维度是否为 2
elif dims == 2:
# 返回 2D 平均池化模块
return nn.AvgPool2d(*args, **kwargs)
# 判断维度是否为 3
elif dims == 3:
# 返回 3D 平均池化模块
return nn.AvgPool3d(*args, **kwargs)
# 抛出不支持的维度错误
raise ValueError(f"unsupported dimensions: {dims}")
# 定义一个 AlphaBlender 类，用于实现不同的混合策略
class AlphaBlender(nn.Module):
# 定义支持的混合策略
strategies = ["learned", "fixed", "learned_with_images"]
# 初始化方法，设置混合因子和重排模式
def __init__(
self,
alpha: float,
merge_strategy: str = "learned_with_images",
rearrange_pattern: str = "b t -> (b t) 1 1",
):
super().__init__()
# 保存混合策略和重排模式
self.merge_strategy = merge_strategy
self.rearrange_pattern = rearrange_pattern
# 确保混合策略在支持的策略中
assert merge_strategy in self.strategies, f"merge_strategy needs to be in {self.strategies}"
# 根据混合策略注册混合因子
if self.merge_strategy == "fixed":
# 注册一个固定的混合因子
self.register_buffer("mix_factor", torch.Tensor([alpha]))
elif self.merge_strategy == "learned" or self.merge_strategy == "learned_with_images":
# 注册一个可学习的混合因子
self.register_parameter("mix_factor", torch.nn.Parameter(torch.Tensor([alpha])))
else:
# 抛出未知混合策略错误
raise ValueError(f"unknown merge strategy {self.merge_strategy}")
# 定义获取 alpha 值的函数，输入为图像指示器，返回一个张量
def get_alpha(self, image_only_indicator: torch.Tensor) -> torch.Tensor:
# 根据合并策略选择 alpha 值的计算方式
if self.merge_strategy == "fixed":
# 如果合并策略为固定值，则直接使用 mix_factor 作为 alpha
alpha = self.mix_factor
elif self.merge_strategy == "learned":
# 如果合并策略为学习得到的值，则对 mix_factor 应用 sigmoid 函数
alpha = torch.sigmoid(self.mix_factor)
elif self.merge_strategy == "learned_with_images":
# 如果合并策略为图像学习，需要确保提供图像指示器
assert image_only_indicator is not None, "need image_only_indicator ..."
# 根据图像指示器选择 alpha 值，真值对应 1，假值则使用 mix_factor 的 sigmoid 结果
alpha = torch.where(
image_only_indicator.bool(), # 将图像指示器转换为布尔值
torch.ones(1, 1, device=image_only_indicator.device), # 为真值创建全 1 的张量
rearrange(torch.sigmoid(self.mix_factor), "... -> ... 1"), # 为假值应用 sigmoid 并调整维度
)
# 根据 rearrange_pattern 重新排列 alpha 的维度
alpha = rearrange(alpha, self.rearrange_pattern)
else:
# 如果合并策略不在已知范围内，抛出未实现错误
raise NotImplementedError
# 返回计算得到的 alpha 值
return alpha
# 定义前向传播函数，输入为空间和时间张量，以及可选的图像指示器
def forward(
self,
x_spatial: torch.Tensor, # 空间输入张量
x_temporal: torch.Tensor, # 时间输入张量
image_only_indicator: Optional[torch.Tensor] = None, # 可选的图像指示器
) -> torch.Tensor:
# 调用 get_alpha 函数获取 alpha 值
alpha = self.get_alpha(image_only_indicator)
# 计算最终输出张量，结合空间和时间输入，使用 alpha 进行加权
x = alpha.to(x_spatial.dtype) * x_spatial + (1.0 - alpha).to(x_spatial.dtype) * x_temporal
# 返回计算得到的输出张量
return x
.\cogvideo-finetune\sat\sgm\modules\diffusionmodules\wrappers.py
# 导入 PyTorch 库
import torch
# 导入 PyTorch 的神经网络模块
import torch.nn as nn
# 从 packaging 库导入 version 模块以便进行版本比较
from packaging import version
# 定义 OpenAIWrapper 的模块路径
OPENAIUNETWRAPPER = "sgm.modules.diffusionmodules.wrappers.OpenAIWrapper"
# 定义身份包装器类，继承自 nn.Module
class IdentityWrapper(nn.Module):
# 初始化函数，接收扩散模型、是否编译模型的标志和数据类型
def __init__(self, diffusion_model, compile_model: bool = False, dtype: torch.dtype = torch.float32):
# 调用父类的初始化函数
super().__init__()
# 判断 PyTorch 版本，选择是否编译模型
compile = (
torch.compile
if (version.parse(torch.__version__) >= version.parse("2.0.0")) and compile_model
else lambda x: x # 如果不满足条件，返回原始输入
)
# 编译扩散模型并赋值给实例变量
self.diffusion_model = compile(diffusion_model)
# 设置数据类型
self.dtype = dtype
# 前向传播函数，接收任意数量的位置参数和关键字参数
def forward(self, *args, **kwargs):
# 调用扩散模型的前向传播，并返回结果
return self.diffusion_model(*args, **kwargs)
# 定义 OpenAIWrapper 类，继承自 IdentityWrapper
class OpenAIWrapper(IdentityWrapper):
# 重写前向传播函数，接收输入张量 x、时间步 t、上下文字典 c 以及其他关键字参数
def forward(self, x: torch.Tensor, t: torch.Tensor, c: dict, **kwargs) -> torch.Tensor:
# 将上下文字典中的每个张量转换为指定的数据类型
for key in c:
c[key] = c[key].to(self.dtype)
# 如果输入张量是 4 维，按维度 1 拼接上下文中的 "concat" 张量
if x.dim() == 4:
x = torch.cat((x, c.get("concat", torch.Tensor([]).type_as(x))), dim=1)
# 如果输入张量是 5 维，按维度 2 拼接上下文中的 "concat" 张量
elif x.dim() == 5:
x = torch.cat((x, c.get("concat", torch.Tensor([]).type_as(x))), dim=2)
# 如果输入张量的维度不符合要求，抛出值错误
else:
raise ValueError("Input tensor must be 4D or 5D")
# 调用扩散模型的前向传播，传入处理后的张量、时间步和上下文等
return self.diffusion_model(
x,
timesteps=t,
context=c.get("crossattn", None),
y=c.get("vector", None),
**kwargs,
)
.\cogvideo-finetune\sat\sgm\modules\diffusionmodules\__init__.py
# 从当前包的 denoiser 模块导入 Denoiser 类
from .denoiser import Denoiser
# 从当前包的 discretizer 模块导入 Discretization 类
from .discretizer import Discretization
# 从当前包的 model 模块导入 Decoder、Encoder 和 Model 类
from .model import Decoder, Encoder, Model
# 从当前包的 openaimodel 模块导入 UNetModel 类
from .openaimodel import UNetModel
# 从当前包的 sampling 模块导入 BaseDiffusionSampler 类
from .sampling import BaseDiffusionSampler
# 从当前包的 wrappers 模块导入 OpenAIWrapper 类
from .wrappers import OpenAIWrapper
.\cogvideo-finetune\sat\sgm\modules\distributions\distributions.py
# 导入 numpy 库并简化为 np
import numpy as np
# 导入 PyTorch 库
import torch
# 定义抽象类 AbstractDistribution
class AbstractDistribution:
# 抽象方法，生成样本
def sample(self):
raise NotImplementedError()
# 抽象方法，返回分布的众数
def mode(self):
raise NotImplementedError()
# 定义 DiracDistribution 类，继承自 AbstractDistribution
class DiracDistribution(AbstractDistribution):
# 初始化 DiracDistribution，接受一个值
def __init__(self, value):
self.value = value
# 重写 sample 方法，返回固定值
def sample(self):
return self.value
# 重写 mode 方法，返回固定值
def mode(self):
return self.value
# 定义 DiagonalGaussianDistribution 类
class DiagonalGaussianDistribution(object):
# 初始化，接受参数和一个确定性标志
def __init__(self, parameters, deterministic=False):
self.parameters = parameters
# 将参数分为均值和对数方差
self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
# 将对数方差限制在指定范围内
self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
self.deterministic = deterministic
# 计算标准差和方差
self.std = torch.exp(0.5 * self.logvar)
self.var = torch.exp(self.logvar)
# 如果是确定性模式，将方差和标准差设为零
if self.deterministic:
self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
# 生成样本
def sample(self):
# 使用均值和标准差生成随机样本
x = self.mean + self.std * torch.randn_like(self.mean).to(device=self.parameters.device)
return x
# 计算 KL 散度
def kl(self, other=None):
if self.deterministic:
return torch.Tensor([0.0]) # 确定性时 KL 散度为 0
else:
if other is None:
# 计算与标准正态分布的 KL 散度
return 0.5 * torch.sum(
torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
dim=[1, 2, 3],
)
else:
# 计算两个分布间的 KL 散度
return 0.5 * torch.sum(
torch.pow(self.mean - other.mean, 2) / other.var
+ self.var / other.var
- 1.0
- self.logvar
+ other.logvar,
dim=[1, 2, 3],
)
# 计算负对数似然
def nll(self, sample, dims=[1, 2, 3]):
if self.deterministic:
return torch.Tensor([0.0]) # 确定性时 NLL 为 0
logtwopi = np.log(2.0 * np.pi) # 计算 2π 的对数
# 计算负对数似然值
return 0.5 * torch.sum(
logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
dim=dims,
)
# 返回均值
def mode(self):
return self.mean
# 定义 normal_kl 函数，计算两个高斯分布之间的 KL 散度
def normal_kl(mean1, logvar1, mean2, logvar2):
"""
source: https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/losses.py#L12
计算两个高斯分布之间的 KL 散度。
形状会自动广播，支持批量比较和标量等用例。
"""
tensor = None
# 找到第一个是 Tensor 的对象
for obj in (mean1, logvar1, mean2, logvar2):
if isinstance(obj, torch.Tensor):
tensor = obj
break
# 确保至少有一个参数是 Tensor
assert tensor is not None, "at least one argument must be a Tensor"
# 强制将方差转换为 Tensor
logvar1, logvar2 = [x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor) for x in (logvar1, logvar2)]
# 计算并返回某种损失值，具体公式由多个部分组成
return 0.5 * ( # 返回整个计算结果的一半
-1.0 + logvar2 - logvar1 + # 计算对数方差差值，并减去常数项
torch.exp(logvar1 - logvar2) + # 计算指数项，表示方差的相对差异
((mean1 - mean2) ** 2) * torch.exp(-logvar2) # 计算均值差的平方乘以方差的倒数
)
.\cogvideo-finetune\sat\sgm\modules\distributions\__init__.py
请提供需要注释的代码。
.\cogvideo-finetune\sat\sgm\modules\ema.py
# 导入 PyTorch 库
import torch
# 从 PyTorch 导入神经网络模块
from torch import nn
# 定义一个名为 LitEma 的类，继承自 nn.Module
class LitEma(nn.Module):
# 初始化函数，接受模型、衰减因子和更新次数的使用标志
def __init__(self, model, decay=0.9999, use_num_upates=True):
# 调用父类构造函数
super().__init__()
# 检查衰减因子是否在有效范围内
if decay < 0.0 or decay > 1.0:
raise ValueError("Decay must be between 0 and 1")
# 初始化模型参数名称到阴影参数名称的映射字典
self.m_name2s_name = {}
# 注册衰减因子为一个缓冲区
self.register_buffer("decay", torch.tensor(decay, dtype=torch.float32))
# 根据是否使用更新次数注册相应的缓冲区
self.register_buffer(
"num_updates",
torch.tensor(0, dtype=torch.int) if use_num_upates else torch.tensor(-1, dtype=torch.int),
)
# 遍历模型的每个命名参数
for name, p in model.named_parameters():
# 只处理需要梯度的参数
if p.requires_grad:
# 将参数名称中的点替换为字符，以便注册为缓冲区
s_name = name.replace(".", "")
# 更新名称映射字典
self.m_name2s_name.update({name: s_name})
# 注册参数的副本为缓冲区
self.register_buffer(s_name, p.clone().detach().data)
# 初始化存储的参数列表
self.collected_params = []
# 重置更新次数的函数
def reset_num_updates(self):
# 删除当前的更新次数缓冲区
del self.num_updates
# 注册新的更新次数缓冲区，初始值为 0
self.register_buffer("num_updates", torch.tensor(0, dtype=torch.int))
# 前向传播函数，更新阴影参数
def forward(self, model):
# 获取当前的衰减因子
decay = self.decay
# 如果更新次数有效，更新衰减因子
if self.num_updates >= 0:
self.num_updates += 1
decay = min(self.decay, (1 + self.num_updates) / (10 + self.num_updates))
# 计算 1 减去衰减因子
one_minus_decay = 1.0 - decay
# 在不计算梯度的情况下执行以下操作
with torch.no_grad():
# 获取模型的参数字典
m_param = dict(model.named_parameters())
# 获取当前缓冲区中的阴影参数字典
shadow_params = dict(self.named_buffers())
# 遍历模型参数
for key in m_param:
# 只处理需要梯度的参数
if m_param[key].requires_grad:
# 获取阴影参数名称
sname = self.m_name2s_name[key]
# 将阴影参数转换为与模型参数相同的类型
shadow_params[sname] = shadow_params[sname].type_as(m_param[key])
# 更新阴影参数
shadow_params[sname].sub_(one_minus_decay * (shadow_params[sname] - m_param[key]))
else:
# 确保该参数不在名称映射中
assert not key in self.m_name2s_name
# 将阴影参数复制回模型的函数
def copy_to(self, model):
# 获取模型参数和阴影参数的字典
m_param = dict(model.named_parameters())
shadow_params = dict(self.named_buffers())
# 遍历模型参数
for key in m_param:
# 只处理需要梯度的参数
if m_param[key].requires_grad:
# 将阴影参数的数据复制到模型参数
m_param[key].data.copy_(shadow_params[self.m_name2s_name[key]].data)
else:
# 确保该参数不在名称映射中
assert not key in self.m_name2s_name
# 存储当前参数以便稍后恢复的函数
def store(self, parameters):
"""
保存当前参数以便稍后恢复。
参数:
parameters: 可迭代的 `torch.nn.Parameter`；要临时存储的参数。
"""
# 将当前参数的副本存储在列表中
self.collected_params = [param.clone() for param in parameters]
# 定义一个恢复参数的方法
def restore(self, parameters):
# 文档字符串，说明此方法的作用及参数
"""
Restore the parameters stored with the `store` method.
Useful to validate the model with EMA parameters without affecting the
original optimization process. Store the parameters before the
`copy_to` method. After validation (or model saving), use this to
restore the former parameters.
Args:
parameters: Iterable of `torch.nn.Parameter`; the parameters to be
updated with the stored parameters.
"""
# 遍历收集的参数和传入的参数，进行一一对应
for c_param, param in zip(self.collected_params, parameters):
# 将收集的参数数据复制到当前参数
param.data.copy_(c_param.data)
.\cogvideo-finetune\sat\sgm\modules\encoders\modules.py
# 导入数学库
import math
# 导入上下文管理器的空上下文
from contextlib import nullcontext
# 导入部分函数的工具
from functools import partial
# 导入类型提示相关的类型
from typing import Dict, List, Optional, Tuple, Union
# 导入 Kornia 图像处理库
import kornia
# 导入 NumPy 库
import numpy as np
# 导入 PyTorch 库
import torch
# 导入 PyTorch 的神经网络模块
import torch.nn as nn
# 导入 einops 库的重排列和重复函数
from einops import rearrange, repeat
# 导入 OmegaConf 库的 ListConfig
from omegaconf import ListConfig
# 导入 PyTorch 的检查点工具
from torch.utils.checkpoint import checkpoint
# 导入 Hugging Face 的 T5 模型和分词器
from transformers import (
T5EncoderModel,
T5Tokenizer,
)
# 从自定义工具模块导入多个函数
from ...util import (
append_dims,
autocast,
count_params,
default,
disabled_train,
expand_dims_like,
instantiate_from_config,
)
# 定义一个抽象的嵌入模型类，继承自 PyTorch 的 nn.Module
class AbstractEmbModel(nn.Module):
# 初始化方法
def __init__(self):
# 调用父类初始化方法
super().__init__()
# 初始化是否可训练的标志
self._is_trainable = None
# 初始化 UCG 速率
self._ucg_rate = None
# 初始化输入键
self._input_key = None
# 返回是否可训练的属性
@property
def is_trainable(self) -> bool:
return self._is_trainable
# 返回 UCG 速率属性，可能是浮点数或张量
@property
def ucg_rate(self) -> Union[float, torch.Tensor]:
return self._ucg_rate
# 返回输入键属性
@property
def input_key(self) -> str:
return self._input_key
# 设置是否可训练的属性
@is_trainable.setter
def is_trainable(self, value: bool):
self._is_trainable = value
# 设置 UCG 速率属性
@ucg_rate.setter
def ucg_rate(self, value: Union[float, torch.Tensor]):
self._ucg_rate = value
# 设置输入键属性
@input_key.setter
def input_key(self, value: str):
self._input_key = value
# 删除是否可训练的属性
@is_trainable.deleter
def is_trainable(self):
del self._is_trainable
# 删除 UCG 速率属性
@ucg_rate.deleter
def ucg_rate(self):
del self._ucg_rate
# 删除输入键属性
@input_key.deleter
def input_key(self):
del self._input_key
# 定义一个通用条件器类，继承自 PyTorch 的 nn.Module
class GeneralConditioner(nn.Module):
# 定义输出维度到键的映射
OUTPUT_DIM2KEYS = {2: "vector", 3: "crossattn", 4: "concat", 5: "concat"}
# 定义键到拼接维度的映射
KEY2CATDIM = {"vector": 1, "crossattn": 2, "concat": 1}
# 初始化方法，接受嵌入模型配置及相关参数
def __init__(self, emb_models: Union[List, ListConfig], cor_embs=[], cor_p=[]):
# 调用父类初始化方法
super().__init__()
# 用于存储嵌入模型实例的列表
embedders = []
# 遍历每个嵌入模型的配置
for n, embconfig in enumerate(emb_models):
# 根据配置实例化嵌入模型
embedder = instantiate_from_config(embconfig)
# 确保实例是 AbstractEmbModel 的子类
assert isinstance(
embedder, AbstractEmbModel
), f"embedder model {embedder.__class__.__name__} has to inherit from AbstractEmbModel"
# 设置嵌入模型是否可训练
embedder.is_trainable = embconfig.get("is_trainable", False)
# 设置嵌入模型的 ucg_rate 参数
embedder.ucg_rate = embconfig.get("ucg_rate", 0.0)
# 如果模型不可训练
if not embedder.is_trainable:
# 禁用训练方法
embedder.train = disabled_train
# 将模型参数的 requires_grad 属性设为 False
for param in embedder.parameters():
param.requires_grad = False
# 将模型设置为评估模式
embedder.eval()
# 打印嵌入模型的初始化信息
print(
f"Initialized embedder #{n}: {embedder.__class__.__name__} "
f"with {count_params(embedder, False)} params. Trainable: {embedder.is_trainable}"
)
# 检查配置中是否有 input_key，并赋值给嵌入模型
if "input_key" in embconfig:
embedder.input_key = embconfig["input_key"]
# 检查配置中是否有 input_keys，并赋值给嵌入模型
elif "input_keys" in embconfig:
embedder.input_keys = embconfig["input_keys"]
# 如果都没有，抛出 KeyError
else:
raise KeyError(f"need either 'input_key' or 'input_keys' for embedder {embedder.__class__.__name__}")
# 设置嵌入模型的 legacy_ucg_value 参数
embedder.legacy_ucg_val = embconfig.get("legacy_ucg_value", None)
# 如果有 legacy_ucg_val，则初始化随机状态
if embedder.legacy_ucg_val is not None:
embedder.ucg_prng = np.random.RandomState()
# 将嵌入模型添加到列表中
embedders.append(embedder)
# 将嵌入模型列表转换为 nn.ModuleList
self.embedders = nn.ModuleList(embedders)
# 如果有 cor_embs，确保 cor_p 的长度正确
if len(cor_embs) > 0:
assert len(cor_p) == 2 ** len(cor_embs)
# 设置相关嵌入和参数
self.cor_embs = cor_embs
self.cor_p = cor_p
# 获取 UCG 值的方法，可能会基于概率进行赋值
def possibly_get_ucg_val(self, embedder: AbstractEmbModel, batch: Dict) -> Dict:
# 确保 legacy_ucg_val 不为 None
assert embedder.legacy_ucg_val is not None
# 获取 ucg_rate 参数
p = embedder.ucg_rate
# 获取 legacy_ucg_val 值
val = embedder.legacy_ucg_val
# 遍历 batch 中的输入数据
for i in range(len(batch[embedder.input_key])):
# 根据概率选择是否替换为 legacy_ucg_val
if embedder.ucg_prng.choice(2, p=[1 - p, p]):
batch[embedder.input_key][i] = val
# 返回更新后的 batch
return batch
# 获取 UCG 值的方法，基于条件进行赋值
def surely_get_ucg_val(self, embedder: AbstractEmbModel, batch: Dict, cond_or_not) -> Dict:
# 确保 legacy_ucg_val 不为 None
assert embedder.legacy_ucg_val is not None
# 获取 legacy_ucg_val 值
val = embedder.legacy_ucg_val
# 遍历 batch 中的输入数据
for i in range(len(batch[embedder.input_key])):
# 如果条件为真，则替换为 legacy_ucg_val
if cond_or_not[i]:
batch[embedder.input_key][i] = val
# 返回更新后的 batch
return batch
# 获取单个嵌入的方法
def get_single_embedding(
self,
embedder,
batch,
output,
cond_or_not: Optional[np.ndarray] = None,
force_zero_embeddings: Optional[List] = None,
):
# 根据 embedder 是否可训练选择适当的上下文管理器
embedding_context = nullcontext if embedder.is_trainable else torch.no_grad
# 进入上下文管理器以控制梯度计算
with embedding_context():
# 检查 embedder 是否有输入键属性并且不为 None
if hasattr(embedder, "input_key") and (embedder.input_key is not None):
# 检查 embedder 的遗留 UCG 值是否不为 None
if embedder.legacy_ucg_val is not None:
# 如果条件为 None，获取可能的 UCG 值
if cond_or_not is None:
batch = self.possibly_get_ucg_val(embedder, batch)
# 否则，确保获取 UCG 值
else:
batch = self.surely_get_ucg_val(embedder, batch, cond_or_not)
# 使用指定的输入键从 batch 中获取嵌入输出
emb_out = embedder(batch[embedder.input_key])
# 检查 embedder 是否有输入键列表
elif hasattr(embedder, "input_keys"):
# 解包 batch 中的输入键以获取嵌入输出
emb_out = embedder(*[batch[k] for k in embedder.input_keys])
# 确保嵌入输出是张量或序列类型
assert isinstance(
emb_out, (torch.Tensor, list, tuple)
), f"encoder outputs must be tensors or a sequence, but got {type(emb_out)}"
# 如果嵌入输出不是列表或元组，则将其转换为列表
if not isinstance(emb_out, (list, tuple)):
emb_out = [emb_out]
# 遍历嵌入输出中的每个嵌入
for emb in emb_out:
# 根据嵌入的维度获取对应的输出键
out_key = self.OUTPUT_DIM2KEYS[emb.dim()]
# 如果 UCG 率大于 0 且没有遗留 UCG 值
if embedder.ucg_rate > 0.0 and embedder.legacy_ucg_val is None:
# 如果条件为 None，随机应用 UCG 率
if cond_or_not is None:
emb = (
expand_dims_like(
torch.bernoulli((1.0 - embedder.ucg_rate) * torch.ones(emb.shape[0], device=emb.device)),
emb,
)
* emb
)
# 否则，基于条件应用 UCG
else:
emb = (
expand_dims_like(
torch.tensor(1 - cond_or_not, dtype=emb.dtype, device=emb.device),
emb,
)
* emb
)
# 如果输入键在强制零嵌入列表中，将嵌入置为零
if hasattr(embedder, "input_key") and embedder.input_key in force_zero_embeddings:
emb = torch.zeros_like(emb)
# 如果输出中已有该输出键，则拼接嵌入
if out_key in output:
output[out_key] = torch.cat((output[out_key], emb), self.KEY2CATDIM[out_key])
# 否则，直接保存嵌入
else:
output[out_key] = emb
# 返回最终输出字典
return output
# 定义一个前向传播函数，接收批次数据和强制零嵌入列表
def forward(self, batch: Dict, force_zero_embeddings: Optional[List] = None) -> Dict:
# 初始化输出字典
output = dict()
# 如果没有强制零嵌入，初始化为空列表
if force_zero_embeddings is None:
force_zero_embeddings = []
# 如果相关嵌入列表不为空
if len(self.cor_embs) > 0:
# 获取批次大小
batch_size = len(batch[list(batch.keys())[0]])
# 根据相关概率随机选择索引
rand_idx = np.random.choice(len(self.cor_p), size=(batch_size,), p=self.cor_p)
# 遍历相关嵌入索引
for emb_idx in self.cor_embs:
# 计算条件是否满足
cond_or_not = rand_idx % 2
# 更新随机索引
rand_idx //= 2
# 获取单个嵌入并更新输出字典
output = self.get_single_embedding(
self.embedders[emb_idx],
batch,
output=output,
cond_or_not=cond_or_not,
force_zero_embeddings=force_zero_embeddings,
)
# 遍历所有嵌入
for i, embedder in enumerate(self.embedders):
# 如果当前索引在相关嵌入列表中，则跳过
if i in self.cor_embs:
continue
# 获取单个嵌入并更新输出字典
output = self.get_single_embedding(
embedder, batch, output=output, force_zero_embeddings=force_zero_embeddings
)
# 返回最终输出字典
return output
# 定义获取无条件条件的函数
def get_unconditional_conditioning(self, batch_c, batch_uc=None, force_uc_zero_embeddings=None):
# 如果没有强制无条件嵌入，初始化为空列表
if force_uc_zero_embeddings is None:
force_uc_zero_embeddings = []
# 初始化无条件生成率列表
ucg_rates = list()
# 遍历所有嵌入，保存其生成率并将生成率设置为零
for embedder in self.embedders:
ucg_rates.append(embedder.ucg_rate)
embedder.ucg_rate = 0.0
# 保存当前相关嵌入和概率
cor_embs = self.cor_embs
cor_p = self.cor_p
# 清空相关嵌入和概率
self.cor_embs = []
self.cor_p = []
# 计算条件输出
c = self(batch_c)
# 计算无条件输出
uc = self(batch_c if batch_uc is None else batch_uc, force_uc_zero_embeddings)
# 恢复每个嵌入的生成率
for embedder, rate in zip(self.embedders, ucg_rates):
embedder.ucg_rate = rate
# 恢复相关嵌入和概率
self.cor_embs = cor_embs
self.cor_p = cor_p
# 返回条件输出和无条件输出
return c, uc
# 定义一个名为 FrozenT5Embedder 的类，继承自 AbstractEmbModel
class FrozenT5Embedder(AbstractEmbModel):
"""使用 T5 变换器编码器处理文本"""
# 初始化方法，设置模型的基本参数
def __init__(
self,
model_dir="google/t5-v1_1-xxl", # 模型目录，默认是 T5 模型
device="cuda", # 设备设置，默认使用 GPU
max_length=77, # 输入文本的最大长度
freeze=True, # 是否冻结模型参数，默认是冻结
cache_dir=None, # 缓存目录，默认无
):
super().__init__() # 调用父类的初始化方法
# 检查模型目录是否为默认 T5 模型
if model_dir is not "google/t5-v1_1-xxl":
# 从指定目录加载 tokenizer 和 transformer 模型
self.tokenizer = T5Tokenizer.from_pretrained(model_dir)
self.transformer = T5EncoderModel.from_pretrained(model_dir)
else:
# 从指定目录加载 tokenizer 和 transformer 模型，同时指定缓存目录
self.tokenizer = T5Tokenizer.from_pretrained(model_dir, cache_dir=cache_dir)
self.transformer = T5EncoderModel.from_pretrained(model_dir, cache_dir=cache_dir)
# 设置设备
self.device = device
# 设置最大输入长度
self.max_length = max_length
# 如果需要冻结模型参数，调用 freeze 方法
if freeze:
self.freeze()
# 定义冻结方法
def freeze(self):
# 将 transformer 设置为评估模式
self.transformer = self.transformer.eval()
# 遍历模型参数，设置为不需要梯度更新
for param in self.parameters():
param.requires_grad = False
# @autocast 注解（注释掉的），用于自动混合精度计算
def forward(self, text):
# 使用 tokenizer 对输入文本进行编码，返回批处理编码结果
batch_encoding = self.tokenizer(
text,
truncation=True, # 超出最大长度时进行截断
max_length=self.max_length, # 设置最大长度
return_length=True, # 返回编码长度
return_overflowing_tokens=False, # 不返回溢出的 tokens
padding="max_length", # 填充到最大长度
return_tensors="pt", # 返回 PyTorch 张量
)
# 将输入 id 转移到指定设备
tokens = batch_encoding["input_ids"].to(self.device)
# 在禁用自动混合精度的上下文中进行前向传播
with torch.autocast("cuda", enabled=False):
# 使用 transformer 进行前向传播，获取输出
outputs = self.transformer(input_ids=tokens)
# 获取 transformer 输出的最后隐藏状态
z = outputs.last_hidden_state
# 返回最后隐藏状态
return z
# 定义编码方法，直接调用 forward 方法
def encode(self, text):
return self(text) # 将输入文本传递给 forward 方法进行编码
标签：None,
CogVideoX,
self,
torch,
---,
源码,
return,
sigma,
def
From： https://www.cnblogs.com/apachecn/p/18494408