Prhub

#22059 [diffusion] fix FLUX[1,2]

原始 PR 作者 yhyang201 合并时间 2026-04-05 16:03 文件变更 11 提交数 14 评论 2 代码增减 +123 / -88

执行摘要

修复 FLUX.1 和 FLUX.2 模型中 guidance_embedder 权重加载和 guidance_scale 缩放问题,对齐 HuggingFace Diffusers。

根据PR body,FLUX.1-dev和FLUX.2-dev检查点包含trained guidance_embedder权重,但SGLang在guidance_embeds=False时丢弃这些权重。设置guidance_embeds=True是正确的修复,但_build_guidance()函数错误地将guidance_scale乘以1000,导致模型的正弦时间步嵌入层接收到分布外的输入。Diffusers直接传递原始guidance_scale值,因此需要移除缩放以匹配其规范。

建议技术管理者和工程师精读denoising.py中的_build_guidance()修改和flux.py中的配置重构,关注如何通过条件判断区分FLUX.1和FLUX.2的逻辑。这些变更体现了在兼容不同模型版本时的设计权衡,值得学习其模块化处理和一致性对齐策略。

讨论亮点

review中没有实质性讨论,仅有gemini-code-assist[bot]的一条评论指出变更移除了1000.0缩放因子以对齐Diffusers规范,防止out-of-distribution embeddings。无其他reviewer参与争议或反馈,变更直接通过。

实现拆解

实现方案按模块拆解:1) 编码器配置模块:新增flux_2.py定义FLUX.2的Mistral文本编码器配置,并在__init__.py中导出;重构flux.py移除冗余的Flux2MistralTextArchConfig和format_text_input函数,改用build_flux2_text_messages。2) 采样参数模块:在flux.py中更新guidance_scale默认值,FluxSamplingParams为3.5,Flux2SamplingParams为4.0,并新增Flux2SamplingParams类;在registry.py中注册配置时使用Flux2SamplingParams。3) 运行时逻辑模块:在denoising.py中修改_build_guidance()函数,仅对FLUX.1模型保持*1000缩放,FLUX.2则使用原始值;在text_encoding.py中调整attention_mask处理,区分FLUX.1和FLUX.2的编码流程。4) 加载器模块:在component_loader.py中将Flux2PipelineConfig的tokenizer加载从AutoTokenizer改为AutoProcessor以对齐Diffusers。5) 测试模块:在test_sampling_params.py中添加测试验证guidance_scale默认值。

文件 模块 状态 重要度
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py multimodal_gen/pipelines_core modified 8.0
python/sglang/multimodal_gen/configs/pipeline_configs/flux.py multimodal_gen/configs modified 7.0
python/sglang/multimodal_gen/configs/sample/flux.py multimodal_gen/configs modified 6.0
python/sglang/multimodal_gen/configs/models/encoders/flux_2.py multimodal_gen/encoders added 6.0

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

关键符号

_build_guidance build_flux2_text_messages flux_2_preprocess_text format_text_input

评论区精华

对齐 Diffusers 规范 正确性

gemini-code-assist[bot] 评论指出变更移除了 1000.0 缩放因子以对齐 HuggingFace Diffusers,防止模型嵌入层接收到 out-of-distribution 输入。

结论:变更被接受,无进一步争议,直接合并。 · 已解决

风险与影响

技术风险包括:1) 回归风险:修改guidance_scale默认值和_build_guidance()逻辑可能影响现有FLUX模型的生成质量,需依赖准确性测试验证;2) 兼容性问题:denoising.py中条件缩放逻辑需正确区分FLUX.1和FLUX.2,若判断错误可能导致嵌入计算偏差;3) 测试覆盖不足:新增的FLUX.2配置和变更逻辑缺乏充分的单元测试,可能隐藏边缘情况bug;4) 依赖变更:component_loader.py中使用AutoProcessor替代AutoTokenizer,需确保processor与tokenizer接口兼容。

影响范围主要限于扩散模型生成管线中的FLUX系列模型,涉及配置、编码和去噪阶段。影响程度中等:修正了权重加载和嵌入计算,应显著改善图像生成的准确性,并提升与HuggingFace Diffusers的一致性。对用户而言,使用FLUX模型的开发者将获得更可靠的生成结果;对系统,增强了多模态生成组件的健壮性;对团队,展示了对齐开源实现时的代码重构模式。

核心路径变更 默认值调整 测试覆盖需验证 兼容性风险

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

执行摘要

本次PR修复了SGLang中FLUX.1和FLUX.2扩散模型的bug,通过启用guidance_embeds=True加载检查点权重并移除guidance_scale的错误缩放,对齐HuggingFace Diffusers实现。影响范围限于多模态生成模块,提升了生成准确性和一致性,建议团队关注条件缩放逻辑的设计权衡。

功能与动机

PR body明确指出:FLUX.1-dev和FLUX.2-dev检查点包含训练好的guidance_embedder权重,但SGLang在guidance_embeds=False时丢弃这些权重。设置guidance_embeds=True是正确的修复,然而_build_guidance()函数错误地将guidance_scale乘以1000,导致模型的正弦时间步嵌入层接收分布外的输入。Diffusers直接传递原始guidance_scale值,因此需要移除缩放以匹配其规范。

实现拆解

  • 编码器配置模块:新增flux_2.py定义FLUX.2的Mistral文本编码器配置,包括build_flux2_text_messages函数和Flux2MistralTextConfig类;在flux.py中移除冗余的Flux2MistralTextArchConfigformat_text_input,改用build_flux2_text_messages
  • 采样参数模块:更新flux.py中的sampling params,修正默认值:FluxSamplingParams.guidance_scale=3.5,新增Flux2SamplingParams.guidance_scale=4.0;在registry.py中注册配置时使用Flux2SamplingParams
  • 运行时逻辑模块:在denoising.py中修改_build_guidance()函数,通过判断pipeline_config类型,仅对FLUX.1模型保持*1000缩放:
    python if isinstance(self.server_args.pipeline_config, FluxPipelineConfig) and not isinstance(self.server_args.pipeline_config, Flux2PipelineConfig): guidance_val = guidance_val * 1000.0
    text_encoding.py中调整attention_mask处理,区分FLUX.1和FLUX.2的编码流程。

  • 加载器模块:在component_loader.py中将Flux2PipelineConfig的tokenizer加载从AutoTokenizer改为AutoProcessor以对齐Diffusers。

  • 测试模块:在test_sampling_params.py中添加测试验证guidance_scale默认值匹配模型需求。

评论区精华

review中仅有gemini-code-assist[bot]的一条评论:

"This pull request removes the 1000.0 scaling factor from the guidance tensor construction in the denoising stage to align with the HuggingFace Diffusers convention and prevent out-of-distribution embeddings."
无其他讨论,变更直接通过,表明团队对对齐Diffusers的共识。

风险与影响

  • 风险:1) 回归风险:修改默认值和缩放逻辑可能影响FLUX模型生成质量,需准确性测试保障;2) 兼容性:条件缩放逻辑需精确区分FLUX.1/FLUX.2,否则可能导致嵌入计算错误;3) 测试覆盖:新增配置和逻辑变更缺乏充分单元测试,可能隐藏bug。
  • 影响:范围限于扩散模型生成管线,程度中等。修正后提升生成准确性,确保与Diffusers一致性,对用户提供更可靠图像生成,对团队展示代码重构和对齐最佳实践。

关联脉络

从近期历史PR看,

  • 22157 "[CI] No diffusers backend in lora case":同属扩散模型修复,涉及CI测试和与Diffusers的一致性,共享多模态生成模块。

  • 22146 "Isolate spec V1 path in decode post-processing":类似结构重构模式,隔离不同版本路径,与本次PR中区分FLUX.1/FLUX.2逻辑有共通设计思路,反映仓库在演进中注重版本管理和一致性维护。

参与讨论