尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

Diffusers 中 PNDMScheduler 深度解析:基于 Runge-Kutta 与线性多步法的扩散模型伪数值求解器

Diffusers 中 PNDMScheduler 深度解析:基于 Runge-Kutta 与线性多步法的扩散模型伪数值求解器 Diffusers 中 PNDMScheduler 深度解析基于 Runge-Kutta 与线性多步法的扩散模型伪数值求解器【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusersPNDMSchedulerPseudo Numerical Methods for Diffusion Models是 Hugging Face Diffusers 调度器家族中一款经典的高阶 ODE 求解器它通过 Runge-Kutta 法与线性多步法两种伪数值积分技术的组合用远少于训练步数的推理步数高质量地逼近扩散模型的反向去噪轨迹。本文以 PNDMScheduler API 参考文档 为主线结合其在 scheduling_pndm.py 中的完整实现与 调度器测试用例逐项拆解其全部构造参数、时间步构造逻辑、PRK/PLMS 两步式步进算法与底层递推公式并给出可直接运行的接入示例帮助你在 Stable Diffusion、ControlNet 等管线中熟练使用与替换该调度器。PNDMScheduler 是什么PNDMScheduler是 Diffusers 对扩散模型伪数值方法Pseudo Numerical Methods for Diffusion Models论文编号 2202.09778的实现。与一阶采样器如 DDPM相比它使用更高级的 ODE 积分技术——Runge-Kutta 方法与线性多步方法linear multi-step method——来提高采样精度与收敛速度。其原始参考实现来自 k-diffusion 项目中的 PLMS 采样器文档中已注明出处而在本仓库中该实现以 scheduling_pndm.py 的形式完整落地代码头部同时声明了它对ermongroup/ddim仓库思路的借鉴。从类继承关系看PNDMScheduler继承自SchedulerMixin与ConfigMixin见 scheduling_pndm.py 第 80 行因此它天然具备所有调度器共有的通用能力配置保存与加载通过save_config/from_pretrained将全部构造参数序列化为调度器配置文件scheduler_config.json测试代码PNDMSchedulerTest.check_over_configs专门验证了保存配置→重新加载→输出一致这一行为与模型、管线组合所有构造参数由register_to_config自动注册可通过scheduler.config随时读取可替换性在 schedulers/init.py 中通过_import_structure[scheduling_pndm] [PNDMScheduler]导出用户可直接from diffusers import PNDMScheduler。在库内PNDMScheduler被广泛用作 Stable Diffusion 系列管线的可替换调度器之一例如 pipeline_stable_diffusion.py 及 ControlNet 系列管线均将其与DDIMScheduler、LMSDiscreteScheduler并列作为scheduler参数的合法取值pipeline_controlnet_blip_diffusion.py 更是直接以PNDMScheduler作为默认调度器类型。此外仓库还提供其简化变体IPNDMScheduler四阶改进伪线性多步调度器见 scheduling_ipndm.py。构造参数详解PNDMScheduler的构造函数签名scheduling_pndm.py 第 119-132 行定义了以下参数全部带有默认值均可通过scheduler_config.json持久化参数默认值合法取值作用num_train_timesteps1000int训练时使用的扩散步数也是 β 序列的长度beta_start0.0001floatβ 序列的起始值beta_end0.02floatβ 序列的结束值beta_schedulelinearlinear/scaled_linear/squaredcos_cap_v2β 调度策略trained_betasNonenp.ndarray或List[float]直接传入 β 数组绕过beta_start/beta_endskip_prk_stepsFalsebool是否跳过 PRK 预热步直接进入 PLMS 阶段set_alpha_to_oneFalsebool最后一步无前一步 α 累积值时是否将其固定为1prediction_typeepsilonepsilon/v_prediction模型输出类型timestep_spacingleadinglinspace/leading/trailing时间步的缩放方式steps_offset0int叠加到推理时间步上的偏移量beta_schedule 三种策略的底层实现源码scheduling_pndm.py 第 133-144 行中 β 序列的生成逻辑如下linear在[beta_start, beta_end]区间上均匀取num_train_timesteps个点即torch.linspace(beta_start, beta_end, num_train_timesteps)这也是 PNDM 测试默认使用的调度scaled_linear先对beta_start、beta_end开平方后再线性插值、最后平方还原linspace(beta_start**0.5, beta_end**0.5, ...) ** 2。源码注释明确指出该调度对潜空间扩散模型latent diffusion model非常特定Stable Diffusion 系列模型常使用它squaredcos_cap_v2即 Glide 余弦调度调用同文件顶部定义的betas_for_alpha_bar辅助函数scheduling_pndm.py 第 27-77 行。该函数将给定的alpha_bar连续函数离散化为 β 序列beta min(1 - alpha_bar(t2)/alpha_bar(t1), max_beta)并支持cosine、exp、laplace三种alpha_transform_type默认max_beta0.999以避免数值不稳定。构造完成后调度器随即派生出一系列内部量alphas 1 - betas、alphas_cumprod cumprod(alphas)α 的累积乘积是后续所有递推公式的核心并根据set_alpha_to_one决定final_alpha_cumprod取1.0还是alphas_cumprod[0]同时初始化init_noise_sigma 1.0初始噪声分布标准差、pndm_order 4四阶方法以及一组运行状态变量cur_model_output、counter、cur_sample、ets。参数语义补充说明trained_betas若传入则完全忽略beta_schedule的分支逻辑直接以torch.tensor(trained_betas, dtypetorch.float32)作为 β 序列适用于从预训练模型中迁移精确噪声调度表的场景skip_prk_steps对某些模型如 Stable Diffusion可以跳过论文规定的 PRK 预热步以获得更好的效果。源码注释指出此时实现基于 crowsonkb 的 PLMS 采样器。set_timesteps中会据此分支构造不同的prk_timesteps/plms_timestepstimestep_spacing三种取值对应论文Common Diffusion Noise Schedules and Sample Steps are Flawed编号 2305.08891表 2 的标注用于解决训练 1000 步、推理仅 20~50 步时时间步如何缩放的歧义prediction_typeepsilon表示模型预测扩散过程的噪声v_prediction表示预测速度场v参考 Imagen Video 论文 2.4 节常用于视频扩散等模型。时间步构造set_timesteps在推理前必须调用set_timesteps(num_inference_steps, deviceNone)scheduling_pndm.py 第 172-230 行为扩散链构建离散时间步序列。若跳过该调用直接执行step_prk/step_plms源码会抛出ValueErrorNumber of inference steps is None, you need to run set_timesteps after creating the scheduler。timestep_spacing的三种模式决定了基础时间步_timesteps的取法linspace在[0, num_train_timesteps - 1]上等间隔取num_inference_steps个点后四舍五入leading默认step_ratio num_train_timesteps // num_inference_steps时间步为arange(0, num_inference_steps) * step_ratio再叠加steps_offset。源码注释特别提醒转为整数以避免推理步数为 3 的幂时出现索引错误——对应测试test_pow_of_3_inference_steps以 27 步验证trailing从num_train_timesteps开始以-step_ratio递减取整后再整体减1使时间步偏向末尾。随后set_timesteps依据skip_prk_steps拆分两阶段时间步不跳过 PRK取最后pndm_order4个基础时间步重复并叠加半程偏移构造出prk_timestepsPRK 预热段其余时间步构成plms_timesteps跳过 PRKprk_timesteps置空数组plms_timesteps直接由_timesteps拼接得到。最终timesteps concatenate([prk_timesteps, plms_timesteps])并转为torch.Tensor移到指定设备同时清空运行状态ets、counter、cur_model_output。测试 test_steps_offset 中给出了steps_offset1、10 步推理时的完整时间步序列断言[901, 851, 851, 801, 801, 751, 751, 701, 701, 651, ...]可直观看到 PRK 段时间步的成对重复特征。两步式步进step 的分发逻辑与 PRK/PLMS 权重调度器的核心入口是step(model_output, timestep, sample, return_dictTrue)scheduling_pndm.py 第 232-263 行它依据内部计数器与 PRK 时间步数量自动分发if self.counter len(self.prk_timesteps) and not self.config.skip_prk_steps: return self.step_prk(...) else: return self.step_plms(...)即推理初期先跑若干轮 Runge-KuttaPRK预热积累历史输出后切换为线性多步PLMS加速。step_prk四阶 Runge-Kutta 预热step_prkscheduling_pndm.py 第 265-322 行以 4 次前向传播为一轮按counter % 4累加加权系数形成经典的 RK4 组合第 1 次% 4 0记录cur_sample累加1/6 * model_output并缓存model_output到ets第 2、3 次分别累加1/3 * model_output第 4 次累加1/6 * model_output后合成最终model_output并清零累加器。其中prev_timestep由timestep - (num_train_timesteps // num_inference_steps // 2)计算即每次推进半个步长体现了 PRK 阶段在粗粒度时间步之间的细粒度探测。step_plms线性多步外推step_plmsscheduling_pndm.py 第 324-392 行维护一个最多 4 项的历史输出窗口ets按下述 Adams 型多步公式外推当前模型输出已缓存历史数外推公式1 项且counter 0直接用model_output1 项且counter 1(model_output ets[-1]) / 22 项(3 * ets[-1] - ets[-2]) / 23 项(23 * ets[-1] - 16 * ets[-2] 5 * ets[-3]) / 124 项及以上(55 * ets[-1] - 59 * ets[-2] 37 * ets[-3] - 9 * ets[-4]) / 24注意step_plms有前置约束若skip_prk_stepsFalse且ets历史不足 3 项会抛出ValueError提示该调度器必须在 PRK 模式运行至少 12 次迭代后才能调用对应测试test_inference_plms_no_past_residuals验证了此异常。step_plms每步推进一个完整步长prev_timestep timestep - num_train_timesteps // num_inference_steps首步有特殊处理以保持轨迹对齐。底层递推公式_get_prev_sample 与 add_noise无论 PRK 还是 PLMS最终都调用私有方法_get_prev_sample(sample, timestep, prev_timestep, model_output)scheduling_pndm.py 第 409-458 行计算前一步样本x_(t-δ)其核心是 PNDM 论文公式 (9)alpha_prod_t alphas_cumprod[t] alpha_prod_t_prev alphas_cumprod[t-δ] t-δ 0 时取 final_alpha_cumprod sample_coeff sqrt(alpha_prod_t_prev / alpha_prod_t) model_output_denom sqrt(alpha_prod_t * (1 - alpha_prod_t_prev)) sqrt(alpha_prod_t * (1 - alpha_prod_t) * alpha_prod_t_prev) prev_sample sample_coeff * sample - (alpha_prod_t_prev - alpha_prod_t) * model_output / model_output_denom源码注释详细推导了sample_coeff的恒等变换(α_(t−δ) − α_t) / (sqrt(α_t) * (sqrt(α_(t−δ)) sqrt(α_t))) sqrt(α_(t−δ) / α_t)。在计算前若prediction_type v_prediction会先把模型输出转换为噪声形式model_output sqrt(alpha_prod_t) * model_output sqrt(beta_prod_t) * sample其他非法类型直接抛出ValueError。另一个可复用方法是add_noisescheduling_pndm.py 第 461-501 行从DDPMScheduler复制而来它按前向扩散公式sqrt(alpha_prod) * original sqrt(1 - alpha_prod) * noise向样本注入指定时间步的噪声并会自动把alphas_cumprod与时间步搬运到与样本相同的设备与数据类型避免推理时反复的 CPU↔GPU 拷贝。此外scale_model_input直接原样返回输入因为 PNDM 无需按时间步缩放模型输入以保证与其他需要缩放的调度器接口兼容。在管线中接入 PNDMScheduler接入方式一直接替换已有管线的调度器。例如在 Stable Diffusion 系列管线中将scheduler换为 PNDMScheduler 实例from diffusers import PNDMScheduler, StableDiffusionPipeline pipe StableDiffusionPipeline.from_pretrained(模型仓库标识) # 替换为 PNDMScheduler显式配置 β 调度与预测类型 pipe.scheduler PNDMScheduler.from_config( pipe.scheduler.config, beta_schedulescaled_linear, prediction_typeepsilon, timestep_spacingleading, skip_prk_stepsTrue, # 跳过 PRK 预热直接使用 PLMS )from_config会继承原调度器配置中兼容的字段如num_train_timesteps未被覆盖的参数沿用默认值是社区实践中替换调度器的标准做法。接入方式二在自定义去噪循环中手动驱动两步式采样与测试full_loop的模式一致参考 test_scheduler_pndm.py 第 95-113 行import torch from diffusers import PNDMScheduler scheduler PNDMScheduler( num_train_timesteps1000, beta_start0.0001, beta_end0.02, beta_schedulelinear, ) num_inference_steps 50 scheduler.set_timesteps(num_inference_steps) # 初始化纯噪声样本init_noise_sigma 1.0 sample torch.randn((1, 4, 64, 64)) # model 为已训练好的去噪网络返回模型输出预测噪声或 v for t in scheduler.timesteps: model_output model(sample, t) sample scheduler.step(model_output, int(t), sample).prev_samplestep返回的SchedulerOutput定义于 scheduling_utils.py 的SchedulerOutput数据类只包含一个字段prev_sample——即上一时间步的样本x_(t-1)应作为下一轮模型输入若传入return_dictFalse则返回(prev_sample,)元组。测试test_step_shape验证了无论走 PRK 还是 PLMS 分支输出张量形状都与输入样本一致。测试与质量保障test_scheduler_pndm.py 是理解 PNDM 行为的最佳辅助材料其覆盖点包括配置往返一致性check_over_configs/check_over_forward验证save_config后from_pretrained重建的调度器其step_prk/step_plms输出与原实例误差小于1e-5确定性端到端结果test_full_loop_no_noise、test_full_loop_with_v_prediction、test_full_loop_with_set_alpha_to_one等对 10 步完整去噪循环的结果做数值快照断言例如 epsilon 模式下结果和约为198.1318一旦算法或系数改动导致输出漂移测试会立即失败——这相当于 PNDM 递推公式的黄金样本回归基准参数组合矩阵test_betas、test_schedules、test_prediction_type、test_timesteps、test_steps_offset、test_time_indices、test_inference_steps覆盖不同 β 范围、调度类型、预测类型、时间步数量与偏移量组合边界与异常test_pow_of_3_inference_steps回归了推理步数为 3 的幂时的索引问题test_inference_plms_no_past_residuals验证未经过 PRK 预热就调用step_plms会正确抛出ValueError。小结PNDMScheduler 是 Diffusers 中兼具历史地位与实用价值的经典调度器它以论文公式 (9) 为递推核心用四阶 Runge-KuttaPRK预热 四阶线性多步PLMS外推的组合在少量推理步数下保持较高的采样质量。其全部行为由 scheduling_pndm.py 中可读的系数与分支逻辑决定并通过 test_scheduler_pndm.py 的数值快照得到严格保障。理解它的参数体系尤其是beta_schedule、prediction_type、timestep_spacing与skip_prk_steps与两步式步进机制能让你在 Stable Diffusion、ControlNet、AnimateDiff 等管线中精准地将其接入、替换与调试若需要更精简的无 PRK 变体还可以关注同目录下的 IPNDMScheduler。【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表