【Bug已解决】A bug on the set_timesteps function of the FlowMatchEulerDiscreteScheduler 解决方案一、现象长什么样FlowMatchEulerDiscreteScheduler是 FLUX 等流匹配模型用的调度器它在[0, 1000]或归一化[0, 1]上定义连续的 sigma/timestep。很多用户在自定义采样循环里手动调set_timesteps或者改了num_inference_steps后复跑会撞到两类异常。第一类最后一步去噪「没生效」输出比默认配置糊一点且没有任何报错from diffusers import FlowMatchEulerDiscreteScheduler scheduler FlowMatchEulerDiscreteScheduler() scheduler.set_timesteps(num_inference_steps28) print(len(scheduler.timesteps)) # 期望 28 print(len(scheduler.sigmas)) # 期望 29含末端点 0 # 实际常见错误timesteps 长度 sigmas 长度末端点 sigma0 被丢掉第二类在自定义step循环里索引越界IndexError: index 28 is out of bounds for dimension 0 with size 28出现在self.sigmas[self.step_index 1]这种取「下一步 sigma」的代码里因为sigmas比预期少一个元素。现象总结set_timesteps生成的timesteps与sigmas长度不一致差一导致 step 内部按step_index1取 sigma 时越界或最末步的sigma0完全去噪被悄悄跳过结果图偏糊。二、背景流匹配调度器的set_timesteps一般要产出两组量timesteps长度为num_inference_steps是模型每一步实际输入的去噪时刻sigmas长度为num_inference_steps 1比timesteps多一个「终点 sigma 0」表示第num_inference_steps步之后噪声已完全移除。step()在计算prev_sample sample (sigma_next - sigma) * ...时需要同时用到「当前 sigma」和「下一步 sigma」所以它依赖sigmas比timesteps多一个元素。一旦set_timesteps把sigmas也只生成num_inference_steps个漏了末尾的 0step_index走到最后一步时sigmas[step_index 1]就越界如果实现改成「不取末尾、用clamp兜底」那最后一步的去噪幅度算错图就糊。此外set_timesteps还涉及device它内部创建timesteps/sigmas张量时若不显式搬到调用方指定的device自定义循环里sample在 GPU、timesteps在 CPU会出 device mismatch。三、根因根因三点sigmas长度比timesteps少一set_timesteps在生成sigmas时把「端点 0」和「起点 1000」都塞进了timesteps的同样长度没有预留1的终点。或者它用sigmas np.linspace(...)只取了num_inference_steps个点漏掉末尾 0。step_index初始化与推进未联动set_timestepsset_timesteps没把self.step_index 0复位复跑时step_index残留在上一次的尾巴配合长度差一就立刻越界。device 未跟随set_timesteps不在签名里接受device内部默认 CPU 创建张量自定义 GPU 循环调用scheduler.scale_model_input(sample, t)时 device 不一致。本质set_timesteps没有保证「len(sigmas) len(timesteps) 1」且「状态step_index、device在每次 set 时干净复位」这条不变量。四、最小可运行复现用 numpy 复现「sigmas 比 timesteps 少一」这个长度错配import numpy as np class BuggyFlowMatchScheduler: def set_timesteps(self, num_inference_steps28, devicecpu): # 错误sigmas 只生成 num_inference_steps 个点漏掉末尾 0 sigmas np.linspace(1.0, 0.0, num_inference_steps) # 长度 28 self.sigmas sigmas self.timesteps (sigmas * 1000).round().astype(long) # 长度 28 # 没有复位 step_index if not hasattr(self, step_index): self.step_index 0 def step(self, model_output, sample): sigma self.sigmas[self.step_index] sigma_next self.sigmas[self.step_index 1] # 最后一步越界 self.step_index 1 return sample (sigma_next - sigma) * model_output s BuggyFlowMatchScheduler() s.set_timesteps(28) try: sample np.zeros(4) for _ in range(28): sample s.step(np.ones(4), sample) except IndexError as e: print(IndexError, e) # index 28 out of bounds for size 28把sigmas np.linspace(1.0, 0.0, num_inference_steps 1)改回长度 29并在set_timesteps开头加self.step_index 0问题消失。五、解决方案第一层最小直接修复最小修复保证sigmas长度 num_inference_steps 1并在set_timesteps入口复位step_index同时接受一个device把张量搬过去import torch import numpy as np class FixedFlowMatchScheduler: def __init__(self): self.step_index 0 def set_timesteps(self, num_inference_steps28, devicecpu, sigma_min0.0, sigma_max1.0): self.step_index 0 # 关键每次 set 复位 # 关键长度多一个含末端点 sigma_min0 sigmas np.linspace(sigma_max, sigma_min, num_inference_steps 1) self.sigmas torch.from_numpy(sigmas).to(devicedevice, dtypetorch.float32) self.timesteps (self.sigmas * 1000).round().long().to(devicedevice) self.num_inference_steps num_inference_steps def step(self, model_output, sample): sigma self.sigmas[self.step_index] sigma_next self.sigmas[self.step_index 1] # 现在安全 prev sample (sigma_next - sigma) * model_output self.step_index 1 return prev这样len(sigmas) len(timesteps) 1恒成立step_index从 0 干净推进最后一步也能正确取到sigma_next 0完成完全去噪。六、解决方案第二层结构性改进把「长度不变量 状态复位 device」收敛成一个 dataclass 单一真源让调度器的set_timesteps永远从它取约束from dataclasses import dataclass, field dataclass(frozenTrue) class FlowMatchTimestepPolicy: FlowMatchEulerDiscreteScheduler 时间步约定的单一真源。 # sigmas 必须比 timesteps 多一个末端点 sigma_length: str num_inference_steps 1 sigma_endpoint_min: float 0.0 sigma_endpoint_max: float 1.0 # set_timesteps 必须复位的内部状态 reset_state_fields: tuple (step_index,) # 是否把张量搬到指定 device respect_device: bool True # 归一化系数timesteps sigmas * scale timestep_scale: float 1000.0 def expected_sigma_len(self, num_steps: int) - int: return num_steps 1 def build_sigmas(self, num_steps: int): import numpy as np, torch return torch.from_numpy( np.linspace(self.sigma_endpoint_max, self.sigma_endpoint_min, self.expected_sigma_len(num_steps)) ).to(dtypetorch.float32) def build_timesteps(self, sigmas: torch.Tensor): return (sigmas * self.timestep_scale).round().long() class FlowMatchSchedulerV2: def __init__(self, policy: FlowMatchTimestepPolicy FlowMatchTimestepPolicy()): self.policy policy self.step_index 0 def set_timesteps(self, num_inference_steps28, devicecpu): # 1) 复位状态 for f in self.policy.reset_state_fields: setattr(self, f, 0) # 2) 按 policy 构造长度恒为 num1 self.sigmas self.policy.build_sigmas(num_inference_steps).to(devicedevice) self.timesteps self.policy.build_timesteps(self.sigmas).to(devicedevice) self.num_inference_steps num_inference_steps任何改动set_timesteps的 PR 都只动FlowMatchTimestepPolicy长度不变量由expected_sigma_len强制。七、解决方案第三层断言 / CI 守护用 pytest 把「长度不变量 状态复位 末步去噪」固化成回归import torch import pytest from mylib.flowmatch import FlowMatchSchedulerV2, FlowMatchTimestepPolicy POLICY FlowMatchTimestepPolicy() def test_sigma_length_invariant(): s FlowMatchSchedulerV2(POLICY) s.set_timesteps(28, devicecpu) assert len(s.sigmas) len(s.timesteps) 1, sigmas 必须比 timesteps 多一个末端点 assert s.sigmas[-1].item() 0.0, 末端点 sigma 必须为 0完全去噪 assert s.sigmas[0].item() 1.0 def test_step_index_reset_on_repeat(): s FlowMatchSchedulerV2(POLICY) s.set_timesteps(4) for _ in range(4): s.step(torch.ones(4), torch.zeros(4)) assert s.step_index 4 s.set_timesteps(4) # 复跑必须复位 assert s.step_index 0 def test_no_index_out_of_range(): s FlowMatchSchedulerV2(POLICY) s.set_timesteps(28, devicecpu) sample torch.zeros(4) for _ in range(28): # 跑满 28 步不应越界 sample s.step(torch.ones(4), sample) assert sample.shape (4,) assert sample.abs().sum() 0 # 确实做了去噪 def test_device_follows_set(): s FlowMatchSchedulerV2(POLICY) s.set_timesteps(4, devicecuda) assert s.timesteps.device.type cudaCI 把test_sigma_length_invariant与test_no_index_out_of_range作为调度器模块的必过项杜绝再次引入差一错误。八、排查清单FlowMatchEulerDiscreteScheduler异常按顺序查len(scheduler.sigmas) len(scheduler.timesteps) 1不是就差一step 必越界。sigmas[-1]是否为 0不为 0 说明末尾完全去噪点丢失输出偏糊。复跑前set_timesteps是否复位step_index没复位会从上次尾巴继续配合长度差一立刻越界。timesteps/sigmas的device和sample是否一致不一致会 device mismatch。自定义step里取sigmas[step_index 1]时是否守卫「最后一步」正确做法是用长度多一的 sigmas而不是 clamp 兜底。num_inference_steps改变后是否重新set_timesteps只改变量不改调度器内部状态会导致 schedule 与步数错配。九、小结FlowMatchEulerDiscreteScheduler.set_timesteps的「Bug」本质是没有守住「len(sigmas) len(timesteps) 1」且不变量且step_index/device 未在每次 set 时干净复位。第一层在set_timesteps里补上「长度多一的 sigmas step_index 复位 device 跟随」让越界和末步丢去噪消失第二层把长度不变量、状态复位字段、device 策略收敛到FlowMatchTimestepPolicy单一真源第三层用 pytest 守住「长度不变量、状态复位、跑满不越界、device 跟随」。通用教训任何生成「采样时间表」的接口都必须显式保证表与步数、与内部游标的一致性否则差一个元素就是越界或质量退化。