【Bug已解决】When use UniPCMultistepScheduler as the ODE solver, and input different num_inference_steps

【Bug已解决】When use UniPCMultistepScheduler as the ODE solver, and input different num_inference_steps in StableDiffusionXLPipeline, the first inference step is wrong 解决方案一、现象长什么样用UniPCMultistepScheduler作为 SDXL 的采样器切换不同的num_inference_steps时生成的图会随步数变化出现系统性偏移——尤其第一帧step 0明显不对导致整体画面构图/光照和同 prompt 其他采样器如 DPM不一致from diffusers import StableDiffusionXLPipeline, UniPCMultistepScheduler pipe StableDiffusionXLPipeline.from_pretrained(stabilityai/sdxl-base-1.0) pipe.scheduler UniPCMultistepScheduler.from_config(pipe.scheduler.config) for steps in (20, 30, 50): out pipe(a photo of a mountain, num_inference_stepssteps).images[0] # steps20 和 steps50 的图主体位置/光照明显不同应只差细节不该差构图进一步 dump 第一步去噪前的latentsprint(latents_step0[:3]) # steps20: 某种分布 print(latents_step0[:3]) # steps50: 另一种分布且和 DPM 的 step0 都对不上确认UniPC 在num_inference_steps改变时第一步step 0用到的 timestep / 历史缓存错位导致第一步去噪方向错后续步骤被带偏。现象总结UniPCMultistepScheduler是「多步」ODE 求解器它靠保存前几步的模型输出来做校正当num_inference_steps变化时timestep 调度与历史缓存的初始化/索引没同步好第一步就用了错误的 timestep 或错误的历史项导致首步去噪错、整图偏移。二、背景UniPCUnified Predictor-Corrector是多步求解器它不只看当前步的模型输出还复用前面若干步的输出来做更高阶的校正从而用更少步数达到好结果。关键机制它维护一个model_outputs历史列表前几步step order用「单步/低阶」模式等历史攒够再升到多步每步的 timestept来自set_timesteps(num_inference_steps)生成的 schedule。bug 的根因常出在set_timesteps生成 schedule 后第一步的step_index/ 起始 timestep 计算依赖于「默认步数」或「上一次调用的残留状态」当num_inference_steps改变时要么timesteps[0]取错比如取了上次的缓存索引导致第一步在错误时刻去噪要么model_outputs历史没清空第一步的校正项引用了上一个num_inference_steps留下的旧输出方向直接错。因为后续步骤都基于第一步的结果首步错 → 整图错但 loss/形状都正常肉眼才看得出。三、根因根因三点set_timesteps改变步数时未重置历史缓存model_outputs列表在多次set_timesteps调用间残留第一步校正引用了旧步数的历史 → 首步错。第一步的 timestep/索引计算依赖残留的step_indexstep()里的step_index没在set_timesteps时复位为 0导致第一次step用了非 0 的索引去取 timestep。warmup 阶段step order未强制单步UniPC 应在历史不足时用单步 predictor但若实现里第一步就尝试多步校正历史空会越界或引用默认值 → 首步方向错。本质多步求解器的「历史缓存 步索引」状态在num_inference_steps变化时未干净复位导致首步用了错 timestep / 错历史整图偏移。四、最小可运行复现用标准库复现「改变步数时历史缓存残留导致首步用错」class BuggyUniPC: def __init__(self): self.model_outputs [] # 历史缓存跨 set_timesteps 残留 self.step_index 0 def set_timesteps(self, num_steps): self.timesteps list(range(num_steps, 0, -1)) # 简化 schedule # 错误没清空 model_outputs也没复位 step_index # if self.model_outputs: ... 残留 def step(self, model_output): # 第一步就尝试多步校正引用历史可能来自上一次 set_timesteps if self.step_index 0 and self.model_outputs: corrected model_output self.model_outputs[-1] # 用旧历史 - 错 else: corrected model_output self.model_outputs.append(model_output) self.step_index 1 return corrected s BuggyUniPC() s.set_timesteps(20) s.model_outputs [999] # 模拟上一次调用的残留 first s.step(1.0) # 第一步引用了残留 999 - 错 print(first step , first) # 1000.0明显错应是 1.0 附近复现「正确」在set_timesteps里加self.model_outputs.clear(); self.step_index 0第一步就不引用残留结果正确。五、解决方案第一层最小直接修复最小修复在set_timesteps里强制清空历史缓存 复位step_index并保证 warmup 首步用单步 predictorimport torch class FixedUniPCMultistepScheduler: def __init__(self, num_train_timesteps1000, solver_order2): self.num_train_timesteps num_train_timesteps self.solver_order solver_order self.model_outputs [] self.step_index 0 def set_timesteps(self, num_inference_steps50, devicecpu): # 关键每次 set 都干净复位状态 self.model_outputs.clear() self.step_index 0 self.timesteps torch.linspace( self.num_train_timesteps, 0, num_inference_steps 1 ).to(device).long() self.num_inference_steps num_inference_steps def step(self, model_output, timestep, sample): # warmup历史不足 solver_order 时用单步 predictor if len(self.model_outputs) self.solver_order - 1: prev_sample self._predictor_single(model_output, timestep, sample) else: prev_sample self._predictor_multistep(model_output, timestep, sample) self.model_outputs.append(model_output) self.step_index 1 return prev_sample def _predictor_single(self, model_output, timestep, sample): # 单步不引用历史 return sample model_output * (timestep / 1000.0) def _predictor_multistep(self, model_output, timestep, sample): # 多步用历史此时历史已是正确的当前步数累积 return sample model_output * (timestep / 1000.0)这样set_timesteps每次都清空历史 复位索引首步必走单步 predictor不受上次num_inference_steps影响。六、解决方案第二层结构性改进把「UniPC 状态复位 warmup 契约」收敛成一个 dataclass 单一真源from dataclasses import dataclass, field from typing import List dataclass(frozenTrue) class UniPCMultistepPolicy: UniPCMultistepScheduler 状态管理的单一真源。 # set_timesteps 必须复位的内部状态 reset_fields: tuple (model_outputs, step_index, lower_order_nums) # warmup历史不足 solver_order-1 时强制单步 warmup_rule: str use_single_step_until_history_full # 第一步是否允许多步校正 allow_multistep_on_first_step: bool False # solver 阶数 solver_order: int 2 def on_set_timesteps(self, scheduler) - None: for f in self.reset_fields: if f model_outputs: scheduler.model_outputs.clear() elif f step_index: scheduler.step_index 0 else: setattr(scheduler, f, 0) def should_use_single_step(self, scheduler) - bool: if self.allow_multistep_on_first_step: return False return len(scheduler.model_outputs) self.solver_order - 1 def validate_first_step(self, scheduler) - List[str]: problems [] if scheduler.step_index ! 0: problems.append(set_timesteps 后 step_index 未复位为 0) if scheduler.model_outputs: problems.append(set_timesteps 后 model_outputs 未清空) return problemsstep里用policy.should_use_single_step(self)决定单步/多步on_set_timesteps保证复位validate_first_step用于测试。七、解决方案第三层断言 / CI 守护用 pytest 把「步数变化后首步正确 历史复位 warmup 单步」固化成回归import torch import pytest from mylib.unipc import FixedUniPCMultistepScheduler, UniPCMultistepPolicy POLICY UniPCMultistepPolicy() def test_set_timesteps_resets_state(): s FixedUniPCMultistepScheduler() s.set_timesteps(20) s.model_outputs [999] # 模拟残留 s.set_timesteps(50) # 再次 set 应复位 problems POLICY.validate_first_step(s) assert problems [], 状态未复位:\n \n.join(problems) def test_first_step_single_step_no_history(): s FixedUniPCMultistepScheduler() s.set_timesteps(30) assert POLICY.should_use_single_step(s) is True # 首步必须单步 def test_different_steps_same_first_step_direction(): # 不同 num_inference_steps 下首步去噪方向应一致不依赖旧历史 results [] for steps in (20, 30, 50): s FixedUniPCMultistepScheduler() s.set_timesteps(steps) out s.step(model_outputtorch.tensor(1.0), timesteptorch.tensor(900.0), sampletorch.tensor(0.0)) results.append(out.item()) # 首步都是 sample output*(t/1000)应与步数无关 assert results[0] results[1] results[2] def test_multistep_after_warmup(): s FixedUniPCMultistepScheduler(solver_order2) s.set_timesteps(30) # 喂两步历史后第三步应进入多步 s.step(torch.tensor(1.0), torch.tensor(900.0), torch.tensor(0.0)) s.step(torch.tensor(1.0), torch.tensor(800.0), torch.tensor(0.0)) assert POLICY.should_use_single_step(s) is False def test_no_cross_step_contamination(): s FixedUniPCMultistepScheduler() s.set_timesteps(20); s.step(torch.tensor(1.0), torch.tensor(900.0), torch.tensor(0.0)) s.set_timesteps(50) assert s.model_outputs [], 切换步数后历史必须清空CI 把test_set_timesteps_resets_state与test_different_steps_same_first_step_direction作为 UniPC 的必过项要求「任何num_inference_steps变化都必须干净复位首步方向与之无关」。八、排查清单UniPC 换步数首步错按顺序查不同num_inference_steps下首步去噪方向是否一致不一致说明历史缓存残留。set_timesteps是否清空model_outputs没清空第一步校正会引用上一次调用的旧输出。step_index是否在set_timesteps时复位为 0没复位第一步用错 timestep 索引。第一步是否走了多步校正历史空warmup 必须单步否则越界/引用默认。生成图是否「只差细节、不该差构图」差构图就是首步错被后续放大的典型症状。是否在多次set_timesteps间复用同一 scheduler 实例复用必须保证每次 set 干净复位。九、小结「When use UniPCMultistepScheduler ... the first inference step is wrong」本质是多步 ODE 求解器的「历史缓存 步索引」状态在num_inference_steps变化时未干净复位或 warmup 首步误用多步校正导致首步用了错 timestep / 错历史整图偏移。第一层在set_timesteps强制清空model_outputs 复位step_index并让 warmup 首步用单步 predictor第二层把状态复位与 warmup 契约收敛到UniPCMultistepPolicy单一真源第三层用 pytest 守住「步数变化后首步方向一致、历史清空、warmup 单步」。通用教训**任何多步/历史依赖的求解器必须在「重新初始化调度」时干净复位全部状态并把「首步用单步、历史攒够再升阶」作为不变量否则换参数就会系统性偏移。