【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_steps=steps).images[0] # steps=20 和 steps=50 的图，主体位置/光照明显不同（应只差细节，不该差构图）进一步 dump 第一步去噪前的latents：
print(latents_step0[:3]) # steps=20: 某种分布 print(latents_step0[:3]) # steps=50: 另一种分布，且和 DPM++ 的 step0 都对不上确认：UniPC 在num_inference_steps改变时，第一步（step 0）用到的 timestep / 历史缓存错位，导致第一步去噪方向错，后续步骤被带偏。
现象总结：UniPCMultistepScheduler是「多步」ODE 求解器，它靠保存前几步的模型输出来做校正；当num_inference_steps变化时，timestep 调度与历史缓存的初始化/索引没同步好，第一步就用了错误的 timestep 或错误的历史项，导致首步去噪错、整图偏移。
二、背景
UniPC（Unified Predictor-Corrector）是多步求解器：它不只看当前步的模型输出，还复用前面若干步的输出来做更高阶的校正，从而用更少步数达到好结果。关键机制：
- 它维护一个
model_outputs历史列表； - 前几步（step < order）用「单步/低阶」模式，等历史攒够再升到多步；
- 每步的 timestep
t来自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_index：step()里的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 首步用单步 predictor：
import torch class FixedUniPCMultistepScheduler: def __init__(self, num_train_timesteps=1000, solver_order=2): 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_steps=50, device="cpu"): # 关键：每次 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(frozen=True) 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_output=torch.tensor(1.0), timestep=torch.tensor(900.0), sample=torch.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_order=2) 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 单步」。通用教训：**任何多步/历史依赖的求解器，必须在「重新初始化调度」时干净复位全部状态，并把「首步用单步、历史攒够再升阶」作为不变量，否则换参数就会系统性偏移。